fix(ds4): isolate auxiliary-stream events across CUDA graphs

tests/models/deepseek_v4/test_attention_stream_events.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import torch
+
+from vllm.models.deepseek_v4 import attention as attention_module
+from vllm.models.deepseek_v4.nvidia.b12x import DeepseekV4B12xMLAAttention
+from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
+    DeepseekV4FlashInferSM120Attention,
+)
+
+
+def test_sm120_backends_enable_post_gemm_aux_streams_for_dspark() -> None:
+    assert DeepseekV4B12xMLAAttention.enable_post_gemm_aux_streams is True
+    assert DeepseekV4FlashInferSM120Attention.enable_post_gemm_aux_streams is True
+
+
+def test_post_gemm_aux_stream_gate_covers_every_attention_path() -> None:
+    streams = [object(), object(), object()]
+    layer = SimpleNamespace(
+        aux_stream_list=streams,
+        enable_post_gemm_aux_streams=False,
+    )
+
+    for index in range(len(streams)):
+        assert (
+            attention_module.DeepseekV4Attention._post_gemm_aux_stream(layer, index)
+            is None
+        )
+
+    layer.enable_post_gemm_aux_streams = True
+    for index, stream in enumerate(streams):
+        assert (
+            attention_module.DeepseekV4Attention._post_gemm_aux_stream(layer, index)
+            is stream
+        )
+
+
+def test_gemm_and_attention_overlap_use_distinct_event_sets(monkeypatch) -> None:
+    calls = []
+
+    def fake_execute_in_parallel(
+        default_fn,
+        aux_fns,
+        start_event,
+        done_events,
+        aux_streams,
+        enable,
+        **kwargs,
+    ):
+        del default_fn, aux_streams, enable, kwargs
+        calls.append((start_event, tuple(done_events)))
+        return torch.empty(1), [None] * len(aux_fns)
+
+    monkeypatch.setattr(
+        attention_module, "execute_in_parallel", fake_execute_in_parallel
+    )
+    monkeypatch.setattr(
+        attention_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(attn_metadata=None),
+    )
+
+    ln_events = [object() for _ in range(4)]
+    attn_events = [object() for _ in range(3)]
+    layer = SimpleNamespace(
+        aux_stream_list=[object(), object(), object()],
+        compressor=object(),
+        indexer=object(),
+        fused_wqa_wkv=object(),
+        ln_events=ln_events,
+        attn_events=attn_events,
+        _post_gemm_event_lease=lambda: nullcontext(attn_events),
+        enqueue_default_before_indexer=True,
+        enable_post_gemm_aux_streams=True,
+        indexer_rotary_emb=object(),
+        rotary_emb=object(),
+        forward_mqa=lambda *args: None,
+    )
+    tensor = torch.empty(1)
+
+    attention_module.DeepseekV4Attention.attn_gemm_parallel_execute(layer, tensor)
+    attention_module.DeepseekV4Attention.attention_impl(
+        layer,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+    )
+
+    assert calls == [
+        (ln_events[0], tuple(ln_events[1:4])),
+        (attn_events[0], tuple(attn_events[1:3])),
+    ]
+    assert set(map(id, ln_events)).isdisjoint(map(id, attn_events))
+
+
+def test_attention_overlap_uses_capture_private_events(monkeypatch) -> None:
+    calls = []
+    captured_events = [object() for _ in range(3)]
+
+    def fake_execute_in_parallel(
+        default_fn,
+        aux_fns,
+        start_event,
+        done_events,
+        aux_streams,
+        enable,
+        **kwargs,
+    ):
+        del default_fn, aux_streams, enable, kwargs
+        calls.append((start_event, tuple(done_events)))
+        return torch.empty(1), [None] * len(aux_fns)
+
+    monkeypatch.setattr(
+        attention_module, "execute_in_parallel", fake_execute_in_parallel
+    )
+    monkeypatch.setattr(
+        attention_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(attn_metadata=None),
+    )
+
+    layer = SimpleNamespace(
+        aux_stream_list=[object(), object(), object()],
+        compressor=object(),
+        indexer=object(),
+        attn_events=[object() for _ in range(3)],
+        attn_event_pool=SimpleNamespace(
+            lease=lambda **_kwargs: nullcontext(captured_events)
+        ),
+        _post_gemm_event_lease=lambda: nullcontext(captured_events),
+        enqueue_default_before_indexer=True,
+        enable_post_gemm_aux_streams=True,
+        indexer_rotary_emb=object(),
+        rotary_emb=object(),
+        forward_mqa=lambda *args: None,
+    )
+    tensor = torch.empty(1)
+
+    attention_module.DeepseekV4Attention.attention_impl(
+        layer,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+    )
+
+    assert calls == [(captured_events[0], tuple(captured_events[1:3]))]

tests/test_multi_stream_utils.py

@@ -6,7 +6,12 @@
 import pytest
 import torch
 
