Account V2 FlashInfer cudagraph workspace

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

Author: lesj0610

tests/v1/worker/test_cudagraph_memory_profiling.py

@@ -505,3 +505,180 @@ def null_context(*args, **kwargs):
     assert warmup_calls[2][1]["profile_seq_lens"] == 1
     assert warmup_calls[3][1]["profile_seq_lens"] is None
     assert cleanup_calls == ["cleanup"]
+
+
+def test_v2_profile_accounts_attention_workspace(monkeypatch):
+    from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    @contextlib.contextmanager
+    def null_context(*args, **kwargs):
+        yield
+
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.vllm_config = object()
+
+    events = []
+
+    monkeypatch.setattr(
+        gpu_model_runner_v2,
+        "set_current_vllm_config",
+        lambda *args, **kwargs: null_context(),
+    )
+
+    def reserve_attention_workspace():
+        events.append("reserve")
+        return 4096
+
+    runner._init_minimal_kv_cache_for_profiling = lambda: events.append("init")
+    runner._reserve_attention_workspace_for_cudagraph_capture = (
+        reserve_attention_workspace
+    )
+    runner._cleanup_profiling_kv_cache = lambda: events.append("cleanup")
+
+    estimate = runner.profile_cudagraph_memory()
+
+    assert estimate == 4096
+    assert runner.cudagraph_memory_persistent_estimate == 4096
+    assert runner.cudagraph_memory_graph_pool_estimate == 0
+    assert events == ["init", "reserve", "cleanup"]
+
+
+def test_v2_cleanup_profiling_kv_cache_releases_builder_refs(monkeypatch):
+    from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    class Builder:
+        pass
+
+    builder = Builder()
+    builder_ref = weakref.ref(builder)
+    layer = SimpleNamespace(
+        kv_cache=torch.empty(1),
+        impl=SimpleNamespace(_k_scale_cache=object(), _v_scale_cache=object()),
+    )
+
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.cache_config = SimpleNamespace(num_gpu_blocks=4)
+    runner.kv_caches = [torch.empty(1)]
+    runner.attn_groups = [[SimpleNamespace(metadata_builders=[builder])]]
+    runner.kv_cache_config = object()
+    runner.block_tables = object()
+    runner.kernel_block_sizes = [16]
+    runner.cudagraph_manager = object()
+    runner.compilation_config = SimpleNamespace(static_forward_context={"layer": layer})
+    del builder
+
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator, "synchronize", lambda: None
+    )
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator, "empty_cache", lambda: None
+    )
+
+    runner._cleanup_profiling_kv_cache()
+    gc.collect()
+
+    assert runner.kv_caches == []
+    assert not hasattr(runner, "attn_groups")
+    assert not hasattr(runner, "kv_cache_config")
+    assert not hasattr(runner, "block_tables")
+    assert not hasattr(runner, "kernel_block_sizes")
+    assert not hasattr(runner, "cudagraph_manager")
+    assert runner.cache_config.num_gpu_blocks is None
+    assert isinstance(layer.kv_cache, torch.Tensor)
+    assert layer.kv_cache.numel() == 0
+    assert layer.impl._k_scale_cache is None
+    assert layer.impl._v_scale_cache is None
+    assert builder_ref() is None
+
+
+def test_v2_capture_reserves_workspace_before_measurement_and_locks(monkeypatch):
+    from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    @contextlib.contextmanager
+    def null_context(*args, **kwargs):
+        yield
+
+    class Builder:
+        reserved = False
+
+        def reserve_workspace_for_cudagraph_capture(self):
+            events.append("builder_reserve")
+            self.reserved = True
+
+    class FakeCudaGraphManager:
+        def needs_capture(self):
+            return True
+
+        def capture(
+            self,
+            model,
+            model_state,
+            input_buffers,
+            intermediate_tensors,
+            block_tables,
+            attn_groups,
+            kv_cache_config,
+            **kwargs,
+        ):
+            events.append("capture")
+            assert attn_groups[0][0].metadata_builders[0].reserved
+            return {}
+
+    events = []
+    builder = Builder()
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.device = torch.device("cpu")
+    runner.cudagraph_manager = FakeCudaGraphManager()
+    runner.lora_config = None
+    runner.maybe_setup_dummy_loras = lambda lora_config: null_context()
+    runner.model = object()
+    runner.model_state = object()
+    runner.input_buffers = object()
+    runner.intermediate_tensors = None
+    runner.block_tables = object()
+    runner.attn_groups = [[SimpleNamespace(metadata_builders=[builder])]]
+    runner.kv_cache_config = object()
+    runner.use_aux_hidden_state_outputs = False
+    runner.speculator = None
+
+    memory_reserved_values = iter([1_000, 1_128])
+    get_memory_info_values = iter([(10_000, 0), (9_000, 0)])
+
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator, "synchronize", lambda: None
+    )
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator, "empty_cache", lambda: None
+    )
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator,
+        "memory_reserved",
+        lambda device: next(memory_reserved_values),
+    )
+
+    def get_memory_info():
+        events.append("memory_info")
+        return next(get_memory_info_values)
+
+    monkeypatch.setattr(
+        gpu_model_runner_v2.torch.accelerator,
+        "get_memory_info",
+        get_memory_info,
+    )
+    monkeypatch.setattr(
+        gpu_model_runner_v2,
+        "lock_workspace",
+        lambda: events.append("lock"),
+    )
+
+    assert runner.capture_model() == 1_000
+    assert events == [
+        "builder_reserve",
+        "memory_info",
+        "capture",
+        "memory_info",
+        "lock",
+    ]

