[Bugfix][Core] Profile and validate kv_cache_memory_bytes instead of skipping the profile

tests/v1/worker/test_cudagraph_memory_profiling.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The minimal KV cache the CUDA graph profiler bootstraps.
+
+SCALE_CACHE  teardown clears quantized scale views, with or without a kv_cache
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.v1.worker.utils as worker_utils
+from vllm.v1.worker import gpu_model_runner as v1
+
+
+def _layer(with_kv_cache):
+    impl = SimpleNamespace(_k_scale_cache=object(), _v_scale_cache=object())
+    layer = SimpleNamespace(impl=impl)
+    if with_kv_cache:
+        layer.kv_cache = torch.empty(4)
+    return layer
+
+
+@pytest.mark.parametrize("with_kv_cache", [True, False])
+def test_scale_cache_is_cleared_with_or_without_a_kv_cache(with_kv_cache):
+    """Scale views can live on a layer that has no ``kv_cache`` attribute."""
+    layer = _layer(with_kv_cache)
+    worker_utils.clear_layer_kv_caches([layer])
+
+    assert layer.impl._k_scale_cache is None
+    assert layer.impl._v_scale_cache is None
+    if with_kv_cache:
+        assert isinstance(layer.kv_cache, torch.Tensor)
+        assert layer.kv_cache.numel() == 0
+    else:
+        assert not hasattr(layer, "kv_cache")
+
+
+def test_scale_cache_shutdown_and_profiling_teardown_share_one_helper(monkeypatch):
+    """V1 clears layers the same way on both paths, so neither can drift."""
+    layer, seen = _layer(with_kv_cache=False), []
+    runner = v1.GPUModelRunner.__new__(v1.GPUModelRunner)
+    runner.cache_config = SimpleNamespace(num_gpu_blocks=4)
+    runner.compilation_config = SimpleNamespace(static_forward_context={"l": layer})
+    runner.kv_caches = []
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+    monkeypatch.setattr(v1, "clear_layer_kv_caches", lambda ls: seen.append(list(ls)))
+
+    runner._cleanup_profiling_kv_cache()
+    assert seen == [[layer]]
+    assert runner.cache_config.num_gpu_blocks is None

vllm/v1/worker/gpu_model_runner.py

@@ -217,6 +217,7 @@
 )
 from vllm.v1.worker.utils import (
     EncoderTimingStats,
+    clear_layer_kv_caches,
     raise_if_nan_logits,
 )
 from vllm.v1.worker.workspace import lock_workspace
@@ -6238,19 +6239,7 @@ def _cleanup_profiling_kv_cache(self) -> None:
         # "runtime" pool; the real initialize_kv_cache rebuilds it.
         self._init_block_sizes = []
 
-        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
+        clear_layer_kv_caches(self.compilation_config.static_forward_context.values())
 
         gc.collect()
         torch.accelerator.empty_cache()

vllm/v1/worker/gpu_worker.py

@@ -9,6 +9,7 @@
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from datetime import timedelta
 from fnmatch import filter as fnmatch_filter
+from functools import partial
 from types import NoneType
 from typing import TYPE_CHECKING, Any, cast
 
@@ -634,6 +635,177 @@ def reload_weights(self, *args, **kwargs) -> None:
         with set_current_vllm_config(self.vllm_config):
             self.model_runner.reload_weights(*args, **kwargs)
 
