Skip to content

Derive the qmv fast path K alignment from bits - #3965

Open
jlee5814 wants to merge 1 commit into
ml-explore:mainfrom
jlee5814:fix/qmv-alignment-from-bits
Open

Derive the qmv fast path K alignment from bits#3965
jlee5814 wants to merge 1 commit into
ml-explore:mainfrom
jlee5814:fix/qmv-alignment-from-bits

Conversation

@jlee5814

@jlee5814 jlee5814 commented Aug 1, 2026

Copy link
Copy Markdown

Motivation

qmv and gather_qmv both gate the fast kernels on K % 512 == 0, but 512 is not a free constant: qmv_fast_impl steps its K loop by

pack_factor<bits, 32>() * (bits == 2 ? 1 : 2) * SIMD_SIZE

which is 512 for bits in {2, 3, 4, 5} and 256 for bits in {6, 8}. At 6 and 8 bits the
dispatcher requires twice the alignment the kernel needs, so any K that is an
odd multiple of 256 takes the slow kernel unnecessarily.

The gate is shared across quantization modes: fp_qmv_fast_impl derives its K step the same way, so MXFP8 weights at 8 bits are affected identically.

This has been true since the fast kernel landed; 6 bit was added later without revisiting the constant.

Modifications

Adds qmv_fast_k_alignment(bits), using the existing get_pack_factor from backend/common/quantized.h, which CPU and CUDA already include. Both dispatch sites call it instead of the literal.

Accuracy Tests

No new test needed: test_qmv already parametrizes bits 6 and 8 with a contraction dimension of 256, so roughly 48 subtests move from the slow kernel to the fast one under this change. python/tests/test_quantized.py passes identically on both builds, 32 tests and 2896 subtests.

MXFP8 at K=768 and K=1280 was checked separately against a dequantized reference on both builds since test_qmv is affine only. Both newly eligible shapes change output within fp16 rounding; the already fast K=1024 control is unchanged.

Speed Tests and Profiling

M4 Pro, N=2048, M=1, group_size=64, affine, fp16. Eight interleaved rounds per build, median, microseconds per call. The 4 bit rows are the control, since the alignment is unchanged there.

bits K main patched ratio
4 1024 4.41 4.39 0.995
4 1280 6.15 6.04 0.983
4 1536 5.73 5.71 0.995
8 512 3.49 3.53 1.011
8 768 5.51 4.52 0.820
8 1024 5.55 5.54 0.997
8 1280 7.30 6.53 0.894
8 1536 7.56 7.52 0.993

Coefficient of variation under 2.5% on both builds for every row above. Two further configurations, 4 bit at K=512 and K=768, were the first measured in each process and too noisy to report.

benchmarks/python/gather_qmm_bench.py is unchanged within noise, as expected since its shapes are 4 bit and already 512 aligned.

Shapes newly reaching the fast path pay a one time kernel compilation on first use, since these variants are not instantiated today. Cached thereafter.

@jlee5814
jlee5814 marked this pull request as ready for review August 1, 2026 19:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant