[II] Define per-cache DCP block-table geometry

tests/v1/core/test_kv_cache_utils.py

@@ -58,6 +58,7 @@
     SlidingWindowMLASpec,
     SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
+    get_kv_cache_dcp_shard_count,
     get_kv_cache_spec_kind,
     get_kv_cache_spec_sliding_window,
 )
@@ -197,6 +198,16 @@ def new_mamba_spec(
     )
 
 
+def test_mamba_cache_has_one_dcp_token_position_shard():
+    spec = new_mamba_spec(block_size=768, num_speculative_blocks=7)
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(mamba_cache_mode="align")
+    )
+
+    assert get_kv_cache_dcp_shard_count(spec, dcp_world_size=16) == 1
+    assert spec.max_num_blocks_per_req(vllm_config, max_len=1_000_000) == 1310
+
+
 def test_unify_kv_cache_spec_page_size_uses_lcm_for_non_divisible_pages():
     mimo_spec = FullAttentionSpec(
         block_size=64,
@@ -1309,8 +1320,11 @@ def test_uniform_type_spec_block_table_width_matches_layer_spec(
     # The runner sizes the block table from the group spec while the metadata
     # builders are constructed from the per-layer spec, so the aggregate must
     # report the same width as the layers it wraps.
-    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
-    vllm_config.parallel_config.decode_context_parallel_size = dcp_size
+    vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size),
+        cache_config=SimpleNamespace(mamba_cache_mode="none"),
+        model_config=SimpleNamespace(max_model_len=1024),
+    )
     if layer_type == "mla":
         layer_spec = new_mla_spec()
     elif layer_type == "replicated":

vllm/v1/kv_cache_interface.py

@@ -148,6 +148,21 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
         """
         return cdiv(max_len, self.block_size)
 
+    def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int:
+        """Return the number of unique token-position shards under DCP.
+
+        Cache types that store recurrent or otherwise rank-local state do not
+        shard that state by token position. Attention cache specifications
+        override this method because their default layout is DCP-sharded.
+        """
+        configured_dcp = int(dcp_world_size)
+        if configured_dcp < 1:
+            raise ValueError(
+                "Configured decode-context-parallel size must be positive: "
+                f"{configured_dcp}"
+            )
+        return 1
+
     def copy_with_new_block_size(self, block_size: int) -> Self:
         """
         Create a new KVCacheSpec from self but replacing the block size.
@@ -240,11 +255,38 @@ def real_page_size_bytes(self) -> int:
 
     def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
         parallel_config = vllm_config.parallel_config
-        kv_shard_count = get_kv_cache_dcp_shard_count(
-            self, parallel_config.decode_context_parallel_size
+        kv_shard_count = self.get_num_dcp_kv_shards(
+            parallel_config.decode_context_parallel_size
         )
         return cdiv(max_len, self.block_size * kv_shard_count)
 
+    def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int:
+        """Return the configured or explicitly overridden attention shard count."""
+        configured_dcp = int(dcp_world_size)
+        if configured_dcp < 1:
+            raise ValueError(
+                "Configured decode-context-parallel size must be positive: "
+                f"{configured_dcp}"
+            )
+        replicated = bool(getattr(self, "dcp_replicated", False))
+        override = getattr(self, "dcp_kv_shard_count", None)
+        if replicated:
+            if override not in (None, 1):
+                raise ValueError(
+                    "dcp_replicated cannot be combined with "
+                    f"dcp_kv_shard_count={override}"
+                )
+            return 1
+        if override is None:
+            return configured_dcp
+        override = int(override)
+        if override < 1 or override > configured_dcp or configured_dcp % override != 0:
+            raise ValueError(
+                "dcp_kv_shard_count must be a positive divisor of the configured "
+                f"DCP size, got shards={override}, DCP={configured_dcp}"
+            )
+        return override
+
 
 @dataclass(frozen=True, kw_only=True)
 class FullAttentionSpec(AttentionSpec):
@@ -398,36 +440,24 @@ def get_kv_cache_dcp_shard_count(
     dcp_world_size: int,
 ) -> int:
     """Return the number of unique DCP token-position shards for a cache group."""
-    configured_dcp = int(dcp_world_size)
-    if configured_dcp < 1:
-        raise ValueError(
-            f"Configured decode-context-parallel size must be positive: "
-            f"{configured_dcp}"
-        )
-    replicated = bool(getattr(spec, "dcp_replicated", False))
-    override = getattr(spec, "dcp_kv_shard_count", None)
-    if replicated:
-        if override not in (None, 1):
-            raise ValueError(
-                f"dcp_replicated cannot be combined with dcp_kv_shard_count={override}"
-            )
-        return 1
-    if override is None:
-        return configured_dcp
-    override = int(override)
-    if override < 1 or override > configured_dcp or configured_dcp % override != 0:
-        raise ValueError(
-            "dcp_kv_shard_count must be a positive divisor of the configured "
-            f"DCP size, got shards={override}, DCP={configured_dcp}"
-        )
-    return override
+    return spec.get_num_dcp_kv_shards(dcp_world_size)
 
 
 def has_nondefault_kv_dcp_layout(
     spec: KVCacheSpec,
     dcp_world_size: int,
 ) -> bool:
-    return get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size)
+    layer_specs = (
+        spec.kv_cache_specs.values()
+        if isinstance(spec, UniformTypeKVCacheSpecs)
+        else (spec,)
+    )
+    is_attention_group = all(
+        isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
+    )
+    return is_attention_group and (
+        get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size)
+    )
 
 
 def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):
@@ -1046,6 +1076,18 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
         )
         return next(iter(widths))
 
+    def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int:
+        shard_counts = {
+            spec.get_num_dcp_kv_shards(dcp_world_size)
+            for spec in self.kv_cache_specs.values()
+        }
+        if len(shard_counts) != 1:
+            raise ValueError(
+                "All layers in a uniform KV cache group must use the same "
+                f"number of DCP KV shards, got {sorted(shard_counts)}."
+            )
+        return next(iter(shard_counts))
+
     @classmethod
     def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
         """

vllm/v1/worker/gpu/model_runner.py

@@ -65,7 +65,6 @@
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
-from vllm.utils.math_utils import cdiv
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -619,19 +618,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         for kv_cache_group in kv_cache_config.kv_cache_groups:
             spec = kv_cache_group.kv_cache_spec
             block_sizes.append(spec.block_size)
-            # One local block covers `block_size * dcp_shard_count` tokens in
-            # the global sequence. Replicated groups keep the full cache on
-            # every rank instead.
             group_cp_size = get_kv_cache_dcp_shard_count(spec, self.dcp_size)
             group_cp_sizes.append(group_cp_size)
-            max_num_blocks = cdiv(
-                block_table_max_model_len, spec.block_size * group_cp_size
+            # Cache specifications own their block-table geometry. Attention
+            # caches account for their token-position DCP shards, while
+            # recurrent state and replicated attention caches remain unscaled.
+            max_num_blocks = spec.max_num_blocks_per_req(
+                self.vllm_config, block_table_max_model_len
             )
-            # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens
             if isinstance(spec, MambaSpec):
-                max_num_blocks = (
-                    max_num_blocks if self.cache_config.enable_prefix_caching else 1
-                ) + spec.num_speculative_blocks
                 max_num_blocks = get_block_table_width(
                     max_num_blocks, spec.block_size, token_alignment=None
                 )