Add tests for CUDA graph profiling cache sizing

Author: s0l

tests/v1/worker/test_gpu_model_runner.py

@@ -42,6 +42,7 @@
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
+    MambaSpec,
 )
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.sample.metadata import SamplingMetadata
@@ -240,6 +241,80 @@ def _make_kv_cache_spec() -> FullAttentionSpec:
     return FullAttentionSpec(block_size=1, num_kv_heads=1, head_size=1, dtype="float16")
 
 
+def test_minimal_cudagraph_profiling_blocks_scales_tokens_to_blocks():
+    attn_spec = FullAttentionSpec(
+        block_size=592,
+        num_kv_heads=4,
+        head_size=128,
+        dtype=torch.float16,
+    )
+    groups = [KVCacheGroupSpec(["attn"], attn_spec)]
+
+    min_blocks, has_mamba_cache = (
+        GPUModelRunner._get_minimal_kv_cache_blocks_for_cudagraph_profiling(
+            groups,
+            max_capture_tokens=8192,
+            max_num_reqs=8,
+            cp_size=3,
+        )
+    )
+
+    assert min_blocks == 5  # ceil(8192 / (592 * 3))
+    assert not has_mamba_cache
+
+
+def test_minimal_cudagraph_profiling_blocks_cover_mamba_cache_lines():
+    attn_spec = FullAttentionSpec(
+        block_size=592,
+        num_kv_heads=4,
+        head_size=128,
+        dtype=torch.float16,
+    )
+    mamba_spec = MambaSpec(
+        shapes=((256, 4),),
+        dtypes=(torch.float16,),
+        block_size=592,
+    )
+    groups = [
+        KVCacheGroupSpec(["attn"], attn_spec),
+        KVCacheGroupSpec(["mamba"], mamba_spec),
+    ]
+
+    min_blocks, has_mamba_cache = (
+        GPUModelRunner._get_minimal_kv_cache_blocks_for_cudagraph_profiling(
+            groups,
+            max_capture_tokens=4096,
+            max_num_reqs=8,
+            cp_size=3,
+        )
+    )
+
+    assert min_blocks == 8  # max(ceil(4096 / (592 * 3)), max_num_reqs)
+    assert has_mamba_cache
+
+
+def test_minimal_cudagraph_profiling_blocks_without_capture_size():
+    attn_spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=128,
+        dtype=torch.float16,
+    )
+    groups = [KVCacheGroupSpec(["attn"], attn_spec)]
+
+    min_blocks, has_mamba_cache = (
+        GPUModelRunner._get_minimal_kv_cache_blocks_for_cudagraph_profiling(
+            groups,
+            max_capture_tokens=None,
+            max_num_reqs=8,
+            cp_size=1,
+        )
+    )
+
+    assert min_blocks == 1
+    assert not has_mamba_cache
+
+
 def test_select_common_block_size_prefers_manager_block_size():
     backend_a = _make_mock_backend_for_kernel_block_size([MultipleOf(32)])
     backend_b = _make_mock_backend_for_kernel_block_size([64, MultipleOf(16)])

vllm/v1/worker/gpu_model_runner.py

@@ -6425,6 +6425,31 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
+    @staticmethod
+    def _get_minimal_kv_cache_blocks_for_cudagraph_profiling(
+        kv_cache_groups: list[KVCacheGroupSpec],
+        max_capture_tokens: int | None,
+        max_num_reqs: int,
+        cp_size: int,
+    ) -> tuple[int, bool]:
+        if max_capture_tokens is None:
+            return 1, False
+
+        block_requirements: list[int] = []
+        has_mamba_cache = False
+        for group in kv_cache_groups:
+            kv_cache_spec = group.kv_cache_spec
+            if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
+                continue
+            if isinstance(kv_cache_spec, MambaSpec):
+                has_mamba_cache = True
+                block_requirements.append(max_num_reqs)
+            else:
+                block_requirements.append(
+                    cdiv(max_capture_tokens, kv_cache_spec.block_size * cp_size)
+                )
+        return max(1, *block_requirements), has_mamba_cache
+
     def _init_minimal_kv_cache_for_profiling(self) -> None:
         from vllm.v1.core.kv_cache_utils import (
             get_kv_cache_config_from_groups,
@@ -6435,27 +6460,16 @@ def _init_minimal_kv_cache_for_profiling(self) -> None:
         KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec)
         kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec)
         max_capture_tokens = self.compilation_config.max_cudagraph_capture_size
-        if max_capture_tokens is None:
-            min_blocks = 1
-        else:
-            cp_size = get_total_cp_world_size()
-            block_requirements: list[int] = []
-            has_mamba_cache = False
-            for group in kv_cache_groups:
-                kv_cache_spec = group.kv_cache_spec
-                if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
-                    continue
-                if isinstance(kv_cache_spec, MambaSpec):
-                    has_mamba_cache = True
-                    block_requirements.append(self.max_num_reqs)
-                else:
-                    block_requirements.append(
-                        cdiv(max_capture_tokens, kv_cache_spec.block_size * cp_size)
-                    )
-            min_blocks = max(
-                1,
-                *block_requirements,
+        cp_size = get_total_cp_world_size()
+        min_blocks, has_mamba_cache = (
+            self._get_minimal_kv_cache_blocks_for_cudagraph_profiling(
+                kv_cache_groups,
+                max_capture_tokens,
+                self.max_num_reqs,
+                cp_size,
             )
+        )
+        if max_capture_tokens is not None:
             logger.info(
                 "Using %d KV blocks for CUDA graph profiling "
                 "(max_capture_tokens=%d, cp_size=%d, has_mamba_cache=%s)",