[Bugfix][Attention] Size FlashInfer sparse MLA workspace for decode-context-parallel

[thegoldenflow]

Purpose

Fixes #50781.

With --decode-context-parallel-size > 1, the FlashInfer sparse MLA backend (FLASHINFER_MLA_SPARSE, B200/SM100) crashes inside FlashInfer's trtllm-gen launcher with a workspace buffer overflow — a single request with a moderately long prompt is enough to bring down every DCP worker and the API server.

Root cause

The workspace handed to trtllm_batch_decode_with_kv_cache_mla is a static default, VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE = 394 MiB (413,138,944 bytes), with no awareness of DCP, head count, or batch size.

Whenever the caller requests LSE, FlashInfer's trtllm-gen launcher carves a softmax-stats slab out of that workspace (csrc/trtllm_fmha_kernel_launcher.cu, identical from v0.6.14 through current main):

softmax_bytes = sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256)
                + 1 MiB guard

vLLM's sparse MQA path (forward_mqa) passes q_len == 1 per token, so round_up(1, 256) == 256 slots per (head, token). Three factors combine under DCP to blow this up:

1. LSE is only requested under DCP. need_to_return_lse_for_decode = dcp_world_size > 1 and can_return_lse_for_decode (the shared DCP reducer needs LSE). Without DCP the slab is never carved, which is why this only reproduces with --decode-context-parallel-size > 1.
2. DCP all-gathers the query in the head dim. Before forward_mqa, the decode query is all-gathered across the DCP group (mla_attention.py), so the kernel sees num_heads_per_rank * dcp_world_size heads per rank — for the reporter's GLM-5.2 (64 q heads) at TP=8, DCP=8 that is 8 × 8 = 64 heads instead of 8.
3. batch_size is the full per-step token count, not just decodes. The sparse backend routes prefill tokens through the same MQA kernel (each token carries its own top-k row, so the sparse mask is fully per-token), which is why "a single request" is enough: the reporter's single ~12K prompt lands as a 12,288-row batch. Whether that routing is the right long-term design is an upstream question and out of scope here — the workspace has to be sized for what the kernel is actually handed today.

Byte-exact validation against the reported crash

The reporter's crashing step scheduled 12,288 tokens (a single ~12K-token prompt). Plugging into the formula:

8 bytes * (8 heads/rank * 8 DCP) * 12288 tokens * 256 + 1 MiB = 1,611,661,312