+    def _profile_attention_for_pinned_kv(self) -> None:
+        """Extend the profile with attention, for a pinned KV size only.
+
+        ``profile_run`` forwards with ``skip_attn=True``: there is no KV cache
+        yet, so the attention kernels never run and neither their activations
+        nor the device memory their first launch takes reach the profile. Under
+        ``gpu_memory_utilization`` that is harmless -- what the profile missed
+        lands in the slice the utilization fraction left unallocated. A pinned
+        ``kv_cache_memory_bytes`` has no such slice, so the same bytes are
+        spent twice and the rank OOMs in the middle of a request.
+
+        Build the same minimal KV cache the CUDA graph profiling uses and run a
+        ladder of batch widths through it, so the surrounding
+        ``memory_profiling`` block sees the attention path. The ladder is not
+        redundant with its widest rung: the caching allocator keeps whole
+        segments per size class, so a served mix of widths costs more than the
+        widest one alone.
+
+        Best effort -- a backend that cannot serve a dummy attention forward
+        leaves the estimate where it was, which is no worse than skipping this
+        entirely.
+        """
+        max_tokens = self.model_runner.max_num_tokens
+        # Widest first, so a backend that refuses the narrow shapes still
+        # contributes the rung that dominates the peak.
+        widths = sorted(
+            {max(1, max_tokens * n // 8) for n in range(1, 9)} | {1}, reverse=True
+        )
+        runner = cast(Any, self.model_runner)
+        build_kv_cache: Callable[[], None]
+        teardown: Callable[[], None]
+        if self.use_v2_model_runner:
+            from vllm.v1.worker.gpu.cudagraph_utils import (
+                _init_minimal_kv_cache_for_profiling,
+                _teardown_profiling_state,
+            )
+
+            build_kv_cache = partial(_init_minimal_kv_cache_for_profiling, runner)
+            teardown = partial(_teardown_profiling_state, runner)
+        else:
+            build_kv_cache = runner._init_minimal_kv_cache_for_profiling
+            teardown = runner._cleanup_profiling_kv_cache
+
+        try:
+            with set_current_vllm_config(self.vllm_config):
+                build_kv_cache()
+        except Exception:
+            logger.warning(
+                "Could not build a profiling KV cache to forward attention "
+                "against, so the profile keeps excluding the attention path. "
+                "The pinned kv_cache_memory_bytes is checked against an "
+                "estimate that is known to be low; leave headroom for it.",
+                exc_info=True,
+            )
+            return
+
+        # Serving rebuilds the allocator pool from scratch: the profiling run's
+        # segments are released before the KV cache is allocated, and what the
+        # pool then has to re-reserve for the same shapes is what this measures.
+        # Against the warm pool profile_run leaves, it would report none of it.
+        gc.collect()
+        torch.accelerator.empty_cache()
+        reserved_before = torch.accelerator.memory_reserved(self.device)
+        allocated_before = torch.accelerator.memory_allocated(self.device)
+        # The reset below also clears the peak ``memory_profiling`` reads on
+        # exit, and profile_run's own peak (encoder cache, sampler) can sit
+        # above the ladder's. Keep it so the pinned check takes the larger one.
+        self._pinned_profile_run_peak_bytes = torch.accelerator.max_memory_allocated(
+            self.device
+        )
+        torch.accelerator.reset_peak_memory_stats(self.device)
+        try:
+            for num_tokens in widths:
+                try:
+                    runner._dummy_run(num_tokens, is_profile=True)
+                except Exception:
+                    logger.warning(
+                        "Attention profiling forward at %d tokens failed; "
+                        "continuing with the other widths.",
+                        num_tokens,
+                        exc_info=True,
+                    )
+            # What the allocator had to hold, not just what was live at once.
+            # ``memory_profiling`` accounts for the live peak alone, which is
+            # the right call where the utilization slice absorbs the difference.
+            self._pinned_allocator_overhead_bytes = max(
+                0,
+                (torch.accelerator.max_memory_reserved(self.device) - reserved_before)
+                - (
+                    torch.accelerator.max_memory_allocated(self.device)
+                    - allocated_before
+                ),
+            )
+        finally:
+            teardown()
+
+    def _apply_pinned_kv_cache_memory(
+        self,
+        kv_cache_memory_bytes: int,
+        profile_result: Any,
+        cudagraph_memory_estimate_applied: int,
+    ) -> int:
+        """Check a pinned KV size against the profile and report it."""
+        allocator_overhead = getattr(self, "_pinned_allocator_overhead_bytes", 0)
+        # The attention ladder reset the peak counter, so add back whatever of
+        # profile_run's peak the final measurement no longer reflects.
+        non_kv_cache_memory = profile_result.non_kv_cache_memory + max(
+            0,
+            getattr(self, "_pinned_profile_run_peak_bytes", 0)
+            - profile_result.after_profile.torch_peak,
+        )
+        # Manual control deliberately ignores gpu_memory_utilization, so the
+        # budget is the whole free device, not the utilization slice.
+        fits_kv_cache_memory_bytes = (
+            self.init_snapshot.free_memory
+            - non_kv_cache_memory
+            - cudagraph_memory_estimate_applied
+            - allocator_overhead
+        )
+        if kv_cache_memory_bytes > fits_kv_cache_memory_bytes:
+            # Refuse rather than quietly serve a smaller cache: a pinned size is
+            # an explicit capacity decision, and halving it behind the
+            # operator's back hides the one number this option exists to
+            # control. Starting anyway is worse -- the bytes are spent either
+            # way, and the rank OOMs once traffic reaches the peak just
+            # measured.
+            raise ValueError(
+                f"kv_cache_memory_bytes={kv_cache_memory_bytes} "
+                f"({format_gib(kv_cache_memory_bytes)} GiB) does not fit on "
+                f"this rank. Of {format_gib(self.init_snapshot.free_memory)} "
+                f"GiB free at startup the model needs "
+                f"{format_gib(non_kv_cache_memory)} GiB for "
+                f"weights, workspaces and its activation peak"
+                + (
+                    f", {format_gib(cudagraph_memory_estimate_applied)} GiB for "
+                    f"CUDA graphs"
+                    if cudagraph_memory_estimate_applied
+                    else ""
+                )
+                + (
+                    f" and {format_gib(allocator_overhead)} GiB the caching "
+                    f"allocator holds beyond that peak"
+                    if allocator_overhead
+                    else ""
+                )
+                + f", leaving {format_gib(fits_kv_cache_memory_bytes)} GiB. Set "
+                f"--kv-cache-memory-bytes={int(fits_kv_cache_memory_bytes)} or "
+                f"lower, or give this rank more room (fewer layers on it, a "
+                f"smaller max_num_batched_tokens or max_model_len). That figure "
+                f"is an upper bound, not a safe setting: the profile runs "
+                f"against a minimal KV cache, so peaks that scale with context "
+                f"length are not in it. Leave headroom below it."
+            )
+
+        self.available_kv_cache_memory_bytes = kv_cache_memory_bytes
+        logger.info(
+            "Initial free memory %s GiB, reserved %s GiB for KV Cache as "
+            "specified by kv_cache_memory_bytes config (profiled non-KV usage "
+            "%s GiB, %s GiB left unused). This does not respect the "
+            "gpu_memory_utilization config.",
+            format_gib(self.init_snapshot.free_memory),
+            format_gib(kv_cache_memory_bytes),
+            format_gib(non_kv_cache_memory),
+            format_gib(fits_kv_cache_memory_bytes - kv_cache_memory_bytes),
+        )
+        return reserve_mm_ipc_gpu_memory(
+            kv_cache_memory_bytes,
+            self.model_config.multimodal_config,
+            getattr(self.parallel_config, "_api_process_count", 1),
+        )
+
     @torch.inference_mode()
     def determine_available_memory(self) -> int:
         """Profiles the peak memory usage of the model to determine how much
@@ -648,9 +820,15 @@ def determine_available_memory(self) -> int:
             by adjusting the `gpu_memory_utilization` parameter.
 
         """
+        pinned_kv_cache_memory_bytes = self.cache_config.kv_cache_memory_bytes
         maybe_apply_startup_plan(self)
 
-        if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
+        # A startup plan restores the size a fully profiled earlier boot of this
+        # exact configuration measured, so it keeps the fast path the plan
+        # exists for. Only a size the operator pinned is measured below.
+        if pinned_kv_cache_memory_bytes is None and (
+            kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes
+        ):
             # still need a profile run which compiles the model for
             # max_num_batched_tokens
             with set_current_vllm_config(self.vllm_config):
@@ -677,24 +855,39 @@ def determine_available_memory(self) -> int:
                 getattr(self.parallel_config, "_api_process_count", 1),
             )
 
+        # Pinning the KV size decides how much memory the cache takes, not how
+        # much the rest of the model needs: weights, workspaces and the
+        # activation peak are still there, and the profiling run below is what
+        # measures them. So it runs either way and kv_cache_memory_bytes is
+        # checked against the result, rather than trusted blind -- skipping the
+        # profile left the operator with no number to size the pin against, and
+        # a pin that does not fit surfaces much later as an OOM mid-request.
+        kv_cache_memory_bytes = self.cache_config.kv_cache_memory_bytes
+
         # Execute a forward pass with dummy inputs to profile the memory usage
         # of the model.
         with (
             memory_profiling(
                 self.init_snapshot,
                 weights_memory=int(self.model_runner.model_memory_usage),
             ) as profile_result,
+            set_current_vllm_config(self.vllm_config),
+        ):
             # Workspaces (e.g. the MoE workspace) grow in steps during this pass,
             # freeing each smaller buffer. Without a split limit, a later small
             # allocation can be carved out of a freed multi-GiB block and pin the
             # whole segment past the empty_cache() in memory_profiling, so it is
             # counted as consumed and taken from the KV cache. Blocks above the
             # limit are never split, so they stay releasable. Exits before
             # memory_profiling measures, restoring the original limit.
-            self._scoped_allocator_max_split(max_split_size_mb=20),
-            set_current_vllm_config(self.vllm_config),
-        ):
-            self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)
+            with self._scoped_allocator_max_split(max_split_size_mb=20):
+                self.model_runner.profile_run(
+                    randomize_inputs=self.randomize_dummy_inputs
+                )
+            if kv_cache_memory_bytes:
+                # Outside the split limit, so the allocator overhead the
+                # ladder measures comes from the settings serving runs with.
+                self._profile_attention_for_pinned_kv()
 
         # Profile CUDA graph memory if graphs will be captured.
         # ROCm is included: #44825 moved the profiler to
