[Bugfix][SparseMLA] Fix piecewise cudagraph capture crash in index group

[JaredforReal]

What

Fix an engine-startup crash when a sparse-MLA model captures decode batches with piecewise cudagraphs (e.g. GLM-5.3-Flash, a hybrid KDA + sparse-MLA model):

File "vllm/v1/attention/backends/mla/index_group.py", line 133, in _convert_once
    self.side_stream.wait_event(self.logical_topk_ready)
torch.AcceleratorError: CUDA error: invalid argument

Full startup traceback (4x GB300, TP4, load_format=dummy; identical with real weights, the capture path does not read weights):

File "vllm/v1/worker/gpu/model_runner.py", line 988, in profile_cudagraph_memory
File "vllm/v1/worker/gpu/cudagraph_utils.py", line 920, in profile_cudagraph_memory
File "vllm/v1/worker/gpu/cudagraph_utils.py", line 456, in capture
File "vllm/compilation/breakable_cudagraph.py", line 206, in add_eager
File "vllm/model_executor/layers/attention/mla_attention.py", line 1483, in unified_mla_attention_with_output
File "vllm/model_executor/layers/attention/sparse_mla_attention.py", line 736, in _convert_logical_to_physical_topk
File "vllm/v1/attention/backends/mla/index_group.py", line 95, in convert_logical_to_physical_topk
File "vllm/v1/attention/backends/mla/index_group.py", line 133, in _convert_once
    self.side_stream.wait_event(self.logical_topk_ready)
torch.AcceleratorError: CUDA error: invalid argument

Root cause

SparseMLAIndexGroup overlaps the logical→physical topk conversion with the compute stream via the logical_topk_ready event: the indexer records it on the compute stream after computing topk, and _convert_once waits on it on the side stream.

In piecewise cudagraph mode the indexer runs inside a captured graph segment, while the consumer runs in the following eager break (unified_mla_attention_with_output is an @eager_break_during_capture point, executed by add_eager() after the segment's capture_end()). An event recorded during stream capture is graph-local; waiting on it eagerly is cudaErrorInvalidValue.

Minimal reproduction of the exact pattern (no vLLM involved):

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    ev.record(torch.cuda.current_stream())   # inside capture
side.wait_event(ev)                          # eager, after capture -> invalid argument

Why only some models: FULL cudagraph capture (non-hybrid sparse MLA, e.g. DeepSeek-V3.2 decode) records and waits in-graph, which is legal; prefill-sized batches take the workspace-overflow branch of _convert_once, which never touches the event. Hybrid models like GLM-5.3-Flash capture decode batches piecewise and hit the crash during profile_cudagraph_memory.

Fix

In piecewise mode, skip the event entirely: set_logical_topk_ready does not record (a record inside a capture segment would be graph-local and unusable anyway), and _convert_once orders the side stream with wait_stream(compute_stream) — identical to the existing has_indexer=False path, and correct because the convert is always invoked after the indexer on the same compute stream. FULL/eager modes keep the fine-grained event. The same guard is applied to the HiSparse copy stream. Adds in_piecewise_cudagraph() to forward_context.py.

Tests

• New regression test (record inside captured segment + convert in eager break under BreakableCUDAGraphCapture); fails pre-fix at index_group.py:133 with the production error, passes post-fix:

  pytest tests/v1/attention/test_sparse_mla_backends.py -k piecewise_capture

• Sparse MLA subset (convert / index_group / piecewise / hisparse): 47 passed

  pytest tests/v1/attention/test_sparse_mla_backends.py -k "convert or index_group or piecewise or hisparse"

• E2E on 4x GB300, GLM-5.3-Flash real weights, TP4: engine starts with cudagraphs enabled (crashed at startup before), generation works.

Model evaluation

GSM8K, 5-shot, 1319 questions, GLM-5.3-Flash TP4 on 4x GB300 (this branch, tests/evals/gsm8k/gsm8k_eval.py, no speculative decoding): 91.8% accuracy (invalid rate 0.0%).

Notes

• Bisect: the crash starts with [3/N] HiSparse: host-resident sparse-MLA decode hot-buffering #53781, which introduced the side-stream + logical_topk_ready event for the topk conversion (before it, the convert ran on the compute stream with no event). Verified empirically: its parent commit starts cleanly, it crashes at engine startup. Purely informational for locating the change — the piecewise interaction was easy to miss.
• Duplicate check: no open PR covers this crash; [Bugfix][Model] Bake real KV caches into captured fused_norm_rope under PIECEWISE (fixes silent garbage output for DSA models on multi-node TP) #54851 is a different piecewise-mode issue (KV-cache writes baked into captured graphs for DSA models on multi-node TP).
• AI assistance: developed with AI assistance (Kimi Code); all changes reviewed and tested by the submitter.

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

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88820 for commit 771908b2197e.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89001 for commit 0605face83e4.

[JaredforReal]

/ci retry

[github-actions]

✅ Queued 3 failed job(s) for retry in Buildkite CI #89001.