[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
+        # cover the per-rank prefill chunk. The DeepGEMM MegaMoE SiTU kernel
-        # DeepGEMM mega kernel via a sentinel activation_clamp with the K3
+        # bakes in the K3 activation constants.
-        # constants baked in.
         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",
+            activation="situ",
-            # 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,
             fast_math=True,
         )
         y = y[:num_tokens]

test/registered/models_e2e/test_kimi_k3_b300.py

@@ -1,6 +1,6 @@
 """B300 per-commit CI coverage for Kimi-K3 serving recipes.
 
-Runs the Low Latency DSPARK recipe and the Balanced DCP/HiCache recipe on
+Runs the Low Latency DSPARK, Balanced DCP/HiCache, and MegaMoE recipes on
 eight B300 GPUs. Each server must preserve basic model quality on GSM8K, and
 the Low Latency recipe must also preserve single-request decode performance.
 """
@@ -18,13 +18,17 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=900, stage="base-c", runner_config="8-gpu-b300")
+register_cuda_ci(est_time=1200, stage="base-c", runner_config="8-gpu-b300")
 
 MODEL_PATH = (
     "/data/radixark/model-cache/hub/models--moonshotai--Kimi-K3/"
     "snapshots/9f62e4e9fffbd0a83ddd60e1c209d828994b3569"
 )
 DSPARK_DRAFT_MODEL = "RadixArk/Kimi-K3-DSpark"
+MEGAMOE_URL = "http://0.0.0.0:30000"
+MEGAMOE_ENV = {
+    "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "8320",
+}
 SERVER_LAUNCH_TIMEOUT = 3600
 GPU_IDLE_TIMEOUT = 120
 
@@ -124,5 +128,53 @@ class TestKimiK3B300Balanced(GSM8KMixin, CustomTestCase):
         _stop_server(getattr(cls, "process", None))
 
 
+class TestKimiK3B300MegaMoE(GSM8KMixin, CustomTestCase):
+    """TP8/EP8/DCP8 MegaMoE recipe with DSPARK speculation."""
+
+    gsm8k_score_threshold = 0.95
+    gsm8k_num_examples = 200
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = MODEL_PATH
+        cls.base_url = MEGAMOE_URL
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=[
+                "--trust-remote-code",
+                "--tp-size",
+                "8",
+                "--moe-a2a-backend",
+                "megamoe",
+                "--ep",
+                "8",
+                "--dcp-size",
+                "8",
+                "--mem-fraction-static",
+                "0.85",
+                "--reasoning-parser",
+                "kimi_k3",
+                "--tool-call-parser",
+                "kimi_k3",
+                "--mamba-full-memory-ratio",
+                "5.13",
+                "--speculative-algorithm",
+                "DSPARK",
+                "--speculative-draft-model-path",
+                DSPARK_DRAFT_MODEL,
+                "--speculative-dspark-block-size",
+                "7",
+                "--enable-linear-replayssm-spec",
+            ],
+            env=MEGAMOE_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        _stop_server(getattr(cls, "process", None))
+
+
 if __name__ == "__main__":
     unittest.main()