launcher/dcp2/diffs/18-v1_kv_cache_interface.py.diff

--- 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 = {