diff --git a/tests/test_cuda_graph.py b/tests/test_cuda_graph.py index fb901c42..4aa05cee 100644 --- a/tests/test_cuda_graph.py +++ b/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 diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index 7cba34fe..db1a7519 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/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()