diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx
index b10ac265e..99bf8c935 100644
--- a/docs/docs/references/environment_variables.mdx
+++ b/docs/docs/references/environment_variables.mdx
@@ -1837,6 +1837,16 @@ SGLang supports various environment variables that can be used to configure its
| Disable linear-layer quantization on ROCm. |
false |
+
+ SGLANG_ROCM_K3_FUSE_KDA_INPROJ |
+ Kimi-K3 on ROCm: fold the KDA [f_a|b] tail into the wide [q,k,v,g] projection so the whole input projection is one GEMM. Applies to unquantized weights only; falls back to the split projection otherwise. |
+ true |
+
+
+ SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS |
+ Token count above which SGLANG_ROCM_K3_FUSE_KDA_INPROJ stops applying and the split projection runs instead. The merged shape is only faster while the projection is bandwidth bound. |
+ 256 |
+
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 7698b1200..ae8a96d61 100644
--- a/python/sglang/srt/environ.py
+++ b/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
diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py
index a815afa7a..9770af47b 100644
--- a/python/sglang/srt/models/kimi_k3.py
+++ b/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()
diff --git a/test/registered/amd/test_kimi_k3_kda_inproj_fusion.py b/test/registered/amd/test_kimi_k3_kda_inproj_fusion.py
new file mode 100644
index 000000000..e6ea31072
--- /dev/null
+++ b/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()