diff --git a/python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh index 509113b4b..c75fa100d 100644 --- a/python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh +++ b/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; } diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 9630872a9..458bc4926 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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(); - # None when the lazy path computes the gate here instead. + # (gate, producer stream); the merged qkv-a GEMM uses None as its + # 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 - self._precompute_output_gate(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( diff --git a/test/registered/unit/models/test_kimi_k3_mla_gate_fusion.py b/test/registered/unit/models/test_kimi_k3_mla_gate_fusion.py new file mode 100644 index 000000000..d0e0dd94b --- /dev/null +++ b/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()