@@ -747,6 +940,13 @@ def determine_available_memory(self) -> int:
             - cudagraph_memory_estimate_applied
         )
 
+        if kv_cache_memory_bytes:
+            return self._apply_pinned_kv_cache_memory(
+                kv_cache_memory_bytes,
+                profile_result,
+                cudagraph_memory_estimate_applied,
+            )
+
         unrequested_memory = self.init_snapshot.free_memory - self.requested_memory
         logger.debug(
             "Initial free memory: %s GiB; Requested memory: %f (util), %s GiB",

vllm/v1/worker/utils.py

@@ -893,17 +893,20 @@ def clear_layer_kv_caches(layers: Iterable[Any]) -> None:
     alone does not release the KV cache memory on teardown paths.
     """
     for layer in layers:
-        if not hasattr(layer, "kv_cache"):
-            continue
-        kv_cache = layer.kv_cache
-        layer.kv_cache = torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
+        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
+        # (int8_per_token_head, fp8_per_token_head). A layer can hold these
+        # without a kv_cache attribute, so they are cleared independently.
+        impl = getattr(layer, "impl", None)
+        if impl is not None:
+            if hasattr(impl, "_k_scale_cache"):
+                impl._k_scale_cache = None
+            if hasattr(impl, "_v_scale_cache"):
+                impl._v_scale_cache = None
 
 
 def copy_kv_cache_blocks_inplace(