diff --git a/python/pyproject.toml b/python/pyproject.toml index 589bfea4e..bf1f6602f 100755 --- a/python/pyproject.toml +++ b/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", diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index b3136cff4..8da877441 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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 diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 1cc190eb1..f85b1a1b4 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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] diff --git a/test/registered/models_e2e/test_kimi_k3_b300.py b/test/registered/models_e2e/test_kimi_k3_b300.py index 33fec0b49..7275546cc 100644 --- a/test/registered/models_e2e/test_kimi_k3_b300.py +++ b/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()