From e226bb711c255ad378ee10ea12772b2ed914de9a Mon Sep 17 00:00:00 2001 From: Chunan Zeng Date: Sun, 9 Aug 2026 23:30:05 -0700 Subject: [PATCH] enable TRT-LLM for MiniMax M3 by preserving SwiGLU params (#33962) --- .../srt/layers/moe/flashinfer_trtllm_moe.py | 18 +++++++++++ .../srt/layers/moe/fused_moe_triton/layer.py | 2 ++ .../sglang/srt/layers/moe/moe_runner/base.py | 1 + .../moe/moe_runner/flashinfer_trtllm.py | 9 ++++++ python/sglang/srt/layers/quantization/fp8.py | 30 +++++++++++++++++++ .../srt/layers/quantization/fp8_utils.py | 6 ++++ .../runner/flashinfer_autotune.py | 3 +- python/sglang/srt/models/minimax_m3.py | 1 + 8 files changed, 68 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py index 7e0ead6eb..c786aeeb3 100644 --- a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py @@ -12,6 +12,9 @@ def _fake_fp8_block_scale_moe_out( hidden_states_scale: torch.Tensor, gemm1_weights: torch.Tensor, gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, gemm2_weights_scale: torch.Tensor, output: torch.Tensor, @@ -45,6 +48,9 @@ def trtllm_fp8_block_scale_moe_out_wrapper( hidden_states_scale: torch.Tensor, gemm1_weights: torch.Tensor, gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, gemm2_weights_scale: torch.Tensor, output: torch.Tensor, @@ -79,6 +85,9 @@ def trtllm_fp8_block_scale_moe_out_wrapper( "hidden_states_scale": hidden_states_scale, "gemm1_weights": gemm1_weights, "gemm1_weights_scale": gemm1_weights_scale, + "gemm1_alpha": gemm1_alpha, + "gemm1_beta": gemm1_beta, + "gemm1_clamp_limit": gemm1_clamp_limit, "gemm2_weights": gemm2_weights, "gemm2_weights_scale": gemm2_weights_scale, "output": output, @@ -116,6 +125,9 @@ def _fake_fp8_block_scale_routed_moe_out( hidden_states_scale: torch.Tensor, gemm1_weights: torch.Tensor, gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, gemm2_weights_scale: torch.Tensor, num_experts: int, @@ -149,6 +161,9 @@ def trtllm_fp8_block_scale_routed_moe_out_wrapper( hidden_states_scale: torch.Tensor, gemm1_weights: torch.Tensor, gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, gemm2_weights_scale: torch.Tensor, num_experts: int, @@ -183,6 +198,9 @@ def trtllm_fp8_block_scale_routed_moe_out_wrapper( "hidden_states_scale": hidden_states_scale, "gemm1_weights": gemm1_weights, "gemm1_weights_scale": gemm1_weights_scale, + "gemm1_alpha": gemm1_alpha, + "gemm1_beta": gemm1_beta, + "gemm1_clamp_limit": gemm1_clamp_limit, "gemm2_weights": gemm2_weights, "gemm2_weights_scale": gemm2_weights_scale, "output": output, 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 b3f1291f4..3ab58a354 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -249,6 +249,7 @@ class FusedMoE(torch.nn.Module): no_combine: bool = False, routed_scaling_factor: Optional[float] = None, gemm1_alpha: Optional[float] = None, + gemm1_beta: Optional[float] = None, gemm1_clamp_limit: Optional[float] = None, swiglu_limit: Optional[float] = None, use_weight_loader_fused: bool = False, @@ -353,6 +354,7 @@ class FusedMoE(torch.nn.Module): no_combine=no_combine, routed_scaling_factor=routed_scaling_factor, gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, swiglu_limit=swiglu_limit, is_gated=is_gated, diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index 49f3ed5c7..e334ab1cb 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -54,6 +54,7 @@ class MoeRunnerConfig: no_combine: bool = False routed_scaling_factor: Optional[float] = None gemm1_alpha: Optional[float] = None + gemm1_beta: Optional[float] = None gemm1_clamp_limit: Optional[float] = None swiglu_limit: Optional[float] = None # Whether gate/up weights are stored interleaved (vs split). Only the diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index f821d095c..e0a39cbae 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -646,6 +646,9 @@ class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo): weight_block_k: int | None = None w13_weight_scale_inv: torch.Tensor | None = None w2_weight_scale_inv: torch.Tensor | None = None + gemm1_alpha: torch.Tensor | None = None + gemm1_beta: torch.Tensor | None = None + gemm1_clamp_limit: torch.Tensor | None = None # Per-tensor path w13_input_scale: torch.Tensor | None = None @@ -744,6 +747,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( hidden_states_scale=a_sf_t, gemm1_weights=quant_info.w13_weight, gemm1_weights_scale=quant_info.w13_weight_scale_inv, + gemm1_alpha=quant_info.gemm1_alpha, + gemm1_beta=quant_info.gemm1_beta, + gemm1_clamp_limit=quant_info.gemm1_clamp_limit, gemm2_weights=quant_info.w2_weight, gemm2_weights_scale=quant_info.w2_weight_scale_inv, num_experts=quant_info.global_num_experts, @@ -779,6 +785,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( hidden_states_scale=a_sf_t, gemm1_weights=quant_info.w13_weight, gemm1_weights_scale=quant_info.w13_weight_scale_inv, + gemm1_alpha=quant_info.gemm1_alpha, + gemm1_beta=quant_info.gemm1_beta, + gemm1_clamp_limit=quant_info.gemm1_clamp_limit, gemm2_weights=quant_info.w2_weight, gemm2_weights_scale=quant_info.w2_weight_scale_inv, output=symm_output, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 0498e6be7..4990d60bd 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -2139,12 +2139,39 @@ class Fp8MoEMethod(FusedMoEMethodBase): align_fp8_moe_weights_for_flashinfer_trtllm(layer) + if ( + get_moe_runner_backend().is_flashinfer_trtllm() + or get_moe_runner_backend().is_flashinfer_trtllm_routed() + ): + self._prepare_flashinfer_trtllm_activation_params(layer) + if get_moe_runner_backend().is_hpc_ops(): self._prepare_hpc_ops_weights(layer) if hasattr(layer, "dispatcher"): layer.dispatcher.set_quant_config({"weight_dtype": layer.w13_weight.dtype}) + def _prepare_flashinfer_trtllm_activation_params(self, layer: Module) -> None: + """Materialize optional TRT-LLM SwiGLU parameters once per expert.""" + num_experts = int(layer.num_local_experts) + device = layer.w13_weight.device + for name, value in ( + ("gemm1_alpha", self.moe_runner_config.gemm1_alpha), + ("gemm1_beta", self.moe_runner_config.gemm1_beta), + ("gemm1_clamp_limit", self.moe_runner_config.gemm1_clamp_limit), + ): + tensor = ( + None + if value is None + else torch.full( + (num_experts,), + float(value), + dtype=torch.float32, + device=device, + ) + ) + setattr(layer, f"_flashinfer_trtllm_{name}", tensor) + def _prepare_hpc_ops_weights(self, layer: Module) -> None: """Precompute the scale layouts consumed by the HPC-Ops fused MoE kernels. @@ -2553,6 +2580,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_weight_scale_inv=( layer.w2_weight_scale_inv if self.block_quant else None ), + gemm1_alpha=layer._flashinfer_trtllm_gemm1_alpha, + gemm1_beta=layer._flashinfer_trtllm_gemm1_beta, + gemm1_clamp_limit=layer._flashinfer_trtllm_gemm1_clamp_limit, w13_input_scale=layer.w13_input_scale if not self.block_quant else None, output1_scales_scalar=( getattr(layer, "output1_scales_scalar", None) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index b772079c5..e45d6a440 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1247,6 +1247,12 @@ def flashinfer_mxfp8_blockscaled_linear( else: output_dtype = torch.bfloat16 + # At small M the persistent CUTLASS kernel is 2-5x slower than the + # CuTe-DSL swap-AB/split-K kernels (both consume the same swizzled + # 1D scales). + if backend == "cutlass" and q_input.shape[0] <= 64: + backend = "cute-dsl" + if backend == "trtllm": weight_scale_t = weight_scale.view(-1) else: diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index d7d1fda99..e656e6c3f 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -33,8 +33,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# TODO: Remove after FlashInfer fixes the mxfp8_gemm autotuning IMA. -FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset({"mxfp8_gemm"}) +FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset() def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]: diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 003c1d75a..6fbca7dcd 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -323,6 +323,7 @@ class MiniMaxM3MoE(nn.Module): activation="silu", is_gated=True, gemm1_alpha=config.swiglu_alpha, + gemm1_beta=1.0, gemm1_clamp_limit=config.swiglu_limit, prefix=add_prefix("experts", prefix), gate_up_interleaved=False,