[Bugfix][Kimi-K3] Declare max_tp_shards on the DSpark MLA KV cache spec

tests/models/test_dspark_mla.py

@@ -336,3 +336,36 @@ def __call__(self, main_x):
     assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2)
     assert calls[0][3] is slot_mappings[0]
     assert calls[1][3] is slot_mappings[2]
+
+
+@pytest.mark.cpu_test
+def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla():
+    """The draft's MLA layers must share a KV cache group with the target's."""
+    from vllm.model_executor.layers.attention.mla_attention import MLAAttention
+    from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
+    from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size
+
+    vllm_config = SimpleNamespace(
+        model_config=None, cache_config=SimpleNamespace(block_size=64)
+    )
+    target_attn = SimpleNamespace(
+        kv_cache_dtype="fp8",
+        head_size=576,
+        sliding_window=None,
+        indexer=None,
+        non_causal_multi_token_decode=False,
+        attn_backend=SimpleNamespace(get_name=lambda: "ROCM_AITER_MLA"),
+        _uses_flat_kv_cache=lambda: False,
+    )
+    draft_attn = SimpleNamespace(
+        kv_cache_dtype="fp8", head_size=576, non_causal_multi_token_decode=True
+    )
+    target_spec = MLAAttention.get_kv_cache_spec(target_attn, vllm_config)
+    draft_spec = MultiHeadLatentAttention.get_kv_cache_spec(draft_attn, vllm_config)
+
+    kv_cache_spec = {f"model.layers.{i}.attn": target_spec for i in range(24)}
+    kv_cache_spec |= {f"draft.layers.{i}.attn": draft_spec for i in range(5)}
+
+    groups = _get_kv_cache_groups_uniform_page_size(kv_cache_spec)
+    assert len(groups) == 1
+    assert len(groups[0].layer_names) == 29

vllm/models/kimi_k3/nvidia/mla.py

@@ -422,6 +422,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         return MLAAttentionSpec(  # type: ignore[call-arg]
             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,