From 8327270c7263fb9f62e03a1e7106f079aa7216f3 Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:28:21 +0300 Subject: [PATCH] [Kernel] Support FlashInfer TRTLLM-Gen fused MoE for non-gated FP4 & FP8 (Nemotron) (#21321) --- .../srt/layers/moe/flashinfer_trtllm_moe.py | 21 ++ .../moe/moe_runner/flashinfer_trtllm.py | 305 ++++++++++++++++-- .../compressed_tensors_w8a8_fp8_moe.py | 6 + python/sglang/srt/layers/quantization/fp8.py | 7 + .../srt/layers/quantization/modelopt_quant.py | 28 +- .../sglang/srt/layers/quantization/utils.py | 15 +- python/sglang/srt/models/nemotron_h.py | 2 + python/sglang/srt/server_args.py | 10 +- 8 files changed, 341 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py index c38126679..cbd43d8b2 100644 --- a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py +++ b/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) 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 b28a8e61e..734883526 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/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] diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index c4139f54a..6e00b51b7 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index b75e78936..ac2c31b6d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 42753f31c..d8ad8a1dd 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/utils.py b/python/sglang/srt/layers/quantization/utils.py index ac9007513..d61e6c36a 100644 --- a/python/sglang/srt/layers/quantization/utils.py +++ b/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) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index b858950c9..0426729ce 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5e3f592d0..1aedd9a73 100644 --- a/python/sglang/srt/server_args.py +++ b/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,