[Bugfix] Collect garbage before CUDA graph capture

tests/test_cuda_graph.py

@@ -1,5 +1,7 @@
 """Graph replay must update KV addresses, request identities and host progress."""
 
+import gc
+
 import pytest
 import torch
 
@@ -101,3 +103,54 @@ def test_graph_cache_limit_fallback_and_abort():
         assert engine.cache_manager.num_used_blocks == 0
     assert engine.model_runner.graphs.captures == 1
     assert engine.model_runner.graphs.fallbacks > 0
+
+
+@pytest.mark.gpu
+def test_capture_survives_collecting_an_earlier_engines_garbage():
+    # torch.cuda.graph does not collect garbage before capture by default. Reference
+    # cycles left by an earlier engine (graphs, streams, pinned buffers) must not be
+    # freed inside a later capture when an allocation there runs the collector.
+    torch.manual_seed(7)
+    model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16)
+    recurrent = model.recurrent_prepared
+
+    def collecting(*args, **kwargs):
+        if torch.cuda.is_current_stream_capturing():
+            gc.collect()
+        return recurrent(*args, **kwargs)
+
+    def run(collect_in_capture):
+        model.recurrent_prepared = collecting if collect_in_capture else recurrent
+        engine = LLMEngine(
+            model,
+            cache_config=CacheConfig(num_blocks=64, block_size=16),
+            scheduler_config=SchedulerConfig(
+                max_num_seqs=3, max_num_batched_tokens=8, prefill_chunk_size=4
+            ),
+            execution_config=ExecutionConfig(cuda_graphs=True, async_scheduling=True),
+            exit_config=ExitConfig("ouro_delayed"),
+            attention_backend="triton",
+        )
+        for i in range(3):
+            engine.add_request(
+                str(i),
+                [2 + i, 3, 4] * (i + 1),
+                SamplingParams(max_tokens=5 + i, min_loops=2, exit_threshold=0.2, ignore_eos=True),
+            )
+        finished = {}
+        while engine.has_unfinished_requests():
+            for out in engine.step():
+                if out.finished:
+                    finished[out.request_id] = (out.token_ids, out.exit_depths)
+        assert engine.model_runner.graphs.captures > 0
+        return finished
+
+    gc.collect()
+    gc.disable()  # Keep the first engine's cycles alive until the second engine captures.
+    try:
+        expected = run(collect_in_capture=False)
+        actual = run(collect_in_capture=True)
+    finally:
+        gc.enable()
+        model.recurrent_prepared = recurrent
+    assert actual == expected

vllm_rlt/worker/cuda_graph.py

@@ -1,5 +1,6 @@
 """Bounded recurrent-core graphs; request bookkeeping and sampling stay eager."""
 
+import gc
 import math
 from dataclasses import dataclass, replace
 from types import SimpleNamespace
@@ -21,6 +22,10 @@ def _capture(hidden, stream, pool, run):
         for _ in range(2):
             run()
     stream.synchronize()
+    # torch.cuda.graph no longer collects garbage before capture. Do it here, so
+    # an allocation inside the capture cannot run the collector and free an earlier
+    # engine's cyclic CUDA state (graphs, streams, pinned buffers) mid-capture.
+    gc.collect()
     graph = torch.cuda.CUDAGraph()
     with torch.cuda.graph(graph, pool=pool, stream=stream):
         output = run()