[Bugfix][SparseMLA] Fix piecewise cudagraph capture crash in index group (#56825)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Kimi Code <noreply@moonshot.ai>
Jared Wen committed
Sep 15, 2026 at 14:51 UTC
51918e252af4dfff381001f42b7ebe6a0454501b
4 files changed
+100
-6
tests/v1/attention/test_sparse_mla_backends.py
+77
-1
@@ -22,10 +22,17 @@ from tests.v1.attention.utils import (
22
create_vllm_config,
23
)
24
from vllm import _custom_ops as ops
25
-from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config
25
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
26
+from vllm.config import (
27
+ CUDAGraphMode,
28
+ HiSparseConfig,
29
+ SpeculativeConfig,
30
+ set_current_vllm_config,
31
+)
32
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
33
HiSparseConnectorWorker,
34
)
35
+from vllm.forward_context import set_forward_context
36
from vllm.model_executor.layers.attention import mla_attention
37
from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha
38
from vllm.model_executor.layers.attention.sparse_mla_attention import (
@@ -897,6 +904,75 @@ def test_triton_convert_req_index_to_global_index_decode_only(
904
torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
905
906
907
+def test_index_group_convert_during_piecewise_capture():
908
+ """In piecewise cudagraph mode the indexer is captured while the convert
909
+ runs in the following eager break. The convert's side stream must not
910
+ wait on ``logical_topk_ready`` there: an event recorded inside a captured
911
+ segment is graph-local, and an eager wait on it raises
912
+ cudaErrorInvalidValue."""
913
+ device = torch.device(DEVICE_TYPE)
914
+ num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16
915
+
916
+ logical_topk_indices = torch.randint(
917
+ 0,
918
+ block_size * blocks_per_req,
919
+ (num_tokens, num_topk),
920
+ dtype=torch.int32,
921
+ device=device,
922
+ )
923
+ builder = SparseMLAIndexGroupBuilder(logical_topk_indices)
924
+ group, layer_index = builder.register_layer(is_index_producing_layer=True)
925
+ assert layer_index == 0 and group.has_indexer
926
+
927
+ req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests
928
+ block_table = torch.arange(
929
+ num_requests * blocks_per_req, dtype=torch.int32, device=device
930
+ ).view(num_requests, blocks_per_req)
931
+ attn_metadata = SimpleNamespace(
932
+ req_id_per_token=req_id, block_table=block_table, block_size=block_size
933
+ )
934
+ expected = triton_convert_req_index_to_global_index(
935
+ req_id,
936
+ block_table,
937
+ logical_topk_indices,
938
+ BLOCK_SIZE=block_size,
939
+ NUM_TOPK_TOKENS=num_topk,
940
+ )
941
+
942
+ vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
943
+ marker = torch.zeros(1, device=device)
944
+ capture_stream = torch.cuda.Stream(device=device)
945
+ # torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream,
946
+ # so current_stream() tracks the capture stream as in production.
947
+ with (
948
+ torch.cuda.stream(capture_stream),
949
+ set_forward_context(
950
+ None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
951
+ ),
952
+ ):
953
+ capture = BreakableCUDAGraphCapture()
954
+ with capture:
955
+ marker.add_(1)
956
+ # Record inside a captured segment, convert in the eager
957
+ # break — mirrors mla.py -> unified_mla_attention_with_output.
958
+ group.set_logical_topk_ready(layer_index)
959
+ capture.add_eager(
960
+ lambda: group.convert_logical_to_physical_topk(
961
+ layer_index,
962
+ logical_topk_indices,
963
+ attn_metadata,
964
+ block_stride_rows=None,
965
+ return_valid_counts=False,
966
+ )
967
+ )
968
+ marker.add_(1)
969
+ capture.replay()
970
+ capture_stream.synchronize()
971
+
972
+ result = group.physical_topk_indices[:num_tokens]
973
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
974
+
975
+
976
@pytest.mark.parametrize("block_size", [16])
977
@pytest.mark.skipif(
978
torch.cuda.get_device_capability() < (9, 0),
vllm/forward_context.py
+9
@@ -209,6 +209,15 @@ def is_forward_context_available() -> bool:
209
return _forward_context is not None
210
211
212
+def in_piecewise_cudagraph() -> bool:
213
+ """Whether the current forward runs in piecewise cudagraph mode (graph
214
+ segments separated by eager breaks), at capture or replay time."""
215
+ return (
216
+ is_forward_context_available()
217
+ and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
218
+ )
219
+
220
+
221
def create_forward_context(
222
attn_metadata: Any,
223
vllm_config: VllmConfig,
vllm/v1/attention/backends/mla/index_group.py
+12
-3
@@ -10,6 +10,7 @@ import torch
10
11
from vllm import _custom_ops as ops
12
from vllm.config import VllmConfig
13
+from vllm.forward_context import in_piecewise_cudagraph
14
from vllm.utils.torch_utils import current_stream
15
from vllm.v1.attention.backends.mla.sparse_utils import (
16
triton_convert_req_index_to_global_index,
@@ -56,8 +57,16 @@ class SparseMLAIndexGroup:
57
return layer_index
58
59
def set_logical_topk_ready(self, layer_index: int) -> None:
59
- if layer_index == 0 and self.has_indexer:
60
- self.logical_topk_ready.record(current_stream())
60
+ if layer_index != 0 or not self.has_indexer:
61
+ return
62
+ # In piecewise cudagraph mode the consumer converts on the side stream
63
+ # from an eager break; an event recorded inside a captured segment is
64
+ # graph-local and any eager wait on it raises cudaErrorInvalidValue.
65
+ # The convert side falls back to wait_stream there instead (and never
66
+ # touches this event), so skip the record as well.
67
+ if in_piecewise_cudagraph():
68
+ return
69
+ self.logical_topk_ready.record(current_stream())
70
71
def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
72
pass
@@ -129,7 +138,7 @@ class SparseMLAIndexGroup:
138
valid_topk_counts = self.valid_topk_counts[:num_tokens]
139
compute_stream = current_stream()
140
if layer_index == 0:
132
- if self.has_indexer:
141
+ if self.has_indexer and not in_piecewise_cudagraph():
142
self.side_stream.wait_event(self.logical_topk_ready)
143
else:
144
self.side_stream.wait_stream(compute_stream)
vllm/v1/hisparse/runtime.py
+2
-2
@@ -16,7 +16,7 @@ import torch
16
17
from vllm.config import CUDAGraphMode, VllmConfig
18
from vllm.distributed import get_tp_group
19
-from vllm.forward_context import get_forward_context
19
+from vllm.forward_context import get_forward_context, in_piecewise_cudagraph
20
from vllm.logger import init_logger
21
from vllm.platforms import current_platform
22
from vllm.utils.math_utils import round_up
@@ -975,7 +975,7 @@ class HiSparseRuntime:
975
) -> None:
976
group = self.index_group
977
compute_stream = current_stream()
978
- if group.logical_topk_ready is not None:
978
+ if group.logical_topk_ready is not None and not in_piecewise_cudagraph():
979
group.copy_stream.wait_event(group.logical_topk_ready)
980
else:
981
group.copy_stream.wait_stream(compute_stream)