[AMD] [Kimi-K3] Fuse the KDA input projection into a single GEMM on ROCm (#35176)

Co-authored-by: HAI <hixiao@gmail.com>

docs/docs/references/environment_variables.mdx

@@ -1837,6 +1837,16 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable linear-layer quantization on ROCm.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ROCM_K3_FUSE_KDA_INPROJ</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Kimi-K3 on ROCm: fold the KDA <code>[f_a|b]</code> tail into the wide <code>[q,k,v,g]</code> projection so the whole input projection is one GEMM. Applies to unquantized weights only; falls back to the split projection otherwise.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Token count above which <code>SGLANG_ROCM_K3_FUSE_KDA_INPROJ</code> stops applying and the split projection runs instead. The merged shape is only faster while the projection is bandwidth bound.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>256</code></td>
+    </tr>
   </tbody>
 </table>
 

python/sglang/srt/environ.py

@@ -857,6 +857,11 @@ class Envs:
     # Enable dual-stream MoE (shared experts vs routed experts) on the
     # ROCm/AITER path. Requires GPU_MAX_HW_QUEUES>=5 to avoid HW-queue serialization.
     SGLANG_ROCM_USE_MULTI_STREAM = EnvBool(False)
+    # Fold the KDA [f_a|b] tail into the wide [q,k,v,g] projection so the whole
+    # in-proj is one GEMM. Decode is bandwidth bound there, so the 144 extra
+    # output columns ride along nearly free.
+    SGLANG_ROCM_K3_FUSE_KDA_INPROJ = EnvBool(True)
+    SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS = EnvInt(256)
     SGLANG_HACK_FLASHMLA_BACKEND = EnvStr("tilelang")
     SGLANG_USE_AITER_FP8_PER_TOKEN = EnvBool(False)
     # Above 8192 tokens of context, aiter's non-static workspace is large enough

python/sglang/srt/models/kimi_k3.py

@@ -10,6 +10,7 @@ import logging
 import os
 from collections.abc import Iterable
 from functools import cached_property
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
@@ -1451,6 +1452,8 @@ class KimiK3DeltaAttention(nn.Module):
             # (6144/rank at TP8). Folding b (12/rank) and f_a (128, replicated)
             # in as well skews the output dim to 6284 and measurably degrades
             # the GEMM kernel selection; they stay as separate tiny GEMVs.
+            # (ROCm reverses this below the token threshold -- see
+            # _merge_kda_inproj_weights_hip.)
             self.fused_qkvg_proj = MergedColumnParallelLinear(
                 self.hidden_size,
                 [
@@ -1500,6 +1503,17 @@ class KimiK3DeltaAttention(nn.Module):
             # _merge_bfa_weights().
             self._bfa_w: Optional[torch.Tensor] = None
             self._bfa_f_b_w: Optional[torch.Tensor] = None
+            if _is_hip:
+                # ROCm only: _merge_kda_inproj_weights_hip() may merge the
+                # whole [q,k,v,g | f_a | b] in-proj instead, making _bfa_w a
+                # tail view of that buffer. _qkvgbfa_sizes is the split of the
+                # buffer, and stays None when the fusion does not apply. These
+                # attributes exist on ROCm only; every reader is _is_hip-gated.
+                self._qkvgbfa_layer: Optional[SimpleNamespace] = None
+                self._qkvgbfa_sizes: Optional[list[int]] = None
+                self._qkvgbfa_bs_limit = (
+                    envs.SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS.get()
+                )
         elif self.do_fuse_qkvbfg:
             self.qkvb_sizes = [
                 projection_size,
@@ -1723,6 +1737,11 @@ class KimiK3DeltaAttention(nn.Module):
             return
         if _is_npu:
             return
+        if _is_hip and self._merge_kda_inproj_weights_hip():
+            # Split-path f_b GEMM still uses this when the fused in-proj
+            # is above the token threshold.
+            self._bfa_f_b_w = self.f_b_proj.weight
+            return
         mods = [self.f_a_proj, self.b_proj]
         if self._bfa_uses_block_fp8:
             weights = [_get_k3_dense_weight(mod) for mod in mods]
@@ -1739,6 +1758,59 @@ class KimiK3DeltaAttention(nn.Module):
             self._bfa_f_b_w = self.f_b_proj.weight
         self._bfa_fa_size, self._bfa_b_size = sizes
 
+    def _merge_kda_inproj_weights_hip(self) -> bool:
+        """ROCm only: append the [f_a | b] tail to the wide [q,k,v,g] buffer so
+        one GEMM covers the whole in-proj, and take _bfa_w as a tail view of
+        that buffer. The merge is view-only, so the wide-only and whole-buffer
+        weights both stay live and forward_qkvbfg_fused picks per batch size.
+
+        Returns False when the fusion does not apply, leaving the caller to do
+        the plain [f_a | b] merge."""
+        if not self._may_fuse_kda_inproj():
+            return False
+
+        # [q,k,v,g | f_a | b | pad]; f_a/b keep the same relative order and the
+        # same pad (both widths are 4 short of a multiple of 8), so the tail
+        # view is byte-identical to the wide-only merge.
+        merged, sizes = _merge_weights_as_views(
+            [self.fused_qkvg_proj, self.f_a_proj, self.b_proj], pad_rows_to=8
+        )
+        self._bfa_fa_size, self._bfa_b_size = sizes[-2:]
+        self._bfa_w = merged[sizes[0] :]
+        # Stand-in "layer" so the fused GEMM goes through the same
+        # quant_method.apply (and therefore the same backend choice) as the
+        # wide projection, whose own .weight stays the 6144-row view for the
+        # above-threshold split path. Not an nn.Module on purpose: this must
+        # not add a duplicate entry to state_dict.
+        self._qkvgbfa_layer = SimpleNamespace(weight=merged)
+        self._qkvgbfa_sizes = [
+            *self.split_sizes,  # q,k,v then g
+            self._bfa_fa_size,
+            self._bfa_b_size,
+            merged.shape[0] - sum(sizes),  # alignment pad
+        ]
+        return True
+
+    def _may_fuse_kda_inproj(self) -> bool:
+        """Whether the [f_a|b] tail can share the wide projection's buffer.
+
+        Needs the wide fused projection to exist and all three weights to be
+        plain unquantized 2-D tensors of one dtype and width -- the checkpoint
+        keeps attention in bf16, but a quantized variant would carry scales
+        that a raw row-cat would silently drop."""
+        if not (_is_hip and envs.SGLANG_ROCM_K3_FUSE_KDA_INPROJ.get()):
+            return False
+        if not (self.do_fuse_qkvbfg and self.use_full_rank_gate):
+            return False
+        # Block-FP8 in-proj needs dequantized BF16 buffers; a raw row-cat
+        # would drop the scales. Leave fusion to the split [f_a|b] path.
+        if self._bfa_uses_block_fp8:
+            return False
+        ws = [m.weight for m in (self.fused_qkvg_proj, self.f_a_proj, self.b_proj)]
+        if not all(type(w.data) is torch.Tensor and w.dim() == 2 for w in ws):
+            return False
+        return len({(w.dtype, w.shape[1]) for w in ws}) == 1
+
     def _prepare_fused_decode(self) -> None:
         """Static inputs for the fused KDA decode kernel
         (kernels/ops/attention/kda_fused_decode): per-segment transposed fp32 conv
@@ -1847,6 +1919,25 @@ class KimiK3DeltaAttention(nn.Module):
                 n_fa, n_b = self._bfa_fa_size, self._bfa_b_size
                 from sglang.kernels.ops.kimi_k3 import kimi_k3_tiny_gemm as gemm
 
+                if (
+                    _is_hip
+                    and self._qkvgbfa_sizes is not None
+                    and 0 < hidden_states.shape[0] <= self._qkvgbfa_bs_limit
+                ):
+                    # ROCm only. One GEMM for the whole in-proj: the [f_a|b]
+                    # tail rides along in the wide projection's bandwidth
+                    # instead of paying its own launch. Worth ~30% of the
+                    # in-proj at decode on gfx950; see
+                    # SGLANG_ROCM_K3_FUSE_KDA_INPROJ.
+                    fused_states = self.fused_qkvg_proj.quant_method.apply(
+                        self._qkvgbfa_layer, hidden_states, None
+                    )
+                    qkv, g_proj_states, f_a, beta, _pad = torch.split(
+                        fused_states, self._qkvgbfa_sizes, dim=-1
+                    )
+                    forget_gate = gemm(f_a, self._bfa_f_b_w)
+                    return qkv, beta, forget_gate, g_proj_states
+
                 if (
                     self._bfa_alt_stream is not None
                     and get_is_capture_mode()

test/registered/amd/test_kimi_k3_kda_inproj_fusion.py

@@ -0,0 +1,122 @@
+"""Layout check for the ROCm fused Kimi-K3 KDA input projection.
+
+Below SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS the whole in-proj is one GEMM
+over ``[q,k,v,g | f_a | b | pad]`` instead of a wide GEMM plus a tiny [f_a|b]
+GEMV. Both layouts are views over the same buffer, so the two paths have to
+agree; this pins the slice offsets, the tail view the split path still reads,
+and the fact that the strided f_a slice is a legal input to the f_b GEMM and
+to the fused decode kernel's shape gate.
+
+The two paths run different GEMM kernels (N=6288 has no tuned aiter config,
+N=6144 does), so they agree to bf16 rounding, not bitwise.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.kimi_k3 import kimi_k3_tiny_gemm
+from sglang.srt.models.kimi_k3 import _merge_weights_as_views
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+HIDDEN = 7168
+HEADS_TP = 12  # num_heads / tp8
+HEAD_DIM = 128
+PROJ_TP = HEADS_TP * HEAD_DIM  # 1536
+WIDE = 4 * PROJ_TP  # 6144, [q,k,v,g]
+MERGED = 6288  # + f_a(128) + b(12) + pad(4)
+# bf16 carries ~8 mantissa bits, so one ULP is ~4e-3 relative.
+TOL = 8e-3
+
+
+class _Fake(torch.nn.Module):
+    """Minimal stand-in for a linear layer: _merge_weights_as_views only
+    touches .weight.data."""
+
+    def __init__(self, rows, device):
+        super().__init__()
+        self.weight = torch.nn.Parameter(
+            torch.randn(rows, HIDDEN, dtype=torch.bfloat16, device=device) * 0.02,
+            requires_grad=False,
+        )
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "no GPU")
+class TestKimiK3KDAInProjFusion(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(0)
+        dev = torch.device("cuda", 0)
+        cls.qkvg = _Fake(WIDE, dev)
+        cls.f_a = _Fake(HEAD_DIM, dev)
+        cls.b = _Fake(HEADS_TP, dev)
+        cls.pre = [m.weight.data.clone() for m in (cls.qkvg, cls.f_a, cls.b)]
+        cls.f_b_w = (
+            torch.randn(PROJ_TP, HEAD_DIM, dtype=torch.bfloat16, device=dev) * 0.02
+        )
+        cls.merged, cls.sizes = _merge_weights_as_views(
+            [cls.qkvg, cls.f_a, cls.b], pad_rows_to=8
+        )
+        cls.split_sizes = [3 * PROJ_TP, PROJ_TP]
+        cls.all_sizes = cls.split_sizes + [
+            HEAD_DIM,
+            HEADS_TP,
+            MERGED - WIDE - HEAD_DIM - HEADS_TP,
+        ]
+
+    def test_merged_layout(self):
+        self.assertEqual(self.sizes, [WIDE, HEAD_DIM, HEADS_TP])
+        self.assertEqual(tuple(self.merged.shape), (MERGED, HIDDEN))
+        self.assertEqual(sum(self.all_sizes), MERGED)
+
+    def test_views_alias_and_preserve_values(self):
+        """The merge must re-point, not reorder or copy-and-drop."""
+        tail = self.merged[WIDE:]
+        self.assertEqual(tail.data_ptr(), self.f_a.weight.data_ptr())
+        self.assertTrue(tail.is_contiguous())
+        self.assertTrue(self.qkvg.weight.is_contiguous())
+        self.assertEqual(self.qkvg.weight.data_ptr(), self.merged.data_ptr())
+        for got, want in zip((self.qkvg, self.f_a, self.b), self.pre):
+            self.assertTrue(torch.equal(got.weight.data, want))
+
+    def test_split_and_fused_paths_agree(self):
+        tail = self.merged[WIDE:]
+        for tokens in (1, 4, 8, 33, 128, 256):
+            with self.subTest(tokens=tokens):
+                x = torch.randn(
+                    tokens, HIDDEN, dtype=torch.bfloat16, device=self.merged.device
+                )
+
+                wide = torch.nn.functional.linear(x, self.qkvg.weight)
+                s_qkv, s_g = torch.split(wide, self.split_sizes, dim=-1)
+                s_bfa = kimi_k3_tiny_gemm(x, tail)
+                s_beta = s_bfa[..., HEAD_DIM : HEAD_DIM + HEADS_TP]
+                s_fg = kimi_k3_tiny_gemm(s_bfa[..., :HEAD_DIM], self.f_b_w)
+
+                allp = torch.nn.functional.linear(x, self.merged)
+                f_qkv, f_g, f_fa, f_beta, _pad = torch.split(
+                    allp, self.all_sizes, dim=-1
+                )
+                # The fused decode kernel's shape gate requires a unit last
+                # stride but tolerates the wider row stride.
+                self.assertEqual(f_fa.stride(-1), 1)
+                self.assertEqual(f_qkv.stride(-1), 1)
+                self.assertEqual(f_fa.stride(0), MERGED)
+                f_fg = kimi_k3_tiny_gemm(f_fa, self.f_b_w)
+
+                for name, got, want in (
+                    ("qkv", f_qkv, s_qkv),
+                    ("g", f_g, s_g),
+                    ("beta", f_beta, s_beta),
+                    ("forget_gate", f_fg, s_fg),
+                ):
+                    scale = want.float().abs().max().clamp_min(1e-6)
+                    rel = ((got.float() - want.float()).abs().max() / scale).item()
+                    self.assertLess(rel, TOL, f"{name} rel err {rel:.2e}")
+
+
+if __name__ == "__main__":
+    unittest.main()