fix(cudagraph): retain profiling resources through teardown

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -82,6 +82,8 @@ 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
 
+    lifecycle: list[str] = []
+
     class FakeCapture:
         def __init__(self, pool):
             self.pool = pool
@@ -102,13 +104,21 @@ def join_after_forward(self):
     resource = object()
 
     def runnable():
+        lifecycle.append("capture")
         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(
+        torch.accelerator,
+        "synchronize",
+        lambda: lifecycle.append("synchronize"),
+    )
+    monkeypatch.setattr(breakable.gc, "collect", lambda: lifecycle.append("collect"))
+    monkeypatch.setattr(
+        torch.accelerator, "empty_cache", lambda: lifecycle.append("empty-cache")
+    )
     monkeypatch.setattr(breakable, "get_offloader", lambda: FakeOffloader())
     monkeypatch.setattr(breakable, "BreakableCUDAGraphCapture", FakeCapture)
     monkeypatch.setattr(breakable, "weak_ref_tensors", lambda value: value)
@@ -121,6 +131,7 @@ def runnable():
     wrapper._capture(entry, (), {})
 
     assert entry.resources == [resource]
+    assert lifecycle == ["synchronize", "collect", "empty-cache", "capture"]
 
 
 @pytest.fixture(autouse=True)

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -47,6 +47,44 @@ def _create_vllm_config() -> MagicMock:
     return vllm_config
 
 
