diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index f080d6aa03e3..517bab5fd65d 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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": diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 14f238b4c735..d3ebdb6fe3ec 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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: """ diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 16685bd7a2a4..53bfa20a77bd 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 )