[Bugfix][Attention] Stabilize sparse-MLA DCP for GLM PCP evals

[khluu]

Purpose

Make the B200 GLM-5.2 PCP evaluation complete reliably across its TP/PCP/DCP configurations.

Exact main build #89033 first exposed two independent memory-budget failures after merged backend repair #56677. Successive literal-head and actual-branch gates then exposed the workspace, empty-rank, warmup-key, packed-pointer, and backend-combination edges described below.

Root causes and fixes

TP2/PCP2 autotuning headroom

FlashInfer CUTLASS autotuning requests 12.22 GiB with only about 7.6 GiB free after the implicit 0.92 budget reserves 33.46 GiB for KV cache. Set --gpu-memory-utilization 0.85 to leave about 12.5 GiB more autotuning headroom.

TP1/PCP4/DCP4 KV-cache budget

CUDA graph profiling accounts for 6.15 GiB and reports total non-KV memory of 161.16 GiB, leaving a negative 3.23 GiB KV budget at 0.92. The runtime calculates 0.9545 as the utilization needed to preserve the old effective KV budget. Set 0.96, which produces a positive KV-cache budget.

FlashInfer sparse-MLA DCP workspace

Under DCP, the TRTLLM-gen launcher requests LSE and carves a softmax-stats slab proportional to gathered heads and scheduled tokens. The old fixed workspace was 413,138,944 bytes; an earlier gate requested 470,810,624 bytes during warmup and failed before evaluation.

This PR carries Jason Yao's byte-derived sizing fix from #50791 with its original author commit, then adapts it to current main's split sparse-MLA backends: preallocation occurs only in FlashInferMLASparseTRTLLMMetadataBuilder (SM100), so the separate SM120 backend keeps its own behavior. Explicit VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE overrides remain respected and warn if undersized; non-DCP allocation is unchanged.

The GLM DCP eval now caps --max-num-batched-tokens at its existing 4,096-token model length. With 64 heads per TP1 rank and DCP4, that bounds the computed workspace to exactly 2,561,671,168 bytes (2.39 GiB) rather than 16.39 GiB at the previous 32,768-token scheduler cap. This preserves the evaluation's context-length and accuracy contract while keeping the up-front buffer inside the measured B200 reserve.

Zero-local-sequence DCP warmup rank

CUDA graph warmup can give one DCP rank no local sequences while its padded LSE buffer still has rows. The eager empty-shard mask previously clamped the derived sequence index to -1 and indexed an empty seq_lens tensor. When the statically shaped sequence count is zero, the mask now fills every padded LSE row with -inf and returns before the generic lookup. A focused CPU regression covers the exact shape.

PCP+DCP first-request JIT keys

Actual-branch gate #89065 showed DCP4 clearing engine initialization and graph capture, then stalling its first real inference batch after an unexpected BuildPrefillChunkMetadataKernel JIT. PCP chunk plans normalize this dispatch to (DCP_RANK=0, DCP_WORLD=1), while startup warmup previously compiled only the configured per-process (rank, world=4) key. Warmup now includes both key families.

Actual-branch gate #89083 narrowed the remaining miss: the PCP plan stores global_cu as the second row of one packed 3 x (rows+1) allocation, so that input pointer is not always 16-byte aligned. Warmup had modeled alignment only for uncompressed_seq_lens and always replayed global_cu as aligned. It now covers the Cartesian product of both runtime alignment classes for both pointers, alongside the normalized DCP key.

Unsupported FlashInfer sparse-MLA PCP+DCP combination

Actual-branch #89087 proved the preceding fixes: TP2/PCP2 and TP1/PCP4 passed; DCP4 initialized with 29.45 GiB of KV cache, completed graph capture with no late JIT, and then hung on its first real batch. Diagnostic #89089 reproduced the stall with NCCL tracing. The first PCP+DCP request had 153 query rows on rank 0 and 160 on ranks 1-3; the final DCP all-gather consequently entered with counts 9,792 versus 10,240 (query_rows x 64 heads).

The selected FLASHINFER_MLA_SPARSE backend filters top-k KV indices to the local DCP shard and returns LSE for every rank-local PCP query, but does not gather the full cross-DCP KV context first. The shared sparse-MLA wrapper then attempts to DCP-combine unequal PCP query-row tensors. FLASHMLA_SPARSE is the implementation that gathers each DCP KV shard before running rank-local PCP prefill queries and returns only the decode LSE that needs DCP combination.