which matches the reported allocation request exactly. The reported "only 404,750,336 bytes available" is also exact: FlashInfer ≤ 0.6.14 first carves an 8 MiB multi-CTA-KV counter slab (8192 * 256 * sizeof(uint32) = 8,388,608) from the same buffer, and 413,138,944 − 8,388,608 = 404,750,336. (FlashInfer ≥ 0.6.15 moved the counter to a separate buffer — flashinfer-ai/flashinfer#3582 — which is the version vLLM currently pins; the fix keeps the 8 MiB inside the preserved baseline either way.)

The overflow threshold for this config is 3,081 tokens in a step (8 * 64 * 3080 * 256 + 1 MiB exactly fills the remaining 404,750,336 bytes), i.e. any prompt longer than ~3K tokens crashes the server.

Fix

Add a pure-Python sizing function compute_trtllm_sparse_mla_workspace_bytes() colocated with the backend, and pre-allocate the shared workspace in FlashInferMLASparseMetadataBuilder.__init__:

required = base (394 MiB default) + 8 * (heads/rank * dcp) * max_num_batched_tokens * 256 + 1 MiB

• The DCP delta is exactly the derived softmax slab — no arbitrary margin. The base stays reserved for trtllm-gen's counter + multi-CTA-KV scratch regions, which are batch/DCP-independent (scratch parallelism is clamped by SM count; multi-CTA-KV is disabled outright once the CTA grid covers the device) and demonstrably fit in today's default — non-DCP serving exercises exactly those regions.
• Sizing happens before the first forward/capture. Cudagraph support here is UNIFORM_BATCH; the buffer address is baked into captured graphs, so the buffer must reach its final size up front — the metadata builder constructor runs before warmup/capture. No lazy regrow, no catch-and-retry.
• max_num_batched_tokens is a hard upper bound for the kernel's batch dim: scheduler steps are capped by it, and cudagraph capture sizes are clamped to it (max_cudagraph_capture_size = min(max_num_tokens, ...)).
• Explicit env override is respected. If VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE is explicitly set, that value is used as-is; if it is below the computed requirement a prominent warning logs both numbers and the exact value to set.
• Sizing is validated against the pinned FlashInfer version. The slab formula is byte-identical in csrc/trtllm_fmha_kernel_launcher.cu at the reporter's 0.6.14.dev20260705, at v0.6.15.post1 (what requirements/cuda.txt pins today), and on FlashInfer main. The only layout change in that range is the multi-CTA-KV counter moving to a caller-supplied buffer in 0.6.15 (Add separate trtllm-gen KV counter buffer flashinfer-ai/flashinfer#3582); since the fix preserves the existing default as the base and adds the slab on top, it is correct on both sides of that change. No version-specific constant is hard-coded beyond the slab's own alignment/guard terms, which are named and commented with their upstream source.

Memory impact

• Non-DCP configs: zero change. Without DCP no LSE is requested, the computed size equals the existing default, and the allocation still happens lazily on first use. The other three readers of VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE (dense FlashInfer backend, MLA prefill backends) are untouched.
• DCP configs: the workspace grows by the exact slab the kernel was already trying to carve, e.g. the reporter's config allocates 394 MiB + 2,148,532,224 bytes ≈ 2.39 GiB per rank (vs. crashing today); a default 8192-token config with the same gathered head count is ≈ 1.07 GiB.
• Like the existing lazy allocation (and the base class's DCP-enlarged chunked-prefill workspace), this buffer lives outside the KV-cache memory profiling budget. Under DCP it is larger than before, so a deployment with very little post-KV-cache headroom would now OOM at builder init instead of crashing mid-serving — the failure moves to startup and becomes actionable (lower --max-num-batched-tokens or gpu_memory_utilization).
• PCP+DCP combinations (dcp ∈ {pcp, tp*pcp} per config validation) gather fewer or differently-grouped heads than plain DCP; the formula upper-bounds those layouts (never under-allocates).

Test plan

New CPU-only tests (no GPU markers) in tests/v1/attention/test_flashinfer_sparse_mla_workspace.py:

• byte-exact reproduction of the reported 1,611,661,312-byte request,
• reporter config coverage (computed ≥ observed request + 8 MiB pre-carve),
• non-DCP size identical to today's default (regression guard),
• env-override semantics (explicit-set respected, warning on undersize, no warning on oversize),
• sync guard between the local default constant and vllm/envs.py.

$ pytest tests/v1/attention/test_flashinfer_sparse_mla_workspace.py -v
...
7 passed, 14 warnings in 6.45s

$ pre-commit run --files vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py \
    tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
ruff check ....... Passed
ruff format ...... Passed
typos ............ Passed
Run mypy for Python 3.10 ... Passed

GPU-only sparse MLA tests (tests/v1/attention/test_sparse_mla_backends.py) were not run — no B200 available to me. This change is sizing-only and does not alter kernel inputs or numerics, so no model-eval results are included; happy to run any suite a reviewer wants that does not need SM100 hardware.

Hardware verification request: @flexwang could you validate on your 8×B200 rig? Expected behavior with this patch on your exact command line: startup logs unchanged, per-rank extra VRAM ≈ 2.0 GiB for the workspace, and the >3K-token single-request crash gone. My prior FlashInfer buffer sizing fix #50022 followed the same size-up-front pattern and was GPU-verified there.

Follow-up (FlashInfer side, not this PR)

trtllm-gen exposes no public workspace-size API for its MLA decode path (the CuteDSL path has _get_split_kv_and_workspace_size; trtllm-gen sizes are only discoverable by reading the C++ launcher). I plan to file a FlashInfer issue requesting a sizing contract mirroring the CuteDSL one, so callers can stop re-deriving these constants.

---

This PR is AI-assisted (analysis and draft authored with Claude Code); I have reviewed every line and the derivation. Duplicate-work check at submission time: #50781 had no comments, no linked PRs, and no open PRs touching sparse-MLA/DCP/workspace sizing.

FIX #50781

[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.

[khluu]

Your byte-exact diagnosis now reproduces in the current B200 PCP CI lane: vllm/ci #89040 reaches the DCP path and requests 470,810,624 bytes from the fixed 413,138,944-byte workspace.

Kevin asked me to make that exact lane pass. I carried your original commit into #55879 with authorship and DCO trailers preserved, then added a small current-main adaptation because sparse MLA now has separate SM100 TRTLLM and SM120 metadata builders; the preallocation is scoped to the SM100 TRTLLM builder. On current main, all 7 focused tests and every applicable pre-commit hook pass locally. I am launching an exact B200 gate and will post the terminal result.

Linking this explicitly so the original investigation and credit remain visible; I am not claiming the underlying derivation as new work.

[mergify]

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

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

[thegoldenflow]

Root-caused and implemented the original fix for vLLM #50781/#50791, deriving the FlashInfer sparse-MLA DCP workspace overflow byte-exactly; the fix was later integrated into main through #55879 and validated on B200 CI.