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