[Perf] mla: Use masked K concatenation for padded head counts

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -730,10 +730,11 @@ def _concat_and_cast_mha_k(
             concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe)
         elif (
             _is_cuda
-            and next_power_of_2(self.num_local_heads) == self.num_local_heads
             and next_power_of_2(self.qk_nope_head_dim) == self.qk_nope_head_dim
             and next_power_of_2(self.qk_rope_head_dim) == self.qk_rope_head_dim
         ):
+            # The Triton helper masks padded heads, including Kimi-K3's 12
+            # local heads at TP8; only the feature dimensions need powers of two.
             # fa3 mha support fp8 inputs
             if (
                 self.current_attention_backend == "fa3"

test/registered/kernels/ops/attention/test_concat_mla.py

@@ -1,5 +1,7 @@
 import itertools
 import sys
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -90,6 +92,100 @@ def jit_concat_mla_absorb_q(
 NUM_TOKENS_LIST = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
 
 
+@pytest.mark.parametrize(
+    "num_tokens,num_heads",
+    [(0, 12), (1, 12), (17, 12), (257, 24), (19, 8), (65536, 12)],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_model_mha_k_concat(num_tokens, num_heads, dtype):
+    """Exercise model dispatch with strided KV projection and latent RoPE views."""
+    from sglang.srt.models.deepseek_common.attention_forward_methods import forward_mha
+
+    attn = SimpleNamespace(
+        num_local_heads=num_heads,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        qk_head_dim=192,
+        current_attention_backend="trtllm_mla",
+        kv_cache_dtype="fp8_e4m3",
+    )
+    kv = torch.randn(num_tokens, num_heads, 256, device="cuda", dtype=dtype)
+    latent = torch.randn(num_tokens, 1, 576, device="cuda", dtype=dtype)
+    k_nope, k_rope = kv[..., :128], latent[..., 512:]
+    expected = torch.empty(num_tokens, num_heads, 192, device="cuda", dtype=dtype)
+    torch_concat_mla_k(expected, k_nope, k_rope)
+    with patch.object(
+        forward_mha,
+        "concat_and_cast_mha_k_triton",
+        wraps=forward_mha.concat_and_cast_mha_k_triton,
+    ) as concat:
+        actual = forward_mha.DeepseekMHAForwardMixin._concat_and_cast_mha_k(
+            attn, k_nope, k_rope, None
+        )
+    concat.assert_called_once()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_model_mha_k_concat_graph():
+    from sglang.srt.models.deepseek_common.attention_forward_methods import forward_mha
+
+    attn = SimpleNamespace(
+        num_local_heads=12,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        qk_head_dim=192,
+        current_attention_backend="trtllm_mla",
+        kv_cache_dtype="fp8_e4m3",
+    )
+    kv = torch.randn(17, 12, 256, device="cuda", dtype=torch.bfloat16)
+    latent = torch.randn(17, 1, 576, device="cuda", dtype=torch.bfloat16)
+    k_nope, k_rope = kv[..., :128], latent[..., 512:]
+    concat = forward_mha.DeepseekMHAForwardMixin._concat_and_cast_mha_k
+    concat(attn, k_nope, k_rope, None)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = concat(attn, k_nope, k_rope, None)
+    for _ in range(2):
+        kv.normal_()
+        latent.normal_()
+        graph.replay()
+        expected = torch.empty_like(actual)
+        torch_concat_mla_k(expected, k_nope, k_rope)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_model_mha_k_concat_fa3_fp8():
+    from sglang.srt.models.deepseek_common.attention_forward_methods import forward_mha
+
+    attn = SimpleNamespace(
+        num_local_heads=12,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        qk_head_dim=192,
+        current_attention_backend="fa3",
+        kv_cache_dtype="fp8_e4m3",
+    )
+    k_nope = torch.randn(17, 12, 128, device="cuda", dtype=torch.bfloat16)
+    k_rope = torch.randn(17, 1, 64, device="cuda", dtype=torch.bfloat16)
+    expected = torch.empty(17, 12, 192, device="cuda", dtype=torch.bfloat16)
+    torch_concat_mla_k(expected, k_nope, k_rope)
+    with patch.object(
+        forward_mha,
+        "get_token_to_kv_pool",
+        return_value=SimpleNamespace(dtype=torch.float8_e4m3fn),
+    ):
+        actual = forward_mha.DeepseekMHAForwardMixin._concat_and_cast_mha_k(
+            attn, k_nope, k_rope, None
+        )
+    assert actual.dtype == torch.float8_e4m3fn
+    torch.testing.assert_close(
+        actual.view(torch.uint8),
+        expected.to(actual.dtype).view(torch.uint8),
+        rtol=0,
+        atol=0,
+    )
+
+
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS_LIST)
 def test_concat_mla_k_jit_vs_torch(num_tokens: int) -> None:
     """Test JIT kernel against PyTorch reference."""