perf(speculative): pad the CUDA MTP verify forward past the qmv dispatch threshold

[inureyes]

Context

PR #733 (issue #638) measured B=1 MTP speculative decoding on GB10 (CUDA, sm_121) as a consistent regression: baseline 14.5 tok/s, 0.77x at K=2, 0.52x at K=4/8, acceptance 35-56%, mean accepted length 0.56-1.05 (docs/benchmark_results/speculative-pairing-gb10-2026-07-10.md). The 2026-07-11 follow-up investigation established the root cause: the MTP verify is a [1, K] sequence-axis forward, and the CUDA quantized dispatch (src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:229-245) routes M*B < 8 to per-row qmv (the sm90 CUTLASS path is Hopper-gated), so a K-wide verify re-reads the full weight matrix K times. Measured: classic forward 69 ms, K=2 round 140 ms, K=4 round ~270 ms, exactly linear in K; the measured speedups equal tokens-per-round divided by round-cost-in-forwards exactly. This is the same root cause as the flat batched-decode throughput tracked in #725.

The Metal speculative path already solves the analogous problem: on M5+ it pads the verify sequence up to a 32-token NA tile boundary to force the tiled qmm_nax GEMM instead of GEMV (src/lib/mlxcel-core/src/speculative/mod.rs:380-425: align_to_na_tile, a padding-aware mask, and KV trim of the padding). The CUDA path has no analogous padding.

The hardware is not the limit: LMSYS measured SGLang on the same DGX Spark hardware with near-linear batched decode scaling (20.5 tok/s at B=1 to 368 tok/s at B=32, llama-3.1-8b) and up to 2x end-to-end EAGLE-3 speculative decoding (https://www.lmsys.org/blog/2025-10-13-nvidia-dgx-spark/), and NVIDIA ships speculative-decoding playbooks for DGX Spark (https://build.nvidia.com/spark/speculative-decoding).

Important caveat

#725 already measured that with today's kernels, qmm_sm80 at M=8 is WORSE than per-row qmv: the Ampere 128-wide tile is ~94% wasted at M=8. So naive padding to M=8 may regress with the current kernel. This issue is therefore two things:

• (a) An A/B experiment to confirm the behavior on the verify shape, including a MLXCEL_QMM_TILE_M tile sweep (the JIT-per-tile env hatch from perf(cuda/quant): sm_120/sm_121-tuned quantized GEMM for Blackwell prefill #637 / PR perf(cuda/quant): sm_120/121 qmm CTA tile (M=128) for Blackwell prefill #723; only tile_m is tunable).
• (b) The integration consumer of perf(cuda/quant): amortizing small-M quantized GEMM for batched decode on Blackwell #725: once an amortizing small-M path lands, the padded verify is how B=1 speculative decoding exploits it.

Impact

Without a padded (or otherwise amortized) verify, B=1 MTP speculative decoding on GB10 stays a strict regression regardless of acceptance rate, and any future small-M kernel improvement from #725 has no consumer on the speculative path.

What's needed

• Implement verify padding on the CUDA MTP path, mirroring the Metal NA-tile approach: pad tokens, padding-aware mask, trim the padding out of the KV cache.
• Gate it behind an env variable (e.g. MLXCEL_MTP_VERIFY_PAD).
• A/B measure on GB10 with speculative_bench --k-values 2,4 on the gemma-4-12b-it-4bit + gemma-4-12b-it-assistant-4bit pairing, sweeping MLXCEL_QMM_TILE_M.
• Verify greedy parity: temperature 0 output byte-identical with padding on and off.
• Document the result in docs/benchmark_results.

Acceptance Criteria

• Padded-vs-unpadded A/B table on GB10 for K=2 and K=4, including the MLXCEL_QMM_TILE_M sweep, documented in docs/benchmark_results.
• Greedy parity confirmed (temperature 0 output byte-identical with padding on/off).
• A documented enable/disable decision wired into the actual MTP verify code path, not a standalone experiment script.
• Re-validation after perf(cuda/quant): amortizing small-M quantized GEMM for batched decode on Blackwell #725 lands: re-run the A/B with the amortizing small-M path and update the enable/disable decision.

Technical Considerations

• Depends on / re-validated by perf(cuda/quant): amortizing small-M quantized GEMM for batched decode on Blackwell #725. Refs perf(speculative): CUDA drafter pairing matrix and continuous-batching integration #638, PR perf(speculative): CUDA pairing matrix and burst HOL observability #733.
• CUDA dispatch: src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:229-245. Metal reference implementation: src/lib/mlxcel-core/src/speculative/mod.rs:380-425.

[inureyes]

Superseded for K in [2,7] by the multirow qmv path in PR #740 (#725): the MTP verify [1, K] forward now amortizes weight reads inside qmv itself, without padding past the dispatch threshold. Measured on GB10: K=4 MTP 0.52x -> 1.46x. Padding would only matter for shapes that must cross into qmm_sm80 (M*B >= 8), which no current MTP pairing reaches (K=8 runs at effective K=4). Recommend closing this after #740 merges, or re-scoping it to the M*B == 8 boundary if such a shape shows up hot.

[inureyes]

Closing as superseded by #740 (merged as a17cbc0): the multirow qmv path amortizes the verify for MB in [2,8) without padding, and every current MTP pairing runs at effective K <= 4. If an MB == 8 boundary shape ever shows up hot, re-open with that scope (the multirow kernel already accepts 8 rows; only the dispatch gate would change).