perf(cuda): GPU-side argmax for batched decode + BM=32 MMQ tile (#205, #203)

[pekkah]

Summary

Pushes multi-user batched decode clearly over the 3-4% bar by stacking two orthogonal levers. Measured on Qwen3-8B Q4_K_M @ RTX 4070 Ti (aggregate t/s vs baseline, all-greedy):

N	baseline	+BM=32	+argmax	+both
6	264.3	271.1 (+2.6%)	273.6 (+3.5%)	281.7 (+6.6%)
7	298.6	306.4 (+2.6%)	310.0 (+3.8%)	317.3 (+6.3%)
8	321.1	328.9 (+2.4%)	332.9 (+3.7%)	342.3 (+6.6%)

The levers target different constraints (host/PCIe tail vs device matmul), so they add.

Lever 1 — GPU-side argmax for batched decode (the larger one, +3.5–3.8% alone)

BatchForwardMulti downloaded the full [N×vocab] logits every step (4.86 MB at N=8, vocab 151936), host-copied it, and allocated N per-row float[]. New BatchForwardMultiArgmax reuses the already-shipped #219 ArgmaxRows kernel to compute the per-sequence argmax on-device and copy back only rows*8 bytes.

• perf(cuda): GPU-side greedy argmax — every decode token downloads the full-vocab logits (0.6–1 MB) with a blocking stream sync #219 found single-user GPU argmax "within noise" because the 1-row download is tiny — but the batched download is N× larger, so eliding it is a real win here.
• The trunk is hoisted into RunBatchedTrunk; BatchForwardMulti keeps the full-logits tail unchanged.
• Bit-exact to a host argmax of the same logits (lowest-index tie-break == Sampler.Greedy).
• ContinuousBatchingEngine takes it only when every active sequence is greedy this step (temp≤0, not force-closing </think>); any sampled/forced seq reverts the whole step to the full path.
• Gated by SupportsBatchedGpuArgmax (the existing SHARPI_GPU_ARGMAX kill-switch; off under the profiling variants). Single-user spec-decode BatchVerify is untouched (stays bit-exact WS).

Lever 2 — BM=32 decode-MMQ tile (#205 item 1, +2.4–2.6% alone)

Emits llm_mmq_q4k_soa_acts_n16 at a second row-tile size from one template; the dispatcher routes grid-starved shapes (ceil(rows/64) < 2·SM: Q/O proj + the Q4_K ffn_down half — 64 blocks on a 60-SM card) to the BM=32/128-thread tile. High-row shapes (gate/up, lm-head) keep BM=64. Bit-identical to BM=64. SM count from cudaDevAttrMultiProcessorCount; SHARPI_DECODE_MMQ_BM32=0 kill-switch.

#205 item 2 (cp.async) dropped — an adversarial design panel plus the #176/#152 prefill precedent rate it ~0–1% / e2e-neutral on decode; not worth the kernel complexity.

Validation

• New Qwen3_8B_BatchForwardMultiArgmax_MatchesBatchForwardMulti_N6: GPU argmax token == host Argmax of the full logits, and the returned logit == that row's value (exact).
• CudaDecodeMmqTests covers both MMQ tiles vs the fp32 CPU reference (SimdKernels.DotQ4K).
• N=6 MMQ oracle + all CudaBatchForwardMultiTests (N=2, two-step, empty-batch, prefill) green — guard the trunk refactor.
• Release build clean (0 warnings, TreatWarningsAsErrors). All 10 relevant CUDA tests green on the 4070 Ti.

Closes #205. Refs #203 (umbrella), #219 (extends GPU argmax to the batched path).

🤖 Generated with Claude Code

[pekkah]

/gemini review

[gemini-code-assist]

Code Review

This pull request optimizes batched decode performance by introducing a BM=32 tile size for decode-MMQ to prevent grid starvation on low-row shapes, and implementing an on-device greedy argmax tail (BatchForwardMultiArgmax) to elide full N×vocab logits downloads when all active sequences are greedy. These changes include kernel generation updates, SM count queries, engine integration, and new unit tests and benchmarks. As there are no review comments, I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Code Review

This pull request optimizes the batched decode path by introducing a smaller BM=32 tile size for decode-MMQ to prevent grid starvation on low-row shapes, and an on-device greedy argmax tail (BatchForwardMultiArgmax) to avoid downloading full logits when all sequences are greedy. Feedback suggests checking the return value of CuBlasInterop.DeviceGetAttribute when querying the SM count to ensure robustness and consistency.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The return value of CuBlasInterop.DeviceGetAttribute is ignored here. While smCount is initialized to 0 and a failure will safely fall back to disabling the BM=32 route (since _smCount > 0 is checked later), it is more robust and consistent with other DeviceGetAttribute calls in this method to check the return value and handle or log any failure.

        int smCount = 0;
        if (CuBlasInterop.DeviceGetAttribute(out smCount, CuBlasInterop.CudaDevAttrMultiProcessorCount, 0) != 0)
        {
            smCount = 0;
        }