Skip to content

Tune M1 quantized matmul dispatch for six-row large shapes - #4629

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
DeadByDawn101:ravenx/m1-qmm-dispatch-evidence
Oct 6, 2026
Merged

zcbenz merged 2 commits into
ml-explore:mainfrom
DeadByDawn101:ravenx/m1-qmm-dispatch-evidence

Conversation

@DeadByDawn101

Copy link
Copy Markdown
Contributor

Summary

For applegpu_g13s, keep large (D > 4096, O > 4096) affine quantized matmul inputs on the QMV path through six rows by changing the matrix-dispatch boundary from 6 to 7.

This is intentionally limited to the measured M1 g13s architecture and leaves every other architecture/shape threshold unchanged. It does not add or modify a Metal kernel.

The PR also adds a reproducible paired benchmark for the row-dispatch boundary.

Why

On an M1 Max with affine 4-bit/group-64 weights at K=5376, N=8192, the existing six-row matrix dispatch was consistently slower than a QMV control. Five independent 40-trial baseline runs measured chunked-QMV/direct ratios:

0.8717, 0.8930, 0.8947, 0.8767, 0.9152
median: 0.8930

The five-row same-route calibration median was 1.0052, which bounds local paired timing bias.

A final 100-trial run produced:

                 M=5 direct/chunk       M=6 direct/chunk
main             1.4492 / 1.4515 ms     1.9305 / 1.6750 ms
this patch       1.4501 / 1.4479 ms     1.6987 / 1.7290 ms

The direct six-row primitive decreased by about 12% in this isolated run. Across seven 200-trial runs, the baseline six-row chunk/direct median was 0.8807; after the threshold change it was 1.0118, showing the direct full-row QMV route aligns with the QMV control.

These are isolated primitive measurements, not end-to-end model or TTFT claims.

Reproduction

python benchmarks/python/qmm_dispatch_boundary.py \
  --rows 5 6 8 12 --warmups 10 --trials 100 --seed 61

Validation

  • source build: pass
  • python/tests/test_quantized.py: 51 passed, 1 skipped
  • pre-commit on changed files: pass
  • exact-shape direct/chunked numeric comparisons: pass (rtol=0.01, atol=0.06)
  • independent review: pass

The full Python suite ran 943 tests; 4 unrelated test_fast_sdpa.py::test_sdpa_vector_head_dim_512 cases fail on this M1 Max because the kernels request 1024 threads while the device reports 640/832 maximum. The same failure reproduces on an unmodified build of the base revision.

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.

2 participants