vllm/v1/worker/gpu/model_runner.py

@@ -29,7 +29,7 @@
 
 import vllm.envs as envs
 from vllm.compilation.counter import compilation_counter
-from vllm.config import VllmConfig
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import (
     get_dcp_group,
@@ -50,6 +50,7 @@
 from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
 from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
@@ -113,6 +114,7 @@
 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 from vllm.v1.worker.utils import KVBlockZeroer
+from vllm.v1.worker.workspace import lock_workspace
 
 logger = init_logger(__name__)
 
@@ -402,7 +404,9 @@ def main_stream(self) -> torch.cuda.Stream:
     def get_kv_cache_spec(self):
         return get_kv_cache_spec(self.vllm_config)
 
-    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+    def initialize_kv_cache(
+        self, kv_cache_config: KVCacheConfig, is_profiling: bool = False
+    ) -> None:
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
 
@@ -475,7 +479,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
         check_attention_cp_compatibility(self.vllm_config)
-        if isinstance(self.speculator, DraftModelSpeculator):
+        if isinstance(self.speculator, DraftModelSpeculator) and not is_profiling:
             # HACK(woosuk)
             self.speculator.set_attn(
                 self.model_state, self.kv_cache_config, self.block_tables
@@ -492,7 +496,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
             self.kernel_block_sizes,
             self.vllm_config,
         )
-        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
+        if not is_profiling:
+            self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
 
     def _init_kv_zero_meta(self) -> None:
         """Build KV-block zeroing metadata; invoked from gpu_worker."""
@@ -682,9 +687,106 @@ def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
         # SP is not supported yet.
         return num_scheduled_tokens
 
+    def _init_minimal_kv_cache_for_profiling(self) -> None:
+        from vllm.v1.core.kv_cache_utils import (
+            get_kv_cache_config_from_groups,
+            get_kv_cache_groups,
+        )
+
+        kv_cache_spec = self.get_kv_cache_spec()
+        KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec)
+        kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec)
+        min_blocks = self.compilation_config.max_cudagraph_capture_size or 1
+
+        # Temporarily allocate just enough KV cache state to instantiate
+        # attention metadata builders for workspace sizing.
+        saved_override = self.cache_config.num_gpu_blocks_override
+        self.cache_config.num_gpu_blocks_override = min_blocks
+        try:
+            minimal_config = get_kv_cache_config_from_groups(
+                self.vllm_config,
+                kv_cache_groups,
+                available_memory=0,
+            )
+        finally:
+            self.cache_config.num_gpu_blocks_override = saved_override
+
+        self.initialize_kv_cache(minimal_config, is_profiling=True)
+        self.cache_config.num_gpu_blocks = minimal_config.num_blocks
+
+        logger.debug("Initialized minimal KV cache for CUDA graph profiling")
+
+    def _cleanup_profiling_kv_cache(self) -> None:
+        torch.accelerator.synchronize()
+
+        if hasattr(self, "kv_caches") and self.kv_caches:
+            for i in range(len(self.kv_caches)):
+                self.kv_caches[i] = None  # type: ignore[assignment]
+            self.kv_caches.clear()
+        if hasattr(self, "attn_groups"):
+            self.attn_groups.clear()
+        for attr in (
+            "attn_groups",
+            "kv_cache_config",
+            "block_tables",
+            "kernel_block_sizes",
+            "cudagraph_manager",
+        ):
+            if hasattr(self, attr):
+                delattr(self, attr)
+        self.cache_config.num_gpu_blocks = None
+
+        for layer in self.compilation_config.static_forward_context.values():
+            if hasattr(layer, "kv_cache"):
+                kv_cache = layer.kv_cache
+                layer.kv_cache = (
+                    torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
+                )
+            # Clean up quantized KV cache scale views
+            # (int8_per_token_head, fp8_per_token_head).
+            if hasattr(layer, "impl"):
+                if hasattr(layer.impl, "_k_scale_cache"):
+                    layer.impl._k_scale_cache = None
+                if hasattr(layer.impl, "_v_scale_cache"):
+                    layer.impl._v_scale_cache = None
+
+        gc.collect()
+        torch.accelerator.empty_cache()
+        logger.debug("Cleaned up profiling KV cache and CUDA graphs")
+
+    def _reserve_attention_workspace_for_cudagraph_capture(self) -> int:
+        if not getattr(self, "attn_groups", None):
+            return 0
+
+        reserved_before = torch.accelerator.memory_reserved(self.device)
+        for groups in self.attn_groups:
+            for attn_group in groups:
+                for builder in attn_group.metadata_builders:
+                    builder.reserve_workspace_for_cudagraph_capture()
+        torch.accelerator.synchronize()
+        torch.accelerator.empty_cache()
+        reserved_after = torch.accelerator.memory_reserved(self.device)
+        return max(reserved_after - reserved_before, 0)
+
     def profile_cudagraph_memory(self) -> int:
