minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Chunan Zeng
2026-08-09 23:30:05 -07:00
committed by GitHub
parent 2969ab3d41
commit e226bb711c
8 changed files with 68 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py
+18
View File
@@ -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
+2
View File
@@ -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
+1
View File
@@ -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
+9
View File
@@ -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
+30
View File
@@ -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
+6
View File
@@ -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
+1 -2
View File
@@ -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]:
python/sglang/srt/models/minimax_m3.py
+1
View File
@@ -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,