-from vllm.utils.multi_stream_utils import execute_in_parallel
+from vllm.utils.multi_stream_utils import (
+    CUDAGraphCaptureEventPool,
+    execute_in_parallel,
+    is_vllm_cudagraph_capture_active,
+    vllm_cudagraph_capture_scope,
+)
 
 
 class _FakeEvent:
@@ -21,6 +26,57 @@ def wait(self) -> None:
         self.calls.append(f"{self.name}.wait")
 
 
+def test_vllm_cudagraph_capture_scope_is_nested_and_exception_safe() -> None:
+    assert not is_vllm_cudagraph_capture_active()
+    with pytest.raises(RuntimeError), vllm_cudagraph_capture_scope():
+        assert is_vllm_cudagraph_capture_active()
+        with vllm_cudagraph_capture_scope():
+            assert is_vllm_cudagraph_capture_active()
+        raise RuntimeError("capture failed")
+    assert not is_vllm_cudagraph_capture_active()
+
+
+def test_cudagraph_capture_event_pool_isolates_capture_generations(monkeypatch):
+    created = []
+
+    def fake_event():
+        event = object()
+        created.append(event)
+        return event
+
+    monkeypatch.setattr(torch.cuda, "Event", fake_event)
+    capturing = False
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing)
+
+    pool = CUDAGraphCaptureEventPool(2)
+    with pool.lease() as default_events:
+        assert default_events is pool.default_events
+
+    with pool.lease(private_eager=True) as eager_a:
+        pass
+    with pool.lease(private_eager=True) as eager_b:
+        assert eager_a is not eager_b
+        assert set(map(id, eager_a)).isdisjoint(map(id, eager_b))
+    assert not pool._captured_event_sets
+
+    with (
+        vllm_cudagraph_capture_scope(),
+        pool.lease(private_eager=True) as scoped_capture,
+    ):
+        assert set(map(id, eager_a)).isdisjoint(map(id, scoped_capture))
+    assert pool._captured_event_sets == [scoped_capture]
+
+    capturing = True
+    with pool.lease() as capture_a:
+        pass
+    with pool.lease() as capture_b:
+        assert capture_a is not capture_b
+        assert set(map(id, capture_a)).isdisjoint(map(id, capture_b))
+        assert set(map(id, pool.default_events)).isdisjoint(map(id, capture_a))
+    assert len(pool._captured_event_sets) == 3
+    assert len(created) == 12
+
+
 @pytest.mark.parametrize("enqueue_default_first", [False, True])
 def test_execute_in_parallel_enqueue_order(monkeypatch, enqueue_default_first):
     calls: list[str] = []

vllm/compilation/breakable_cudagraph.py

@@ -48,6 +48,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.offloader.base import get_offloader
 from vllm.platforms import current_platform
+from vllm.utils.multi_stream_utils import vllm_cudagraph_capture_scope
 from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors
 
 logger = init_logger(__name__)
@@ -384,8 +385,7 @@ def _capture(
 
         if b12x_cuda_graph_wrapper_prewarm_enabled(
             is_piecewise=(
-                get_forward_context().cudagraph_runtime_mode
-                == CUDAGraphMode.PIECEWISE
+                get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
             )
         ):
             prewarm_output = self.runnable(*args, **kwargs)
@@ -394,9 +394,13 @@ def _capture(
             get_offloader().sync_prev_onload()
 
         capture = BreakableCUDAGraphCapture(pool=self.graph_pool)
-        with guard_b12x_kernel_resolution(
-            "vLLM BreakableCUDAGraphWrapper capture after B12X eager warmup"
-        ), capture:
+        with (
+            guard_b12x_kernel_resolution(
+                "vLLM BreakableCUDAGraphWrapper capture after B12X eager warmup"
+            ),
+            vllm_cudagraph_capture_scope(),
+            capture,
+        ):
             output = self.runnable(*args, **kwargs)
             # Join the offloader's copy stream while we still hold the last
             # segment open, so the join is captured into the graph (otherwise

vllm/compilation/cuda_graph.py

@@ -28,6 +28,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.offloader.base import get_offloader
 from vllm.platforms import current_platform
+from vllm.utils.multi_stream_utils import vllm_cudagraph_capture_scope
 from vllm.utils.torch_utils import current_stream, weak_ref_tensors
 
 logger = init_logger(__name__)
@@ -328,12 +329,16 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
                     get_offloader().sync_prev_onload()
 
                 # mind-exploding: carefully manage the reference and memory.
-                with guard_b12x_kernel_resolution(
-                    "vLLM CUDAGraphWrapper capture after B12X eager warmup"
-                ), torch.cuda.graph(
-                    cudagraph,
-                    pool=self.graph_pool,
-                    stream=current_stream(),
+                with (
+                    guard_b12x_kernel_resolution(
+                        "vLLM CUDAGraphWrapper capture after B12X eager warmup"
+                    ),
+                    vllm_cudagraph_capture_scope(),
+                    torch.cuda.graph(
+                        cudagraph,
+                        pool=self.graph_pool,
+                        stream=current_stream(),
+                    ),
                 ):
                     # `output` is managed by pytorch's cudagraph pool
                     output = self.runnable(*args, **kwargs)