Route affine qmv_wide by bit-width and batch size#5
Merged
Conversation
Fused greedy speculative-decoding verify as mx.fast.spec_decode_verify: inputs draft_tokens [B,K] + target_logits [B,K+1,V]; outputs n_accepted [B] and committed [B,K+1] (accepted draft prefix ++ corrected token). v1 is a pure op composition (argmax -> longest-match prefix -> corrected), so it needs no new Primitive/Metal/stubs and runs on any backend; it is the intended fallback for a future fused Metal kernel. Greedy only (temperature==0); caller prepends the seed 'last' token. Validated vs a Swift-derived oracle across fp32/fp16/bf16.
Promote the v1 composition to a fast::Custom primitive (SpecDecodeVerify). The op reuses mlx argmax on-device (result stays on GPU), then the fused primitive does only the trivial per-row prefix-match in one thread/row -> no large-V reduction, no tie-break risk. use_fallback routes CPU to the op composition (the correctness oracle) and GPU to the kernel. Adds metal eval_gpu + spec_decode.metal, CMake wiring, and no_gpu/cuda link stubs. Validated on a Metal build: kernel(gpu) and fallback(cpu) both match the oracle across fp32/fp16/bf16. Perf to be judged on an e2e decode loop, not this microbench (the win is the host round-trip/launch collapse, not FLOPs).
The qmv_wide kernel dequantizes a full group of weights into registers and reuses them across the M input rows. That amortization only pays off for 2-bit once three or more rows share a group; at M=2 it breaks even, and for 1-bit the weight traffic is small enough that the per-row dequant dominates and the specialized qmv is faster. Gate the dispatch accordingly: 1-bit and 2-bit M<2 keep the specialized qmv; 2-bit routes to qmv_wide only at M>=3 on gen-15+; fp modes are unchanged. Add 1-bit affine coverage (full sweep + tiny shapes) to test_qmv_wide. Measured on affine 2-bit matvecs across a range of projection shapes: qmv_wide wins at widths >=3 (more so as N grows), with no regression at width <=2 or for 1-bit.
Author
|
@copilot review |
There was a problem hiding this comment.
Pull request overview
Adds selective qmv_wide dispatch alongside the stacked speculative-decoding operation and upstream wide-kernel implementation.
Changes:
- Routes affine
qmv_wideby bit width, batch size, and GPU generation. - Adds affine/FP wide kernels and quantized-matmul coverage.
- Introduces
mx.fast.spec_decode_verifywith Metal and fallback implementations.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
test_spec_verify.py |
Adds standalone verification tests. |
python/tests/test_quantized.py |
Tests wide quantized matvec paths. |
python/src/fast.cpp |
Exposes the Python fast API. |
mlx/fast.h |
Declares the C++ API. |
mlx/fast.cpp |
Implements validation and fallback composition. |
mlx/fast_primitives.h |
Defines the verification primitive. |
mlx/backend/no_gpu/primitives.cpp |
Enables non-GPU fallback. |
mlx/backend/cuda/primitives.cpp |
Enables CUDA fallback. |
mlx/backend/metal/spec_decode.cpp |
Dispatches the Metal kernel. |
mlx/backend/metal/kernels/spec_decode.metal |
Implements fused verification. |
mlx/backend/metal/quantized.cpp |
Adds selective wide dispatch. |
mlx/backend/metal/kernels/quantized.h |
Implements affine wide matvec. |
mlx/backend/metal/kernels/quantized.metal |
Instantiates affine kernels. |
mlx/backend/metal/kernels/fp_quantized.h |
Implements FP wide matvec. |
mlx/backend/metal/kernels/fp_quantized.metal |
Instantiates FP kernels. |
mlx/backend/metal/kernels/CMakeLists.txt |
Builds the new Metal kernel. |
mlx/backend/metal/CMakeLists.txt |
Builds the Metal dispatcher. |
docs/src/python/fast.rst |
Adds API documentation entry. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused | ||
| // primitive does only the tiny per-row prefix-match. Primitive inputs are | ||
| // int32 token arrays, so the kernel needs no large-V reduction. | ||
| auto dft = astype(draft_tokens, int32, s); |
Comment on lines
+993
to
+995
| auto d = inputs[0]; // [B, K] int32 | ||
| auto t = inputs[1]; // [B, K+1] int32 | ||
| auto t_pref = slice(t, Shape{0, 0}, Shape{B, K}, s); |
| @@ -0,0 +1,111 @@ | |||
| """Validate mx.fast.spec_decode_verify against a pure-Python oracle mirroring the | |||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Contents of this branch vs
prismspec_decode_verifyfast op (commitsa29368ea,67f3651d,8e6c9feb,c8eaa12e) — CPU/composition v1 + fused Metal kernel v2,mx.fastplumbing, docs, andtest_spec_verify.py. The core MLX-side primitive DSpark needs.qmv_widecherry-pick (c9a59b06, Add small-batch quantized matvec kernel (qmv_wide) ml-explore/mlx#3764) — small-batch quantized matvec kernel.qmv_widedispatch (bff994ff, this PR's intended change) — see below.The intended change (selective dispatch)
Make the affine
qmv_widedispatch bit-width and batch-size aware instead of routing every affine matvec withM >= 2to the wide kernel.qmv_widedequantizes a full group of weights into registers and reuses them across theMinput rows. That reuse only amortizes when enough rows share a group:M = 2, wins fromM >= 3(largerNbenefits more).qmvis faster — routing 1-bit throughqmv_wideis a measurable regression.use_qmv_widenow takesbitsandM: 1-bit and 2-bitM < 2stay on the specializedqmv; 2-bit routes toqmv_wideonly atM >= 3on gen-15+; fp modes unchanged. Adds 1-bit affine coverage totest_qmv_wide.Testing
test_qmv_widepasses, including new 1-bit cases.qmv_widewins at widths>= 3, no regression at width<= 2or for 1-bit.Planned split before merge
spec_decode_verifyfast op ->prism.qmv_widecherry-pick + selective dispatch, stacked on A.