FlashInferMLASparseTRTLLMBackend.supports_combination now rejects combined PCP+DCP. Automatic selection therefore falls back to FLASHMLA_SPARSE; an explicit incompatible backend request receives an actionable error. CPU regressions cover both rejection of the combined mode and continued support for PCP-only and DCP-only modes.

DCP decode rank scalar

Actual-branch gate #89093 confirmed the backend fix: DCP selected FLASHMLA_SPARSE, completed graph capture and warmup, and began serving mixed prefill/decode traffic. The first decode then hit the runtime sync guard because get_dcp_local_seq_lens constructed the static integer DCP rank as a new device tensor inside metadata build().

The metadata builder now materializes that scalar once on its device during initialization and passes the tensor through the localization helper. The helper accepts and directly reuses a correctly typed, colocated scalar tensor; a CPU regression replaces torch.tensor with a hard failure during this path to prove that no per-build scalar construction remains.

Existing work / duplicate check

• FlashInfer MLA decode workspace buffer overflow with decode-context-parallel #50781 is the exact workspace issue. Open [Bugfix][Attention] Size FlashInfer sparse MLA workspace for decode-context-parallel #50791 contains the original diagnosis and implementation, but its head predates the separate TRTLLM/SM120 metadata builders. Original authorship and DCO trailers are preserved here; the current-main adaptation is a separate signed-off commit.
• Searches found no open PR or issue for the empty-rank mask, normalized PCP+DCP warmup key, packed-pointer alignment miss, or unsupported FlashInfer combined-mode selection.
• [CI] Select the supported DCP backend for PCP eval #56677 is merged and the exact gates confirm its ag_rs backend repair is active.
• Open [Core] Size the KV cache from this process's own device memory, not the device-wide delta #55828 addresses device-wide versus per-process memory attribution under co-tenancy. These builds use one worker per B200 and report identical per-rank accounting, so it does not own these deterministic config budgets.

Validation

• Negative control: main #89033 reached both configs; TP2/PCP2 failed in CUTLASS autotuning and TP1/PCP4/DCP4 failed with a negative KV-cache budget.
• Intermediate exact gate #89040 at a4149f66f: both PCP configurations passed their corrected budget stages, then reproduced FlashInfer MLA decode workspace buffer overflow with decode-context-parallel #50781 at the fixed workspace boundary.
• Diagnostic literal-head #89050: TP2/PCP2 and TP1/PCP4 passed, DCP4 received a positive 29.45-GiB KV cache, and warmup isolated the zero-local-sequence failure. It is not counted as a PR-branch gate because it was attributed to main.
• Actual-branch #89065 at 2f9bdae78: both non-DCP configs passed; DCP4 cleared sizing, initialization, and graph warmup, then exposed the normalized JIT-key miss.
• #89076 at 2d93cf0ef was infrastructure-only: its Buildkite pre-commit proxy received a GitHub API 504, so no B200 job ran.
• Actual-branch #89083 at 2d93cf0ef: both non-DCP configs passed; DCP4 reached its first real request and isolated the remaining rank-0 unaligned-pointer JIT key.
• Actual-branch #89087 at 5d2d5c905: both non-DCP configs passed; DCP4 completed sizing, initialization, graph capture, and all warmups with no late JIT, then isolated the incompatible backend collective.
• Diagnostic #89089 at 5d2d5c905: NCCL tracing recorded the exact unequal DCP all-gather counts (9,792 vs 10,240) on the first PCP+DCP request.
• pytest -q tests/v1/attention/test_flashinfer_sparse_mla_workspace.py: 10 passed.
• Focused zero-local-sequence regression: passed.
• pytest tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py -k indexer_warmup: 2 passed.
• Both YAML files parse; server_args contain exactly the intended memory-utilization and scheduler-cap values.
• All applicable pre-commit hooks (including Ruff, mypy, SPDX, and forbidden-import checks) and git diff --check: passed.
• Actual-branch #89093 at 7b2f95468: both non-DCP configs passed; DCP4 selected FLASHMLA_SPARSE, completed startup, served mixed prefill/decode traffic, and then exposed the runtime DCP-rank scalar construction in decode localization.
• pytest -q tests/v1/attention/test_indexer_dcp_localize.py -k get_dcp_local_seq_lens: 13 passed.
• Final actual-branch exact #89097 passed terminally at 05d607a8b: bootstrap, GitHub pre-commit, CUDA image, and the sole B200 PCP lane all passed without retry. TP2/PCP2 scored 0.9121, TP1/PCP4 scored 0.9454, and TP1/PCP4/DCP4 scored 0.8832; each completed all 1,319 questions with invalid rate 0.000. The exact job finished 3 passed in 17m54s.

