--- baked-ours/v1/kv_cache_interface.py 2026-09-10 19:56:14.000000000 +0200 +++ baked/v1/kv_cache_interface.py 2026-09-11 10:13:47.000000000 +0200 @@ -240,8 +240,11 @@ f"Unsupported KV cache spec type: {type(self)}. " "Please register it using @register_kv_cache_spec decorator." ) + dcp_replicated = getattr(self, "dcp_replicated", None) return all( - isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() + isinstance(spec, uniform_type_base_spec) + and getattr(spec, "dcp_replicated", None) == dcp_replicated + for spec in kv_cache_specs.values() ) @@ -400,6 +403,13 @@ """Tokens covered by one stored state. Ints > 1 compress multiple tokens into one state (DSv4 sparse MLA); fractions < 1 store multiple states per token (Whisper block pooling: ``Fraction(1, block_pool_size)``).""" + dcp_replicated: bool = False + """DCP: every rank holds this layer's full KV instead of a 1/dcp shard. + A replicated group keeps its plain block geometry under DCP (no block-size + scaling, no per-rank memory divide, block table sized for the whole + context) and every rank writes every position (DeepSeek-V4.1 SWA cache and + compressor ring). Sharded groups (the default) round-robin positions across + ranks at ``cp_kv_cache_interleave_size`` granularity.""" def __post_init__(self): if self.head_size_v is None: @@ -438,7 +448,9 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = parallel_config.decode_context_parallel_size + kv_shard_count = ( + 1 if self.dcp_replicated else parallel_config.decode_context_parallel_size + ) return cdiv(max_len, self.block_size * kv_shard_count) @@ -471,7 +483,7 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not self.dcp_replicated: max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @@ -519,6 +531,7 @@ num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, tokens_per_state=specs[0].tokens_per_state, + dcp_replicated=specs[0].dcp_replicated, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -602,6 +615,7 @@ state_content_bytes=specs[0].state_content_bytes, cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), + dcp_replicated=specs[0].dcp_replicated, model_version=model_version_set.pop(), storage_block_size=storage_block_size_set.pop(), non_causal_multi_token_decode=non_causal_mtd_set.pop(), @@ -657,6 +671,7 @@ num_head_slots=base.num_head_slots, state_content_bytes=base.state_content_bytes, tokens_per_state=base.tokens_per_state, + dcp_replicated=base.dcp_replicated, sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, @@ -700,6 +715,7 @@ return all( isinstance(spec, ChunkedLocalAttentionSpec) and spec.attention_chunk_size == self.attention_chunk_size + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -743,9 +759,10 @@ return cdiv(num_tokens, self.block_size) + 1 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - assert vllm_config.parallel_config.decode_context_parallel_size == 1, ( - "DCP not support sliding window." - ) + assert ( + vllm_config.parallel_config.decode_context_parallel_size == 1 + or self.dcp_replicated + ), "DCP not support sliding window (sharded); mark the spec dcp_replicated." max_blocks = self.max_admission_blocks_per_request( max_in_flight_tokens=vllm_config.max_in_flight_tokens, max_model_len=vllm_config.model_config.max_model_len, @@ -758,6 +775,7 @@ return all( isinstance(spec, SlidingWindowSpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -786,7 +804,9 @@ self, kv_cache_specs: dict[str, KVCacheSpec] ) -> bool: return all( - isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values() + isinstance(spec, CircularBufferSpec) + and spec.dcp_replicated == self.dcp_replicated + for spec in kv_cache_specs.values() ) @property @@ -847,6 +867,7 @@ extra_retained_tokens=extra_retained_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), + dcp_replicated=specs[0].dcp_replicated, model_version=model_version_set.pop(), ) @@ -856,6 +877,7 @@ return all( isinstance(spec, SlidingWindowMLASpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -1013,6 +1035,7 @@ page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, + dcp_replicated=specs[0].dcp_replicated, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), @@ -1147,6 +1170,28 @@ ) +def is_dcp_replicated_spec(kv_cache_spec: KVCacheSpec) -> bool: + """Whether a KV cache group spec is (or wraps) DCP-replicated attention + layers: every DCP rank holds the full KV, so the group keeps its plain + block geometry and every rank writes every position.""" + layer_specs = iter_layer_specs(kv_cache_spec) + return len(layer_specs) > 0 and all( + isinstance(spec, AttentionSpec) and spec.dcp_replicated + for spec in layer_specs + ) + + +def is_dcp_sharded_spec(kv_cache_spec: KVCacheSpec) -> bool: + """Whether a KV cache group spec is (or wraps) DCP-sharded full attention + (incl. MLA): positions are round-robined across DCP ranks, so manager block + size, prefix hashing and per-rank memory scale by the DCP size.""" + layer_specs = iter_layer_specs(kv_cache_spec) + return len(layer_specs) > 0 and all( + isinstance(spec, FullAttentionSpec) and not spec.dcp_replicated + for spec in layer_specs + ) + + def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind: if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): inner_kinds = {