Fix V2 CUDA graph pool memory sampling

Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>

Author: lesj0610

tests/v1/worker/test_cudagraph_memory_profiling.py

@@ -573,7 +573,7 @@ def test_v2_profile_is_noop_on_non_cuda(monkeypatch):
     assert runner.cudagraph_memory_graph_pool_estimate == 0
 
 
-def test_v2_cuda_graph_pool_sample_uses_peak(monkeypatch):
+def test_v2_cuda_graph_pool_sample_separates_residual_and_peak(monkeypatch):
     from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
     from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
@@ -623,18 +623,77 @@ def test_v2_cuda_graph_pool_sample_uses_peak(monkeypatch):
         lambda device: 530,
     )
 
-    assert (
-        runner._measure_cuda_graph_pool_sample(lambda: events.append("capture")) == 120
-    )
+    sample = runner._measure_cuda_graph_pool_sample(lambda: events.append("capture"))
+    assert sample.residual == 50
+    assert sample.peak == 120
     assert events == [
         "sync",
         "empty_cache",
         "reset_peak",
         "capture",
         "sync",
+        "empty_cache",
     ]
 
 
+def test_v2_graph_pool_profile_does_not_multiply_transient_peak(monkeypatch):
+    from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
+    from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+    from vllm.v1.worker.gpu.model_runner import (
+        CudaGraphMemorySample,
+        CUDAGraphMode,
+        GPUModelRunner,
+    )
+
+    descs = [
+        BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 128, None),
+        BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 64, None),
+        BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 32, None),
+    ]
+
+    class FakeCudaGraphManager:
+        pool = "manager-original"
+        graphs = {}
+        _graphs_captured = False
+        hidden_states = None
+        aux_hidden_states = []
+        intermediate_tensors = None
+        use_aux_hidden_state_outputs = False
+        use_breakable_cg = False
+        breakable_cg_runner = None
+
+        def get_capture_descs(self):
+            return [(CUDAGraphMode.PIECEWISE, descs)]
+
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.cudagraph_manager = FakeCudaGraphManager()
+    samples = iter(
+        [
+            CudaGraphMemorySample(residual=10 << 20, peak=200 << 20),
+            CudaGraphMemorySample(residual=2 << 20, peak=180 << 20),
+        ]
+    )
+
+    monkeypatch.setattr(
+        gpu_model_runner_v2.current_platform,
+        "graph_pool_handle",
+        lambda: "profile-pool",
+    )
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator,
+        "empty_cache",
+        lambda: None,
+    )
+    monkeypatch.setattr(
+        runner,
+        "_measure_cuda_graph_pool_sample",
+        lambda capture_fn: capture_fn() or next(samples),
+    )
+    monkeypatch.setattr(runner, "_capture_model_cudagraphs", lambda **kwargs: None)
+
+    assert runner._profile_cudagraph_memory_graph_pool() == 204 << 20
+
+
 def test_v2_graph_pool_profile_restores_capture_state(monkeypatch):
     from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
     from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
@@ -680,7 +739,16 @@ def init_breakable_cg_runner(self, model):
     runner = GPUModelRunner.__new__(GPUModelRunner)
     runner.model = "model"
     runner.cudagraph_manager = FakeCudaGraphManager()
-    sample_values = iter([1_000, 2_000_000, 800, 3_000_000])
+    from vllm.v1.worker.gpu.model_runner import CudaGraphMemorySample
+
+    sample_values = iter(
+        [
+            CudaGraphMemorySample(residual=1_000, peak=10_000),
+            CudaGraphMemorySample(residual=2_000_000, peak=200_000_000),
+            CudaGraphMemorySample(residual=800, peak=8_000),
+            CudaGraphMemorySample(residual=3_000_000, peak=300_000_000),
+        ]
+    )
     capture_calls = []
 
     monkeypatch.setattr(
@@ -727,7 +795,7 @@ def capture_model_cudagraphs(**kwargs):
 
     estimate = runner._profile_cudagraph_memory_graph_pool()
 
-    assert estimate == 7_001_000
+    assert estimate == 7_010_000
     assert [call["capture_descs"] for call in capture_calls] == [
         {CUDAGraphMode.PIECEWISE: [piecewise_descs[0]]},
         {CUDAGraphMode.PIECEWISE: [piecewise_descs[1]]},

vllm/v1/worker/gpu/model_runner.py

@@ -121,6 +121,11 @@
 logger = init_logger(__name__)
 
 
+class CudaGraphMemorySample(NamedTuple):
+    residual: int
+    peak: int
+
+
 class GPUModelRunner(LoRAModelRunnerMixin):
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.vllm_config = vllm_config
@@ -796,7 +801,9 @@ def _capture_model_cudagraphs(
             if capture_speculator and self.speculator is not None:
                 self.speculator.capture(attn_states)
 
-    def _measure_cuda_graph_pool_sample(self, capture_fn: Callable[[], None]) -> int:
+    def _measure_cuda_graph_pool_sample(
+        self, capture_fn: Callable[[], None]
+    ) -> CudaGraphMemorySample:
         torch.accelerator.synchronize()
         torch.accelerator.empty_cache()
         free_before = torch.accelerator.get_memory_info()[0]
@@ -807,8 +814,6 @@ def _measure_cuda_graph_pool_sample(self, capture_fn: Callable[[], None]) -> int
         capture_fn()
         torch.accelerator.synchronize()
 
-        free_after = torch.accelerator.get_memory_info()[0]
-        residual_delta = max(free_before - free_after, 0)
         peak_reserved_delta = max(
             torch.accelerator.max_memory_reserved(self.device) - reserved_before,
             0,
@@ -817,7 +822,11 @@ def _measure_cuda_graph_pool_sample(self, capture_fn: Callable[[], None]) -> int
             torch.accelerator.max_memory_allocated(self.device) - allocated_before,
             0,
         )
-        return max(residual_delta, peak_reserved_delta, peak_allocated_delta)
+        torch.accelerator.empty_cache()
+        free_after = torch.accelerator.get_memory_info()[0]
+        residual_delta = max(free_before - free_after, 0)
+        peak_delta = max(residual_delta, peak_reserved_delta, peak_allocated_delta)
+        return CudaGraphMemorySample(residual=residual_delta, peak=peak_delta)
 
     def _profile_cudagraph_memory_graph_pool(self) -> int:
         assert self.cudagraph_manager is not None
@@ -858,7 +867,7 @@ def _profile_cudagraph_memory_graph_pool(self) -> int:
 
             for mode, descs in capture_descs:
                 profile_descs = descs[:2]
-                mem_samples: list[int] = []
+                mem_samples: list[CudaGraphMemorySample] = []
 
                 for desc in profile_descs:
                     sample_descs = {mode: [desc]}
@@ -878,8 +887,11 @@ def capture_sample(
                         self._measure_cuda_graph_pool_sample(capture_sample)
                     )
 
-                first_capture = mem_samples[0]
-                per_graph = max(mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20)
+                first_capture = mem_samples[0].peak
+                per_graph = max(
+                    mem_samples[1].residual if len(mem_samples) > 1 else 0,
+                    1 << 20,
+                )
                 shared_memory_estimate[mode] = first_capture
                 per_graph_estimate[mode] = per_graph * (len(descs) - 1)
                 logger.debug(