[drakosha]

Fixes vllm-project/vllm#59305 on SM90 (regression of vllm-project/vllm#53413). The SM100 side is #42.

What happens

sparse_attn_decode_interface allocates o_accum and lse_accum, (b + num_sm_parts) x s_q x h_q x d_v in fp32, whenever enable_split_kv is set, and on SM90 it always is (must_split_kv = arch.is_sm90a() || ...). That holds for the scheduler, not for the accumulators. With num_sm_parts == 1 the scheduler puts every request into the single partition and marks nothing as split (get_decoding_sched_meta.cu), so the SM90 kernel takes its is_no_split branch and writes out / lse directly (sm90/decode/sparse/splitkv_mla.cuh), and the combine kernel returns for every request (combine.cu). The accumulators are never read.

#19 skipped them in exactly this case. The V4.1 sync (#20) replaced sparse_decode.h with upstream's sparse_decode.cpp and the guard did not survive. vLLM picked that up with the pin bump in vllm-project/vllm#56893: in its mixed-batch layout (b=1, s_q=num_tokens) the dead workspace is 2 x s_q x h_q x d_v x 4 bytes, 8 GiB at 32k tokens with 64 heads, and the runtime OOM from #53413 is back.

Change

• needs_split_workspace = enable_split_kv && !(arch.is_sm90a() && impl_meta.num_sm_parts == 1) gates the accumulator allocation and the combine launch. The scheduler metadata is still built and passed whenever enable_split_kv is set, so kernel selection does not change.
• tests/test_flash_mla_sparse_decoding.py: test_no_split_workspace_allocation from [Bugfix] Avoid unused sparse decode split-KV workspace #19 is back, SM90 only, over (h_q, d_qk) in {(64, 576), (128, 576), (128, 512)} with top-k 2048.

SM100 is not touched. Since #30 its kernels live in the new csrc/cuda_kernels/sm100/ tree, which builds the o_accum tensor map and offsets the accumulator pointers unconditionally, so skipping the allocation there would need the null guards from #19 ported into that tree. #42 takes the other route for SM100: it turns enable_split_kv off when it is not mandatory and there is one partition. That condition never fires on SM90 (must_split_kv is always true there) and this PR's gate never fires off SM90, so the two do not overlap and merge cleanly in either order.

The first revision of this PR also restored #19's null guards on the legacy SM100 head128 path. After #30 that path is no longer the one that runs, so this revision drops them.

Validation

H200 NVL (sm_90a, 132 SMs), vLLM nightly 0bfc7a15d, _flashmla_C built from the first revision of this branch and swapped into the image. Baseline is the pinned 0eee43b. Mixed-batch layout (b=1, s_q=T), fp8 V3.2 cache, top-k 2048, 20% of the indices masked out. On SM90 the first revision's gate (num_sm_parts > 1) and this one select the same calls.

Transient allocation during one flash_mla_with_kvcache call on top of inputs and outputs:

h_q	T	num_sm_parts	before	after
64	1	132	16.7 MiB	16.7 MiB
64	8	16	18.0 MiB	18.0 MiB
64	128	1	33.1 MiB	1.0 MiB
64	1024	1	256.5 MiB	0
64	8192	1	2052 MiB	0
64	32768	1	8208 MiB	0
128	1	66	16.8 MiB	16.8 MiB
128	8	8	18.0 MiB	18.0 MiB
128	1024	1	513 MiB	0
128	8192	1	4104 MiB	0

out and lse are bit-identical to the pinned build for h_q in {64, 128} and T in {1, 8, 128, 132, 1024, 8192}, for both split and no-split schedules. In the no-split case both builds write the output from the kernel; the old one also allocated the buffers and launched a combine kernel that returned immediately. Per-call time is unchanged within 0.8% over the same grid (64 heads / 8192 tokens: 7512 us before, 7474 us after; 128 heads / 8192 tokens: 14158 us before, 14078 us after).

The restored test fails on the pinned build (No-split decode (h_q=64, d_qk=576) allocated 42.35 MiB for 8.28 MiB of outputs) and passes with the branch. The FlashMLA test suite needs kernelkit, which is not on PyPI, so it was run against vLLM's build through a small stand-in for the helpers the test touches; the full correctness suite was not run.

Not yet rebuilt after the rebase onto 847ef30 (#38, #39, #41 touch the fused RoPE kernels, the SM100 head64 decode and the SM90 dense decode, not this file's path). @cjackal reports in vllm-project/vllm#59305 that the first revision of this branch together with #42, both rebased onto 847ef30, fixes the memory peak for Mistral Large 4.

AI assistance (Claude) was used for the analysis and the draft of this change.

[cjackal]

Seems a rebase is needed after #30. And the rebase would send the scratch-buffer guard to legacy codepath, we also need to replicate it on the new csrc/cuda_kernels/sm100/ path.

[drakosha]

Rebased onto 847ef30. I narrowed it to SM90 instead of porting the guards into csrc/cuda_kernels/sm100/: your #42 already covers SM100 by turning split-KV off before scheduling, and that is less code than guarding the new tree's tensor map and pointer offsets. #42 never fires on SM90 (must_split_kv is always true there) and this gate only fires on SM90, so they merge cleanly in either order. Thanks for testing the pair on Mistral Large 4.

AI assistance was used for this comment.