[Perf][DSv4.1] Native-head Triton sparse MLA for SM90 few-head prefill

[positive666]

Purpose

SM90 FlashMLA sparse prefill (flash_mla_sparse_fwd) tiles WGMMA with B_H=64, so the kernel only accepts h_q in {64, 128}. DeepSeek-V4.1-Flash under TP8 has 8 local Q heads, and prefill currently zero-pads 8 → 64. Decode FP8 still needs that width. BF16 prefill does not: most of the QK/PV work is dummy heads.

This PR adds a Triton sparse MLA kernel that runs at native h_q for those prefill chunks. Decode stays on padded FP8 FlashMLA.

The path is on by default, and only used when:

• 0 < n_local_heads < padded_heads (the pad-64 / pad-128 case)
• chunk s_q >= VLLM_DSV41_FEWHEAD_MIN_SQ (default 2048)

Short prefills stay on FlashMLA. VLLM_DSV41_FEWHEAD_PREFILL=0 restores the old path.

This is not covered by #47629 (SM80/SM121 Triton MLA backend), #54929 (SM12x fallback), or #54592 (Q8KV8 sparse prefill). None of those skip SM90 pad-64 FlashMLA prefill for DSV4.1 TP8.

Test Plan

pytest tests/models/test_fewhead_prefill_gate.py -v
pytest tests/kernels/attention/test_fewhead_sparse_prefill.py -v

Kernel tests compare native h_q=8 against pad-64 FlashMLA at s_q in {64, 2048, 8192} using the production layout: non-contiguous q[:, :8] / out[:, :8] views into a pad-64 buffer, out= in-place, and topk_length=512. The wrapper test also uses a dim-0 slice of a larger token buffer, matching _forward_prefill.

Serving A/B on 8×H20, TP8, DSpark 5, same launch flags, uncached prompts. Those numbers were collected on frozen image dsv41-flash-h20:d-bb550774-pr56697-r2 with this kernel overlaid; the in-tree tests above are the checks for this PR SHA.

• closed-loop tok/s: C16 × 8k × 48 and C4 × 32k × 12
• cold TTFT: 32k and 128k (n=3)
• 512-token decode before/after an 8k prefill (MIN_SQ=2048, so 512 stays on FlashMLA)
• GSM8K complete test (1319) with VLLM_DSV41_FEWHEAD_MIN_SQ=0 so every prefill chunk, including 15-token turns, uses the kernel

Test Result

Kernel vs pad-64 FlashMLA (h_q=8, d=512, combined topk 640, attn_sink + topk_length). In-tree tests on H20 at s_q 64/2048/8192 with non-contiguous q[:, :8] / out[:, :8] and out=: mean_abs 3.4e-5, max_abs 3.91e-3. Isolated 8192-token chunk ~1.51× vs pad-64.

Closed-loop throughput (MIN_SQ=2048, 0 errors). Token counts are identical on both arms, so the relative gain is wall-clock only. tok/s below is (prompt + completion) / wall; output tok/s is completion / wall.

load	total tok/s	output tok/s	mean TTFT	p99 TTFT
C16 × 8k × 48	8939 → 9460 (+5.8%)	137.5 → 145.5 (+5.8%)	4.77s → 4.23s (−11.2%)	13.75s → 13.97s
C4 × 32k × 12	9094 → 9728 (+7.0%)	35.4 → 37.9 (+7.0%)	6.94s → 5.94s (−14.4%)	13.85s → 12.78s

DSpark accept: C16 0.355 → 0.377; C4 0.412 → 0.395.

Cold TTFT (n=3, uncached): 32k 3.49s → 3.22s (−7.8%); 128k 15.38s → 14.30s (−7.0%).

512-token short requests at the production gate (MIN_SQ=2048) stay on FlashMLA. After an 8k prefill, 512 e2e was 0.304s → 0.293s (n=5). The before-8k 512 set moved with DSpark accept noise (one sample), not the Triton path.

GSM8K, MIN_SQ=0: 1275/1319 → 1277/1319. Spec accept rate 0.589 → 0.594. Five items flipped baseline-correct / few-head-wrong and seven the other way. This is a no-collapse check, not a claim of exact quality equivalence.

AI assistance was used for coding and testing.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.