AI assistance was used for investigation, current-main adaptation, CI configuration, and validation. No human review is claimed.

[tlrmchlsmth]

merged main to pick up #55499, which I think was breaking the CI for this PR

[khluu]

Expanded this draft after the exact merged-main lane exposed a second config-specific capacity failure.

• Negative control: #89033. [CI] Select the supported DCP backend for PCP eval #56677's ag_rs backend fix is active. TP2/PCP2 now reaches and OOMs in CUTLASS autotuning; TP1/PCP4 now reaches graph profiling and gets Available KV cache memory: -3.23 GiB at 0.92, with the runtime recommending 0.9545.
• New literal head: a4149f66fa8abe77ff269a60a4a04e89f06a5699. The PCP4 config uses 0.96; the existing PCP2 config remains 0.85.
• Static validation: both YAML/config argument parses, all applicable pre-commit hooks, and git diff --check pass.
• Exact four-B200 lane: #89040 is starting with only lm-eval-pcp-4xb200 selected.

The new commit was added after the prior approval, so please treat that approval as stale until the exact lane is terminal and the PCP4 line is re-reviewed.

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

The prior #89050 run was launched against main, so I am not treating it as PR-branch gate evidence. Its literal-head logs did expose one more deterministic DCP warmup edge case after both PCP configurations passed their memory-budget stages: a rank with zero local sequences indexed an empty seq_lens tensor while masking padded LSE rows.

Head 2f9bdae78ca8133f95312e30a4e3475736271362 now treats all padded LSE rows as empty when a DCP rank has no local sequences, with a focused regression test. Locally: 1/1 focused test passed, and all applicable pre-commit hooks (including ruff, mypy, SPDX, and forbidden-import checks) passed.

Replacement exact gate #89065 is running on the actual PR branch codex/ci-pcp-autotune-headroom, carries PR metadata for #55879, and selects only lm-eval-pcp-4xb200: https://buildkite.com/vllm/ci/builds/89065

[khluu]

Actual-branch gate #89065 reached all three GLM configurations at head 2f9bdae78:

• TP2/PCP2 passed (GSM8K accuracy 0.9143, 0 invalid).
• TP1/PCP4 passed (accuracy 0.9378, 0 invalid).
• TP1/PCP4/DCP4 now clears workspace allocation, KV-cache sizing, engine initialization, and CUDA-graph warmup. Its first real inference batch then exposed a distinct warmup-key miss: the runtime PCP chunk plan dispatches BuildPrefillChunkMetadataKernel with already-localized rows and normalized (DCP_RANK=0, DCP_WORLD=1), while startup had compiled only the configured per-process DCP key (rank, world=4). Rank-local late JIT then stalled sample_tokens until the 300-second RPC timeout.

Head 2d93cf0ef adds the normalized PCP+DCP key to startup warmup and a CPU regression test that verifies both configured and normalized key families. Focused tests pass (2/2), as do all applicable pre-commit hooks, mypy, and git diff --check.

Replacement actual-branch exact gate: Buildkite #89076, with PR #55879 metadata and only lm-eval-pcp-4xb200 selected.

[khluu]

#89076 did not execute the B200 lane. Its Buildkite pre-commit proxy received a transient GitHub API HTTP 504 after five successful polls and exited 1; both actual GitHub pre-commit check runs on 2d93cf0ef are terminal success, and the exact image build passed.

Rebuilt the same actual branch/head and PR metadata as #89083, still selecting only lm-eval-pcp-4xb200. No source change or blind B200 retry was made; this only retries the failed CPU proxy after its upstream check is confirmed green.

[khluu]

Exact actual-branch gate #89083 at 2d93cf0ef narrowed the remaining DCP4 stall to one missing Triton pointer-alignment specialization:

