launcher/dcp2/diffs/17-v1_core_kv_cache_utils.py.diff

--- baked-ours/v1/core/kv_cache_utils.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/v1/core/kv_cache_utils.py	2026-09-11 10:13:47.000000000 +0200
@@ -37,6 +37,7 @@
     SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
     compute_layout_strides,
+    is_dcp_sharded_spec,
     iter_layer_specs,
     replace_as,
 )
@@ -651,13 +652,15 @@
 
 
 def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
-    """Return the token span of a cache block under DCP."""
-    layer_specs = iter_layer_specs(spec)
-    if len(layer_specs) > 0 and all(
-        isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
-    ):
-        return spec.block_size * dcp_world_size
-    return spec.block_size
+    """Return the token span of a cache block under DCP.
+
+    Only DCP-sharded groups span ``dcp_world_size`` blocks of tokens per
+    manager block; replicated attention groups (``dcp_replicated``), Mamba and
+    other per-rank state keep their plain block size. Must agree with
+    ``dcp_world_size_for_kv_cache_spec`` (the coordinator scales manager block
+    sizes with that helper).
+    """
+    return spec.block_size * dcp_world_size_for_kv_cache_spec(spec, dcp_world_size)
 
 
 def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
@@ -686,15 +689,13 @@
     must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
 
     Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
-    ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
-    draft group would need a different spec, not this helper.
+    ``MLAAttentionSpec`` and therefore keep the process DCP size. Attention
+    groups whose layers set ``dcp_replicated`` (DeepSeek-V4.1 SWA cache and
+    compressor ring) hold the full KV on every rank and keep ``1``.
     """
     if dcp_world_size <= 1:
         return 1
-    inner = spec
-    if isinstance(spec, UniformTypeKVCacheSpecs):
-        inner = next(iter(spec.kv_cache_specs.values()))
-    if isinstance(inner, FullAttentionSpec):
+    if is_dcp_sharded_spec(spec):
         return dcp_world_size
     return 1