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 c75fa100d..509113b4b 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,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; } diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 458bc4926..9630872a9 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -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( 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 deleted file mode 100644 index d0e0dd94b..000000000 --- a/test/registered/unit/models/test_kimi_k3_mla_gate_fusion.py +++ /dev/null @@ -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()