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

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

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()