[Bugfix][DCP] Handle sparse MLA metadata after DCP Manager refactor

[cjackal]

Purpose

DCP Manager refactoring #50484 is focused on DCP for dense MLA, which does not handle sparse MLA properly.

• MLASparseMetadata does not inherit MLACommonMetadata, thus attn_metadata.decode attribute access should be guarded by getattr
• DCP query-gather assumes a decode query, so its workspace (direct_dcp_q_gather_workspace) is too small for prefill queries (which is passed via forced MQA codepath).
• Forced MQA codepath must be considered in DCP Manager combine as well.

Note that the two sparse MLA backend (flashinfer_mla_sparse, flashmla_sparse) shares the same codepath and thus this PR covers both.

Test Plan

Cherry-pick #46514 atop of this PR and run GLM-5.2 with DCP.

# On H100 x 8 worker node
vllm serve zai-org/GLM-5.2-FP8 \
  --kv-cache-dtype fp8_ds_mla \
  --tensor-parallel-size 8 \
  --decode-context-parallel-size 8 \
  --enable-auto-tool-choice \
  --tool-call-parser glm47 \
  --reasoning-parser glm45

1. Pass a sufficiently long (topk_tokens=2k for GLM-5.2) input and check that model server runs okay.
2. No accuracy regression

Test Result

1. vllm bench serve --dataset-name random --random-input-len 4096 --random-output-len 512 --random-range-ratio 0 --num-prompts 1

============ Serving Benchmark Result ============
Successful requests:                     1
Failed requests:                         0
Benchmark duration (s):                  9.60
Total input tokens:                      4096
Total generated tokens:                  512
Request throughput (req/s):              0.10
Output token throughput (tok/s):         52.26
Peak output token throughput (tok/s):    44.00
Peak concurrent requests:                1.00
Total token throughput (tok/s):          470.33
--------------Time to First Token----------------
Mean TTFT (ms):                          795.75
Median TTFT (ms):                        795.75
P99 TTFT (ms):                           795.75
-----Time per Output Token (excl. 1st token)-----
Mean TPOT (ms):                          17.61
Median TPOT (ms):                        17.61
P99 TPOT (ms):                           17.61
---------------Inter-token Latency---------------
Mean ITL (ms):                           22.90
Median ITL (ms):                         22.89
P99 ITL (ms):                            24.36
=================================================

2. GSM-8K, max_tokens=4096

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9466|±  |0.0044|
|     |       |strict-match    |     5|exact_match|↑  |0.9473|±  |0.0044|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[cjackal]

Also cc @GirasoleY in case this PR may not befit the DCP Manager refactor.

[Leoyzen]

Production validation from the GLM-5.2 / DCP pipeline this PR fixes.

Context

We run GLM-5.2-NVFP4 in production (TP8, DCP4, MTP spec decode x5, fp8_ds_mla KV, Mooncake KV store) on a rebuilt stack that includes #46514 + this PR, and hit exactly the failure this PR addresses: after #50484 landed, long-context decode collapsed into repetition with ~100% acceptance — matching the per-token seq_lens localization problem noted in #46514.

