[Kernel] Support FlashInfer TRTLLM-Gen fused MoE for non-gated FP4 & FP8 (Nemotron) (#21321)

python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py

@@ -28,6 +28,7 @@ def _fake_fp8_block_scale_moe(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -58,6 +59,7 @@ def trtllm_fp8_block_scale_moe_wrapper(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     try:
         from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
@@ -94,6 +96,11 @@ def trtllm_fp8_block_scale_moe_wrapper(
 
         kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_block_scale_moe(**kwargs)
 
 
@@ -120,6 +127,7 @@ def _fake_fp8_block_scale_routed_moe(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -150,6 +158,7 @@ def trtllm_fp8_block_scale_routed_moe_wrapper(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     try:
         from flashinfer.fused_moe import trtllm_fp8_block_scale_routed_moe
@@ -186,6 +195,11 @@ def trtllm_fp8_block_scale_routed_moe_wrapper(
 
         kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_block_scale_routed_moe(**kwargs)
 
 
@@ -210,6 +224,7 @@ def _fake_fp8_per_tensor_scale_moe(
     routing_method_type: int = 0,
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -238,6 +253,7 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper(
     routing_method_type: int = 0,
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     # lazy import
     try:
@@ -271,4 +287,9 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper(
         "tune_max_num_tokens": tune_max_num_tokens,
     }
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_per_tensor_scale_moe(**kwargs)

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -34,6 +34,14 @@ from sglang.srt.utils.common import (
     next_power_of_2,
 )
 
+logger = __import__("logging").getLogger(__name__)
+
+
+def round_up_to_multiple(x: int, m: int) -> int:
+    """Round up *x* to the nearest multiple of *m*."""
+    return (x + m - 1) // m * m
+
+
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import (
         StandardCombineInput,
@@ -52,6 +60,49 @@ _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8: dict[
 ] = {}
 
 
+def _is_gated(layer: Module) -> bool:
+    """Return whether the MoE layer uses a gated activation (default True)."""
+    is_gated = (
+        getattr(layer, "moe_runner_config", None) and layer.moe_runner_config.is_gated
+    )
+    return True if is_gated is None else is_gated
+
+
+def _align_fp8_moe_weights(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    is_gated: bool,
+    min_alignment: int = 16,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Pad intermediate size so FlashInfer TRTLLM FP8 kernels' alignment holds.
+
+    Returns (w13, w2, padded_intermediate).
+    """
+    num_experts, hidden_size, intermediate = w2.shape
+
+    padded_intermediate = round_up_to_multiple(intermediate, min_alignment)
+    if padded_intermediate == intermediate:
+        return w13, w2, intermediate
+
+    logger.info(
+        "FP8 MoE: padding intermediate size from %d to %d (alignment=%d)",
+        intermediate,
+        padded_intermediate,
+        min_alignment,
+    )
+
+    up_mult = 2 if is_gated else 1
+    padded_gate_up = up_mult * padded_intermediate
+
+    padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2]))
+    padded_w13[:, : w13.shape[1], :] = w13
+
+    padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
+    padded_w2[:, :, :intermediate] = w2
+
+    return padded_w13, padded_w2, padded_intermediate
+
+
 def align_fp8_moe_weights_for_flashinfer_trtllm(
     layer: Module, swap_w13_halves: bool = False
 ) -> None:
@@ -63,32 +114,47 @@ def align_fp8_moe_weights_for_flashinfer_trtllm(
             This is needed for ModelOpt FP8 checkpoints which store weights in
             [Up, Gate] order, while regular FP8 checkpoints store them in [Gate, Up].
     """
-    from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
+    from flashinfer import shuffle_matrix_a
+
+    is_gated = _is_gated(layer)
 
     w13_weight = cast(torch.Tensor, layer.w13_weight)
     w2_weight = cast(torch.Tensor, layer.w2_weight)
-    num_experts, two_n, hidden = w13_weight.shape
+    num_experts, gate_up_dim, hidden = w13_weight.shape
 
-    # Optionally swap W13 halves: [Up, Gate] -> [Gate, Up]
-    if swap_w13_halves:
-        inter = two_n // 2
+    # Optionally swap W13 halves: [Up, Gate] -> [Gate, Up] (only for gated)
+    if swap_w13_halves and is_gated:
+        inter = gate_up_dim // 2
         w13_weight = (
             w13_weight.reshape(num_experts, 2, inter, hidden)
             .flip(dims=[1])
-            .reshape(num_experts, two_n, hidden)
+            .reshape(num_experts, gate_up_dim, hidden)
         )
 
-    w13_interleaved_list = [
-        reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts)
-    ]
-    w13_interleaved: torch.Tensor = torch.stack(w13_interleaved_list).reshape(
-        num_experts, two_n, hidden
+    # Pad for kernel alignment (non-gated needs 128, gated needs 16)
+    min_alignment = 16 if is_gated else 128
+    w13_weight, w2_weight, _ = _align_fp8_moe_weights(
+        w13_weight, w2_weight, is_gated, min_alignment
     )
+    num_experts, gate_up_dim, hidden = w13_weight.shape
+
+    epilogue_tile_m = 128
+
+    if is_gated:
+        from flashinfer import reorder_rows_for_gated_act_gemm
+
+        w13_interleaved_list = [
+            reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts)
+        ]
+        w13_processed: torch.Tensor = torch.stack(w13_interleaved_list).reshape(
+            num_experts, gate_up_dim, hidden
+        )
+    else:
+        w13_processed = w13_weight
 
     # Shuffle weights for transposed MMA output (both W13, W2)
-    epilogue_tile_m = 128
     w13_shuffled = [
-        shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
+        shuffle_matrix_a(w13_processed[i].view(torch.uint8), epilogue_tile_m)
         for i in range(num_experts)
     ]
     w2_shuffled = [
@@ -117,7 +183,16 @@ def align_fp8_moe_weights_for_flashinfer_trtllm(
     w13_weight_scale = cast(torch.Tensor, layer.w13_weight_scale).to(torch.float32)
     w2_weight_scale = cast(torch.Tensor, layer.w2_weight_scale).to(torch.float32)
 
-    output1_scales_scalar = w13_weight_scale * input_scale * (1.0 / activation_scale)
+    # For gated (SwiGLU): g1_alphas = w1_scale * a1_scale, g1_scale_c = g1_alphas / a2_scale
+    # For non-gated (Relu2): g1_scale_c = 1 / a2_scale (no gate dequant contribution)
+    if is_gated:
+        output1_scales_scalar = (
+            w13_weight_scale * input_scale * (1.0 / activation_scale)
+        )
+    else:
+        output1_scales_scalar = torch.ones_like(w13_weight_scale) * (
+            1.0 / activation_scale
+        )
     output1_scales_gate_scalar = w13_weight_scale * input_scale
     output2_scales_scalar = activation_scale * w2_weight_scale
 
@@ -128,6 +203,55 @@ def align_fp8_moe_weights_for_flashinfer_trtllm(
     layer.output2_scales_scalar = Parameter(output2_scales_scalar, requires_grad=False)
 
 
+def _align_mxfp8_moe_weights(
+    w13: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2: torch.Tensor,
+    w2_scale: torch.Tensor,
+    is_gated: bool,
+    min_alignment: int = 16,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Pad intermediate size so FlashInfer TRTLLM MXFP8 kernels' alignment holds.
+
+    Returns (w13, w13_scale, w2, w2_scale, padded_intermediate).
+    """
+    num_experts, hidden_size, intermediate = w2.shape
+
+    padded_intermediate = round_up_to_multiple(intermediate, min_alignment)
+    if padded_intermediate == intermediate:
+        return w13, w13_scale, w2, w2_scale, intermediate
+
+    logger.info(
+        "MXFP8 MoE: padding intermediate size from %d to %d (alignment=%d)",
+        intermediate,
+        padded_intermediate,
+        min_alignment,
+    )
+
+    up_mult = 2 if is_gated else 1
+    padded_gate_up = up_mult * padded_intermediate
+
+    padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2]))
+    padded_w13[:, : w13.shape[1], :] = w13
+
+    padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
+    padded_w2[:, :, :intermediate] = w2
+
+    padded_w13_scale = w13_scale.new_zeros(
+        (num_experts, padded_gate_up, w13_scale.shape[2])
+    )
+    padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale
+
+    # Scale's last dim tracks intermediate / block_size (MXFP8 block_size = 32)
+    scale_block_k = intermediate // w2_scale.shape[2] if w2_scale.shape[2] > 0 else 32
+    padded_w2_scale = w2_scale.new_zeros(
+        (num_experts, hidden_size, padded_intermediate // scale_block_k)
+    )
+    padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale
+
+    return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
+
+
 def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     """Prepare MXFP8 MoE weights/scales for FlashInfer TRT-LLM kernels."""
     from flashinfer import block_scale_interleave
@@ -139,6 +263,8 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
         get_shuffle_matrix_sf_a_row_indices,
     )
 
+    is_gated = _is_gated(layer)
+
     w13_weight = cast(torch.Tensor, layer.w13_weight).contiguous()
     w2_weight = cast(torch.Tensor, layer.w2_weight).contiguous()
     w13_scale = cast(torch.Tensor, layer.w13_weight_scale_inv).contiguous()
@@ -147,7 +273,13 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     assert w13_scale.dtype == torch.uint8
     assert w2_scale.dtype == torch.uint8
 
-    num_experts, two_n, _ = w13_weight.shape
+    # Pad for kernel alignment (non-gated needs 128, gated needs 16)
+    min_alignment = 16 if is_gated else 128
+    w13_weight, w13_scale, w2_weight, w2_scale, _ = _align_mxfp8_moe_weights(
+        w13_weight, w13_scale, w2_weight, w2_scale, is_gated, min_alignment
+    )
+
+    num_experts, gate_up_dim, _ = w13_weight.shape
     _, hidden_size, _ = w2_weight.shape
     epilogue_tile_m = 128
 
@@ -155,7 +287,7 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     w13_weight_u8 = w13_weight.view(torch.uint8)
     w2_weight_u8 = w2_weight.view(torch.uint8)
     cache_key = (
-        two_n,
+        gate_up_dim,
         hidden_size,
         w2_weight.shape[-1],
         w13_scale.shape[-1],
@@ -168,9 +300,14 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     )
     cache = _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8.get(cache_key)
     if cache is None:
-        reorder_row_indices = get_reorder_rows_for_gated_act_gemm_row_indices(
-            w13_weight_u8[0]
-        ).to(w13_weight.device)
+        if is_gated:
+            reorder_row_indices = get_reorder_rows_for_gated_act_gemm_row_indices(
+                w13_weight_u8[0]
+            ).to(w13_weight.device)
+        else:
+            reorder_row_indices = torch.arange(
+                gate_up_dim, device=w13_weight.device, dtype=torch.long
+            )
         w13_shuffle_row_indices = get_shuffle_matrix_a_row_indices(
             w13_weight_u8[0], epilogue_tile_m
         ).to(w13_weight.device)
@@ -178,7 +315,7 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
             w2_weight_u8[0], epilogue_tile_m
         ).to(w2_weight.device)
         w13_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices(
-            w13_scale[0].reshape(two_n, -1), epilogue_tile_m
+            w13_scale[0].reshape(gate_up_dim, -1), epilogue_tile_m
         ).to(w13_scale.device)
         w2_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices(
             w2_scale[0].reshape(hidden_size, -1), epilogue_tile_m
@@ -212,7 +349,7 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
         )
         w2_shuffled_u8[i].copy_(w2_weight_u8[i].index_select(0, w2_shuffle_row_indices))
 
-        w13_scale_linear = w13_scale_interleaved.reshape(two_n, -1)
+        w13_scale_linear = w13_scale_interleaved.reshape(gate_up_dim, -1)
         w13_scale_shuffled[i].copy_(
             block_scale_interleave(
                 w13_scale_linear.index_select(0, w13_scale_shuffle_row_indices)
@@ -243,10 +380,59 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     layer.w2_weight_scale_inv.format_ue8m0 = True
 
 
+def _align_fp4_moe_weights(
+    w13: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2: torch.Tensor,
+    w2_scale: torch.Tensor,
+    is_gated: bool,
+    min_alignment: int = 16,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Pad intermediate size so FlashInfer TRTLLM FP4 kernels' alignment holds.
+
+    Returns (w13, w13_scale, w2, w2_scale, padded_intermediate).
+    """
+    num_experts, hidden_size, intermediate_packed = w2.shape
+    intermediate = intermediate_packed * 2  # FP4 packs 2 values per byte
+
+    padded_intermediate = round_up_to_multiple(intermediate, min_alignment)
+    if padded_intermediate == intermediate:
+        return w13, w13_scale, w2, w2_scale, intermediate
+
+    logger.info(
+        "FP4 MoE: padding intermediate size from %d to %d (alignment=%d)",
+        intermediate,
+        padded_intermediate,
+        min_alignment,
+    )
+
+    up_mult = 2 if is_gated else 1
+    padded_gate_up = up_mult * padded_intermediate
+
+    padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2]))
+    padded_w13[:, : w13.shape[1], :] = w13
+
+    padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate // 2))
+    padded_w2[:, :, : w2.shape[2]] = w2
+
+    padded_w13_scale = w13_scale.new_zeros(
+        (num_experts, padded_gate_up, w13_scale.shape[2])
+    )
+    padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale
+
+    padded_w2_scale = w2_scale.new_zeros(
+        (num_experts, hidden_size, padded_intermediate // 16)
+    )
+    padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale
+
+    return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
+
+
 def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     """Prepare FP4 MoE weights/scales for FlashInfer TRT-LLM kernels.
 
     This function handles the weight transformation needed for FP4 TRTLLM MoE:
+    - Pads intermediate dimension for kernel alignment constraints
     - Reorders weights for gated activation GEMM
     - Shuffles weights and scales for transposed MMA output
     - Computes the output scale factors
@@ -260,6 +446,21 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     w13_weight_scale = cast(torch.Tensor, layer.w13_weight_scale)
     w2_weight_scale = cast(torch.Tensor, layer.w2_weight_scale)
 
+    is_gated = layer.moe_runner_config.is_gated
+    min_alignment = 16 if is_gated else 128
+
+    # Pad for kernel alignment before shuffle/reorder
+    w13_weight, w13_weight_scale, w2_weight, w2_weight_scale, intermediate_size = (
+        _align_fp4_moe_weights(
+            w13_weight,
+            w13_weight_scale,
+            w2_weight,
+            w2_weight_scale,
+            is_gated,
+            min_alignment,
+        )
+    )
+
     (
         gemm1_weights_fp4_shuffled,
         gemm1_scales_fp4_shuffled,
@@ -271,8 +472,9 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
         w13_weight_scale,
         w2_weight_scale,
         w2_weight.size(-2),  # hidden_size
-        w13_weight.size(-2) // 2,  # intermediate_size
+        intermediate_size,  # padded intermediate_size
         w13_weight.size(0),  # num_experts
+        is_gated=is_gated,
     )
 
     # Set flashinfer parameters in-place
@@ -285,14 +487,39 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
         layer, "w2_weight_scale", gemm2_scales_fp4_shuffled.contiguous()
     )
 
-    # Compute additional scaling factor needed for TRT-LLM
+    # Compute additional scaling factor needed for TRT-LLM.
+    # For gated (SwiGLU): g1_scale_c = g1_alphas * a2_gscale
+    # For non-gated (Relu2): g1_scale_c = a2_gscale (no gate dequant contribution)
     w2_input_scale_quant = cast(torch.Tensor, layer.w2_input_scale_quant)
     g1_alphas = cast(torch.Tensor, layer.g1_alphas)
-    copy_or_rebind_param(
-        layer,
-        "g1_scale_c",
-        (w2_input_scale_quant * g1_alphas).to(torch.float32),
-    )
+    if layer.moe_runner_config.is_gated:
+        g1_scale_c = (w2_input_scale_quant * g1_alphas).to(torch.float32)
+    else:
+        num_experts = g1_alphas.shape[0]
+        g1_scale_c = (
+            w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous()
+        )
+    copy_or_rebind_param(layer, "g1_scale_c", g1_scale_c)
+
+    # Update intermediate_size_per_partition to reflect any padding applied
+    layer.intermediate_size_per_partition = intermediate_size
+
+
+def get_activation_type(activation: str) -> int:
+    """Map SGLang activation string to FlashInfer ActivationType int value."""
+    from flashinfer.fused_moe.core import ActivationType
+
+    _ACTIVATION_STR_TO_TYPE = {
+        "silu": ActivationType.Swiglu,
+        "relu2": ActivationType.Relu2,
+    }
+    act = _ACTIVATION_STR_TO_TYPE.get(activation)
+    if act is None:
+        raise ValueError(
+            f"Unsupported activation '{activation}' for TRTLLM MoE. "
+            f"Expected one of {list(_ACTIVATION_STR_TO_TYPE.keys())}."
+        )
+    return act.value
 
 
 @dataclass
@@ -325,6 +552,9 @@ class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo):
     output2_scales_scalar: torch.Tensor | None = None
     use_routing_scales_on_input: bool = False
 
+    # Activation type (None = kernel default / Swiglu)
+    activation_type: int | None = None
+
 
 def _pack_topk_for_flashinfer_routed(
     topk_ids: torch.Tensor, topk_weights: torch.Tensor
@@ -348,7 +578,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
     from sglang.srt.layers.moe.topk import TopKOutputChecker
     from sglang.srt.layers.moe.utils import RoutingMethodType
 
-    assert runner_config.activation == "silu", "Only silu is supported."
+    _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"}
+    assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, (
+        f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for FP8 MoE, "
+        f"got '{runner_config.activation}'."
+    )
     assert not runner_config.no_combine, "no_combine is not supported for flashinfer."
 
     hidden_states = dispatch_output.hidden_states
@@ -446,6 +680,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
                 use_shuffled_weight=use_shuffled_weight,
                 tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
                 fp8_quantization_type=int(fp8_quantization_type),
+                activation_type=quant_info.activation_type,
             )
         else:
             assert TopKOutputChecker.format_is_bypassed(topk_output)
@@ -479,6 +714,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
                 use_shuffled_weight=use_shuffled_weight,
                 tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
                 fp8_quantization_type=int(fp8_quantization_type),
+                activation_type=quant_info.activation_type,
             )
         # TODO: Once https://github.com/flashinfer-ai/flashinfer/issues/2703 is fixed, pass output to moe kernel and remove this copy.
         symm_output.copy_(output)
@@ -540,6 +776,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
             use_routing_scales_on_input=False,
             routing_method_type=routing_method_type,
             tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+            activation_type=quant_info.activation_type,
         )
         symm_output.copy_(output)
         output = symm_output
@@ -622,8 +859,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
     from sglang.srt.layers.moe.topk import TopKOutputChecker
     from sglang.srt.layers.moe.utils import RoutingMethodType
 
-    assert runner_config.activation == "silu", "Only silu is supported for FP4 MoE."
-    assert runner_config.is_gated, "Only gated MoEs are supported for FP4 MoE."
+    _SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2"}
+    assert runner_config.activation in _SUPPORTED_FP4_ACTIVATIONS, (
+        f"Only {_SUPPORTED_FP4_ACTIVATIONS} are supported for FP4 MoE, "
+        f"got '{runner_config.activation}'."
+    )
 
     hidden_states = dispatch_output.hidden_states
     topk_output = dispatch_output.topk_output
@@ -635,6 +875,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
     hs_scale = hs_scale_linear.view(torch.float8_e4m3fn).reshape(
         *hs_scale_linear.shape[:-1], -1
     )
+    activation_type = get_activation_type(runner_config.activation)
 
     with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
         num_tokens = hs_fp4.shape[0]
@@ -678,6 +919,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
             routed_scaling_factor=None,
             routing_method_type=1,  # Unused, but must be 1 to pass validation.
             do_finalize=True,
+            activation_type=activation_type,
             tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]),
             output=symm_output,
         )[0]
@@ -728,6 +970,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
                 else RoutingMethodType.Default
             ),
             do_finalize=True,
+            activation_type=activation_type,
             tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]),
             output=symm_output,
         )[0]

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py

@@ -398,6 +398,11 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
             return StandardCombineInput(hidden_states=output)
         elif self.weight_quant.strategy == QuantizationStrategy.BLOCK:
             if self.use_flashinfer_trtllm:
+                from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+                    get_activation_type,
+                )
+
+                activation_type = get_activation_type(moe_runner_config.activation)
                 quant_info = FlashInferTrtllmFp8MoeQuantInfo(
                     w13_weight=layer.w13_weight,
                     w2_weight=layer.w2_weight,
@@ -410,6 +415,7 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
                     weight_block_k=self.weight_block_size[1],
                     w13_weight_scale_inv=layer.w13_weight_scale,
                     w2_weight_scale_inv=layer.w2_weight_scale,
+                    activation_type=activation_type,
                 )
             else:
                 quant_info = TritonMoeQuantInfo(

python/sglang/srt/layers/quantization/fp8.py

@@ -1688,6 +1688,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             num_local_experts = int(getattr(layer, "num_local_experts"))
             moe_ep_rank = int(getattr(layer, "moe_ep_rank"))
 
+            from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+                get_activation_type,
+            )
+
+            activation_type = get_activation_type(self.moe_runner_config.activation)
+
             quant_info = FlashInferTrtllmFp8MoeQuantInfo(
                 w13_weight=layer.w13_weight,
                 w2_weight=layer.w2_weight,
@@ -1728,6 +1734,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                     if not self.block_quant
                     else None
                 ),
+                activation_type=activation_type,
             )
         elif self.runner.runner_backend.is_triton():
             quant_info = self.get_triton_quant_info(layer)

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -1008,20 +1008,19 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
 
             topk_config = topk_output.topk_config
 
-            # Constraints for ModelOpt FP8 MoE
-            assert (
-                self.moe_runner_config.activation == "silu"
-            ), "Only silu is supported for flashinfer fp8 moe"
+            from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+                get_activation_type,
+            )
 
-            # Enforce Llama4 routing for ModelOpt FP8 MoE for now.
-            # TODO(brayden): support other routing methods
-            assert topk_config.top_k == 1, "ModelOpt FP8 MoE requires top_k==1"
-            assert (
-                not topk_config.num_expert_group
-            ), "ModelOpt FP8 MoE does not support expert grouping"
-            assert (
-                not topk_config.topk_group
-            ), "ModelOpt FP8 MoE does not support grouped top-k"
+            _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"}
+            assert self.moe_runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, (
+                f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for "
+                f"flashinfer trtllm fp8 moe, got '{self.moe_runner_config.activation}'"
+            )
+
+            routing_method_type = getattr(
+                layer, "routing_method_type", RoutingMethodType.Llama4
+            )
 
             quant_info = FlashInferTrtllmFp8MoeQuantInfo(
                 w13_weight=layer.w13_weight,
@@ -1030,13 +1029,14 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
                 local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
                 local_num_experts=layer.num_local_experts,
                 intermediate_size=layer.w2_weight.shape[2],
-                routing_method_type=RoutingMethodType.Llama4,
+                routing_method_type=routing_method_type,
                 block_quant=False,
                 w13_input_scale=layer.w13_input_scale,
                 output1_scales_scalar=layer.output1_scales_scalar,
                 output1_scales_gate_scalar=layer.output1_scales_gate_scalar,
                 output2_scales_scalar=layer.output2_scales_scalar,
                 use_routing_scales_on_input=True,
+                activation_type=get_activation_type(self.moe_runner_config.activation),
             )
 
             return fused_experts_none_to_flashinfer_trtllm_fp8(

python/sglang/srt/layers/quantization/utils.py

@@ -652,6 +652,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(
     hidden_size,
     intermediate_size,
     num_experts,
+    is_gated: bool = True,
 ):
     from flashinfer import nvfp4_block_scale_interleave
     from flashinfer.fused_moe.core import (
@@ -663,14 +664,16 @@ def prepare_static_weights_for_trtllm_fp4_moe(
     _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
     epilogue_tile_m = 128  # FIXME: this depends on the kernel internals
 
+    gemm1_rows = (2 if is_gated else 1) * intermediate_size
+
     # Convert quantized weights to proper formats
     gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
-        num_experts, 2 * intermediate_size, hidden_size // 2
+        num_experts, gemm1_rows, hidden_size // 2
     )  # packed fp4
     gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
         torch.float8_e4m3fn
     ).reshape(
-        num_experts, 2 * intermediate_size, hidden_size // 16
+        num_experts, gemm1_rows, hidden_size // 16
     )  # fp8 scaling factors
 
     gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
@@ -687,14 +690,11 @@ def prepare_static_weights_for_trtllm_fp4_moe(
     gemm2_weights_fp4_shuffled = []
     gemm2_scales_fp4_shuffled = []
     for i in range(num_experts):
-        # Calculate the permute indices for the following:
-        # 1. Reorder rows of W1 and scales for fused gated activation
-        # 2. Shuffle weights and scaling factors for transposed mma output
-        # for both w3_w1 and w2 weights and scale factors
         permute_indices = _maybe_get_cached_w3_w1_permute_indices(
             _cache_permute_indices,
             gemm1_weights_fp4[i].view(torch.uint8),
             epilogue_tile_m,
+            is_gated_act_gemm=is_gated,
         )
         gemm1_weights_fp4_shuffled.append(
             gemm1_weights_fp4[i]
@@ -707,6 +707,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(
             gemm1_scales_linear_fp4[i].view(torch.uint8),
             epilogue_tile_m,
             num_elts_per_sf=16,
+            is_gated_act_gemm=is_gated,
         )
         gemm1_scales_fp4_shuffled.append(
             nvfp4_block_scale_interleave(
@@ -750,7 +751,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(
     gemm1_scales_fp4_shuffled = (
         torch.stack(gemm1_scales_fp4_shuffled)
         .view(torch.float8_e4m3fn)
-        .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
+        .reshape(num_experts, gemm1_rows, hidden_size // 16)
     )
 
     gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)

python/sglang/srt/models/nemotron_h.py

@@ -51,6 +51,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
 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.utils import RoutingMethodType
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
@@ -190,6 +191,7 @@ class NemotronHMoE(nn.Module):
             activation=config.mlp_hidden_act,
             layer_id=layer_idx,
             is_gated=False,
+            routing_method_type=RoutingMethodType.DeepSeekV3,
         )
         if config.n_shared_experts:
             self.shared_experts = NemotronHMLP(

python/sglang/srt/server_args.py

@@ -2132,7 +2132,15 @@ class ServerArgs:
                         )
                 else:
                     self.quantization = model_config.quantization
-                self.moe_runner_backend = "flashinfer_cutlass"
+                if self.moe_runner_backend == "auto":
+                    if is_sm100_supported() and self.moe_a2a_backend == "none":
+                        self.moe_runner_backend = "flashinfer_trtllm"
+                        logger.info(
+                            "Use flashinfer_trtllm as MoE runner backend on sm100 for "
+                            f"{model_arch}"
+                        )
+                    else:
+                        self.moe_runner_backend = "flashinfer_cutlass"
 
             self._handle_mamba_radix_cache(
                 model_arch=model_arch,