[KV-Offloading][TP] : Expand replicated_layout detection to multi-group MLA (#57652)

tests/v1/core/test_kv_cache_utils.py

@@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import (
     hash_block_tokens,
     init_none_hash,
     is_kv_cache_spec_uniform,
+    kv_cache_groups_tp_replicas,
     make_block_hash_with_group_id,
     tensor_data,
     to_event_extra_keys,
@@ -3274,6 +3275,7 @@ def new_mla_spec(cache_dtype_str=None, block_size: int = 16):
     return MLAAttentionSpec(
         block_size=block_size,
         num_kv_heads=1,
+        max_tp_shards=1,
         head_size=576,
         dtype=torch.float32,
         cache_dtype_str=cache_dtype_str,
@@ -3284,6 +3286,7 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):
     return SlidingWindowMLASpec(
         block_size=16,
         num_kv_heads=1,
+        max_tp_shards=1,
         head_size=head_size,
         dtype=torch.float32,
         sliding_window=sliding_window,
@@ -4637,3 +4640,97 @@ def test_trailing_layer_fallback_requires_exact_partition():
     _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)
 
     assert not any(g.is_eagle_group for g in trimmed)
+
+
+_GQA_SPEC = FullAttentionSpec(
+    block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32
+)
+_GQA_SWA_SPEC = SlidingWindowSpec(
+    block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128
+)
+
+
+@pytest.mark.parametrize(
+    "specs,expected",
+    [
+        pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"),
+        pytest.param(
+            {
+                "l.0": new_mla_spec(),
+                "l.1": new_mla_spec(),
+                "s.0": new_swa_mla_spec(),
+                "s.1": new_swa_mla_spec(),
+            },
+            4,
+            id="mla-and-swa-mla",
+        ),
+        pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"),
+        pytest.param(
+            {"l.0": new_mla_spec(), "l.1": _GQA_SPEC},
+            1,
+            id="mla-and-gqa-uniform-group",
+        ),
+        pytest.param(
+            {
+                "l.0": new_mla_spec(),
+                "l.1": new_mla_spec(),
+                "s.0": _GQA_SWA_SPEC,
+                "s.1": _GQA_SWA_SPEC,
+            },
+            1,
+            id="mla-and-gqa-swa",
+        ),
+    ],
+)
+def test_kv_tp_replicas(monkeypatch, specs, expected):
+    """Resolved per layer before scheduler flattening hides mixed groups."""
+    from vllm.config import ParallelConfig
+    from vllm.platforms import current_platform
+
+    monkeypatch.setattr(current_platform, "device_count", lambda: 4)
+    parallel_config = ParallelConfig(tensor_parallel_size=4)
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config
+    )
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
+    mem = sum(s.page_size_bytes for s in specs.values()) * 10
+    configs = get_kv_cache_configs(vllm_config, [specs], [mem])
+    assert configs[0].kv_tp_replicas == expected
+    scheduler = generate_scheduler_kv_cache_config(configs)
+    assert scheduler.kv_tp_replicas == expected
+
+
+@pytest.mark.parametrize(
+    "specs,tp_size,dcp_size,expected",
+    [
+        pytest.param([new_mla_spec()], 8, 1, 8, id="mla"),
+        pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"),
+        pytest.param(
+            [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset"
+        ),
+        pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"),
+        pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"),
+        pytest.param(
+            [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)],
+            8,
+            1,
+            2,
+            id="mla-and-gqa-partial",
+        ),
+        pytest.param(
+            [
+                HiddenStateCacheSpec(
+                    block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32
+                )
+            ],
+            8,
+            1,
+            1,
+            id="hidden-state",
+        ),
+    ],
+)
+def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected):
+    """Replicas are the gcd of each layer's tp_size // max_tp_shards."""
+    groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)]
+    assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected

tests/v1/kv_connector/unit/offloading_connector/test_config.py

@@ -17,7 +17,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     SchedulerOffloadConfig,
 )
 from vllm.platforms import current_platform
