minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-08-13 08:35:38 +08:00
committed by GitHub
parent bbda7f32b1
commit aefe2d0207
3 changed files with 4 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh
-1
View File
@@ -604,7 +604,6 @@ 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
+4 -62
View File
@@ -15,10 +15,6 @@ 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 (
@@ -77,7 +73,6 @@ 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 (
@@ -1927,8 +1922,6 @@ 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
@@ -1948,8 +1941,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); the merged qkv-a GEMM uses None as its
# producer stream, while the fallback may issue on the alt stream.
# (gate, producer stream) issued on the alt stream by forward();
# None when the lazy path computes the gate here instead.
self._gate_precomputed = None
self._gate_alt_stream = gate_alt_stream
# Above this token count the attention-core kernels fill the SMs
@@ -1967,7 +1960,7 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
self._gate_hidden_states = None
precomputed = self._gate_precomputed
self._gate_precomputed = None
if precomputed is not None and precomputed[1] is not None:
if precomputed 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.
@@ -1990,52 +1983,6 @@ 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
@@ -2068,10 +2015,7 @@ 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
self._precompute_output_gate(hidden_states)
return super().forward(
positions, hidden_states, forward_batch, zero_allocator, **kwargs
)
@@ -3118,8 +3062,6 @@ 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
-118
View File
@@ -1,118 +0,0 @@
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()