perf(cuda/quant): size the qmm_naive tile from the device shared-memory budget, not an sm80 flag #1541

Part of #1536. Phase 1. Baseline from #1538. Independent of #1539; both can land in parallel.

Small, self-contained prefill win that is worth taking on its own merits and stays valuable regardless of whether #1543 ever lands — it improves the scalar path that is the fallback in either case.

Context

mlx/backend/cuda/quantized/qmm/qmm_naive.cu:17-23 sizes the CTA tile:

inline auto make_cta_tiler(int itemsize, int m, int group_size, bool sm80) {
  bool enough_smem = sm80 && itemsize <= 2 && group_size <= 64;
  int tile_m = std::max(16, std::min(64, next_power_of_2(m)));
  int tile_n = enough_smem ? 128 : 64;
  int tile_k = std::max(64, group_size);

sm80 is compute_capability_major() >= 8 (:42), so a Volta part is forced to tile_n = 64 purely by architecture, and enough_smem is the only reason given. The premise does not hold on V100:

• The 128-wide tile at itemsize = 2, group_size = 64 needs roughly (tile_m * tile_k + tile_n * tile_k) * 2 bytes — about 18 KB at tile_m = 16, and ~36 KB in the wider tile_k = 128 case.
• V100 provides 96 KB of shared memory per block through cudaFuncAttributeMaxDynamicSharedMemorySize, and 64 KB even without the opt-in.

So the tile is halved on Volta with roughly 2.5x the required shared memory available. qmm_naive is 82.0% of GPU time on the #1538 baseline (994 launches, 29.3 ms average), so its N-tile is not a marginal knob — halving tile_n doubles the number of CTAs and doubles the A-tile traffic per output element.

The launch site already computes smem_bytes from the resolved layouts (qmm_naive.cu:96-98) and passes it as dynamic shared memory, so the plumbing for a larger tile is present.

Scope

Make the tile decision a function of the actual per-block shared-memory budget of the target device rather than a binary sm80 flag, and opt into V100's 96 KB where the chosen tile needs more than the 48 KB default.

qmm_naive.cu and device/gemm_sm70.cuh are not currently in the overlay (src/lib/mlx-cpp/patches/mlx/backend/cuda/ — see the inventory in #1536), so this issue introduces qmm_naive.cu as a new overlay file. Coordinate with #1543, which will overlay gemm_sm70.cuh and touch the same launch path; whichever lands first establishes the overlay, and the second rebases onto it.

Out of scope: the MMA atom (#1543), and any change to tile_m/tile_k selection beyond what a larger tile_n forces.

Implementation plan

1. Replace the sm80 boolean input to make_cta_tiler with the device's cudaDevAttrMaxSharedMemoryPerBlockOptin, and choose the largest tile_n in {128, 64} whose computed smem_bytes fits with margin. Keep the sm_80+ decision numerically identical to today so the change is provably a no-op above Volta.
2. Where the selected tile exceeds the 48 KB static default, call cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes) before launch. Note this is required on Volta for anything above 48 KB and is a silent launch failure if omitted.
3. Occupancy check with -Xptxas -v: a 128-wide tile at 128 threads may reduce blocks-per-SM enough to offset the win on an 80-SM part. Sweep tile_n in {64, 128} across prefill lengths 512/2048/8192 and pick per-length if the crossover is real; a documented "128 always wins" or "64 wins below N" outcome is equally acceptable as long as it is measured.
4. gather_gemm.cu:16,54-55 has the same sm80 flag feeding an analogous tiler. Check whether it has the same headroom and fix it in the same change if so.

Acceptance criteria

• Prefill tok/s on qwen3.8-27B-4bit at ~600, 2048, and 8192 prompt tokens, before and after, against the test(benchmark): Volta (sm_70) baseline record and build coverage #1538 baseline of ~7.7 tok/s.
• The tile choice is derived from a queried shared-memory attribute, not from a hardcoded arch comparison, and the sm_80+ selection is unchanged for every (itemsize, group_size, m) combination — assert this with a host-side unit test over the tiler function.
• cudaFuncSetAttribute is called whenever the selection exceeds 48 KB, and a deliberately oversized tile is rejected at selection time rather than failing at launch.
• Occupancy data (-Xptxas -v registers and blocks/SM for both tile widths) recorded in the sweep, with the crossover verdict stated.
• Quantized-matmul parity unchanged: qmm_naive output identical across the tile change for bits 4/8 and group sizes 32/64/128. Tile width must not alter results.
• No regression on GB10; cargo test --features cuda green on sm_70 and sm_121.

Validation

MLX_CUDA_ARCHITECTURES=70 make release-cuda
P=$(python3 -c "print(' '.join(['The quick brown fox jumps over the lazy dog.']*60))")
./target/release/mlxcel generate -m ./models/qwen3.8-27B-4bit -p "$P Summarize." -n 8
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o tile_after \
  ./target/release/mlxcel generate -m ./models/qwen3.8-27B-4bit -p "$P Summarize." -n 8

References

• Tiler and its sm80 input: mlx/backend/cuda/quantized/qmm/qmm_naive.cu:17-23,42-43.
• Launch-side shared-memory computation already in place: same file, :96-98.
• Same flag in the MoE gather path: mlx/backend/cuda/gemms/gather_gemm.cu:16,54-55.
• Overlay inventory this issue extends: src/lib/mlx-cpp/patches/mlx/backend/cuda/.

Outcome

The wide tile was measured and refused. The issue's shared-memory premise holds (the 128-wide tile needs 24 KB of a 96 KB per-block budget) but shared memory was never the binding constraint on Volta. Registers are: the 128-wide instantiation needs 255 registers and spills 128 bytes per thread against the 64-wide one's 224 and none, both reach the same 2 blocks per SM, and halving the CTA count on a grid already smaller than an 80-SM part makes qmm_naive 1.50x slower at a 106-token prompt, 1.06x at 516, 1.02x at 1,906 and 1.02x at 4,106, at identical launch counts. There is no crossover: prefill is chunked at 2,048 tokens, so the M extent of the grid is capped and the wide tile's best case is bounded rather than approached. Upstream's exclusion is correct; the name it gave the predicate, enough_smem, is not, and that is what changed. Full record: docs/benchmark_results/qmm-naive-tile-v100-2026-08-31.md.

Notes on two criteria that are ticked with a qualification:

• The prefill ladder ran at 106, 516, 1,906 and 4,106 prompt tokens rather than 600 / 2,048 / 8,192. 8,192 is not measurable on this host for this checkpoint: cudaMallocAsync failed: out of memory inside the third chunk, on a 32 GB card holding a 16 GB checkpoint, which is the chunked-prefill memory profile of fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672. The 4,106 rung replaces it and answers the same question, because prefill chunking caps m at 2,048 and therefore caps how large the grid can ever get.
• The host-side tiler test asserts more than the criterion asked for. It enumerates every (itemsize, group_size, m) combination against every real per-block shared-memory budget and both MMA paths, so it pins the selected tile as identical to upstream's on every architecture, not only on sm_80 and later.

The last criterion stays unticked. No Ampere-or-later device exists on the implementation host, and a full cargo test --features cuda needs --test-threads=1 per #1048, which does not fit this host's run budget; the targeted CUDA and host-side suites for this change are green on sm_70. See the pull request's ## Deferred to GB10 section.