-from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
+from vllm.v1.core.kv_cache_utils import (
+    generate_scheduler_kv_cache_config,
+    kv_cache_groups_tp_replicas,
+)
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
@@ -164,6 +167,7 @@ def _mla_spec(
     return MLAAttentionSpec(
         block_size=block_size,
         num_kv_heads=1,
+        max_tp_shards=1,
         head_size=head_size,
         dtype=dtype,
         tokens_per_state=tokens_per_state,
@@ -321,7 +325,184 @@ def _replicated_layout(
     config.parallel_config.nnodes = nnodes
     if world_size is not None:
         config.parallel_config.world_size = world_size
-    return build_offloading_config(config, kv_cache_config).replicated_layout
+    kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
+        kv_cache_config.kv_cache_groups,
+        config.parallel_config.tensor_parallel_size,
+        config.parallel_config.decode_context_parallel_size,
+    )
+
+    # Scheduler and Worker offload config
+    worker_offload_config = build_offloading_config(config, kv_cache_config)
+    scheduler_offload_config = build_offloading_config(
+        config, generate_scheduler_kv_cache_config([kv_cache_config])
+    )
+    # Verify that the Worker and Scheduler agree on replication.
+    # and worker-bytes-per-block
+    assert (
+        worker_offload_config.replicated_layout
+        == scheduler_offload_config.replicated_layout
+    )
+    assert (
+        worker_offload_config.worker_kv_bytes_per_block
+        == scheduler_offload_config.worker_kv_bytes_per_block
+    )
+
+    return worker_offload_config.replicated_layout
+
+
+_SWA_MLA_PAGE = SlidingWindowMLASpec(
+    block_size=16,
+    num_kv_heads=1,
+    max_tp_shards=1,
+    head_size=512,
+    dtype=torch.float32,
+    sliding_window=128,
+).page_size_bytes
+
+
+def _make_swa_mla_kv_cache_config(
+    num_blocks: int = 4,
+) -> KVCacheConfig:
+    """Single SlidingWindowMLASpec group — DSV4 sliding-window layers."""
+    spec = SlidingWindowMLASpec(
+        block_size=16,
+        num_kv_heads=1,
+        max_tp_shards=1,
+        head_size=512,
+        dtype=torch.float32,
+        sliding_window=128,
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=_SWA_MLA_PAGE * num_blocks,
+                layers=["layer"],
+                layer_stride=_SWA_MLA_PAGE * num_blocks,
+                block_stride=_SWA_MLA_PAGE,
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)],
+    )
+
+
+def _make_dsv4_flash_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig:
+    """Minimal fake of DeepSeek-V4-Flash's five UniformTypeKVCacheSpecs groups.
+
+    The real model has the following groups (layer counts from runtime logs);
+    2 per group is enough to exercise the structural properties that govern
+    replicated_layout:
+        group[0]: 22 SlidingWindowMLASpec layers, sliding_window=128
+        group[1]: 21 SlidingWindowMLASpec layers, sliding_window=128
+        group[2]: 62 MLAAttentionSpec layers  (MLA + indexer.k_cache)
+        group[3]: 42 SlidingWindowMLASpec layers, sliding_window=8
+        group[4]: 20 SlidingWindowMLASpec layers, sliding_window=128
+    All inner specs have num_kv_heads=1 → every group is TP-replicated.
+    """
+    block_size = 16
+    dtype = torch.float32
+    head_size = 512
+
+    def _swa_group(tag: str, sw: int) -> KVCacheGroupSpec:
+        names = [f"{tag}_0", f"{tag}_1"]
+        spec = SlidingWindowMLASpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            max_tp_shards=1,
+            head_size=head_size,
+            dtype=dtype,
+            sliding_window=sw,
+        )
+        return KVCacheGroupSpec(
+            names,
+            UniformTypeKVCacheSpecs(
+                block_size=block_size,
+                kv_cache_specs={name: spec for name in names},
+            ),
+        )
+
+    def _mla_group(tag: str) -> KVCacheGroupSpec:
+        names = [f"{tag}_0", f"{tag}_1"]
+        spec = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            max_tp_shards=1,
+            head_size=head_size,
+            dtype=dtype,
+        )
+        return KVCacheGroupSpec(
+            names,
+            UniformTypeKVCacheSpecs(
+                block_size=block_size,
+                kv_cache_specs={name: spec for name in names},
+            ),
+        )
+
+    groups = [
+        _swa_group("swa_128a", 128),  # group[0]: SWA-MLA, sw=128
+        _swa_group("swa_128b", 128),  # group[1]: SWA-MLA, sw=128
+        _mla_group("mla"),  # group[2]: MLA + indexer layers
+        _swa_group("swa_8", 8),  # group[3]: SWA-MLA, sw=8
+        _swa_group("swa_128c", 128),  # group[4]: SWA-MLA, sw=128
+    ]
+    # All specs share the same page size (num_kv_heads=1, head_size=512, fp32).
+    page = MLAAttentionSpec(
+        block_size=block_size, num_kv_heads=1, head_size=head_size, dtype=dtype
+    ).page_size_bytes
+    all_layers = [name for g in groups for name in g.layer_names]
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=page * len(all_layers) * num_blocks,
+                layers=all_layers,
+                layer_stride=page * num_blocks,
+                block_stride=page,
+            )
+        ],
+        kv_cache_groups=groups,
+    )
+
+
+def _make_dsv3_2_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig:
+    """Minimal fake of DeepSeek-V3.2-Exp's single UniformTypeKVCacheSpecs group.
+
+    The real model has 1 group (layer counts from runtime logs):
+        group[0]: 122 MLAAttentionSpec layers (self_attn.attn + indexer.k_cache)
+    All inner specs have num_kv_heads=1 → the group is TP-replicated.
+
+    GLM-5.3 has the same topology (one UniformType group of MLAAttentionSpec
+    layers, num_kv_heads=1) with 99 layers instead of 122, so this config
+    covers both models.
+    """
+    block_size = 16
+    names = [f"mla_{i}" for i in range(4)]  # 4 layers; real model has 122
+    spec = MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        max_tp_shards=1,
+        head_size=512,
+        dtype=torch.float32,
+    )
+    group = KVCacheGroupSpec(
+        names,
+        UniformTypeKVCacheSpecs(
+            block_size=block_size,
+            kv_cache_specs={name: spec for name in names},
+        ),
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=spec.page_size_bytes * len(names) * num_blocks,
+                layers=names,
+                layer_stride=spec.page_size_bytes * num_blocks,
+                block_stride=spec.page_size_bytes,
+            )
+        ],
+        kv_cache_groups=[group],
+    )
 
 
 @pytest.mark.parametrize("packed", [False, True])
@@ -588,35 +769,64 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
     )
 
 
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_replicated_layout_enabled_for_swa_mla(world_size: int):
+    """SlidingWindowMLASpec groups are also replicated and qualify."""
+    assert _replicated_layout(
+        _make_swa_mla_kv_cache_config(),
+        tensor_parallel_size=world_size,
+    )
+
+
+@pytest.mark.parametrize("world_size", [2, 4, 8])
+def test_replicated_layout_dsv4_flash(world_size: int):
+    """DeepSeek-V4-Flash: all 5 UniformType groups are MLA/SWA-MLA with
+    num_kv_heads=1, so replicated_layout must be True under TP/mp single-node."""
+    assert _replicated_layout(
+        _make_dsv4_flash_kv_cache_config(),
+        tensor_parallel_size=world_size,
+    )
+
+
+@pytest.mark.parametrize("world_size", [2, 4, 8])
+def test_replicated_layout_dsv3_2(world_size: int):
+    """DeepSeek-V3.2 / GLM-5.3: single UniformType group of MLAAttentionSpec
+    layers (122 and 99 respectively), num_kv_heads=1 → replicated_layout must
+    be True under TP/mp single-node."""
+    assert _replicated_layout(
+        _make_dsv3_2_kv_cache_config(),
+        tensor_parallel_size=world_size,
+    )
+
+
+def test_replicated_layout_enabled_for_uniform_wrapper_mla():
+    """UniformTypeKVCacheSpecs wrapping MLA layers of different sizes qualifies:
+    each layer is replicated and the page accounting sums correctly."""
+    mla_layers = ["layer0", "layer1"]
+    spec = UniformTypeKVCacheSpecs(
+        block_size=16,
+        kv_cache_specs={"layer0": _mla_spec(), "layer1": _mla_spec(head_size=256)},
+    )
+    total_page = spec.page_size_bytes
+    num_blocks = 4
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=total_page * num_blocks,
+                layers=mla_layers,
+                layer_stride=total_page * num_blocks,
+                block_stride=total_page,
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(mla_layers, spec)],
+    )
+    assert _replicated_layout(kv_cache_config)
+
+
 @pytest.mark.parametrize(
     ("kv_cache_config", "case"),
     [
-        (
-            KVCacheConfig(
-                num_blocks=4,
-                kv_cache_tensors=[
-                    KVCacheTensor(
-                        size=_MLA_PAGE * 4,
-                        layers=["layer"],
-                        layer_stride=_MLA_PAGE * 4,
-                        block_stride=_MLA_PAGE,
-                    )
-                ],
-                kv_cache_groups=[
-                    KVCacheGroupSpec(
-                        ["layer"],
-                        SlidingWindowMLASpec(
-                            block_size=16,
-                            num_kv_heads=1,
-                            head_size=512,
-                            dtype=torch.float32,
-                            sliding_window=128,
-                        ),
-                    )
-                ],
-            ),
-            "sliding-window-mla",
-        ),
         (
             KVCacheConfig(
                 num_blocks=4,
@@ -642,40 +852,6 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
             ),
             "hidden-state",
         ),
-        (
-            # One group, two page sizes: layer1's run starts past layer0's region.
-            KVCacheConfig(
-                num_blocks=4,
-                kv_cache_tensors=[
-                    KVCacheTensor(
-                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
-                        layers=["layer0"],
-                        layer_stride=_MLA_PAGE * 4,
-                        block_stride=_MLA_PAGE,
-                    ),
-                    KVCacheTensor(
-                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
-                        layers=["layer1"],
-                        layer_stride=_HALF_MLA_PAGE * 4,
-                        block_stride=_HALF_MLA_PAGE,
-                        offset=_MLA_PAGE * 4,
-                    ),
-                ],
-                kv_cache_groups=[
-                    KVCacheGroupSpec(
-                        ["layer0", "layer1"],
-                        UniformTypeKVCacheSpecs(
-                            block_size=16,
-                            kv_cache_specs={
-                                "layer0": _mla_spec(),
-                                "layer1": _mla_spec(head_size=256),
-                            },
-                        ),
-                    )
-                ],
-            ),
-            "uniform-wrapper",
-        ),
         (
             # Overlaid groups with different page sizes: a block is a window of
             # the largest group's packing.
@@ -726,6 +902,84 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
             ),
             "mla-mamba-hybrid",
         ),
+    ],
+    ids=[
+        "hidden-state",
+        "mla-full-hybrid",
+        "mla-mamba-hybrid",
+    ],
+)
+def test_replicated_layout_excludes_unproven_cache_shapes(
+    kv_cache_config: KVCacheConfig,
+    case: str,
+):
+    assert not _replicated_layout(kv_cache_config), case
+
+
+@pytest.mark.parametrize(
+    ("kv_cache_config", "case"),
+    [
+        (
+            KVCacheConfig(
+                num_blocks=4,
+                kv_cache_tensors=[
+                    KVCacheTensor(
+                        size=_MLA_PAGE * 4,
+                        layers=["layer"],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
+                    )
+                ],
+                kv_cache_groups=[
+                    KVCacheGroupSpec(
+                        ["layer"],
+                        SlidingWindowMLASpec(
+                            block_size=16,
+                            num_kv_heads=1,
+                            max_tp_shards=1,
+                            head_size=512,
+                            dtype=torch.float32,
+                            sliding_window=128,
+                        ),
+                    )
+                ],
+            ),
+            "sliding-window-mla",
+        ),
+        (
+            # One group, two page sizes: layer1's run starts past layer0's region.
+            KVCacheConfig(
+                num_blocks=4,
+                kv_cache_tensors=[
+                    KVCacheTensor(
+                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
+                        layers=["layer0"],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
+                    ),
+                    KVCacheTensor(
+                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
+                        layers=["layer1"],
+                        layer_stride=_HALF_MLA_PAGE * 4,
+                        block_stride=_HALF_MLA_PAGE,
+                        offset=_MLA_PAGE * 4,
+                    ),
+                ],
+                kv_cache_groups=[
+                    KVCacheGroupSpec(
+                        ["layer0", "layer1"],
+                        UniformTypeKVCacheSpecs(
+                            block_size=16,
+                            kv_cache_specs={
+                                "layer0": _mla_spec(),
+                                "layer1": _mla_spec(head_size=256),
+                            },
+                        ),
+                    )
+                ],
+            ),
+            "uniform-wrapper",
+        ),
         (
             KVCacheConfig(
                 num_blocks=4,
@@ -748,31 +1002,54 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
     ],
     ids=[
         "sliding-window-mla",
-        "hidden-state",
         "uniform-wrapper",
-        "mla-full-hybrid",
-        "mla-mamba-hybrid",
         "multi-group-mla",
     ],
 )
-def test_replicated_layout_excludes_unproven_cache_shapes(
+def test_replicated_layout_supported_cache_shapes(
     kv_cache_config: KVCacheConfig,
     case: str,
 ):
-    assert not _replicated_layout(kv_cache_config), case
+    assert _replicated_layout(kv_cache_config), case
 
 
-def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
+@pytest.mark.parametrize(
+    ("main_spec", "add_spec", "expected"),
+    [
+        (
+            _mla_spec(head_size=512),
+            _mla_spec(head_size=128, dtype=torch.uint8),  # indexer.k_cache spec
+            True,
+        ),
+        (
+            _mla_spec(head_size=512),
+            _full_attention_spec(),
+            False,
+        ),
+    ],
+    ids=["uniform-mla-indexer", "uniform-mla-fa"],
+)
+def test_replicated_layout_for_dsa_style_groups(
+    main_spec: KVCacheSpec, add_spec: KVCacheSpec, expected: bool
+):
+    """KV Cache groups with packed-specs inside a UniformTypeKVCacheSpec."""
     num_blocks = 4
-    main_spec = _mla_spec(head_size=512)
-    indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
     main_layers = [f"main_{i}" for i in range(61)]
-    indexer_layers = [f"indexer_{i}" for i in range(61)]
-    # A DSA-style group: the main pages and the smaller indexer pages are packed one
-    # after the other, so a block holds more than 61 MLA pages.
+    add_layers = [f"add_{i}" for i in range(61)]
+
     main_bytes = main_spec.page_size_bytes * len(main_layers)
-    indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers)
-    size = (main_bytes + indexer_bytes) * num_blocks
+    add_bytes = add_spec.page_size_bytes * len(add_layers)
+    size = (main_bytes + add_bytes) * num_blocks
+
+    kv_cache_specs = {
+        **{layer: main_spec for layer in main_layers},
+        **{layer: add_spec for layer in add_layers},
+    }
+    layer_names = main_layers + add_layers
+    group_spec: KVCacheSpec = UniformTypeKVCacheSpecs(
+        block_size=16, kv_cache_specs=kv_cache_specs
+    )
+
     kv_cache_config = KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[
@@ -784,16 +1061,16 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
             ),
             KVCacheTensor(
                 size=size,
-                layers=indexer_layers,
-                layer_stride=indexer_spec.page_size_bytes * num_blocks,
-                block_stride=indexer_spec.page_size_bytes,
+                layers=add_layers,
+                layer_stride=add_spec.page_size_bytes * num_blocks,
+                block_stride=add_spec.page_size_bytes,
                 offset=main_bytes * num_blocks,
             ),
         ],
-        kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names, group_spec)],
     )
 