+@pytest.mark.parametrize("requires_raw_tokens", [False, True])
+def test_capture_token_inputs_match_runtime_embedding_contract(requires_raw_tokens):
+    """Prepared embeddings replace token IDs unless the model needs both."""
+
+    class Model:
+        requires_raw_input_tokens = requires_raw_tokens
+
+    input_ids = torch.arange(4, dtype=torch.int32)
+    inputs_embeds = torch.zeros((4, 8), dtype=torch.bfloat16)
+    model_inputs = {
+        "input_ids": input_ids,
+        "positions": torch.arange(4),
+        "inputs_embeds": inputs_embeds,
+    }
+
+    gpu_cudagraph_utils.normalize_model_token_inputs(Model(), model_inputs)
+
+    if requires_raw_tokens:
+        assert model_inputs["input_ids"] is input_ids
+    else:
+        assert model_inputs["input_ids"] is None
+    assert model_inputs["inputs_embeds"] is inputs_embeds
+
+
+def test_capture_token_inputs_keep_ids_without_embeddings():
+    """Text-only token input remains present when no embeddings are supplied."""
+    input_ids = torch.arange(4, dtype=torch.int32)
+    model_inputs = {
+        "input_ids": input_ids,
+        "positions": torch.arange(4),
+        "inputs_embeds": None,
+    }
+
+    gpu_cudagraph_utils.normalize_model_token_inputs(object(), model_inputs)
+
+    assert model_inputs["input_ids"] is input_ids
+
+
 def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
     """FULL capture must set graph_pool_id before entering torch.cuda.graph().
 

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py

@@ -51,12 +51,33 @@ def __init__(
         self._capture_mem_samples: list[int] | None = None
         self.use_breakable_cg = False
         self.graphs: dict[Any, Any] = {}
+        self.graph_capture_resources: dict[Any, list[Any]] = {}
         self._graphs_captured = False
 
     def needs_capture(self) -> bool:
         return self._needs_capture
 
 
+class _RecordingGraph:
+    def __init__(self, lifecycle: list[str], name: str) -> None:
+        self.lifecycle = lifecycle
+        self.name = name
+
+    def reset(self) -> None:
+        self.lifecycle.append(f"reset-{self.name}")
+
+
+class _RecordingDict(dict[Any, Any]):
+    def __init__(self, lifecycle: list[str], name: str) -> None:
+        super().__init__()
+        self.lifecycle = lifecycle
+        self.name = name
+
+    def clear(self) -> None:
+        self.lifecycle.append(f"clear-{self.name}")
+        super().clear()
+
+
 def _make_profiling_runner(
     cudagraph_mode: CUDAGraphMode,
     *,
@@ -108,6 +129,7 @@ def _fake_set_current_vllm_config(_cfg):
     # The profiler reads free GPU memory before/after to compute what it
     # retained; default to a constant (nothing retained).
     monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None)
+    monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None)
     monkeypatch.setattr(
         cgu.torch.accelerator, "get_memory_info", lambda: (1 << 30, 1 << 30)
     )
@@ -282,23 +304,86 @@ def test_profile_cudagraph_memory_clears_captured_graphs(monkeypatch):
     _patch_module(monkeypatch)
     runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)
 
-    cleared: list[str] = []
+    lifecycle: list[str] = []
+    monkeypatch.setattr(
+        cgu.torch.accelerator,
+        "synchronize",
+        lambda: lifecycle.append("synchronize"),
+    )
+    monkeypatch.setattr(
+        cgu.CUDAGraphWrapper,
+        "reset_all_graphs",
+        classmethod(lambda cls: lifecycle.append("reset-piecewise")),
+    )
+    monkeypatch.setattr(
+        cgu.BreakableCUDAGraphWrapper,
+        "reset_all_graphs",
+        classmethod(lambda cls: lifecycle.append("reset-breakable")),
+    )
     monkeypatch.setattr(
         cgu.CUDAGraphWrapper,
         "clear_all_graphs",
-        classmethod(lambda cls: cleared.append("piecewise")),
+        classmethod(lambda cls: lifecycle.append("clear-piecewise")),
     )
     monkeypatch.setattr(
         cgu.BreakableCUDAGraphWrapper,
         "clear_all_graphs",
-        classmethod(lambda cls: cleared.append("breakable")),
+        classmethod(lambda cls: lifecycle.append("clear-breakable")),
     )
+    runner.cudagraph_manager.graphs = _RecordingDict(lifecycle, "full")
+    runner.cudagraph_manager.graphs["profile"] = _RecordingGraph(lifecycle, "full")
+    runner.cudagraph_manager.graph_capture_resources = _RecordingDict(
+        lifecycle, "resources"
+    )
+    runner.cudagraph_manager.graph_capture_resources["profile"] = [object()]
 
     cgu.profile_cudagraph_memory(runner)
 
-    # Profiling captures are discarded so the real capture re-captures them
-    # against the KV cache.
-    assert cleared == ["piecewise", "breakable"]
+    # CUDA graph executables must be destroyed and synchronized before their
+    # B12X channel checkpoints and tensor workspaces are released.
+    assert lifecycle == [
+        "synchronize",
+        "reset-piecewise",
+        "reset-breakable",
+        "reset-full",
+        "synchronize",
+        "clear-piecewise",
+        "clear-breakable",
+        "clear-full",
+        "clear-resources",
+    ]
+
+
+def test_cuda_graph_wrappers_reset_executables_without_releasing_resources():
+    lifecycle: list[str] = []
+    piecewise = object.__new__(cgu.CUDAGraphWrapper)
+    piecewise_entry = SimpleNamespace(
+        cudagraph=_RecordingGraph(lifecycle, "piecewise"),
+        output=object(),
+    )
+    piecewise.concrete_cudagraph_entries = {"profile": piecewise_entry}
+
+    class _RecordingCapture:
+        def reset(self) -> None:
+            lifecycle.append("reset-breakable")
+
+    breakable = object.__new__(cgu.BreakableCUDAGraphWrapper)
+    breakable_entry = SimpleNamespace(
+        capture=_RecordingCapture(),
+        resources=[object()],
+    )
+    breakable.entries = {"profile": breakable_entry}
+
+    piecewise.reset_graphs()
+    breakable.reset_graphs()
+
+    assert lifecycle == ["reset-piecewise", "reset-breakable"]
+    assert piecewise_entry.cudagraph is None
+    assert piecewise_entry.output is not None
+    assert piecewise.concrete_cudagraph_entries == {"profile": piecewise_entry}
+    assert breakable_entry.capture is None
+    assert breakable_entry.resources
+    assert breakable.entries == {"profile": breakable_entry}
 
 
 def test_profile_cudagraph_memory_redirects_wrapper_pools(monkeypatch):
@@ -319,6 +404,9 @@ def __init__(self) -> None:
         def clear_graphs(self) -> None:
             pass
 
+        def reset_graphs(self) -> None:
+            pass
+
     wrapper = _FakeWrapper()
     cgu.CUDAGraphWrapper._all_instances.add(wrapper)
     try:
@@ -351,6 +439,9 @@ def __init__(self) -> None:
         def clear_graphs(self) -> None:
             pass
 
+        def reset_graphs(self) -> None:
+            pass
+
     wrapper: _FakeWrapper | None = None
     capture_model = runner.capture_model
 
@@ -390,8 +481,9 @@ def test_profile_cudagraph_memory_redirects_speculator_managers(monkeypatch):
     def _capture_model() -> int:
         nonlocal pools_during_capture
         pools_during_capture = (prefill_manager.pool, decode_manager.pool)
-        prefill_manager.graphs["profile"] = object()
-        decode_manager.graphs["profile"] = object()
+        lifecycle: list[str] = []
+        prefill_manager.graphs["profile"] = _RecordingGraph(lifecycle, "prefill")
+        decode_manager.graphs["profile"] = _RecordingGraph(lifecycle, "decode")
         prefill_manager._graphs_captured = True
         decode_manager._graphs_captured = True
         return capture_model()
@@ -455,7 +547,7 @@ def reset_attn(self) -> None:
     assert runner.kv_caches == []
     assert runner.attn_groups == []
     assert runner.cudagraph_manager is None
-    assert runner.block_tables is None
+    assert not hasattr(runner, "block_tables")
     assert runner.pcp_manager is None
     assert runner.adaptive_verification is None
     assert not hasattr(runner, "kv_cache_config")

vllm/compilation/breakable_cudagraph.py

@@ -151,6 +151,7 @@ def is_active(cls) -> bool:
     def __init__(self, pool: Any | None = None) -> None:
         self.pool = pool
         self.segments: list[Callable[[], Any]] = []
+        self._graphs: list[torch.cuda.CUDAGraph] = []
         self._num_graphs: int = 0
         self._num_eager_breaks: int = 0
         self._current_graph: torch.cuda.CUDAGraph | None = None
@@ -188,6 +189,7 @@ def _end_segment(self) -> None:
             return
         assert self._current_graph is not None
         self._current_graph.capture_end()
+        self._graphs.append(self._current_graph)
         self.segments.append(self._current_graph.replay)
         self._num_graphs += 1
         self._current_graph = None
@@ -214,6 +216,15 @@ def replay(self) -> None:
         for r in self.segments:
             r()
 
+    def reset(self) -> None:
+        """Destroy every graph segment after its pending work has completed."""
+        if self._capturing:
+            raise RuntimeError("Cannot reset an active breakable CUDA graph capture.")
+        for graph in self._graphs:
+            graph.reset()
+        self._graphs.clear()
+        self.segments.clear()
+
     # --- introspection ---------------------------------------------------
 
     @property
@@ -267,6 +278,12 @@ def clear_all_graphs(cls) -> None:
         for instance in list(cls._all_instances):
             instance.clear_graphs()
 
+    @classmethod
+    def reset_all_graphs(cls) -> None:
+        """Destroy graph segments without releasing entry-owned resources."""
+        for instance in list(cls._all_instances):
+            instance.reset_graphs()
+
     def __init__(
         self,
         runnable: Callable[..., Any],
@@ -307,6 +324,14 @@ def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
     def clear_graphs(self) -> None:
         self.entries.clear()
 
+    def reset_graphs(self) -> None:
+        """Destroy graph segments while retaining their captured resources."""
+        for entry in self.entries.values():
+            capture = entry.capture
+            if capture is not None:
+                capture.reset()
+                entry.capture = None
+
     # --- dispatch --------------------------------------------------------
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
@@ -367,13 +392,16 @@ def _capture(
         else:
             set_graph_pool_id(current_platform.graph_pool_handle())
 
-        # Match torch.cuda.graph()'s pre-capture cleanup once per descriptor.
+        # Match torch.cuda.graph()'s pre-capture barrier and cleanup once per
+        # descriptor. The warmup immediately before this call may use shared
+        # communication scratch. Starting capture before that work completes
+        # lets the captured kernels race the warmup on the same storage.
         # We drive capture_begin/end directly and bypass torch.cuda.graph(),
-        # so its built-in gc + empty_cache never fire. Run them here once
-        # per _capture call -- NOT inside _begin_segment, since this capture
-        # session may issue many begin/end pairs (one per layer's break),
-        # and repeated gc would tank capture time the way it did for the
-        # pre-`gc_disable` piecewise path.
+        # so its synchronize + gc + empty_cache sequence never runs. Run it
+        # here once per _capture call -- NOT inside _begin_segment, since this
+        # capture session may issue many begin/end pairs (one per layer's
+        # break), and repeated cleanup would dominate capture time.
+        torch.accelerator.synchronize()
         gc.collect()
         torch.accelerator.empty_cache()
         # Sync the offloader's copy stream before capture so any in-flight

vllm/compilation/cuda_graph.py

@@ -175,6 +175,12 @@ def clear_all_graphs(cls) -> None:
         for instance in list(cls._all_instances):
             instance.clear_graphs()
 
+    @classmethod
+    def reset_all_graphs(cls) -> None:
+        """Destroy captured graph executables without releasing their entries."""
+        for instance in list(cls._all_instances):
+            instance.reset_graphs()
+
     def __init__(
         self,
         runnable: Callable[..., Any],
@@ -230,6 +236,14 @@ def cudagraph_wrapper(self) -> "CUDAGraphWrapper":
     def clear_graphs(self) -> None:
         self.concrete_cudagraph_entries.clear()
 
+    def reset_graphs(self) -> None:
+        """Destroy graph executables while retaining entry-owned tensors."""
+        for entry in self.concrete_cudagraph_entries.values():
+            graph = entry.cudagraph
+            if graph is not None:
+                graph.reset()
+                entry.cudagraph = None
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
         if not is_forward_context_available():
             # No forward context means we are outside the normal