[NickLucche]

Overview

The FlashMLA sparse fp8_ds_mla mixed-batch path now writes the head-padded query and the kernel output into the workspace buffers that FlashMLASparseImpl already reserves for max_num_batched_tokens, so it no longer allocates them on every call. This removes about 4.25 GiB of per-call allocation at 32K tokens per step. That allocation isn't counted by the startup memory profile, so it can cause runtime OOMs.

Related: #59305, #57878. Complements vllm-project/FlashMLA#29, which removes the kernel-internal fp32 split-KV o_accum buffer (the other 8 GiB at T=32K). This PR only covers the wrapper side.

Claims

• On the fp8 mixed-batch path, the padded q (T x 64 x 576 bf16) and the output (T x 64 x 512 bf16) are written into the reserved workspace, and the output buffer is passed to the kernel via out=. At T=32768 this saves 4.25 GiB per call, measured. The workspace was already reserved, so KV cache capacity does not change.
• Outputs and LSE are bit-identical to the allocating path.
• GLM-5.3-Flash FP8 on 8xH100 (TP8, max_num_batched_tokens=32768, gpu_memory_utilization=0.85) previously hit a runtime OOM in flash_mla_with_kvcache during warmup. With this change it serves the same workload without errors, and GSM8K accuracy matches.

Validation

Per-call peak allocation. Single H100. GLM-5.3-Flash TP8 shapes: 8 heads/rank padded to 64, topk 2048, s_q = T. Measured through FlashMLASparseImpl._fp8_flash_mla_kernel:

T	before (allocating path)	after (workspace path)	saved	outputs equal
1024	0.384 GiB	0.251 GiB	0.133 GiB	yes
8192	3.068 GiB	2.006 GiB	1.063 GiB	yes
16384	6.137 GiB	4.012 GiB	2.125 GiB	yes
32768	12.273 GiB	8.023 GiB	4.250 GiB	yes

The remaining ~8 GiB at T=32768 is the kernel's o_accum, which is addressed by vllm-project/FlashMLA#29.

End-to-end, previously crashing config. GLM-5.3-Flash FP8, 8xH100, --tensor-parallel-size 8 --max-num-seqs 256 --attention-backend FLASHMLA_SPARSE --kv-cache-dtype fp8_ds_mla --max-num-batched-tokens 32768 --gpu-memory-utilization 0.85 --max-model-len 131072.

• Before this change: during the warmup with 8K input and concurrency 256, the server failed with torch.OutOfMemoryError (tried to allocate 8 GiB with 5.29 GiB free), and 255 of 256 requests failed.
• After this change: warmups at 8K/c256 and 64K/c64 completed with no errors, and there are no OOM lines in the server log. Peak nvidia-smi memory.used was 80923 / 81559 MiB. This is tight because o_accum is still allocated per call until FlashMLA#29 lands.

vllm bench serve results, random dataset, --ignore-eos. The comparison is against the existing MNBT=16384 config, which has a larger KV pool (2.66M vs 1.88M tokens), so this is a sanity check rather than an A/B:

ISL/OSL, concurrency	completed	total tok/s, MNBT 32K + this PR	total tok/s, MNBT 16K baseline
8192/1024, c64	192/192	16717	16424
8192/1024, c256	768/768	20189	21522 (KV cache reached 100% at 32K)
65536/1024, c16	48/48	23357	23141
65536/1024, c64	192/192	24263	25101

Accuracy. GSM8K, 5-shot, all 1319 questions, using tests/evals/gsm8k/gsm8k_eval.py against the server:

config	accuracy	invalid
main, MNBT 8192	0.919	0.000
this PR, MNBT 32768	0.914	0.000

Tests.

• New test: pytest tests/v1/attention/test_sparse_mla_backends.py -k test_flashmla_fp8_mixed_batch_stages_in_workspace (H100). 2 passed.
• Existing tests: pytest tests/v1/attention/test_sparse_mla_backends.py -k "test_sparse_backend_decode_correctness and FlashMLA and fp8_ds_mla" (H100). 84 passed, 84 skipped. These include the TP>1 cases that use the mixed-batch path, run with the workspace locked.

Details

Root cause. For fp8_ds_mla with fewer than MIN_HEADS_FOR_BF16_PREFILL heads per rank (fp8_use_mixed_batch), every step makes a single flash_mla_with_kvcache call with s_q = num_tokens. The wrapper pads q to 64 heads with new_zeros, and the kernel allocates out, lse and o_accum, all O(T * 64). The profile run skips attention, so none of this is counted when the KV cache is sized. At large max_num_batched_tokens, a full step therefore OOMs after startup.

Fix. __init__ already reserves (max_tokens, padded_heads, 576) and (max_tokens, padded_heads, 512) bf16 workspace slots for the BF16 prefill path, which is not used in mixed-batch mode. _fp8_flash_mla_kernel(use_workspace=True) takes flat views of those slots sized for the FP8 decode padding (the reserved head padding is at least as wide). It zeroes only the padding heads and passes out= to the kernel.

• Only the mixed-batch caller turns this on. Pure-decode, HiSparse and separate prefill/decode callers keep the allocating path, because their T is small.
• The returned output aliases the workspace. It is consumed by the v_up projection or the DCP merge before any other workspace user runs. The separate prefill/decode path already returns workspace-backed output the same way.
• The q input comes from workspace slot 0, the q-concat buffer, so it doesn't overlap the padded-q slot.
• No new CPU/GPU synchronization: the added checks are host-side shape asserts.
• lse is still allocated by the kernel (T x 64 x 4 B = 8 MiB at 32K).

Limitations. The kernel-internal o_accum (2 x T x 64 x 512 x 4 B = 8 GiB at T=32K on SM90) is still allocated per call and isn't counted by the profile run. That is fixed in vllm-project/FlashMLA#29. This PR removes the wrapper-side share.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[drakosha]

#59319 (ours, open) stops reserving the bf16 prefill workspace on the fp8 mixed-batch path. Its first revision skipped the whole block, including the three max_num_batched_tokens buffers this PR stages q and out in. The two PRs merge without a conflict, but together workspace_specs on that path holds only q_concat, and the get_simultaneous unpack fails on the first forward with ValueError: not enough values to unpack (expected at least 3, got 1).

The current revision of #59319 skips only the 5 * max_model_len buffer (plus the PCP+DCP gather buffer) and keeps the three max_num_batched_tokens buffers, so either merge order works. On one H200 (nightly af7f9488c with both files swapped in), test_sparse_backend_decode_correctness for FlashMLA + fp8_ds_mla fails in all 14 TP=8 cases with the first revision plus this PR, and passes with the current one: 72 passed, including your two staging tests and the #59319 unit tests.

test_flashmla_fp8_mixed_batch_stages_in_workspace builds workspace_specs by hand, so it won't catch a change in what __init__ reserves. The TP=8 cases #59319 adds go through the real __init__ and do.

AI assistance was used for this comment.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @NickLucche.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork