[II] Warm Kimi B12X MLA DCP transport

tests/model_executor/test_flashinfer_autotune_cache.py

@@ -244,6 +244,66 @@ def test_b12x_dcp_warmup_finds_generic_mla_attention(monkeypatch) -> None:
     ]
 
 
+def test_b12x_dcp_warmup_finds_kimi_dense_mla_attention(monkeypatch) -> None:
+    from vllm.distributed import parallel_state
+    from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
+    from vllm.v1.attention.ops import dcp_alltoall
+
+    monkeypatch.setenv("VLLM_USE_B12X_DCP_A2A", "1")
+    attention = MultiHeadLatentAttention.__new__(MultiHeadLatentAttention)
+    torch.nn.Module.__init__(attention)
+    attention.register_parameter(
+        "device_probe",
+        torch.nn.Parameter(torch.empty(1)),
+    )
+    attention.attn_backend = SimpleNamespace(get_name=lambda: "B12X_MLA")
+    attention.impl = SimpleNamespace(dcp_world_size=16)
+    attention.num_local_heads = 6
+    attention.head_size = 576
+    attention.kv_lora_rank = 512
+
+    model = torch.nn.Module()
+    model.add_module("attention", attention)
+    worker = SimpleNamespace(
+        get_model=lambda: model,
+        model_runner=SimpleNamespace(),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4096),
+        vllm_config=SimpleNamespace(
+            parallel_config=SimpleNamespace(
+                decode_context_parallel_size=16,
+                dcp_comm_backend="a2a",
+            ),
+            compilation_config=SimpleNamespace(
+                static_forward_context={"model.layers.0.attn": attention}
+            ),
+        ),
+    )
+    group = object()
+    calls = []
+    monkeypatch.setattr(parallel_state, "get_dcp_group", lambda: group)
+    monkeypatch.setattr(
+        dcp_alltoall,
+        "warmup_b12x_dcp_a2a",
+        lambda *args, **kwargs: calls.append((args, kwargs)),
+    )
+
+    assert kernel_warmup._warmup_b12x_dcp_a2a(worker) == 1
+    assert calls == [
+        (
+            (group,),
+            {
+                "device": torch.device("cpu"),
+                "dtype": torch.bfloat16,
+                "max_batch_size": 4096,
+                "total_heads": 96,
+                "head_dim": 512,
+                "query_head_dim": 576,
+            },
+        )
+    ]
+
+
 def test_kernel_warmup_runs_b12x_mxfp8_linear_warmup(monkeypatch) -> None:
     calls = []
     model = torch.nn.Linear(2, 2)

vllm/model_executor/warmup/kernel_warmup.py

@@ -210,6 +210,9 @@ def _warmup_b12x_dcp_a2a(worker: "Worker") -> int:
     from vllm.models.deepseek_v4.nvidia.b12x import (
         DeepseekV4B12xMLAAttention,
     )
+    from vllm.models.kimi_k3.nvidia.mla import (
+        MultiHeadLatentAttention as KimiK3MLAAttention,
+    )
     from vllm.v1.attention.ops.dcp_alltoall import warmup_b12x_dcp_a2a
 
     model = worker.get_model()
@@ -235,6 +238,17 @@ def _warmup_b12x_dcp_a2a(worker: "Worker") -> int:
             total_heads = int(module.num_heads) * dcp_world_size
             query_head_dim = int(module.kv_lora_rank + module.qk_rope_head_dim)
             output_head_dim = int(module.kv_lora_rank)
+        elif (
+            isinstance(module, KimiK3MLAAttention)
+            and module.attn_backend.get_name() == "B12X_MLA"
+        ):
+            module_dcp_world_size = int(module.impl.dcp_world_size)
+            if module_dcp_world_size <= 1:
+                continue
+            device = next(module.parameters()).device
+            total_heads = int(module.num_local_heads) * module_dcp_world_size
+            query_head_dim = int(module.head_size)
+            output_head_dim = int(module.kv_lora_rank)
         else:
             continue