Derive the qmv fast path K alignment from bits

[jlee5814]

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.

[zcbenz]

Good job finding this!

[zcbenz]

I tried to do some benchmarking but results of GEMV with a small K is too variant to give a definite conclusion to be honest. But it is a bug that we are not using qmv_fast for the said shape so I'm merging this.