-        # NOTE(woosuk): It is TBD whether we keep this API or not.
-        return 0
+        self.cudagraph_memory_persistent_estimate = 0
+        self.cudagraph_memory_graph_pool_estimate = 0
+
+        try:
+            with set_current_vllm_config(self.vllm_config):
+                self._init_minimal_kv_cache_for_profiling()
+            persistent_estimate = (
+                self._reserve_attention_workspace_for_cudagraph_capture()
+            )
+        finally:
+            self._cleanup_profiling_kv_cache()
+
+        self.cudagraph_memory_persistent_estimate = int(persistent_estimate)
+        logger.info(
+            "Estimated CUDA graph persistent workspace memory: %.2f GiB",
+            persistent_estimate / (1 << 30),
+        )
+        return int(persistent_estimate)
 
     @torch.inference_mode()
     def capture_model(self) -> int:
@@ -700,6 +802,8 @@ def capture_model(self) -> int:
 
         start_time = time.perf_counter()
         gc.collect()
+        self._reserve_attention_workspace_for_cudagraph_capture()
+        torch.accelerator.synchronize()
         torch.accelerator.empty_cache()
         start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
 
@@ -729,6 +833,7 @@ def capture_model(self) -> int:
             elapsed_time,
             cuda_graph_size / (1 << 30),
         )
+        lock_workspace()
         return cuda_graph_size
 
     def _remove_request(self, req_id: str) -> bool: