[#57652 suggestion] AttentionSpec.max_tp_shards + KVCacheConfig.kv_tp_replicas for TP-replication detection

tests/v1/core/test_kv_cache_utils.py

@@ -52,6 +52,7 @@
     hash_block_tokens,
     init_none_hash,
     is_kv_cache_spec_uniform,
+    kv_cache_groups_tp_replicas,
     make_block_hash_with_group_id,
     tensor_data,
 )
@@ -3182,6 +3183,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,
@@ -3192,6 +3194,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,
@@ -4547,120 +4550,95 @@ def test_trailing_layer_fallback_requires_exact_partition():
     assert not any(g.is_eagle_group for g in trimmed)
 
 
-# ---------------------------------------------------------------------------
-# Tests for KVCacheConfig.all_groups_are_tp_replicated
-# ---------------------------------------------------------------------------
-
-
-def _make_tp_vllm_config(tp_size: int, dp_size: int = 1) -> "VllmConfig":
-    from vllm.config import ParallelConfig
-
-    cfg = VllmConfig(
-        model_config=ModelConfig(max_model_len=16),
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tp_size,
-            data_parallel_size=dp_size,
-        ),
-    )
-    cfg.cache_config.kv_cache_layout = "LBNHC"
-    return cfg
+_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,num_groups",
+    "specs,expected",
     [
-        pytest.param(
-            {"l.0": new_mla_spec(), "l.1": new_mla_spec()},
-            1,
-            id="single-mla-group",
-        ),
+        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(),
             },
-            2,
-            id="mla-and-swa-mla-groups",
+            4,
+            id="mla-and-swa-mla",
         ),
-    ],
-)
-def test_all_groups_tp_replicated_true(specs, num_groups):
-    mem = sum(s.page_size_bytes for s in specs.values()) * 10
-    configs = get_kv_cache_configs(_make_tp_vllm_config(tp_size=2), [specs], [mem])
-    assert len(configs[0].kv_cache_groups) == num_groups
-    assert configs[0].all_groups_are_tp_replicated is True
-
-
-@pytest.mark.parametrize(
-    "specs,num_groups",
-    [
+        pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"),
         pytest.param(
-            {
-                "l.0": FullAttentionSpec(
-                    block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32
-                ),
-                "l.1": FullAttentionSpec(
-                    block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32
-                ),
-            },
+            {"l.0": new_mla_spec(), "l.1": _GQA_SPEC},
             1,
-            id="non-mla-only",
+            id="mla-and-gqa-uniform-group",
         ),
         pytest.param(
             {
                 "l.0": new_mla_spec(),
                 "l.1": new_mla_spec(),
-                "s.0": SlidingWindowSpec(
-                    block_size=16,
-                    num_kv_heads=8,
-                    head_size=64,
-                    dtype=torch.float32,
-                    sliding_window=128,
-                ),
-                "s.1": SlidingWindowSpec(
-                    block_size=16,
-                    num_kv_heads=8,
-                    head_size=64,
-                    dtype=torch.float32,
-                    sliding_window=128,
-                ),
+                "s.0": _GQA_SWA_SPEC,
+                "s.1": _GQA_SWA_SPEC,
             },
             1,
-            id="mla-and-non-mla-groups",
+            id="mla-and-gqa-swa",
         ),
     ],
 )
-def test_all_groups_tp_replicated_false_spec(specs, num_groups):
-    """Non-MLA specs disqualify TP replication regardless of parallelism."""
+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(_make_tp_vllm_config(tp_size=2), [specs], [mem])
-    assert len(configs[0].kv_cache_groups) == num_groups
-    assert configs[0].all_groups_are_tp_replicated is False
+    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(
-    "tp_size,extra_kwargs,expected",
+    "specs,tp_size,dcp_size,expected",
     [
-        pytest.param(2, {}, True, id="tp2-only"),
-        pytest.param(2, {"data_parallel_size": 2}, True, id="tp2-dp2"),
-        pytest.param(1, {}, False, id="tp1"),
-        pytest.param(1, {"data_parallel_size": 2}, False, id="tp1-dp2"),
-        pytest.param(2, {"pipeline_parallel_size": 2}, False, id="pp"),
-        pytest.param(2, {"prefill_context_parallel_size": 2}, False, id="pcp"),
+        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_all_groups_tp_replicated_parallelism(tp_size, extra_kwargs, expected):
-    """Verify TP replication gating across parallelism configurations."""
-    from vllm.config import ParallelConfig
-
-    vllm_config = VllmConfig(
-        model_config=ModelConfig(max_model_len=16),
-        parallel_config=ParallelConfig(tensor_parallel_size=tp_size, **extra_kwargs),
-    )
-    vllm_config.cache_config.kv_cache_layout = "LBNHC"
-    specs = {"l.0": new_mla_spec(), "l.1": new_mla_spec()}
-    mem = new_mla_spec().page_size_bytes * 2 * 10
-    configs = get_kv_cache_configs(vllm_config, [specs], [mem])
-    assert configs[0].all_groups_are_tp_replicated is expected
+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

@@ -18,8 +18,8 @@
 )
 from vllm.platforms import current_platform
 from vllm.v1.core.kv_cache_utils import (
-    _all_kv_groups_tp_replicated,
     generate_scheduler_kv_cache_config,
+    kv_cache_groups_tp_replicas,
 )
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
@@ -167,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,
@@ -324,8 +325,10 @@ def _replicated_layout(
     config.parallel_config.nnodes = nnodes
     if world_size is not None:
         config.parallel_config.world_size = world_size
-    kv_cache_config.all_groups_are_tp_replicated = _all_kv_groups_tp_replicated(
-        kv_cache_config.kv_cache_groups, config
+    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
@@ -350,6 +353,7 @@ def _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,
@@ -363,6 +367,7 @@ def _make_swa_mla_kv_cache_config(
     spec = SlidingWindowMLASpec(
         block_size=16,
         num_kv_heads=1,
+        max_tp_shards=1,
         head_size=512,
         dtype=torch.float32,
         sliding_window=128,
@@ -403,6 +408,7 @@ def _swa_group(tag: str, sw: int) -> KVCacheGroupSpec:
         spec = SlidingWindowMLASpec(
             block_size=block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=head_size,
             dtype=dtype,
             sliding_window=sw,
@@ -420,6 +426,7 @@ def _mla_group(tag: str) -> KVCacheGroupSpec:
         spec = MLAAttentionSpec(
             block_size=block_size,
             num_kv_heads=1,
+            max_tp_shards=1,
             head_size=head_size,
             dtype=dtype,
         )
@@ -473,6 +480,7 @@ def _make_dsv3_2_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig:
     spec = MLAAttentionSpec(
         block_size=block_size,
         num_kv_heads=1,
+        max_tp_shards=1,
         head_size=512,
         dtype=torch.float32,
     )
@@ -928,6 +936,7 @@ def test_replicated_layout_excludes_unproven_cache_shapes(
                         SlidingWindowMLASpec(
                             block_size=16,
                             num_kv_heads=1,
+                            max_tp_shards=1,
                             head_size=512,
                             dtype=torch.float32,
                             sliding_window=128,
@@ -1109,6 +1118,7 @@ def test_parallelism_agnostic_for_single_full_attention_group():
 _SWA_MLA_SPEC = SlidingWindowMLASpec(
     block_size=16,
     num_kv_heads=1,
+    max_tp_shards=1,
     head_size=576,
     dtype=torch.float32,
     sliding_window=128,
@@ -1201,8 +1211,10 @@ def _get_offloading_config(canonical: bool):
             ],
         )
 
-        kv_config.all_groups_are_tp_replicated = _all_kv_groups_tp_replicated(
-            kv_config.kv_cache_groups, config
+        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)

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

@@ -157,8 +157,14 @@ def build_offloading_config(
     )
     replicated_layout = (
         vllm_config.model_config.use_mla
-        and kv_cache_config.all_groups_are_tp_replicated
+        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
+        and parallel_config.world_size == parallel_config.tensor_parallel_size
         # Shared /dev/shm mmap layout is single-node mp only.
         and parallel_config.distributed_executor_backend == "mp"
         and parallel_config.nnodes_within_dp == 1

vllm/model_executor/layers/attention/mla_attention.py

@@ -1350,6 +1350,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         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 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         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 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         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 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         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 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         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 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             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,