From dc164643897a0b43b5b388551c1ea118fbaca726 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 08:55:54 +0000 Subject: [PATCH] fix: isolate DS4 stream events across CUDA graphs --- .../test_attention_stream_events.py | 160 ++++++++++++++++++ tests/test_multi_stream_utils.py | 58 ++++++- vllm/compilation/breakable_cudagraph.py | 14 +- vllm/compilation/cuda_graph.py | 17 +- vllm/models/deepseek_v4/attention.py | 118 ++++++++----- vllm/models/deepseek_v4/nvidia/b12x.py | 1 + .../deepseek_v4/nvidia/flashinfer_sparse.py | 3 + vllm/models/deepseek_v4/nvidia/model.py | 16 +- vllm/utils/multi_stream_utils.py | 76 ++++++++- vllm/v1/worker/gpu/cudagraph_utils.py | 8 +- 10 files changed, 407 insertions(+), 64 deletions(-) create mode 100644 tests/models/deepseek_v4/test_attention_stream_events.py diff --git a/tests/models/deepseek_v4/test_attention_stream_events.py b/tests/models/deepseek_v4/test_attention_stream_events.py new file mode 100644 index 000000000000..3611777e2d2a --- /dev/null +++ b/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]))] diff --git a/tests/test_multi_stream_utils.py b/tests/test_multi_stream_utils.py index cdfd03b828ac..c76183871743 100644 --- a/tests/test_multi_stream_utils.py +++ b/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] = [] diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index e5e157d26068..aba7c4a7d21d 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/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 diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 4ac705926004..c108d879745a 100644 --- a/vllm/compilation/cuda_graph.py +++ b/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) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 1a421071b207..8c86b36ace3d 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -6,6 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from typing import TYPE_CHECKING, Any, ClassVar, cast import torch @@ -51,6 +52,7 @@ from vllm.models.deepseek_v4.compressor import DeepseekCompressor from vllm.utils.math_utils import cdiv from vllm.utils.multi_stream_utils import ( + CUDAGraphCaptureEventPool, execute_in_parallel, maybe_execute_in_parallel, ) @@ -126,6 +128,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # launch loop can enqueue the independent default-stream Q branch first. # The stream/event dependency graph remains identical. enqueue_default_before_indexer: ClassVar[bool] = False + # Some custom attention kernels are captured as part of a larger FULL graph + # and cannot safely replay the C4 post-GEMM work from auxiliary streams. + enable_post_gemm_aux_streams: ClassVar[bool] = True # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather # workspace allocated in _forward_prefill and is also read by the dummy-run # path to pre-reserve that workspace. @@ -163,6 +168,21 @@ def _uses_fp8_ds_mla_layout(self) -> bool: """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" return self.use_fp8_ds_mla_layout + def _post_gemm_aux_stream(self, index: int) -> torch.cuda.Stream | None: + if not self.enable_post_gemm_aux_streams or self.aux_stream_list is None: + return None + return self.aux_stream_list[index] + + def _post_gemm_event_lease( + self, + ) -> AbstractContextManager[list[torch.cuda.Event]]: + pool = getattr(self, "attn_event_pool", None) + # attention_impl can execute eagerly between CUDA graph segments. Its + # event handles must not be recycled into another graph artifact. + if pool is None: + return nullcontext(self.attn_events) + return pool.lease(private_eager=True) + def __init__( self, vllm_config: VllmConfig, @@ -278,15 +298,16 @@ def __init__( self.indexer_rotary_emb = self.rotary_emb self.topk_indices_buffer = topk_indices_buffer + # Will be None on ROCm for now. + self.aux_stream_list = aux_stream_list + self.indexer = None if self.compress_ratio == 4: # Only C4A uses sparse attention and hence has indexer. # aux_stream_list[2] is free here (outer GEMMs joined) for the inner # overlap of wq_b+fused_indexer_q_rope_quant vs compressor. None on # ROCm, where aux_stream_list is None. - indexer_aux_stream = ( - aux_stream_list[2] if aux_stream_list is not None else None - ) + indexer_aux_stream = self._post_gemm_aux_stream(2) self.indexer = DeepseekV4Indexer( vllm_config, config=config, @@ -301,12 +322,12 @@ def __init__( topk_scores_buffer=topk_scores_buffer, ) - # Will be None on ROCm for now. - self.aux_stream_list = aux_stream_list - # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; - # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins - # before post-GEMM starts. + # Keep the GEMM and post-GEMM event generations independent. The first + # join only enqueues waits on the default stream; re-recording those + # events for cache/indexer overlap before the waits execute can race. self.ln_events = [torch.cuda.Event() for _ in range(4)] + self.attn_event_pool = CUDAGraphCaptureEventPool(3) + self.attn_events = self.attn_event_pool.default_events assert cache_config is not None, "DeepseekV4 attention requires cache_config" # ---- Attention / KV-cache setup ---- @@ -656,30 +677,33 @@ def wq_b_kv_insert() -> torch.Tensor: # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the # MLA compressor. Slot [2] is reserved for the indexer's inner # overlap. ROCm (aux_streams is None) falls back to sequential. - q, _ = execute_in_parallel( - wq_b_kv_insert, - [ - lambda: indexer( - hidden_states, - qr, - indexer_kv_score, - indexer_weights, - positions, - self.indexer_rotary_emb, + with self._post_gemm_event_lease() as attn_events: + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + attn_events[0], + [attn_events[1], attn_events[2]], + [aux_streams[0], aux_streams[1]] + if aux_streams is not None + else None, + enable=( + aux_streams is not None and self.enable_post_gemm_aux_streams ), - lambda: compressor(kv_score, positions, self.rotary_emb), - ], - self.ln_events[0], - [self.ln_events[1], self.ln_events[2]], - [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, - enable=aux_streams is not None, - enqueue_default_first=self.enqueue_default_before_indexer, - ) + enqueue_default_first=self.enqueue_default_before_indexer, + ) elif self.compressor is not None: # wq_b + kv_insert on default, compressor on aux. - aux_stream = ( - self.aux_stream_list[0] if self.aux_stream_list is not None else None - ) + aux_stream = self._post_gemm_aux_stream(0) compressor = self.compressor def wq_b_kv_insert() -> torch.Tensor: @@ -687,13 +711,14 @@ def wq_b_kv_insert() -> torch.Tensor: q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) return q - q, _ = maybe_execute_in_parallel( - wq_b_kv_insert, - lambda: compressor(kv_score, positions, self.rotary_emb), - self.ln_events[0], - self.ln_events[1], - aux_stream, - ) + with self._post_gemm_event_lease() as attn_events: + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + attn_events[0], + attn_events[1], + aux_stream, + ) else: # SWA-only layer: no compressor, no overlap. q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) @@ -1075,10 +1100,8 @@ def __init__( # None on ROCm — maybe_execute_in_parallel falls back to sequential. self.aux_stream = aux_stream - self.ln_events: list[torch.cuda.Event] = [ - torch.cuda.Event(), - torch.cuda.Event(), - ] + self.event_pool = CUDAGraphCaptureEventPool(2) + self.ln_events = self.event_pool.default_events def forward( self, @@ -1107,11 +1130,12 @@ def wq_b_and_q_quant(): # compressor returns None and writes K to the indexer KV cache; the # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), k = maybe_execute_in_parallel( - wq_b_and_q_quant, - lambda: compressor(compressed_kv_score, positions, rotary_emb), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) + with self.event_pool.lease(private_eager=True) as events: + (q_quant, weights), k = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + events[0], + events[1], + self.aux_stream, + ) return self.indexer_op(hidden_states, q_quant, k, weights) diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index bf2cf25a0e32..b51fe2cdcdfc 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -437,6 +437,7 @@ class DeepseekV4B12xMLAAttention(DeepseekV4FlashMLAAttention): DeepseekV4B12xMLASparseBackend ) enqueue_default_before_indexer: ClassVar[bool] = True + enable_post_gemm_aux_streams: ClassVar[bool] = True @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 1848c1930db0..afb8d6c94e80 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -529,6 +529,9 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention): """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels.""" backend_cls = DeepseekV4FlashInferMLASparseBackend + # The outer attention path joins the indexer and compressor events before + # this backend consumes their cache writes and top-k output. + enable_post_gemm_aux_streams: ClassVar[bool] = True use_fp8_ds_mla_layout: ClassVar[bool] = True @staticmethod diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index aabcd92431e9..f3535bfddb3c 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1585,9 +1585,19 @@ def forward( ) if idx + 1 in self.aux_hidden_state_layers: # Reconstruct the aux hidden state for draft models - aux_recon = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) + if layer._should_run_b12x_mhc(int(hidden_states.shape[0])): + from b12x.integration.residual import b12x_mhc_post + + aux_recon = b12x_mhc_post( + hidden_states, + residual, + post_mix, + res_mix, + ) + else: + aux_recon = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) aux_hidden_states.append(aux_recon.mean(dim=1)) final_aux_recon = aux_recon if layer is not None: diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index 148b87664d7e..cd5ea7d60d0c 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -1,12 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from contextvars import ContextVar from enum import Enum from typing import Any import torch +_vllm_cudagraph_capture_depth: ContextVar[int] = ContextVar( + "vllm_cudagraph_capture_depth", default=0 +) + + +@contextmanager +def vllm_cudagraph_capture_scope() -> Iterator[None]: + """Mark Python execution as belonging to a vLLM CUDA graph capture. + + CUDA reports capture status per current stream. Python custom ops can run + on auxiliary streams joined to a multi-stream graph, where querying only + the current stream is insufficient to establish event-handle lifetime. + """ + token = _vllm_cudagraph_capture_depth.set(_vllm_cudagraph_capture_depth.get() + 1) + try: + yield + finally: + _vllm_cudagraph_capture_depth.reset(token) + + +def is_vllm_cudagraph_capture_active() -> bool: + return _vllm_cudagraph_capture_depth.get() > 0 + class AuxStreamType(Enum): Attention = 1 @@ -17,6 +42,55 @@ class EventType(Enum): Attention = 1 +class CUDAGraphCaptureEventPool: + """Keep CUDA event generations private to each graph or eager call. + + Reusing one event handle across independently captured graphs is unsafe + when those graphs can be replayed at different shapes. A replay may record + a new generation while another graph still has waits bound to the same + handle. Some custom ops also execute eagerly between CUDA graph segments, + where capture-state detection is false even though adjacent graph shapes + can still overlap event generations. + + Every real capture gets a retained private set embedded only in that graph. + Eager graph-break callers request a fresh set for every Python invocation; + the wrappers stay alive through enqueue and are then released. CUDA event + destruction is asynchronous for pending work, which is also the lifetime + pattern used by :meth:`torch.cuda.Stream.wait_stream`. Event handles are + never recycled into another graph artifact. + """ + + def __init__(self, num_events: int) -> None: + if num_events < 1: + raise ValueError("num_events must be at least one") + self.num_events = num_events + self.default_events = [torch.cuda.Event() for _ in range(num_events)] + self._captured_event_sets: list[list[torch.cuda.Event]] = [] + + @contextmanager + def lease(self, *, private_eager: bool = False) -> Iterator[list[torch.cuda.Event]]: + if ( + is_vllm_cudagraph_capture_active() + or torch.cuda.is_current_stream_capturing() + ): + events = [torch.cuda.Event() for _ in range(self.num_events)] + # CUDA graphs retain the event handles, and this list keeps the + # wrappers alive for the same lifetime as the owning module. + self._captured_event_sets.append(events) + yield events + return + + if not private_eager: + yield self.default_events + return + + # Keep these wrappers alive until the caller has enqueued every record + # and wait. cudaEventDestroy then defers resource release until pending + # device work completes, so no Python-side retention is required. + events = [torch.cuda.Event() for _ in range(self.num_events)] + yield events + + def maybe_execute_in_parallel( fn0: Callable[[], Any], fn1: Callable[[], Any], diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index f54bfcc797b9..1a5f1991fe04 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -33,6 +33,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import round_up +from vllm.utils.multi_stream_utils import vllm_cudagraph_capture_scope from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables @@ -360,7 +361,11 @@ def capture( capture. """ attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {} - with graph_capture(device=self.device): + # Keep event handles created by descriptor warmups alive together with + # the graph artifacts captured below. Some multi-stream custom ops run + # on joined auxiliary streams where CUDA's per-current-stream capture + # query is false even though later graph nodes retain those handles. + with graph_capture(device=self.device), vllm_cudagraph_capture_scope(): # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger # activations so FULL activations should fit in already allocated # buffers in the graph pool. @@ -412,6 +417,7 @@ def capture( guard_b12x_kernel_resolution( "vLLM full CUDA graph capture after B12X warmup" ), + vllm_cudagraph_capture_scope(), torch.cuda.graph(graph, self.pool), ): forward_fn(CUDAGraphMode.NONE)