• TP2/PCP2 and TP1/PCP4 passed.
• DCP4 got a positive 29.45-GiB KV cache, completed initialization and CUDA-graph capture, and reached the first real 623-token request.
• All four ranks built the same PCP batch, but only rank 0 emitted Triton kernel JIT compilation during inference: BuildPrefillChunkMetadataKernel.kernel; sample_tokens then timed out after 300 seconds.
• The PCP plan's global_cu is row 1 of a packed 3 x (num_rows + 1) allocation. With two PCP rows in this request its pointer is offset by 12 bytes, so it has Triton's unaligned pointer specialization. Warmup always replayed that input as aligned.

New signed head 5d2d5c9059ef43e31c851006baeb0edcd6cb6f7e covers aligned and unaligned global_cu independently of the existing uncompressed_seq_lens variants, while retaining the normalized (DCP_RANK=0, DCP_WORLD=1) key. The focused regression asserts all four pointer-alignment combinations and both DCP key families.

Validation:

• pytest tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py -k indexer_warmup: 2 passed
• applicable pre-commit hooks, including Ruff and mypy: passed
• git diff --check: passed

No duplicate PR was opened; this extends the existing canonical #55879 repair. A replacement actual-branch B200 PCP-only gate is next.

[khluu]

Root cause for the post-warmup DCP hang is now isolated and patched at 7b2f95468.

Diagnostic #89089 reproduced the stall on the first PCP+DCP request with no late JIT. NCCL recorded rank 0 entering the LSE all-gather with count 9,792 while ranks 1-3 entered with 10,240—the exact 153 vs 160 local query rows x 64 heads mismatch.

FLASHINFER_MLA_SPARSE was being auto-selected even though it does not gather full cross-DCP KV before rank-local PCP prefill. It returns LSE for each rank's unequal PCP rows, which the shared DCP combine then tries to all-gather. FLASHMLA_SPARSE is the backend that implements the required KV gather and excludes prefill LSE from DCP combination.

The new capability check rejects combined PCP+DCP for FlashInfer sparse MLA, so auto-selection falls back to FlashMLA sparse; explicit incompatible selection now gets an actionable error. Focused CPU coverage is 10/10, including PCP-only and DCP-only controls, and all applicable pre-commit hooks pass.

Replacement exact actual-branch gate: https://buildkite.com/vllm/ci/builds/89093 (PR metadata present; only lm-eval-pcp-4xb200). I canceled #89091 before test execution because my first launch omitted Buildkite's required top-level PR metadata fields; it is not validation evidence.

[khluu]

#89093 confirmed that the backend-selection repair works: the DCP configuration chose FLASHMLA_SPARSE, completed initialization/graph capture/warmup, and began processing real mixed prefill/decode traffic.

It then exposed the next deterministic edge on the first decode: get_dcp_local_seq_lens called torch.tensor(dcp_rank, device=seq_lens.device) from metadata build(), and the runtime guard rejected that per-batch host-to-device scalar construction as GPU<->CPU sync detected.

Head 05d607a8b now materializes the static rank scalar once during metadata-builder initialization and reuses that device tensor in decode localization. The regression test makes torch.tensor raise during the tensor-rank path, proving the hot path does not construct a replacement. Focused helper tests are 13/13 and all applicable pre-commit hooks, mypy, SPDX, and diff checks pass.

Replacement actual-branch exact gate: https://buildkite.com/vllm/ci/builds/89097 — complete PR metadata, literal head, and only lm-eval-pcp-4xb200.

[khluu]

Final exact validation is green: https://buildkite.com/vllm/ci/builds/89097

At literal head 05d607a8b9aa17e001f921c6fbdb405abeaecb5c and actual PR branch metadata, bootstrap, GitHub pre-commit, CUDA image, and the sole B200 LM Eval PCP job all passed without retry. The lane finished 3 passed in 17m54s:

• TP2 / PCP2: 1,319/1,319, invalid rate 0.000, metric 0.9121
• TP1 / PCP4: 1,319/1,319, invalid rate 0.000, metric 0.9454
• TP1 / PCP4 / DCP4: 1,319/1,319, invalid rate 0.000, metric 0.8832

The combined DCP case selected FLASHMLA_SPARSE, completed mixed prefill/decode traffic end to end, and exited cleanly. This closes the exact gate that had successively exposed the memory, workspace, zero-rank, warmup-key, pointer-alignment, incompatible-backend, and device-scalar edges.