[Kimi K3] Fuse MLA gate projection into QKV-A GEMM (#33623)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh

@@ -604,6 +604,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
 constexpr int pick_tile_m(int hd_in, int hd_out) {
   if (hd_out == 2624 && hd_in == 6144) return 32;
   if (hd_out == 4096 && hd_in == 2048) return 32;
+  if (hd_out == 3648 && hd_in == 7168) return 32;
   return 16;
 }
 

python/sglang/srt/models/kimi_k3.py

@@ -15,6 +15,10 @@ import torch
 from torch import nn
 
 from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated
+from sglang.kernels.ops.gemm.fused_a_gemm import (
+    dsv3_fused_a_gemm,
+    fused_a_gemm_weight_eligible,
+)
 from sglang.srt.configs.kimi_k3 import KimiK3Config
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.distributed import (
@@ -73,6 +77,7 @@ from sglang.srt.layers.moe.utils import (
     get_moe_runner_backend,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import (
@@ -1922,6 +1927,8 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
             # is the wrong group at attn_tp>1 and deadlocks against idle DP
             # ranks.
             self.o_proj.use_dp_attention_reduce = True
+        self._qkv_a_g_proj_weight = None
+        self._qkv_a_g_proj_sizes = None
         if self.use_output_gate:
             projection_size = config.num_attention_heads * config.v_head_dim
             # Shard by attn-TP to match the attention output (DSV2 MLA shards
@@ -1941,8 +1948,8 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
             # cores, so wrap its forward at the instance level; the module
             # itself (weights, reduce_results, loading path) is untouched.
             self._gate_hidden_states = None
-            # (gate, producer stream) issued on the alt stream by forward();
+            # (gate, producer stream); the merged qkv-a GEMM uses None as its
-            # None when the lazy path computes the gate here instead.
+            # producer stream, while the fallback may issue on the alt stream.
             self._gate_precomputed = None
             self._gate_alt_stream = gate_alt_stream
             # Above this token count the attention-core kernels fill the SMs
@@ -1960,7 +1967,7 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
                 self._gate_hidden_states = None
                 precomputed = self._gate_precomputed
                 self._gate_precomputed = None
-                if precomputed is not None:
+                if precomputed is not None and precomputed[1] is not None:
                     # Use wait_stream rather than an explicit event so the
                     # breakable-CUDA-graph runner can track the side-stream
                     # join across graph-segment boundaries.
@@ -1983,6 +1990,52 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
 
             self.o_proj.forward = _gated_o_proj_forward
 
+    def _merge_qkv_a_g_proj_weights(self) -> None:
+        """Merge the same-input MLA qkv-a and TP-local output-gate weights."""
+        if not self.use_output_gate:
+            return
+        mods = [self.fused_qkv_a_proj_with_mqa, self.g_proj]
+        # K3's global MXFP4 config ignores attention; inspect the resolved
+        # methods instead of treating a non-None quant_config as quantized.
+        if any(
+            not isinstance(mod.quant_method, UnquantizedLinearMethod) for mod in mods
+        ):
+            return
+        dtypes = {mod.weight.dtype for mod in mods}
+        if len(dtypes) != 1 or dtypes.pop() not in (torch.bfloat16, torch.float16):
+            return
+        self._qkv_a_g_proj_weight, self._qkv_a_g_proj_sizes = _merge_weights_as_views(
+            mods
+        )
+
+    def prepare_qkv_latent(
+        self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
+    ):
+        weight = self._qkv_a_g_proj_weight
+        if (
+            weight is None
+            or not isinstance(hidden_states, torch.Tensor)
+            or getattr(self.fused_qkv_a_proj_with_mqa, "set_lora", False)
+            or getattr(self.g_proj, "set_lora", False)
+        ):
+            return super().prepare_qkv_latent(hidden_states, forward_batch)
+
+        if self._use_min_latency_fused_a_gemm is None:
+            self._use_min_latency_fused_a_gemm = (
+                not get_exec().deterministic.enable_deterministic_inference
+                and weight.shape[0] % 16 == 0
+                and fused_a_gemm_weight_eligible(self.fused_qkv_a_proj_with_mqa)
+            )
+        if self._use_min_latency_fused_a_gemm and 1 <= hidden_states.shape[0] <= 16:
+            fused = dsv3_fused_a_gemm(
+                hidden_states, weight.T, backend=self.fused_a_gemm_backend
+            )
+        else:
+            fused = _k3_bf16_gemm(hidden_states, weight)
+        qkv_latent, gate = torch.split(fused, self._qkv_a_g_proj_sizes, dim=-1)
+        self._gate_precomputed = (gate, None)
+        return qkv_latent
+
     def _precompute_output_gate(self, hidden_states: torch.Tensor) -> None:
         """Issue the output-gate GEMM on the alt stream so it overlaps the
         attention core; the lazy path in the o_proj wrap otherwise computes
@@ -2015,7 +2068,10 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
     ):
         if self.use_output_gate:
             self._gate_hidden_states = hidden_states
+            if self._qkv_a_g_proj_weight is None:
                 self._precompute_output_gate(hidden_states)
+            else:
+                self._gate_precomputed = None
         return super().forward(
             positions, hidden_states, forward_batch, zero_allocator, **kwargs
         )
@@ -3062,6 +3118,8 @@ class KimiK3LinearForCausalLM(nn.Module):
             if isinstance(layer.self_attn, KimiK3DeltaAttention):
                 layer.self_attn._merge_bfa_weights()
                 layer.self_attn._prepare_fused_decode()
+            elif isinstance(layer.self_attn, KimiK3MLAAttention):
+                layer.self_attn._merge_qkv_a_g_proj_weights()
 
         for layer in self.model.layers:
             if isinstance(layer, PPMissingLayer) or not isinstance(

test/registered/unit/models/test_kimi_k3_mla_gate_fusion.py

@@ -0,0 +1,118 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.models.kimi_k3 import KimiK3MLAAttention
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestKimiK3MlaGateFusion(unittest.TestCase):
+    def test_merged_projection_matches_separate_projections(self):
+        torch.manual_seed(0)
+        qkv_proj = SimpleNamespace(
+            weight=torch.nn.Parameter(torch.randn(12, 16, dtype=torch.bfloat16)),
+            quant_method=UnquantizedLinearMethod(),
+        )
+        g_proj = SimpleNamespace(
+            weight=torch.nn.Parameter(torch.randn(8, 16, dtype=torch.bfloat16)),
+            quant_method=UnquantizedLinearMethod(),
+        )
+        attn = SimpleNamespace(
+            use_output_gate=True,
+            quant_config=object(),
+            fused_qkv_a_proj_with_mqa=qkv_proj,
+            g_proj=g_proj,
+            _qkv_a_g_proj_weight=None,
+            _qkv_a_g_proj_sizes=None,
+            _use_min_latency_fused_a_gemm=False,
+            _gate_precomputed=None,
+        )
+        x = torch.randn(3, 16, dtype=torch.bfloat16)
+        expected_qkv = torch.nn.functional.linear(x, qkv_proj.weight)
+        expected_gate = torch.nn.functional.linear(x, g_proj.weight)
+
+        KimiK3MLAAttention._merge_qkv_a_g_proj_weights(attn)
+        qkv = KimiK3MLAAttention.prepare_qkv_latent(attn, x, None)
+        gate, stream = attn._gate_precomputed
+
+        torch.testing.assert_close(qkv, expected_qkv)
+        torch.testing.assert_close(gate, expected_gate)
+        self.assertIsNone(stream)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_fused_a_cuda_graph_replay(self):
+        if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+            self.skipTest("SM90+ required")
+
+        # K3 TP8: replicated qkv-a [2112, 7168] plus local gate [1536, 7168].
+        qkv_proj = SimpleNamespace(
+            weight=torch.nn.Parameter(
+                torch.randn(2112, 7168, dtype=torch.bfloat16, device="cuda")
+            ),
+            quant_method=UnquantizedLinearMethod(),
+        )
+        g_proj = SimpleNamespace(
+            weight=torch.nn.Parameter(
+                torch.randn(1536, 7168, dtype=torch.bfloat16, device="cuda")
+            ),
+            quant_method=UnquantizedLinearMethod(),
+        )
+        attn = SimpleNamespace(
+            use_output_gate=True,
+            quant_config=object(),
+            fused_qkv_a_proj_with_mqa=qkv_proj,
+            g_proj=g_proj,
+            _qkv_a_g_proj_weight=None,
+            _qkv_a_g_proj_sizes=None,
+            _use_min_latency_fused_a_gemm=None,
+            _gate_precomputed=None,
+            fused_a_gemm_backend="jit",
+        )
+        KimiK3MLAAttention._merge_qkv_a_g_proj_weights(attn)
+
+        exec_config = SimpleNamespace(
+            deterministic=SimpleNamespace(enable_deterministic_inference=False)
+        )
+        for num_tokens in (1, 8, 16):
+            with self.subTest(num_tokens=num_tokens):
+                static_x = torch.randn(
+                    num_tokens, 7168, dtype=torch.bfloat16, device="cuda"
+                )
+                with patch(
+                    "sglang.srt.models.kimi_k3.get_exec", return_value=exec_config
+                ):
+                    KimiK3MLAAttention.prepare_qkv_latent(attn, static_x, None)
+                self.assertTrue(attn._use_min_latency_fused_a_gemm)
+
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph):
+                    qkv = KimiK3MLAAttention.prepare_qkv_latent(attn, static_x, None)
+                gate = attn._gate_precomputed[0]
+
+                static_x.copy_(torch.randn_like(static_x))
+                graph.replay()
+                torch.cuda.synchronize()
+
+                torch.testing.assert_close(
+                    qkv,
+                    torch.nn.functional.linear(static_x, qkv_proj.weight),
+                    rtol=1e-2,
+                    atol=1e-3,
+                )
+                torch.testing.assert_close(
+                    gate,
+                    torch.nn.functional.linear(static_x, g_proj.weight),
+                    rtol=1e-2,
+                    atol=1e-3,
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()