-    assert not _replicated_layout(kv_cache_config)
+    assert _replicated_layout(kv_cache_config) == expected
 
 
 @pytest.mark.parametrize(
@@ -841,6 +1118,7 @@ _SWA_SPEC = SlidingWindowSpec(
 _SWA_MLA_SPEC = SlidingWindowMLASpec(
     block_size=16,
     num_kv_heads=1,
+    max_tp_shards=1,
     head_size=576,
     dtype=torch.float32,
     sliding_window=128,
@@ -904,6 +1182,52 @@ def test_canonical_layout_gate(kv_cache_groups, certified):
     assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified
 
 
+@pytest.mark.parametrize(
+    ("spec"),
+    [_SWA_MLA_SPEC, _mla_spec(tokens_per_state=2)],
+    ids=["swa", "compressed-mla"],
+)
+def test_uncertifiable_canonical_layout_resets_replication(spec: KVCacheSpec):
+    def _get_offloading_config(canonical: bool):
+        config = _make_vllm_config(
+            tensor_parallel_size=2,
+            extra_config={"canonical_layout": True} if canonical else None,
+        )
+        config.model_config.use_mla = True
+
+        PAGE = spec.page_size_bytes
+        kv_config = KVCacheConfig(
+            num_blocks=4,
+            kv_cache_tensors=[
+                KVCacheTensor(
+                    size=PAGE * 4,
+                    layers=["layer"],
+                    layer_stride=PAGE * 4,
+                    block_stride=PAGE,
+                )
+            ],
+            kv_cache_groups=[
+                KVCacheGroupSpec(["layer"], spec),
+            ],
+        )
+
+        kv_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
+            kv_config.kv_cache_groups,
+            config.parallel_config.tensor_parallel_size,
+            config.parallel_config.decode_context_parallel_size,
+        )
+
+        return build_offloading_config(config, kv_config)
+
+    offload_config = _get_offloading_config(canonical=False)
+    assert offload_config.replicated_layout
+    assert not offload_config.parallel.is_parallelism_agnostic
+
+    canon_offload_config = _get_offloading_config(canonical=True)
+    assert not canon_offload_config.replicated_layout
+    assert not canon_offload_config.parallel.is_parallelism_agnostic
+
+
 def test_canonical_layout_certifies_v2_model_runner():
     """Canonical bytes are certified per layer against live tensor strides at
     registration, so the static gate must not depend on the model-runner

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

@@ -157,15 +157,10 @@ def build_offloading_config(
     )
     replicated_layout = (
         vllm_config.model_config.use_mla
-        # Exact type: fail closed on wrappers and sliding-window variants.
-        and type(single_group_spec) is MLAAttentionSpec
-        # Page accounting: one MLA page per layer, no packed/mixed rows.
-        and worker_kv_bytes_per_block > 0
-        and worker_kv_bytes_per_block
-        == single_group_spec.page_size_bytes
-        * len(kv_cache_config.kv_cache_groups[0].layer_names)
-        # Safe MVP boundary: TP-only, no other parallel axes.
         and parallel_config.tensor_parallel_size > 1
+        and kv_cache_config.kv_tp_replicas == parallel_config.tensor_parallel_size
+        and worker_kv_bytes_per_block > 0
+        # Safe MVP boundary: TP-only, no other parallel axes.
         and parallel_config.pipeline_parallel_size == 1
         and parallel_config.prefill_context_parallel_size == 1
         and parallel_config.decode_context_parallel_size == 1
@@ -233,9 +228,10 @@ def build_offloading_config(
             and parallel_config.world_size == tp_size
         )
 
-    if canonical_layout and is_parallelism_agnostic:
+    if canonical_layout:
         replicated_layout = (
-            all(
+            is_parallelism_agnostic
+            and all(
                 type(spec) is MLAAttentionSpec
                 for _, group in selected_groups
                 for spec in iter_layer_specs(group.kv_cache_spec)

vllm/model_executor/layers/attention/mla_attention.py

@@ -1350,6 +1350,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         common_kwargs = dict(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.head_size,
             dtype=kv_cache_dtype,
             cache_dtype_str=self.kv_cache_dtype,

vllm/model_executor/models/deepseek_v2.py

@@ -658,6 +658,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             cache_role=SparseCacheRole.INDEXER,

vllm/models/deepseek_v4/attention.py

@@ -865,6 +865,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         return MLAAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.head_dim,
             dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
             tokens_per_state=self.compress_ratio,
@@ -910,6 +911,7 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             tokens_per_state=self.compress_ratio,

vllm/models/deepseek_v4/compressor.py

@@ -184,6 +184,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
         return SlidingWindowMLASpec(  # only has one vector instead of K + V
             block_size=self.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.state_dim,
             dtype=self.dtype,
             sliding_window=self.sliding_window,

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -124,6 +124,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
             bounded_replay=self.bounded_replay,
             block_size=self.block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             sliding_window=self.window_size,

vllm/v1/core/kv_cache_utils.py

@@ -9,7 +9,7 @@ import os
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field, replace
-from functools import partial
+from functools import partial, reduce
 from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
 
 from vllm import envs
@@ -2409,6 +2409,26 @@ def get_kv_cache_groups(
     return groups
 
 
+def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int:
+    if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None:
+        return 1
+    if spec.dcp_sharded and dcp_size > 1:
+        return 1
+    return max(1, tp_size // spec.max_tp_shards)
+
+
+def kv_cache_groups_tp_replicas(
+    groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1
+) -> int:
+    """Consecutive TP ranks holding identical KV for every layer."""
+    specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)]
+    if not specs:
+        return 1
+    return reduce(
+        math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size
+    )
+
+
 def generate_scheduler_kv_cache_config(
     kv_cache_configs: list[KVCacheConfig],
 ) -> KVCacheConfig:
@@ -2823,6 +2843,13 @@ def get_kv_cache_configs(
             vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
         )
 
+    for kv_cache_config in kv_cache_configs:
+        kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
+            kv_cache_config.kv_cache_groups,
+            vllm_config.parallel_config.tensor_parallel_size,
+            vllm_config.parallel_config.decode_context_parallel_size,
+        )
+
     return kv_cache_configs
 
 

vllm/v1/kv_cache_interface.py

@@ -484,6 +484,9 @@ class HiSparseResidentSpec(KVCacheSpec):
 @dataclass(frozen=True, kw_only=True)
 class AttentionSpec(KVCacheSpec):
     dcp_sharded: bool = True
+    max_tp_shards: int | None = None
+    """Distinct shards this cache splits into across TP; TP ranks beyond this
+    hold replicas. None: unknown, treated as not replicated."""
     num_kv_heads: int
     head_size: int
     dtype: torch.dtype
@@ -618,6 +621,7 @@ class FullAttentionSpec(AttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -705,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -768,6 +773,7 @@ class RSWASpec(FullAttentionSpec):
             dtype=base.dtype,
             kv_quant_mode=base.kv_quant_mode,
             dcp_sharded=base.dcp_sharded,
+            max_tp_shards=base.max_tp_shards,
             page_size_padded=base.page_size_padded,
             num_head_slots=base.num_head_slots,
             state_content_bytes=base.state_content_bytes,
@@ -960,6 +966,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         bounded_replay_set = set(spec.bounded_replay for spec in specs)
         block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
         assert len({spec.dcp_sharded for spec in specs}) == 1
+        assert len({spec.max_tp_shards for spec in specs}) == 1
         assert (
             len(cache_dtype_str_set) == 1
             and len(tokens_per_state_set) == 1
@@ -980,6 +987,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
             head_size=specs[0].head_size,
             dtype=specs[0].dtype,
             dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -1197,6 +1205,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
             dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -1483,6 +1492,9 @@ class KVCacheConfig:
     hisparse_shared_host_pool: bool = False
     """Whether local TP ranks share one physical HiSparse host pool."""
 
+    kv_tp_replicas: int = 1
+    """Consecutive TP ranks holding identical KV for every layer (1: none)."""
+
     @cached_property
     def transfer_group_ids(self) -> tuple[int, ...]:
         """IDs of cache groups that participate in external KV transfer."""