From 9cc7da2ab0b305f7aecd45f7b0854c7e20535640 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:00:38 -0700 Subject: [PATCH] [MegaMoE] Wire Qwen MoE blocks to DeepGEMM MegaMoE (MXFP4 and NVFP4 experts) (#38080) --- .../sglang/srt/arg_groups/deepseek_v4_hook.py | 85 --------- python/sglang/srt/arg_groups/mega_moe_hook.py | 172 ++++++++++++++++++ python/sglang/srt/arg_groups/model_hook.py | 2 - python/sglang/srt/arg_groups/pipeline.py | 6 + .../sglang/srt/arg_groups/resolution_hooks.py | 1 + python/sglang/srt/configs/model_config.py | 1 + .../srt/layers/moe/fused_moe_triton/layer.py | 3 + python/sglang/srt/layers/moe/mega_moe.py | 128 ++++++++++--- python/sglang/srt/layers/moe/mega_moe_sm90.py | 23 +-- .../srt/layers/quantization/modelopt_quant.py | 68 +++++++ .../sglang/srt/layers/quantization/mxfp4.py | 10 +- python/sglang/srt/models/qwen2_moe.py | 71 ++++++++ python/sglang/srt/models/qwen3_moe.py | 61 ++++++- .../layers/moe/test_mega_moe_deepgemm_api.py | 94 +++++++++- .../unit/server_args/test_server_args.py | 88 +++++++++ 15 files changed, 686 insertions(+), 127 deletions(-) diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 22e098c26..6f901263e 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -20,91 +20,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def validate_deepseek_v4_mega_moe_token_budget( - server_args: ServerArgs, -) -> None: - """Ensure the DSV4 prefill budget fits MegaMoE's per-rank buffer.""" - cfg = resolving_view(server_args) - mega_moe_enabled = cfg.moe_a2a_backend == "megamoe" - if not mega_moe_enabled or cfg.disaggregation_mode == "decode": - # decode node will skip the check because decode bs is not relevant with --chunk-prefill-size - return - - if cfg.pp_size > 1 and cfg.enable_dynamic_chunking: - return - - if cfg.chunked_prefill_size is None or cfg.chunked_prefill_size <= 0: - raise ValueError( - "DeepSeekV4 with MegaMoE requires chunked prefill to be enabled. " - "Set --chunked-prefill-size to a positive value; " - "--chunked-prefill-size=-1 is unsafe because MegaMoE's per-rank " - "token requirement would not have a strict prefill-forward bound." - ) - - if cfg.enable_prefill_cp: - token_partition_size = cfg.attn_cp_size - token_partition_name = "attn_cp_size" - token_alignment = 1 - local_chunked_prefill_size = ( - cfg.chunked_prefill_size + token_partition_size - 1 - ) // token_partition_size - elif cfg.enable_dp_attention: - token_partition_size = cfg.dp_size - token_partition_name = "dp_size" - token_alignment = max( - cfg.tp_size // cfg.dp_size // cfg.attn_cp_size, - 1, - ) - local_chunked_prefill_size = cfg.chunked_prefill_size // token_partition_size - else: - # Pure TP and PP with static chunking are handled here. - token_partition_size = 1 - token_partition_name = "none" - # global_num_tokens will ceil_align to attn_tp_size so the validation needs to do alignment as well - token_alignment = max( - cfg.tp_size // token_partition_size // cfg.attn_cp_size, - 1, - ) - local_chunked_prefill_size = cfg.chunked_prefill_size - - if local_chunked_prefill_size <= 0: - raise ValueError( - "DeepSeekV4 with MegaMoE requires a positive effective per-rank " - "chunked prefill size. " - f"Current values: chunked_prefill_size=" - f"{cfg.chunked_prefill_size}, " - f"token_partition={token_partition_name}, " - f"token_partition_size={token_partition_size}." - ) - - required_tokens_per_rank = ( - (local_chunked_prefill_size + token_alignment - 1) - // token_alignment - * token_alignment - ) - max_tokens_per_rank = ( - envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() - ) - if max_tokens_per_rank < required_tokens_per_rank: - raise ValueError( - "DeepSeekV4 with MegaMoE requires " - "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to " - "cover each rank's effective prefill token budget. " - f"Current values: chunked_prefill_size=" - f"{cfg.chunked_prefill_size}, " - f"token_partition={token_partition_name}, " - f"token_partition_size={token_partition_size}, " - f"token_alignment={token_alignment}, " - f"required_per_rank={required_tokens_per_rank}, " - "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" - f"{max_tokens_per_rank}. Set " - "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to at " - f"least {required_tokens_per_rank}, or lower " - "--chunked-prefill-size until the effective per-rank budget fits. " - "Otherwise MegaMoE falls back to the fused MoE path at runtime." - ) - - def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None: """Residual imperative arm of the DeepSeek V4 defaults. diff --git a/python/sglang/srt/arg_groups/mega_moe_hook.py b/python/sglang/srt/arg_groups/mega_moe_hook.py index 148f8a3ce..d10ec4e52 100644 --- a/python/sglang/srt/arg_groups/mega_moe_hook.py +++ b/python/sglang/srt/arg_groups/mega_moe_hook.py @@ -8,14 +8,186 @@ if TYPE_CHECKING: from sglang.srt.arg_groups.overrides import ( declare_resolution, + model_config_of, resolving_view, ) +from sglang.srt.connector import ConnectorType +from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_platform +from sglang.srt.utils.common import parse_connector_type logger = logging.getLogger(__name__) def handle_mega_moe(server_args: ServerArgs) -> None: handle_moe_runner_backend_alias(server_args) + check_mega_moe_compat(server_args) + + +def check_mega_moe_compat(server_args: ServerArgs) -> None: + cfg = resolving_view(server_args) + if cfg.moe_a2a_backend != "megamoe": + return + if cfg.enable_two_batch_overlap or cfg.enable_single_batch_overlap: + # The fused kernel has no dispatch_a / combine_a split for TBO / SBO. + raise ValueError( + "--moe-a2a-backend megamoe has no two-batch / single-batch overlap " + "decomposition; disable --enable-two-batch-overlap and " + "--enable-single-batch-overlap." + ) + platform = get_platform() + if not (platform.is_cuda and (platform.is_sm90 or platform.is_sm100)): + raise ValueError( + "--moe-a2a-backend megamoe needs a CUDA SM90 GPU (block-FP8 experts) " + "or an SM100-class GPU (MXFP4 / NVFP4 experts); it runs DeepGEMM " + "kernels over CUDA symmetric memory." + ) + + +# MoE blocks with no fused-MoE fallback under megamoe. +MEGA_MOE_NO_FALLBACK_ARCHS = frozenset( + { + "DeepseekV4ForCausalLM", + "InternS2PreviewForConditionalGeneration", + "MellumForCausalLM", + "Qwen2MoeForCausalLM", + "Qwen3MoeForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + } +) + + +def mega_moe_needs_token_budget( + model_arch: str, hf_quant_config: dict, quantization: str | None +) -> bool: + # NVFP4 experts are repacked into the mega layout at load: no fallback. + nvfp4_experts = quantization == "modelopt_fp4" or "FP4" in str( + hf_quant_config.get("quant_algo", "") + ) + return model_arch in MEGA_MOE_NO_FALLBACK_ARCHS or nvfp4_experts + + +def validate_mega_moe_token_budget_for_model(server_args: ServerArgs) -> None: + # Runs after speculative resolution so the decode / verify bound is final. + cfg = resolving_view(server_args) + if cfg.moe_a2a_backend != "megamoe": + return + if parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE: + return + model_config = model_config_of(server_args) + model_arch = model_config.hf_config.architectures[0] + if mega_moe_needs_token_budget( + model_arch, model_config.hf_quant_config, cfg.quantization + ): + validate_mega_moe_token_budget(server_args, model_arch) + + +def mega_moe_decode_tokens_per_rank(cfg) -> int: + cg_config = cfg.cuda_graph_config + decode_max_bs = (cg_config.decode.max_bs if cg_config is not None else 0) or 0 + num_tokens_per_req = ( + (cfg.speculative_num_draft_tokens or 1) if cfg.speculative_algorithm else 1 + ) + return decode_max_bs * num_tokens_per_req + + +def validate_mega_moe_token_budget(server_args: ServerArgs, model_label: str) -> None: + cfg = resolving_view(server_args) + if cfg.moe_a2a_backend != "megamoe": + return + + max_tokens_per_rank = ( + envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() + ) + if cfg.disaggregation_mode != "prefill": + decode_tokens = mega_moe_decode_tokens_per_rank(cfg) + if max_tokens_per_rank < decode_tokens: + raise ValueError( + f"{model_label} with MegaMoE requires " + "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to cover the " + "largest decode / verify forward on one rank. Current values: " + f"decode cuda graph max bs x tokens per request = {decode_tokens}, " + f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" + f"{max_tokens_per_rank}. Raise the env var to at least " + f"{decode_tokens} or lower --cuda-graph-max-bs / " + "--speculative-num-draft-tokens." + ) + if cfg.disaggregation_mode == "decode": + # decode node will skip the check because decode bs is not relevant with --chunk-prefill-size + return + + if cfg.pp_size > 1 and cfg.enable_dynamic_chunking: + return + + if cfg.chunked_prefill_size is None or cfg.chunked_prefill_size <= 0: + raise ValueError( + f"{model_label} with MegaMoE requires chunked prefill to be enabled. " + "Set --chunked-prefill-size to a positive value; " + "--chunked-prefill-size=-1 is unsafe because MegaMoE's per-rank " + "token requirement would not have a strict prefill-forward bound." + ) + + if cfg.enable_prefill_cp: + token_partition_size = cfg.attn_cp_size + token_partition_name = "attn_cp_size" + token_alignment = 1 + local_chunked_prefill_size = ( + cfg.chunked_prefill_size + token_partition_size - 1 + ) // token_partition_size + elif cfg.enable_dp_attention: + token_partition_size = cfg.dp_size + token_partition_name = "dp_size" + token_alignment = max( + cfg.tp_size // cfg.dp_size // cfg.attn_cp_size, + 1, + ) + local_chunked_prefill_size = cfg.chunked_prefill_size // token_partition_size + else: + # Pure TP and PP with static chunking are handled here. + token_partition_size = 1 + token_partition_name = "none" + # global_num_tokens will ceil_align to attn_tp_size so the validation needs to do alignment as well + token_alignment = max( + cfg.tp_size // token_partition_size // cfg.attn_cp_size, + 1, + ) + local_chunked_prefill_size = cfg.chunked_prefill_size + + if local_chunked_prefill_size <= 0: + raise ValueError( + f"{model_label} with MegaMoE requires a positive effective per-rank " + "chunked prefill size. " + f"Current values: chunked_prefill_size=" + f"{cfg.chunked_prefill_size}, " + f"token_partition={token_partition_name}, " + f"token_partition_size={token_partition_size}." + ) + + required_tokens_per_rank = ( + (local_chunked_prefill_size + token_alignment - 1) + // token_alignment + * token_alignment + ) + if max_tokens_per_rank < required_tokens_per_rank: + raise ValueError( + f"{model_label} with MegaMoE requires " + "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to " + "cover each rank's effective prefill token budget. " + f"Current values: chunked_prefill_size=" + f"{cfg.chunked_prefill_size}, " + f"token_partition={token_partition_name}, " + f"token_partition_size={token_partition_size}, " + f"token_alignment={token_alignment}, " + f"required_per_rank={required_tokens_per_rank}, " + "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" + f"{max_tokens_per_rank}. Set " + "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to at " + f"least {required_tokens_per_rank}, or lower " + "--chunked-prefill-size until the effective per-rank budget fits." + ) def handle_moe_runner_backend_alias(server_args: ServerArgs) -> None: diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 5867fb299..c4505d0ad 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -421,14 +421,12 @@ def handle_model_specific_adjustments(server_args: Any): ]: from sglang.srt.arg_groups.deepseek_v4_hook import ( validate_deepseek_v4_cp, - validate_deepseek_v4_mega_moe_token_budget, validate_deepseek_v41_features, ) # Before the CP validation: V4.1 rejects CP outright, the actionable message. validate_deepseek_v41_features(server_args) validate_deepseek_v4_cp(server_args) - validate_deepseek_v4_mega_moe_token_budget(server_args) if get_platform().is_sm120: # FP8 wo_a stays opt-in on SM120: only recent DeepGEMM builds ship diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 496f37c9d..acdb83921 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -337,6 +337,12 @@ def run_resolution_pipeline(server_args: Any) -> None: # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final. run_hook(validate_cutedsl_a2a_token_budget, server_args) + from sglang.srt.arg_groups.mega_moe_hook import ( + validate_mega_moe_token_budget_for_model, + ) + + run_hook(validate_mega_moe_token_budget_for_model, server_args) + # Handle model loading format. run_hook(handle_load_format, server_args) diff --git a/python/sglang/srt/arg_groups/resolution_hooks.py b/python/sglang/srt/arg_groups/resolution_hooks.py index 71c4c81fa..c94cc9873 100644 --- a/python/sglang/srt/arg_groups/resolution_hooks.py +++ b/python/sglang/srt/arg_groups/resolution_hooks.py @@ -107,6 +107,7 @@ _OVERRIDABLE_HOOKS: FrozenSet[str] = frozenset( "handle_speculative_decoding", "handle_layernorm_sp", "validate_cutedsl_a2a_token_budget", + "validate_mega_moe_token_budget_for_model", "handle_load_format", "handle_encoder_disaggregation", "handle_tokenizer_batching", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index baa1ac2bf..96161ab8a 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -563,6 +563,7 @@ class ModelConfig: _quant_config_to_dict(getattr(self.hf_config, "quantization_config", None)) or {} ) + self.hf_quant_config: dict = quantization_config routed_experts_quant_method = quantization_config.get( "routed_experts_quant_method" ) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 188964489..739ad7958 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -398,6 +398,9 @@ class FusedMoE(torch.nn.Module): self._num_local_routed = self._num_global_routed // storage_ep_size self.num_local_experts = self._num_local_routed + num_fused_shared_experts self._has_fused_shared = num_fused_shared_experts > 0 + # Set by the quant method when it repacks experts for MegaMoE. + self._mega_moe_weights_built = False + self._mega_moe_nvfp4 = False self._pending_fp8_shared_weights: dict[tuple[int, str], torch.Tensor] = {} self._pending_fp8_shared_scales: dict[tuple[int, str], torch.Tensor] = {} diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 4adb7826d..783ac819e 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -45,7 +45,9 @@ if TYPE_CHECKING: _MEGA_MOE_SYMM_BUFFER: dict = {} -def _mega_moe_mma_type() -> str: +def _mega_moe_mma_type(experts=None) -> str: + if experts is not None and experts._mega_moe_nvfp4: + return "nvfp4xnvfp4" return "mxf4xmxf4" if get_exec().moe.enable_w4a4_mxfp4_megamoe else "fp8xfp4" @@ -86,6 +88,20 @@ def _configure_mega_moe_deep_gemm_num_sms(deep_gemm): deep_gemm.set_num_sms(current_num_sms) +def check_mega_moe_shapes(hidden: int, intermediate: int, mma_type: str) -> None: + # DeepGEMM keeps one scale row per token and needs 16-byte TMA alignment + # on it (layout/mega_moe.cuh), so both dims must be multiples of 16 * group. + scale_group = 16 if mma_type == "nvfp4xnvfp4" else 32 + align = 16 * scale_group + if hidden % align != 0 or intermediate % align != 0: + raise ValueError( + f"DeepGEMM MegaMoE ({mma_type}) needs hidden_size and " + f"moe_intermediate_size to be multiples of {align}; got " + f"hidden_size={hidden}, moe_intermediate_size={intermediate}. " + "Use another --moe-a2a-backend for this model." + ) + + def _get_mega_moe_symm_buffer( group, num_experts: int, @@ -93,10 +109,12 @@ def _get_mega_moe_symm_buffer( num_topk: int, hidden: int, intermediate_hidden: int, + mma_type: Optional[str] = None, ) -> SymmBuffer: import deep_gemm - mma_type = _mega_moe_mma_type() + if mma_type is None: + mma_type = _mega_moe_mma_type() with _configure_mega_moe_deep_gemm_num_sms(deep_gemm): key = ( id(group), @@ -124,14 +142,20 @@ def _get_mega_moe_symm_buffer( return buf +def is_mega_moe_experts_ready(experts) -> bool: + if not experts._mega_moe_weights_built: + return False + if _device_sm == 90: + return is_sm90_fp8_mega_moe_available(experts) + # The SM100 mega kernels exist for compute capability 10.x only. + return _device_sm // 10 == 10 + + def should_use_mega_moe(moe: DeepseekV2MoE, hidden_states: torch.Tensor) -> bool: if not get_moe_a2a_backend().is_megamoe(): return False - if not getattr(moe.experts, "_mega_moe_weights_built", False): + if not is_mega_moe_experts_ready(moe.experts): return False - if _device_sm == 90: - if not is_sm90_fp8_mega_moe_available(moe.experts): - return False if get_is_capture_mode(): return True @@ -188,10 +212,6 @@ def _run_mega_routed( input_ids_global: Optional[torch.Tensor], num_tokens: int, ) -> torch.Tensor: - import deep_gemm - - from sglang.srt.runtime_context import get_parallel - hidden_size = moe.config.hidden_size if num_tokens > 0: @@ -216,10 +236,44 @@ def _run_mega_routed( topk_ids = None topk_weights = None + return run_mega_routed_experts( + moe.experts, + hidden_states, + topk_ids, + topk_weights, + hidden_size=hidden_size, + intermediate_size=moe.config.moe_intermediate_size, + top_k=moe.config.num_experts_per_tok + moe.num_fused_shared_experts, + num_tokens=num_tokens, + activation_clamp=moe.experts.moe_runner_config.swiglu_limit, + routed_scaling_factor=( + 1.0 + if moe.experts.should_fuse_routed_scaling_factor_in_topk + else float(moe.routed_scaling_factor) + ), + ) + + +def run_mega_routed_experts( + experts, + hidden_states: torch.Tensor, + topk_ids: Optional[torch.Tensor], + topk_weights: Optional[torch.Tensor], + *, + hidden_size: int, + intermediate_size: int, + top_k: int, + num_tokens: int, + activation_clamp: Optional[float] = None, + routed_scaling_factor: float = 1.0, +) -> torch.Tensor: + # Rows are this rank's tokens; the returned rows are fully combined. + import deep_gemm + + from sglang.srt.runtime_context import get_parallel + ep_group = get_parallel().moe_ep_group.device_group - num_experts = moe.experts.num_experts - top_k = moe.config.num_experts_per_tok + moe.num_fused_shared_experts - intermediate_size = moe.config.moe_intermediate_size + num_experts = experts.num_experts num_max_tokens_per_rank = ( envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() ) @@ -230,6 +284,7 @@ def _run_mega_routed( f"--cuda-graph-max-bs-decode / --chunked-prefill-size accordingly" ) + mma_type = _mega_moe_mma_type(experts) buf = _get_mega_moe_symm_buffer( ep_group, num_experts=num_experts, @@ -237,6 +292,7 @@ def _run_mega_routed( num_topk=top_k, hidden=hidden_size, intermediate_hidden=intermediate_size, + mma_type=mma_type, ) if num_tokens > 0: @@ -248,16 +304,42 @@ def _run_mega_routed( if _device_sm == 90: return run_sm90_mega_routed( - moe, + experts, hidden_states, topk_ids_in, topk_weights_in, buf, num_tokens, + hidden_size=hidden_size, + activation_clamp=activation_clamp, + routed_scaling_factor=routed_scaling_factor, ) - mma_type = _mega_moe_mma_type() - if mma_type == "mxf4xmxf4": + mega_kwargs = {"recipe": (1, 1, 32)} + if mma_type == "nvfp4xnvfp4": + # Per-token outer scales go to buf.x_scales; the kernel folds them with + # the per-expert alphas in the L1 / L2 epilogues. + deep_gemm.mega_moe_pre_dispatch( + hidden_states, + topk_ids_in, + topk_weights_in, + buf.x, + buf.x_sf, + buf.topk_idx, + buf.topk_weights, + num_tokens=num_tokens, + group_size=16, + mma_type=mma_type, + buf_x_scales=buf.x_scales, + ) + mega_kwargs = { + "recipe": (1, 1, 16), + "use_x_scales": True, + "l1_alphas": experts.mega_l1_alphas, + "l2_alphas": experts.mega_l2_alphas, + "l2_act_scales": experts.mega_l2_act_scales, + } + elif mma_type == "mxf4xmxf4": # FP4 path goes through DeepGEMM's mega_moe_pre_dispatch which # handles the E2M1 packing variant. The jit implementation # only emits FP8. @@ -292,22 +374,21 @@ def _run_mega_routed( dtype=torch.bfloat16, device=hidden_states.device, ) - swiglu_limit = getattr(moe.config, "swiglu_limit", None) with _configure_mega_moe_deep_gemm_num_sms(deep_gemm): deep_gemm.fp8_fp4_mega_moe( y, - moe.experts.mega_l1_weights, - moe.experts.mega_l2_weights, + experts.mega_l1_weights, + experts.mega_l2_weights, buf, - recipe=(1, 1, 32), activation="swiglu", - activation_clamp=swiglu_limit, + activation_clamp=activation_clamp, fast_math=True, + **mega_kwargs, ) y = y[:num_tokens] - if not moe.experts.should_fuse_routed_scaling_factor_in_topk: - y.mul_(moe.routed_scaling_factor) + if routed_scaling_factor != 1.0: + y.mul_(routed_scaling_factor) return y @@ -366,6 +447,7 @@ def build_mega_moe_experts_weights(experts) -> None: num_groups, n1, half_k1 = w13.shape k1 = half_k1 * 2 + check_mega_moe_shapes(hidden=k1, intermediate=n1 // 2, mma_type=mma_type) _, n2, half_k2 = w2.shape k2 = half_k2 * 2 diff --git a/python/sglang/srt/layers/moe/mega_moe_sm90.py b/python/sglang/srt/layers/moe/mega_moe_sm90.py index 2da005dff..02417bd97 100644 --- a/python/sglang/srt/layers/moe/mega_moe_sm90.py +++ b/python/sglang/srt/layers/moe/mega_moe_sm90.py @@ -15,7 +15,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch @@ -24,8 +24,6 @@ from sglang.srt.models.deepseek_common.utils import _device_sm if TYPE_CHECKING: from deep_gemm import SymmBuffer - from sglang.srt.models.deepseek_v2 import DeepseekV2MoE - def is_sm90_fp8_mega_moe_available(experts) -> bool: if _device_sm != 90: @@ -42,20 +40,19 @@ def is_sm90_fp8_mega_moe_available(experts) -> bool: def run_sm90_mega_routed( - moe: DeepseekV2MoE, + experts, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, buf: SymmBuffer, num_tokens: int, + *, + hidden_size: int, + activation_clamp: Optional[float] = None, + routed_scaling_factor: float = 1.0, ) -> torch.Tensor: import deep_gemm - if moe.experts.should_fuse_routed_scaling_factor_in_topk: - routed_scaling_factor = 1.0 - else: - routed_scaling_factor = float(moe.routed_scaling_factor) - deep_gemm.mega_moe_pre_dispatch_sm90( hidden_states, topk_ids, @@ -70,18 +67,18 @@ def run_sm90_mega_routed( ) y = torch.empty( - (max(num_tokens, 1), moe.config.hidden_size), + (max(num_tokens, 1), hidden_size), dtype=torch.bfloat16, device=hidden_states.device, ) deep_gemm.fp8_mega_moe( y, - moe.experts.mega_l1_weights, - moe.experts.mega_l2_weights, + experts.mega_l1_weights, + experts.mega_l2_weights, buf, recipe=(128, 128, 128), activation="swiglu", - activation_clamp=getattr(moe.config, "swiglu_limit", None), + activation_clamp=activation_clamp, fast_math=True, ) y = y[:num_tokens] diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 1add64f5c..24cc6adec 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2539,6 +2539,59 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): w2_input_scale._sglang_require_global_experts = True layer.register_parameter("w2_input_scale", w2_input_scale) + def _build_mega_moe_weights(self, layer: torch.nn.Module) -> None: + # Activations are quantized per token at dispatch, so w13_input_scale + # is not used. + import deep_gemm + from deep_gemm.utils.math import transform_ue4m3_sf_into_required_layout + + from sglang.srt.layers.moe.mega_moe import check_mega_moe_shapes + + assert layer.moe_runner_config.is_gated, "MegaMoE NVFP4 needs a gated MLP" + n1 = layer.w13_weight.shape[1] + n2 = layer.w2_weight.shape[1] + check_mega_moe_shapes( + hidden=layer.w13_weight.shape[2] * 2, + intermediate=n1 // 2, + mma_type="nvfp4xnvfp4", + ) + w13_sf = transform_ue4m3_sf_into_required_layout( + layer.w13_weight_scale.data.view(torch.float8_e4m3fn), n1 + ) + w2_sf = transform_ue4m3_sf_into_required_layout( + layer.w2_weight_scale.data.view(torch.float8_e4m3fn), n2 + ) + l1_pair, l2_pair = deep_gemm.transform_weights_for_mega_moe( + (layer.w13_weight.data.view(torch.int8), w13_sf), + (layer.w2_weight.data.view(torch.int8), w2_sf), + mma_type="nvfp4xnvfp4", + ) + + num_local = layer.num_local_experts + ones = torch.ones(num_local, dtype=torch.float32, device=l1_pair[0].device) + g1_gate, g1_up = _compute_gemm1_alphas(layer.w13_weight_scale_2, ones, True) + l1_alphas = torch.stack([g1_gate, g1_up], dim=1).contiguous() + w2_input_scale = _input_scale_to_local_experts( + layer.w2_input_scale, num_local, layer.num_experts, layer.moe_ep_rank + ) + l2_act_scales = (1.0 / w2_input_scale).contiguous() + l2_alphas = ( + w2_input_scale * layer.w2_weight_scale_2.to(torch.float32) + ).contiguous() + + layer.mega_l1_weights = l1_pair + layer.mega_l2_weights = l2_pair + layer.mega_l1_alphas = l1_alphas + layer.mega_l2_alphas = l2_alphas + layer.mega_l2_act_scales = l2_act_scales + # Free the checkpoint layout. + layer.w13_weight.data = l1_pair[0] + layer.w13_weight_scale.data = l1_pair[1] + layer.w2_weight.data = l2_pair[0] + layer.w2_weight_scale.data = l2_pair[1] + layer._mega_moe_nvfp4 = True + layer._mega_moe_weights_built = True + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """Transform packed FP4 MoE weights and scales for the selected backend.""" if getattr(layer, "inference_moe_w13_interleaved", False) and not getattr( @@ -2553,6 +2606,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) layer._w13_deinterleaved = True + if get_moe_a2a_backend().is_megamoe(): + self._build_mega_moe_weights(layer) + return + # GEMM1 scale processing is deferred until the input scale is known; # see _compute_gemm1_alphas, which splits w13's gate/up weight scales. moe_runner_backend = getattr( @@ -2922,6 +2979,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): self._moe_runner_backend = moe_runner_backend + if get_moe_a2a_backend().is_megamoe(): + # FusedMoE.forward is never reached under megamoe. + self.runner = None + return + if moe_runner_backend.is_flashinfer_cutedsl(): import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func @@ -2972,6 +3034,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): layer: FusedMoE, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + if layer._mega_moe_nvfp4: + raise RuntimeError( + "NVFP4 MegaMoE experts cannot fall back to the fused MoE path; " + "the model block must route every forward through " + "run_mega_routed_experts (check the MegaMoE token budget)." + ) # Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output # attribute -- topk_ids/topk_weights live directly on the dispatch # tuple). Defer per-attribute access to the branches that actually diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 98787d599..12cbc62c0 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -677,9 +677,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # (keeping both layouts OOMs: 92 layers double the experts). from deep_gemm import transform_weights_for_mega_moe - from sglang.srt.layers.moe.mega_moe import _mega_moe_mma_type + from sglang.srt.layers.moe.mega_moe import ( + _mega_moe_mma_type, + check_mega_moe_shapes, + ) mma_type = _mega_moe_mma_type() + check_mega_moe_shapes( + hidden=layer.w13_weight.shape[2] * 2, + intermediate=layer.w13_weight.shape[1] // 2, + mma_type=mma_type, + ) l1_pair, l2_pair = transform_weights_for_mega_moe( (layer.w13_weight.data, layer.w13_weight_scale.data), (layer.w2_weight.data, layer.w2_weight_scale.data), diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8c99eba64..7ebb60e9a 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -377,6 +377,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): or get_moe_a2a_backend().is_deepep_v2() or get_moe_a2a_backend().is_flashinfer() or get_moe_a2a_backend().is_flashinfer_megamoe() + or get_moe_a2a_backend().is_megamoe() ) else {} ), @@ -406,6 +407,10 @@ class Qwen2MoeSparseMoeBlock(nn.Module): ) self.top_k = config.num_experts_per_tok self.is_nextn = is_nextn + self._use_mega_moe = get_moe_a2a_backend().is_megamoe() + self._mega_top_k = config.num_experts_per_tok + self.num_fused_shared_experts + self._mega_intermediate_size = config.moe_intermediate_size + self._mega_hidden_size = config.hidden_size def get_moe_weights(self): return [ @@ -625,6 +630,69 @@ class Qwen2MoeSparseMoeBlock(nn.Module): return final_hidden_states + def _forward_mega_moe( + self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] + ) -> torch.Tensor: + # Same contract as _forward_deepep: combined rows, no TP all-reduce. + from sglang.srt.layers.moe.mega_moe import ( + is_mega_moe_experts_ready, + run_mega_routed_experts, + ) + + if not is_mega_moe_experts_ready(self.experts): + raise RuntimeError( + "moe_a2a_backend=megamoe needs MegaMoE expert weights on this " + "model: on SM100 load a checkpoint with MXFP4 or NVFP4 routed " + "experts; on SM90 load a block-FP8 checkpoint with a DeepGEMM " + "that ships fp8_mega_moe." + ) + + num_tokens = hidden_states.shape[0] + shared_output = None + topk_ids = None + topk_weights = None + if num_tokens > 0: + router_logits, _ = self.gate(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=( + forward_batch.num_token_non_padded + if forward_batch is not None + else None + ), + expert_location_dispatch_info=( + ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) + if not self.is_nextn + else None + ), + ) + if self.enable_shared_expert_fusion: + topk_output = self._append_shared_to_topk_output( + topk_output, hidden_states + ) + assert TopKOutputChecker.format_is_standard(topk_output), ( + "MegaMoE pre-dispatch consumes raw topk ids/weights; " + "pick a MoE runner backend that emits standard TopK output" + ) + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + + final_hidden_states = run_mega_routed_experts( + self.experts, + hidden_states, + topk_ids, + topk_weights, + hidden_size=self._mega_hidden_size, + intermediate_size=self._mega_intermediate_size, + top_k=self._mega_top_k, + num_tokens=num_tokens, + ) + if shared_output is not None: + final_hidden_states.add_(shared_output) + return final_hidden_states + @property def supports_deferred_finalize(self) -> bool: return bool( @@ -743,6 +811,9 @@ class Qwen2MoeSparseMoeBlock(nn.Module): if defer_finalize and num_tokens == 0: raise RuntimeError("Qwen deferred finalize does not support M=0") + if self._use_mega_moe: + return self._forward_mega_moe(hidden_states, forward_batch) + if ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_deepep_v2() diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 8eed5d729..f4825ac4f 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -43,7 +43,7 @@ from sglang.srt.layers.moe import ( ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE -from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.topk import TopK, TopKOutputChecker from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, @@ -291,12 +291,19 @@ class Qwen3MoeSparseMoeBlock(nn.Module): ) self.top_k = config.num_experts_per_tok + self._use_mega_moe = get_moe_a2a_backend().is_megamoe() + self._mega_top_k = config.num_experts_per_tok + self._mega_intermediate_size = config.moe_intermediate_size + self._mega_hidden_size = config.hidden_size + def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, ) -> torch.Tensor: + if self._use_mega_moe: + return self._forward_mega_moe(hidden_states, forward_batch) if ( not is_deepep_class_backend() and not get_moe_a2a_backend().is_ascend_fuseep() @@ -356,6 +363,58 @@ class Qwen3MoeSparseMoeBlock(nn.Module): ) return final_hidden_states + def _forward_mega_moe( + self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] + ) -> torch.Tensor: + # Same contract as forward_deepep: combined rows, no TP all-reduce. + from sglang.srt.layers.moe.mega_moe import ( + is_mega_moe_experts_ready, + run_mega_routed_experts, + ) + + if not is_mega_moe_experts_ready(self.experts): + raise RuntimeError( + "moe_a2a_backend=megamoe needs MegaMoE expert weights on this " + "model: on SM100 load a checkpoint with MXFP4 or NVFP4 routed " + "experts; on SM90 load a block-FP8 checkpoint with a DeepGEMM " + "that ships fp8_mega_moe." + ) + + num_tokens = hidden_states.shape[0] + topk_ids = None + topk_weights = None + if num_tokens > 0: + router_logits, _ = self.gate(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=( + forward_batch.num_token_non_padded + if forward_batch is not None + else None + ), + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + assert TopKOutputChecker.format_is_standard(topk_output), ( + "MegaMoE pre-dispatch consumes raw topk ids/weights; " + "pick a MoE runner backend that emits standard TopK output" + ) + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + + return run_mega_routed_experts( + self.experts, + hidden_states, + topk_ids, + topk_weights, + hidden_size=self._mega_hidden_size, + intermediate_size=self._mega_intermediate_size, + top_k=self._mega_top_k, + num_tokens=num_tokens, + ) + def op_gate(self, state): if is_non_idle_and_non_empty( state.forward_batch.forward_mode, state.hidden_states_mlp_input diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py index e345d00f1..e7ca95663 100644 --- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py +++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py @@ -203,6 +203,8 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): with ( patch.dict(sys.modules, {"deep_gemm": deep_gemm}), patch.object(mega_moe, "_mega_moe_mma_type", return_value="mxf4xmxf4"), + # Toy shapes. + patch.object(mega_moe, "check_mega_moe_shapes"), ): method.process_weights_after_loading(layer) @@ -224,6 +226,9 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): mega_l1_weights=object(), mega_l2_weights=object(), should_fuse_routed_scaling_factor_in_topk=True, + moe_runner_config=SimpleNamespace(swiglu_limit=None), + _mega_moe_weights_built=True, + _mega_moe_nvfp4=False, ) topk_output = SimpleNamespace( topk_ids=torch.tensor([[0, 1]]), @@ -259,14 +264,18 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): "_configure_mega_moe_deep_gemm_num_sms", return_value=nullcontext(), ), + # Toy shapes. + patch.object(mega_moe, "check_mega_moe_shapes"), patch.object( mega_moe.ExpertLocationDispatchInfo, "init_new", return_value=object(), ), patch( - "sglang.srt.distributed.parallel_state.get_moe_ep_group", - return_value=SimpleNamespace(device_group=object()), + "sglang.srt.runtime_context.get_parallel", + return_value=SimpleNamespace( + moe_ep_group=SimpleNamespace(device_group=object()) + ), ), ): mega_moe._run_mega_routed( @@ -282,6 +291,87 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): self.assertEqual(call.kwargs.get("mma_type"), "mxf4xmxf4") self.assertNotIn("use_fp4_acts", call.kwargs) + def test_run_mega_routed_experts_generic_entry(self): + deep_gemm = ModuleType("deep_gemm") + deep_gemm.fp8_fp4_mega_moe = MagicMock() + buffer = SimpleNamespace( + x=object(), + x_sf=object(), + topk_idx=object(), + topk_weights=object(), + ) + experts = SimpleNamespace( + num_experts=8, + mega_l1_weights=object(), + mega_l2_weights=object(), + _mega_moe_weights_built=True, + _mega_moe_nvfp4=False, + ) + hidden_states = torch.zeros((3, 4), dtype=torch.bfloat16) + topk_ids = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int64) + topk_weights = torch.full((3, 2), 0.5, dtype=torch.bfloat16) + + with ( + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + patch.object(mega_moe, "_device_sm", 100), + patch.object(mega_moe, "_mega_moe_mma_type", return_value="fp8xfp4"), + patch.object(mega_moe, "mega_moe_pre_dispatch") as pre_dispatch, + patch.object( + mega_moe, "_get_mega_moe_symm_buffer", return_value=buffer + ) as get_buffer, + patch.object( + mega_moe, + "_configure_mega_moe_deep_gemm_num_sms", + return_value=nullcontext(), + ), + # Toy shapes. + patch.object(mega_moe, "check_mega_moe_shapes"), + patch( + "sglang.srt.runtime_context.get_parallel", + return_value=SimpleNamespace( + moe_ep_group=SimpleNamespace(device_group=object()) + ), + ), + ): + out = mega_moe.run_mega_routed_experts( + experts, + hidden_states, + topk_ids, + topk_weights, + hidden_size=4, + intermediate_size=8, + top_k=2, + num_tokens=3, + activation_clamp=7.0, + routed_scaling_factor=1.0, + ) + + self.assertEqual(out.shape, (3, 4)) + self.assertEqual(out.dtype, torch.bfloat16) + buf_call = get_buffer.call_args + self.assertEqual(buf_call.kwargs.get("num_topk"), 2) + self.assertEqual(buf_call.kwargs.get("hidden"), 4) + self.assertEqual(buf_call.kwargs.get("intermediate_hidden"), 8) + # The kernel wants int32 ids and fp32 weights regardless of the router dtype. + ids_arg, weights_arg = pre_dispatch.call_args.args[1:3] + self.assertEqual(ids_arg.dtype, torch.int32) + self.assertEqual(weights_arg.dtype, torch.float32) + mega_call = deep_gemm.fp8_fp4_mega_moe.call_args + self.assertIs(mega_call.args[1], experts.mega_l1_weights) + self.assertIs(mega_call.args[2], experts.mega_l2_weights) + self.assertEqual(mega_call.kwargs.get("activation_clamp"), 7.0) + + def test_shape_check_rejects_unaligned_intermediate(self): + # Qwen3-30B-A3B: intermediate 768 leaves a 24-byte scale row. + with self.assertRaisesRegex(ValueError, "multiples of 512"): + mega_moe.check_mega_moe_shapes(2048, 768, "fp8xfp4") + mega_moe.check_mega_moe_shapes(4096, 1536, "fp8xfp4") + mega_moe.check_mega_moe_shapes(4096, 1024, "mxf4xmxf4") + # 768 is a multiple of 256, so the NVFP4 (g16) rule accepts it. + mega_moe.check_mega_moe_shapes(2048, 768, "nvfp4xnvfp4") + with self.assertRaisesRegex(ValueError, "multiples of 256"): + mega_moe.check_mega_moe_shapes(2048, 384, "nvfp4xnvfp4") + def test_mxf4_l1_uses_packed_gate_up_interleave(self): source = torch.arange(32).reshape(1, 32) expected = torch.tensor( diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 60da7dcd5..5972c2529 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -189,6 +189,94 @@ class TestPrepareServerArgs(CustomTestCase): self.assertEqual(os.environ["DG_USE_FP4_ACTS"], "0") self.assertEqual(os.environ["DG_USE_MXF4_KIND"], "0") + def test_megamoe_rejects_two_batch_overlap(self): + # The fused kernel has no dispatch/combine split for the TBO ops to call. + with override_platform(is_cuda=True, is_sm90=False, is_sm100=True): + args = ServerArgs( + model_path="dummy", + moe_a2a_backend="megamoe", + enable_two_batch_overlap=True, + ) + with self.assertRaisesRegex(ValueError, "overlap"): + args.resolve_once() + + def test_megamoe_requires_sm90_or_sm100(self): + with override_platform(is_cuda=True, is_sm90=False, is_sm100=False): + args = ServerArgs(model_path="dummy", moe_a2a_backend="megamoe") + with self.assertRaisesRegex(ValueError, "SM90"): + args.resolve_once() + with override_platform(is_cuda=False, is_sm90=False, is_sm100=False): + args = ServerArgs(model_path="dummy", moe_a2a_backend="megamoe") + with self.assertRaisesRegex(ValueError, "CUDA"): + args.resolve_once() + with override_platform(is_cuda=True, is_sm90=False, is_sm100=True): + ServerArgs(model_path="dummy", moe_a2a_backend="megamoe").resolve_once() + + def test_megamoe_token_budget_must_cover_chunked_prefill(self): + from sglang.srt.arg_groups.mega_moe_hook import validate_mega_moe_token_budget + from sglang.srt.environ import envs + + with override_platform(is_cuda=True, is_sm90=False, is_sm100=True): + args = ServerArgs( + model_path="dummy", + moe_a2a_backend="megamoe", + chunked_prefill_size=16384, + ) + args.resolve_once() + with envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override( + 8192 + ): + with self.assertRaisesRegex(ValueError, "required_per_rank=16384"): + validate_mega_moe_token_budget(args, "Qwen3MoeForCausalLM") + with envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override( + 16384 + ): + validate_mega_moe_token_budget(args, "Qwen3MoeForCausalLM") + + def test_megamoe_token_budget_gate_by_arch_and_nvfp4(self): + from sglang.srt.arg_groups.mega_moe_hook import mega_moe_needs_token_budget + + for arch in ( + "InternS2PreviewForConditionalGeneration", + "MellumForCausalLM", + "Qwen3MoeForCausalLM", + "DeepseekV4ForCausalLM", + ): + self.assertTrue(mega_moe_needs_token_budget(arch, {}, None), arch) + # MXFP4 DeepSeek-family models keep their runtime fallback. + self.assertFalse(mega_moe_needs_token_budget("DeepseekV3ForCausalLM", {}, None)) + # NVFP4 experts are repacked at load: every model is checked. + self.assertTrue( + mega_moe_needs_token_budget( + "DeepseekV3ForCausalLM", {"quant_algo": "NVFP4"}, None + ) + ) + self.assertTrue( + mega_moe_needs_token_budget("DeepseekV3ForCausalLM", {}, "modelopt_fp4") + ) + + def test_megamoe_decode_tokens_per_rank_follows_graph_bs_and_draft_tokens(self): + # cuda_graph_config only exists on a GPU host; use a stand-in view. + from sglang.srt.arg_groups.mega_moe_hook import mega_moe_decode_tokens_per_rank + + def view(max_bs, algorithm=None, draft_tokens=None, cg=True): + return SimpleNamespace( + cuda_graph_config=( + SimpleNamespace(decode=SimpleNamespace(max_bs=max_bs)) + if cg + else None + ), + speculative_algorithm=algorithm, + speculative_num_draft_tokens=draft_tokens, + ) + + self.assertEqual(mega_moe_decode_tokens_per_rank(view(16384)), 16384) + self.assertEqual(mega_moe_decode_tokens_per_rank(view(256, "EAGLE", 4)), 1024) + # Draft tokens only count under a speculative algorithm. + self.assertEqual(mega_moe_decode_tokens_per_rank(view(256, None, 4)), 256) + self.assertEqual(mega_moe_decode_tokens_per_rank(view(None)), 0) + self.assertEqual(mega_moe_decode_tokens_per_rank(view(0, cg=False)), 0) + def test_w4a4_mxfp4_megamoe_disabled_preserves_deepgemm_env(self): deepgemm_env = { "DG_USE_FP4_ACTS": "0",