perf(cuda/quant): sm_120/121 qmm CTA tile (M=128) for Blackwell prefill#723
Merged
Conversation
qmm_sm80 used the Ampere CTA tile (M=64), which underutilizes Blackwell SMs: ncu shows 35-47% SM / 52-68% memory throughput on GB10 (sm_121), and 4bit prefill ran at 0.55x the bf16 ceiling. Raising the CTA tile M cap to 128 on sm_120/121 (cc major >= 12) fills the idle SMs. The overlay make_cta_tiler gates the tile_m cap on compute_capability_major(): 128 on Blackwell consumer parts, 64 (stock) on sm_90/sm_80. Small-M decode is unaffected (min(cap, next_power_of_2(m)) stays small; m=1 takes the qmv path). A tile sweep confirmed tile_m is the only safely-tunable axis: wider tile_n or deeper tile_k break the fixed-MMA shared-memory layout (JIT failure). GB10 measurements (arch-gated default vs forced tile_m=64): - llama-3.1-8b-4bit prefill @8192: 2213 -> 3054 tok/s (+38%) - qwen2.5-7b-4bit prefill @8192: 2410 -> 3168 tok/s (+31%) - greedy parity identical; decode unaffected. Findings, dispatch map, and ncu evidence in docs/benchmark_results/. Small-M batched decode still uses per-row qmv (no amortization) and is out of scope here; it is the gap behind #714's --parallel 4 CUDA regression and needs a dedicated small-M kernel upstream.
inureyes
added
type:performance
…(review) pr-reviewer LOW nits: the findings doc said MLX 0.32.0 but the pin is 0.32.1 (57c66cac), and the prefill deltas differed across the PR body, doc, and code comment. Standardize on the same-session paired measurement (arch-gated default tile_m=128 vs forced tile_m=64): +38% llama-3.1-8b-4bit, +31% qwen2.5-7b-4bit @8192, and note the ~5% run-to-run variance of the tile_m=64 baseline. No code behavior change.
This was referenced Jul 10, 2026
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
1 participant
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.
Summary
Raises the CUDA quantized-matmul (
qmm_sm80) CTA tile M cap from 64 to 128 on Blackwell consumer GPUs (sm_120/121,cc_major >= 12), recovering +31-38% prefill throughput on GB10. Part of epic #623.Closes #637
Root cause
qmm_sm80'smake_cta_tilercapstile_mat 64 (Ampere-tuned). On GB10 (sm_121) the Ampere tile leaves the SMs idle: ncu shows 35-47% SM / 52-68% memory throughput, and 4bit prefill runs at 0.55x the bf16 ceiling. (The sm_90 Hopper kernel isn't compiled for arch121, so Blackwell falls to the Ampereqmm_sm80.)Fix
Arch-gated
make_cta_tiler(overlayqmm_sm80.cu):tile_mcap = 128 whendevice.compute_capability_major() >= 12, else 64 (stock). A tile sweep confirmedtile_mis the only safely-tunable axis (widertile_n/ deepertile_kbreak the fixed-MMA shared-memory layout -> JIT failure). AddsMLXCEL_QMM_TILE_M/N/Kenv overrides as a tuning hatch.Benchmark (GB10, full detail committed under
docs/benchmark_results/)qmm_sm80).Scope / not addressed
M*Bin [2,8)) still uses per-row qmv (no weight amortization); thetile_mcap doesn't raise small-M tiles. This is the gap behind perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714's--parallel 4CUDA regression and needs a dedicated small-M kernel (upstream MLX). perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714's--paralleldefault should be made backend-aware until then.Validation
make release-cudabuilds clean. The arch gate uses runtimecompute_capability_major(), so it works for both121and121abuilds (both cc 12.1 at runtime).Acceptance criteria (#637)
docs/benchmark_results/(achieved-vs-ceiling, sm_121 dispatch map, ncu roofline, go/no-go).121/121abuild (arch-gate is runtime cc, not compile-time arch string).