diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 1227669e39ab..01dd3c5f2597 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -839,6 +839,7 @@ def run_inv_rope(*args, **kwargs): def test_b12x_mhc_uses_public_plan_bind_run(monkeypatch) -> None: calls: dict[str, Any] = {} + retained_bindings: list[Any] = [] def make_caps(**kwargs): calls["caps"] = kwargs @@ -878,6 +879,11 @@ def run_post(*args): ) monkeypatch.setattr(b12x_mla, "_require_b12x_mhc", lambda: module) monkeypatch.setattr(b12x_mla, "current_workspace_manager", lambda: _Workspace()) + monkeypatch.setattr( + b12x_mla, + "retain_cuda_graph_capture_resource", + retained_bindings.append, + ) mhc = b12x_mla.B12xMHCResidual( hidden_size=256, @@ -918,6 +924,10 @@ def run_post(*args): assert calls["bind"][1]["scratch"].dtype == torch.uint8 assert calls["pre"][1]["binding"].expected_m == 3 assert calls["post_pre"][1]["expected_m"] == 3 + assert retained_bindings == [ + calls["pre"][1]["binding"], + calls["post_pre"][1]["binding"], + ] assert residual_out.shape == (3, 4, 256) assert layer_input.shape == (3, 256) assert final is next_outputs[0] diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index dc9c1ffb3067..a0667e3895c7 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -78,6 +78,51 @@ def forward_fn(cg_mode): assert create_calls[0][1] is not create_calls[1][1] +def test_breakable_wrapper_retains_capture_resources(monkeypatch): + import vllm.compilation.breakable_cudagraph as breakable + from vllm.v1.worker.workspace import retain_cuda_graph_capture_resource + + class FakeCapture: + def __init__(self, pool): + self.pool = pool + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return None + + class FakeOffloader: + def sync_prev_onload(self): + pass + + def join_after_forward(self): + pass + + resource = object() + + def runnable(): + assert retain_cuda_graph_capture_resource(resource) + return object() + + monkeypatch.setattr(breakable, "validate_cudagraph_capturing_enabled", lambda: None) + monkeypatch.setattr(breakable, "set_graph_pool_id", lambda _pool: None) + monkeypatch.setattr(breakable.gc, "collect", lambda: None) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + monkeypatch.setattr(breakable, "get_offloader", lambda: FakeOffloader()) + monkeypatch.setattr(breakable, "BreakableCUDAGraphCapture", FakeCapture) + monkeypatch.setattr(breakable, "weak_ref_tensors", lambda value: value) + + wrapper = object.__new__(breakable.BreakableCUDAGraphWrapper) + wrapper.runnable = runnable + wrapper.graph_pool = object() + entry = breakable._BreakableEntry(batch_descriptor=object()) + + wrapper._capture(entry, (), {}) + + assert entry.resources == [resource] + + @pytest.fixture(autouse=True) def _reset_breakable_tls(): """Defensively clear thread-local capture state between tests so a diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index a909ea9af89d..8bbe8e9cc0d1 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -100,3 +100,21 @@ def test_workspace_lane_validation(monkeypatch) -> None: with pytest.raises(ValueError, match="at least one"): workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) + + +def test_cuda_graph_capture_resources_are_scoped_to_collector() -> None: + outside = object() + first = object() + nested = object() + second = object() + + assert not workspace.retain_cuda_graph_capture_resource(outside) + with workspace.collect_cuda_graph_capture_resources() as resources: + assert workspace.retain_cuda_graph_capture_resource(first) + with workspace.collect_cuda_graph_capture_resources() as nested_resources: + assert workspace.retain_cuda_graph_capture_resource(nested) + assert workspace.retain_cuda_graph_capture_resource(second) + + assert resources == [first, second] + assert nested_resources == [nested] + assert not workspace.retain_cuda_graph_capture_resource(outside) diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 6da3ec717861..f4fe203714fb 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -45,6 +45,7 @@ from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors +from vllm.v1.worker.workspace import collect_cuda_graph_capture_resources logger = init_logger(__name__) @@ -241,6 +242,7 @@ class _BreakableEntry: capture: BreakableCUDAGraphCapture | None = None output: Any = None input_addresses: list[int] | None = None + resources: list[Any] | None = None class BreakableCUDAGraphWrapper: @@ -379,7 +381,7 @@ def _capture( get_offloader().sync_prev_onload() capture = BreakableCUDAGraphCapture(pool=self.graph_pool) - with capture: + with collect_cuda_graph_capture_resources() as resources, 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 @@ -392,6 +394,7 @@ def _capture( output = weak_ref_tensors(output) entry.capture = capture + entry.resources = resources entry.output = weak_ref_tensors(output) logger.debug( diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 6e4c40030728..649b50289091 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -38,7 +38,10 @@ from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) -from vllm.v1.worker.workspace import current_workspace_manager +from vllm.v1.worker.workspace import ( + current_workspace_manager, + retain_cuda_graph_capture_resource, +) if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import ( @@ -170,7 +173,7 @@ def _binding( raise ValueError("B12x mHC scratch plan did not provide any buffers.") scratch: torch.Tensor | tuple[torch.Tensor, ...] scratch = buffers[0] if len(buffers) == 1 else tuple(buffers) - return self._bind( + binding = self._bind( plan, scratch=scratch, tokens=tokens, @@ -180,6 +183,8 @@ def _binding( out=out, expected_m=expected_m, ) + retain_cuda_graph_capture_resource(binding) + return binding def run_pre( self, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 6203b443104e..e17d1e787b7a 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -38,6 +38,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.utils import AttentionGroup, unbind_kv_cache +from vllm.v1.worker.workspace import collect_cuda_graph_capture_resources if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -133,6 +134,7 @@ def __init__( self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map() self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} + self.graph_capture_resources: dict[BatchExecutionDescriptor, list[Any]] = {} self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None self._graphs_captured = False @@ -377,18 +379,20 @@ def capture( if self._capture_mem_samples is not None: torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] - with torch.cuda.graph(graph, self.pool): + with ( + collect_cuda_graph_capture_resources() as resources, + torch.cuda.graph(graph, self.pool), + ): forward_fn(CUDAGraphMode.NONE) - # Join offloader's copy stream after forward to avoid - # unjoined stream error. The last layer's start_prefetch - # forks copy_stream, but wait_prefetch only happens in - # the next forward pass. + # Join the offloader copy stream because the last layer + # can leave a prefetch pending at capture end. get_offloader().join_after_forward() if self._capture_mem_samples is not None: torch.accelerator.synchronize() free_after = torch.accelerator.get_memory_info()[0] self._capture_mem_samples.append(free_before - free_after) self.graphs[desc] = graph + self.graph_capture_resources[desc] = resources compilation_counter.num_cudagraph_captured += 1 self._graphs_captured = True @@ -817,6 +821,7 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: BreakableCUDAGraphWrapper.clear_all_graphs() for graph_manager in graph_managers: graph_manager.graphs.clear() + graph_manager.graph_capture_resources.clear() graph_manager._graphs_captured = False graph_manager.pool = original_manager_pools[id(graph_manager)] live_wrappers = list(CUDAGraphWrapper._all_instances) + list( diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 39b3b10349c3..e98a05c99e4e 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -8,6 +8,7 @@ from contextvars import ContextVar from itertools import accumulate from math import prod +from typing import Any import torch @@ -30,6 +31,9 @@ def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: # Global workspace manager instance _manager: "WorkspaceManager | None" = None _workspace_lane: ContextVar[int] = ContextVar("vllm_workspace_lane", default=0) +_cuda_graph_capture_resources: ContextVar[list[Any] | None] = ContextVar( + "vllm_cuda_graph_capture_resources", default=None +) @contextmanager @@ -44,6 +48,40 @@ def use_workspace_lane(lane: int) -> Iterator[None]: _workspace_lane.reset(token) +@contextmanager +def collect_cuda_graph_capture_resources() -> Iterator[list[Any]]: + """Collect objects whose storage is referenced by one CUDA graph. + + A CUDA graph records device pointers, but it does not retain the Python + objects that own those allocations. Callers that allocate custom-op output + or scratch tensors during capture can register their owner with + :func:`retain_cuda_graph_capture_resource`. The graph manager keeps the + returned list alive for exactly as long as the captured graph. + """ + resources: list[Any] = [] + token = _cuda_graph_capture_resources.set(resources) + try: + yield resources + finally: + _cuda_graph_capture_resources.reset(token) + + +def retain_cuda_graph_capture_resource(resource: Any) -> bool: + """Retain an object whose storage is referenced by a CUDA graph. + + Args: + resource: Python owner that must remain alive while the graph exists. + + Returns: + ``True`` when a capture resource collector retained the object. + """ + resources = _cuda_graph_capture_resources.get() + if resources is None: + return False + resources.append(resource) + return True + + class WorkspaceManager: """Manager for workspace allocation.