Repository navigation
metal: 1-bit qmv_fast use 1 pack/thread for occupancy (~9% on M5) - #3
Conversation
affine qmv_fast set packs_per_thread=2 for all bits except 2-bit, so 1-bit got values_per_thread=64 (x_thread[64], ~256B/thread) -> low occupancy -> 1-bit decode saturates only ~75% of M5 DRAM BW vs ~90/96% for 2/4-bit. Use 1 pack/thread for bits<=2 (values_per_thread=32), matching 2-bit's register footprint. Measured on M5 Pro (distinct-weight DRAM-bound, 2-bit as drift control): 1-bit 24.0 -> 21.9 us/matvec (~9%, 75->82% BW), 2-bit control 33.2->32.9 (0.8% drift); correct, rel_err 2.7e-4. Also makes scale_step_per_thread (=group_size/values_per_thread) well-defined for group_size=32 at 1-bit.
|
End-to-end validation on the test 1-bit model (MLX, M5 Pro; 128 tok, temp 0, 3 reps each):
Both stable across reps. The occupancy fix lifts every projection matvec in the decode path, so the end-to-end gain slightly exceeds the isolated qmv micro-bench (~9%). |
|
Context-window scaling (test 1-bit model, M5 Pro, decode t/s after prompt prefill):
The win persists across context, diluting only slightly (10.7%→8.7%) as the attention KV cache grows to a larger share of decode (the hybrid-attention design keeps that share small — decode falls only ~19% over 30× context). |
|
Losslessness confirmed — bit-exact on the test 1-bit model. Teacher-forced prefill over a fixed passage (33 positions × full vocab), before vs after: The occupancy change (vpt 64→32) reorders the K-loop tiling but yields byte-identical logits, so output is unchanged while decode is +10%. (Note: the 1-bit checkpoint itself currently generates incoherent text — a separate, pre-existing conversion-coherence issue unrelated to this kernel; identical before/after.) |
|
On m4 got 20% speed up so seems great :) |
affine qmv_fast set packs_per_thread=2 for all bits except 2-bit, so 1-bit got values_per_thread=64 (x_thread[64], ~256B/thread) -> low occupancy -> 1-bit decode saturates only ~75% of M5 DRAM BW vs ~90/96% for 2/4-bit. Use 1 pack/thread for bits<=2 (values_per_thread=32), matching 2-bit's register footprint. Measured on M5 Pro (distinct-weight DRAM-bound, 2-bit as drift control): 1-bit 24.0 -> 21.9 us/matvec (~9%, 75->82% BW), 2-bit control 33.2->32.9 (0.8% drift); correct, rel_err 2.7e-4. Also makes scale_step_per_thread (=group_size/values_per_thread) well-defined for group_size=32 at 1-bit.
affine qmv_fast set packs_per_thread=2 for all bits except 2-bit, so 1-bit got values_per_thread=64 (x_thread[64], ~256B/thread) -> low occupancy -> 1-bit decode saturates only ~75% of M5 DRAM BW vs ~90/96% for 2/4-bit. Use 1 pack/thread for bits<=2 (values_per_thread=32), matching 2-bit's register footprint. Measured on M5 Pro (distinct-weight DRAM-bound, 2-bit as drift control): 1-bit 24.0 -> 21.9 us/matvec (~9%, 75->82% BW), 2-bit control 33.2->32.9 (0.8% drift); correct, rel_err 2.7e-4. Also makes scale_step_per_thread (=group_size/values_per_thread) well-defined for group_size=32 at 1-bit.
…eader Ports PrismML-Eng/mlx#3 into the pre-generated metal kernel header that mlx-swift compiles into its metallib: qmv_fast_impl uses packs_per_thread = bits <= 2 ? 1 : 2, so 1-bit (Q1_0) matches 2-bit (vpt=32) for better occupancy. ~+11% 1-bit decode on M5 (bit-exact, validated). Follow-up: bump the Cmlx/mlx submodule to an mlx rev including #3 to make this permanent.
2 participants
What
affine_qmv_fast(the decode/token-gen quantized matvec) setspacks_per_thread = bits == 2 ? 1 : 2, so 1-bit falls into the: 2case and getsvalues_per_thread = pack_factor * 2 = 64→x_thread[64](~256 B/thread). That registerfootprint drops occupancy, so 1-bit decode can't hide DRAM latency and under-saturates
bandwidth (~75% on M5 Pro) while 2-bit/4-bit saturate (~90/96%).
Use 1 pack/thread for
bits <= 2so 1-bit matches 2-bit'svalues_per_thread = 32footprint. (2-bit already used 1 pack, so it is unchanged — it serves as a control below.)
Why it helps
x_threadpressure for 1-bit → better latency hiding.scale_step_per_thread = group_size / values_per_threadwas0forgroup_size = 32at 1-bit (vpt=64); with vpt=32 it is well-defined for g32.Measured (Apple M5 Pro, distinct-weight DRAM-bound matvec, best-of-10)
2-bit is the drift control (untouched by this change):
Correct:
quantized_matmulvs dequantize+matmul rel_err 2.7e-4. The control moving only 0.8%confirms the 1-bit gain is real, not machine noise.
Notes
control, but a quiet machine would tighten the headline number.
bytes at group 64) — a follow-up (int8/shared scales), not this PR.