[Kimi-K3] Use explicit SiTU activation for MegaMoE (#34883)

python/pyproject.toml

@@ -69,7 +69,7 @@ dependencies = [
   "sentencepiece",
   "setproctitle",
   "sgl-deep-ep==0.1.0",
-  "sgl-deep-gemm==0.1.5.post2",
+  "sgl-deep-gemm==0.1.5.post3",
   "sglang-kernel==0.4.6.post1",
   "smg-grpc-servicer>=0.5.0",
   "soundfile==0.13.1",

python/sglang/srt/layers/quantization/mxfp4.py

@@ -342,6 +342,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         # Weights stay MXFP4 (e2m1 + ue8m0 g32, zero requantization);
         # activations are quantized to fp8 per-token-group-128.
         self.use_deep_gemm = get_moe_runner_backend().is_deep_gemm()
+        # MegaMoE calls DeepGEMM directly instead of going through the selected
+        # MoE runner, so it needs the same unpadded checkpoint layout and
+        # DeepGEMM scale preparation even when the runner is flashinfer_mxfp4.
+        self.use_mega_moe = get_moe_a2a_backend().is_megamoe()
         self.flashinfer_mxfp4_moe_precision = (
             get_exec().moe.flashinfer_mxfp4_moe_precision
         )
@@ -392,7 +396,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         # pad the intermediate size to be a multiple of 2 * mxfp4_block
         # for to hold non-uniform sharded tensor as well as swizzling
         intermediate_size_per_partition_after_pad = intermediate_size_per_partition
-        if self.use_marlin:
+        if self.use_marlin and not self.use_mega_moe:
             intermediate_size_per_partition_after_pad = round_up(
                 intermediate_size_per_partition, 128
             )
@@ -402,7 +406,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 intermediate_size_per_partition_after_pad
                 - layer.intermediate_size_per_partition
             )
-        elif self.use_deep_gemm:
+        elif self.use_deep_gemm or self.use_mega_moe:
             # DeepGEMM fp8_fp4 grouped GEMM consumes the checkpoint layout
             # directly (packed e2m1 K-major + ue8m0 g32 scales); no padding.
             pass
@@ -544,7 +548,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             set_weight_attrs(w2_weight_bias, extra_weight_attrs)
 
     def process_weights_after_loading(self, layer):
-        if self.use_marlin:
+        if self.use_marlin and not self.use_mega_moe:
             from sglang.srt.layers.quantization.marlin_utils import (
                 check_moe_marlin_supports_layer,
             )
@@ -572,7 +576,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             layer._mxfp4_backend = "marlin"
             return
 
-        if self.use_deep_gemm:
+        if self.use_deep_gemm or self.use_mega_moe:
             from deep_gemm import transform_sf_into_required_layout
 
             # Packed fp4 (e2m1 x2 per byte) weights: DeepGEMM expects int8.
@@ -597,7 +601,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                     num_groups=num_experts,
                     disable_ue8m0_cast=False,
                 )
-            if get_moe_a2a_backend().is_megamoe():
+            if self.use_mega_moe:
                 # MegaMoE consumes the same transformed sf, plus its own
                 # interleaved/UTCCP weight layout. K3 routes EVERY batch
                 # through mega (the megamoe backend has no a2a fallback), so

python/sglang/srt/models/kimi_k3.py

@@ -141,14 +141,6 @@ def _cdiv(a: int, b: int) -> int:
     return (a + b - 1) // b
 
 
-# MegaMoE SiTU sentinel: DeepGEMM 0.1.5.post1+ selects the K3 SiTU
-# activation when activation_clamp == 0.03125 (2^-5: exactly representable and
-# unused by any legitimate swiglu clamp; the host asserts clamp >= 0 so a
-# negative sentinel is impossible). beta=4.0 / linear_beta=25.0 are baked into
-# the DeepGEMM kernel.
-_K3_MEGA_SITU_SENTINEL_CLAMP = 0.03125
-
-
 def _k3_bf16_gemm(
     x: torch.Tensor,
     weight: torch.Tensor,
@@ -487,9 +479,8 @@ class KimiK3MoE(nn.Module):
         # through it when enabled — the megamoe backend's non-mega fallback is
         # a StandardDispatcher without a2a, which is wrong for scattered
         # tokens — so SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK must
-        # cover the per-rank prefill chunk. SiTU is selected inside the
-        # DeepGEMM mega kernel via a sentinel activation_clamp with the K3
-        # constants baked in.
+        # cover the per-rank prefill chunk. The DeepGEMM MegaMoE SiTU kernel
+        # bakes in the K3 activation constants.
         self._use_mega_moe = get_moe_a2a_backend().is_megamoe()
         self._mega_intermediate_size = moe_intermediate_size
         self._mega_top_k = config.num_experts_per_token
@@ -499,7 +490,7 @@ class KimiK3MoE(nn.Module):
                 config.activation_situ_beta,
                 config.activation_situ_linear_beta,
             ) == (4.0, 25.0), (
-                "mega SiTU kernel patch bakes beta=4.0/linear_beta=25.0; "
+                "MegaMoE SiTU kernel bakes beta=4.0/linear_beta=25.0; "
                 "got a checkpoint with different constants"
             )
 
@@ -802,10 +793,7 @@ class KimiK3MoE(nn.Module):
             self.experts.mega_l2_weights,
             buf,
             recipe=(1, 1, 32),
-            activation="swiglu",
-            # Sentinel: selects the K3 SiTU branch in the DeepGEMM mega kernel
-            # (beta=4.0 / linear_beta=25.0 baked in).
-            activation_clamp=_K3_MEGA_SITU_SENTINEL_CLAMP,
+            activation="situ",
             fast_math=True,
         )
         y = y[:num_tokens]