fix(cudagraph): retain custom-op capture resources

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]

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

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)

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(

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,

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(

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.