Before (with #46514, without this PR)

• 8+2 -> reasoning output degenerated into jumbled fragments, no answer
• consecutive requests accumulated corruption
• SpecDecoding acceptance collapsed to ~100% (per-position: 1.000, 1.000, 0.260, ...), i.e. the rejection sampler lost its signal because the target model itself was corrupted
• Counting 1-200 broke mid-sequence (repeated tokens)

After (this PR cherry-picked on top of #46514)

• 8+2 -> 10; 3+4 repeated 3x -> 7 each time (no accumulation)
• translate hello -> bonjour
• Counting 1-200 -> 1..170 consecutive, no repetition
• Acceptance back to a healthy reject curve: per-position 0.97, 0.83, 0.40, 0.39, 0.38 (and 0.89, 0.54, 0.39, ... on harder prompts), 0 crashes over the test session

Notes

• The getattr(attn_metadata, "decode", None) guard, the _direct_workspace_query_gather fallback for forced-MQA, and the use_forced_mqa seq_lens/query_start_loc handling all apply cleanly on top of [Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP #46514.
• dcp_comm_backend=ag_rs (not a2a) in our config; the combine path goes through cp_lse_ag_out_rs.

Thanks for the fix — this unblocks GLM-5.2 + DCP + MTP production serving.

[Leoyzen]

Crash in lse_reduce: "token count exceeds symmetric buffer capacity" under forced-MQA prefill

Found this while testing GLM-5.2-NVFP4 in production (TP8, decode-context-parallel-size 4, dcp-comm-backend a2a, MTP-5 speculative decoding, max-num-batched-tokens 8192, chunked prefill enabled). A long-input request consistently kills the engine:

RuntimeError: Worker failed with error 'direct_dcp_a2a_lse_reduce, .../dcp_direct_a2a_lse_reduce.cu:254, token count exceeds symmetric buffer capacity'

Root cause

get_dcp_workspace_max_num_tokens (vllm/v1/attention/ops/dcp_utils.py:97-116) sizes the symmetric buffer for decode-only traffic:

tokens_per_seq = 1 + 2*5 = 11 # MTP-5, parallel drafting
capacity = min(max_num_batched_tokens, max(max_num_seqs * tokens_per_seq, cudagraph_capture))
= min(8192, max(128 * 11, ...))
= 1408

However, on the forced-MQA path (the one this PR already fixes on the query-gather side), mla_attention.py:824-826 sets num_mqa_tokens = q.size(0), i.e. all batch tokens including prefill. A prefill-heavy step can pass up to max_num_batched_tokens (8192) tokens into dcp_manager.combine → lse_reduce → the .cu:254 check fails (8192 >> 1408).

The q_gather fallback added in this PR covers _direct_workspace_query_gather, but the lse_reduce side has no fallback: _init_combine still returns direct_workspace.lse_reduce directly (dcp_utils.py:656-658).

Suggested fixes

1. Fallback, mirroring the q_gather approach: in DirectDCPA2AWorkspace.lse_reduce (dcp_utils.py:249), when partial_output.shape[0] > self.max_num_tokens, fall back to the NCCL dcp_a2a_lse_reduce (buffers are allocated dynamically there, so there is no capacity check). Zero memory cost; only oversized prefill steps pay extra latency.
2. Formula fix: change min(...) to max(...) in get_dcp_workspace_max_num_tokens so the capacity also covers the all-token case. Cost: ~6x symmetric buffer memory (1408 → 8192 tokens).

Workaround used in production

VLLM_USE_DIRECT_DCP_A2A=0 falls back to the NCCL path (dynamic buffers, no capacity check); we measured roughly +15-30% DCP latency on prefill, no impact on decode.

Happy to test a patch or provide further details.

[cjackal]

Crash in lse_reduce: "token count exceeds symmetric buffer capacity" under forced-MQA prefill

Found this while testing GLM-5.2-NVFP4 in production (TP8, decode-context-parallel-size 4, dcp-comm-backend a2a, MTP-5 speculative decoding, max-num-batched-tokens 8192, chunked prefill enabled). A long-input request consistently kills the engine:

Thank you for pointing it out! I just have added the corresponding fallback to A2A attention combine branch, I had slipped it out when testing the PR as the default dcp comm backend is "ag_rs" which does not hit the direct workspace a2a code branch.

I haven't tested the code change yet, will do in this workday but it would be grateful if you can test it as well. I have taken the fallback approach, so the overall latency drop for the problematic queries are the same as VLLM_USE_DIRECT_DCP_A2A=0 that you have already measured.

[Leoyzen]

@cjackal Follow-up on the fallback choice: for the formula side, I've been running the max() version of get_dcp_workspace_max_num_tokens in production (GLM-5.2-NVFP4, TP8, DCP4, MTP-5, max-num-batched-tokens 8192) since yesterday — 0 capacity crashes, health OK, no OOM.

The memory cost is negligible, worth quantifying for the record: with heads_per_rank=8 (64 heads / TP8), head_dim=512 (kv_lora_rank), the symmetric buffer scales linearly with max_num_tokens — going 1408 → 8192 adds ~446 MB/GPU for the A2A lse_reduce workspace (output 2×4×8×512×2B × ΔT), plus ~70-280 MB for q_gather (padded heads), so ~0.5-0.7 GB/GPU total ≈ 0.4% of H200 memory.

Not proposing to change your fix (the fallback is correct and strictly better latency-wise for the common case) — just documenting that taking max() as well would remove the NCCL fallback divergence at negligible cost, if you'd prefer a single code path. Either way, happy to test a patch.

[mergify]

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

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

[LucasWilkinson]

Thanks for the contribution, overall looks good, one nit

[LucasWilkinson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85249 for commit 5c6854acc06f.

[github-actions]

✅ @cjackal, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[LucasWilkinson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85251 for commit 80314a18c0ce.

[cjackal]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #85251.