enable TRT-LLM for MiniMax M3 by preserving SwiGLU params (#33962)

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,

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,

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

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,

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)

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:

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()
-FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset({"mxfp8_gemm"})
 
 
 def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]:

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,