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 839463518..a1b5f4e7d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -1146,14 +1146,14 @@ class FlashInferFusedMoE(FusedMoE): else: - final_hidden_states = self.quant_method.apply_with_router_logits( + final_hidden_states = self.quant_method.apply( layer=self, dispatch_output=StandardDispatchOutput( hidden_states=hidden_states, hidden_states_scale=None, topk_output=topk_output, ), - ) + ).hidden_states # NOTE for symmetric memory tagging: # We do not create the context in this function. diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py new file mode 100644 index 000000000..160ffa068 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import torch +from torch.nn import Module +from torch.nn.parameter import Parameter + +from sglang.srt.distributed import get_tp_group +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + register_fused_func, +) +from sglang.srt.layers.quantization.fp8_kernel import ( + per_token_group_quant_fp8, + scaled_fp8_quant, +) +from sglang.srt.utils.common import next_power_of_2 + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + StandardCombineInput, + StandardDispatchOutput, + ) + + +def align_fp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: + """Prepare FP8 MoE weights/scales for FlashInfer TRT-LLM kernels.""" + from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a + + # Note: No need to swap W13 halves, they are already in the correct order: + # [Gate, Up] + w13_weight = cast(torch.Tensor, layer.w13_weight) + w2_weight = cast(torch.Tensor, layer.w2_weight) + num_experts, two_n, hidden = w13_weight.shape + + 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 + ) + + # 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) + for i in range(num_experts) + ] + w2_shuffled = [ + shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m) + for i in range(num_experts) + ] + + layer.w13_weight = Parameter( + torch.stack(w13_shuffled).view(torch.float8_e4m3fn), + requires_grad=False, + ) + layer.w2_weight = Parameter( + torch.stack(w2_shuffled).view(torch.float8_e4m3fn), + requires_grad=False, + ) + + # Precompute and register per-expert output scaling factors for FI MoE. + # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction. + assert hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None + assert hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None + assert hasattr(layer, "w13_weight_scale") and layer.w13_weight_scale is not None + assert hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None + + input_scale = cast(torch.Tensor, layer.w13_input_scale).to(torch.float32) + activation_scale = cast(torch.Tensor, layer.w2_input_scale).to(torch.float32) + 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) + output1_scales_gate_scalar = w13_weight_scale * input_scale + output2_scales_scalar = activation_scale * w2_weight_scale + + layer.output1_scales_scalar = Parameter(output1_scales_scalar, requires_grad=False) + layer.output1_scales_gate_scalar = Parameter( + output1_scales_gate_scalar, requires_grad=False + ) + layer.output2_scales_scalar = Parameter(output2_scales_scalar, requires_grad=False) + + +@dataclass +class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo): + """Quantization payload consumed by FlashInfer TRT-LLM FP8 MoE kernels.""" + + # Weights + w13_weight: torch.Tensor + w2_weight: torch.Tensor + + # Expert-parallel metadata + global_num_experts: int + local_expert_offset: int + local_num_experts: int + + routing_method_type: int + + # Block-quant path + block_quant: bool + weight_block_k: int | None = None + w13_weight_scale_inv: torch.Tensor | None = None + w2_weight_scale_inv: torch.Tensor | None = None + + # Per-tensor path + w13_input_scale: torch.Tensor | None = None + output1_scales_scalar: torch.Tensor | None = None + output1_scales_gate_scalar: torch.Tensor | None = None + output2_scales_scalar: torch.Tensor | None = None + + +@register_fused_func("none", "flashinfer_trtllm") +def fused_experts_none_to_flashinfer_trtllm_fp8( + dispatch_output: StandardDispatchOutput, + quant_info: FlashInferTrtllmFp8MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + from flashinfer.fused_moe import ( + trtllm_fp8_block_scale_moe, + trtllm_fp8_per_tensor_scale_moe, + ) + + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + 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." + assert not runner_config.no_combine, "no_combine is not supported for flashinfer." + + hidden_states = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + assert TopKOutputChecker.format_is_bypassed(topk_output) + + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + correction_bias = ( + None + if topk_config.correction_bias is None + else topk_config.correction_bias.to(hidden_states.dtype) + ) + + routing_method_type = quant_info.routing_method_type + + if quant_info.block_quant: + assert quant_info.weight_block_k is not None + assert quant_info.w13_weight_scale_inv is not None + assert quant_info.w2_weight_scale_inv is not None + + a_q, a_sf = per_token_group_quant_fp8(hidden_states, quant_info.weight_block_k) + a_sf_t = a_sf.t().contiguous() + + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + # FIXME: there is a bug in the trtllm_fp8_block_scale_moe. + # It ignored the `output` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325 + # so we put the whole function under the ``use_symmetric_memory`` context manager. + # If the bug is fixed, we can only put the output tensor allocation under the context manager. + output = trtllm_fp8_block_scale_moe( + routing_logits=( + router_logits.to(torch.float32) + if routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ), + routing_bias=correction_bias, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale_inv, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale_inv, + num_experts=quant_info.global_num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=int(quant_info.w2_weight.shape[2]), + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), + tile_tokens_dim=None, + routing_method_type=routing_method_type, + use_shuffled_weight=False, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + ) + else: + assert quant_info.w13_input_scale is not None + assert quant_info.output1_scales_scalar is not None + assert quant_info.output1_scales_gate_scalar is not None + assert quant_info.output2_scales_scalar is not None + + a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale) + routing_bias_cast = ( + None if correction_bias is None else correction_bias.to(torch.bfloat16) + ) + + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + output = trtllm_fp8_per_tensor_scale_moe( + routing_logits=router_logits.to(torch.bfloat16), + routing_bias=routing_bias_cast, + hidden_states=a_q, + gemm1_weights=quant_info.w13_weight, + output1_scales_scalar=quant_info.output1_scales_scalar, + output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar, + gemm2_weights=quant_info.w2_weight, + output2_scales_scalar=quant_info.output2_scales_scalar, + num_experts=quant_info.global_num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=int(quant_info.w2_weight.shape[2]), + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), + use_routing_scales_on_input=False, + routing_method_type=routing_method_type, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + ) + + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index fde68df94..8b58cd311 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -39,6 +39,8 @@ class MoeRunner: self.runner_core = DeepGemmRunnerCore(config) elif runner_backend.is_marlin(): self.runner_core = None # Marlin only supports fused path + elif runner_backend.is_flashinfer_trtllm(): + self.runner_core = None # FlashInfer TRT-LLM only supports fused path else: raise NotImplementedError(f"Unsupported runner backend: {runner_backend}") @@ -50,6 +52,12 @@ class MoeRunner: a2a_backend_name, runner_backend_name ) + if self.runner_core is None and self.fused_func is None: + raise NotImplementedError( + f"Runner backend {runner_backend} requires a fused func for a2a backend " + f"{a2a_backend_name}, but none is registered." + ) + self.down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None self.meta_overlap_args: Optional[dict] = None @@ -69,6 +77,7 @@ class MoeRunner: if self.fused_func is not None: return self.fused_func(dispatch_output, quant_info, self.config) + assert self.runner_core is not None dispatch_format = dispatch_output.format.value runner_format = self.runner_core.runner_backend.value self.pre_permute_func = PermuteMethodPool.get_pre_permute( diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index e80e818a0..9bf872720 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -18,8 +18,11 @@ from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo +from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmFp8MoeQuantInfo, +) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import get_moe_runner_backend +from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend from sglang.srt.layers.parameter import ( BlockQuantScaleParameter, ModelWeightParameter, @@ -69,18 +72,13 @@ from sglang.srt.utils import ( is_sm90_supported, is_sm100_supported, log_info_on_rank0, - next_power_of_2, print_warning_once, set_weight_attrs, use_intel_amx_backend, ) if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - DispatchOutput, - StandardDispatchOutput, - ) + from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config @@ -947,83 +945,11 @@ class Fp8MoEMethod(FusedMoEMethodBase): # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled if get_moe_runner_backend().is_flashinfer_trtllm(): - from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a - - # Note: No need to swap W13 halves, they are already in the correct order: [Gate, Up] - num_experts, two_n, hidden = layer.w13_weight.shape - - # 2) Reorder rows for fused gated activation (W13) - w13_interleaved = [ - reorder_rows_for_gated_act_gemm(layer.w13_weight[i]) - for i in range(num_experts) - ] - w13_interleaved = torch.stack(w13_interleaved).reshape( - num_experts, two_n, hidden + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + align_fp8_moe_weights_for_flashinfer_trtllm, ) - # 3) 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 - ) - for i in range(num_experts) - ] - w2_shuffled = [ - shuffle_matrix_a( - layer.w2_weight[i].view(torch.uint8), epilogue_tile_m - ) - for i in range(num_experts) - ] - - layer.w13_weight = Parameter( - torch.stack(w13_shuffled).view(torch.float8_e4m3fn), - requires_grad=False, - ) - layer.w2_weight = Parameter( - torch.stack(w2_shuffled).view(torch.float8_e4m3fn), - requires_grad=False, - ) - - # Precompute and register per-expert output scaling factors for FI MoE - # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction - assert ( - hasattr(layer, "w13_input_scale") - and layer.w13_input_scale is not None - ) - assert ( - hasattr(layer, "w2_input_scale") - and layer.w2_input_scale is not None - ) - assert ( - hasattr(layer, "w13_weight_scale") - and layer.w13_weight_scale is not None - ) - assert ( - hasattr(layer, "w2_weight_scale") - and layer.w2_weight_scale is not None - ) - - input_scale = layer.w13_input_scale.to(torch.float32) - activation_scale = layer.w2_input_scale.to(torch.float32) - w13_weight_scale = layer.w13_weight_scale.to(torch.float32) - w2_weight_scale = layer.w2_weight_scale.to(torch.float32) - - output1_scales_scalar = ( - w13_weight_scale * input_scale * (1.0 / activation_scale) - ) - output1_scales_gate_scalar = w13_weight_scale * input_scale - output2_scales_scalar = activation_scale * w2_weight_scale - - layer.output1_scales_scalar = Parameter( - output1_scales_scalar, requires_grad=False - ) - layer.output1_scales_gate_scalar = Parameter( - output1_scales_gate_scalar, requires_grad=False - ) - layer.output2_scales_scalar = Parameter( - output2_scales_scalar, requires_grad=False - ) + align_fp8_moe_weights_for_flashinfer_trtllm(layer) return def process_weights_hip_int4(self, layer: Module): @@ -1119,7 +1045,11 @@ class Fp8MoEMethod(FusedMoEMethodBase): moe_runner_backend = MoeRunnerBackend.DEEP_GEMM else: moe_runner_backend = MoeRunnerBackend.TRITON - if moe_runner_backend.is_deep_gemm() or moe_runner_backend.is_triton(): + if ( + moe_runner_backend.is_deep_gemm() + or moe_runner_backend.is_triton() + or moe_runner_backend.is_flashinfer_trtllm() + ): self.runner = MoeRunner(moe_runner_backend, moe_runner_config) else: # TODO(cwan): refactor other backends @@ -1245,6 +1175,51 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_scale=w2_scale, block_shape=block_shape, ) + elif self.runner.runner_backend.is_flashinfer_trtllm(): + # FlashInfer TRT-LLM backend only supports fused execution and consumes + # router logits directly (no separate apply_with_router_logits needed). + global_num_experts = int(getattr(layer, "num_experts")) + num_local_experts = int(getattr(layer, "num_local_experts")) + moe_ep_rank = int(getattr(layer, "moe_ep_rank")) + + quant_info = FlashInferTrtllmFp8MoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + global_num_experts=global_num_experts, + local_expert_offset=moe_ep_rank * num_local_experts, + local_num_experts=num_local_experts, + routing_method_type=int( + getattr(layer, "routing_method_type", RoutingMethodType.DeepSeekV3) + ), + block_quant=self.block_quant, + weight_block_k=( + None + if self.quant_config.weight_block_size is None + else self.quant_config.weight_block_size[1] + ), + w13_weight_scale_inv=( + layer.w13_weight_scale_inv if self.block_quant else None + ), + w2_weight_scale_inv=( + layer.w2_weight_scale_inv if self.block_quant else None + ), + w13_input_scale=layer.w13_input_scale if not self.block_quant else None, + output1_scales_scalar=( + getattr(layer, "output1_scales_scalar", None) + if not self.block_quant + else None + ), + output1_scales_gate_scalar=( + getattr(layer, "output1_scales_gate_scalar", None) + if not self.block_quant + else None + ), + output2_scales_scalar=( + getattr(layer, "output2_scales_scalar", None) + if not self.block_quant + else None + ), + ) elif self.runner.runner_backend.is_triton(): quant_info = TritonMoeQuantInfo( w13_weight=layer.w13_weight, @@ -1316,123 +1291,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): self._cutlass_buffers_ready = True - def apply_with_router_logits( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> torch.Tensor: - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - activation = self.moe_runner_config.activation - routed_scaling_factor = self.moe_runner_config.routed_scaling_factor - - from flashinfer.fused_moe import ( - trtllm_fp8_block_scale_moe, - trtllm_fp8_per_tensor_scale_moe, - ) - - from sglang.srt.layers.moe.topk import TopKOutputChecker - from sglang.srt.layers.moe.utils import RoutingMethodType - - assert TopKOutputChecker.format_is_bypassed(topk_output) - router_logits = topk_output.router_logits - topk_config = topk_output.topk_config - assert ( - activation == "silu" - ), "Only silu is supported for flashinfer blockscale fp8 moe" - - if self.block_quant: - a_q, a_sf = per_token_group_quant_fp8( - x, self.quant_config.weight_block_size[1] - ) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() - else: - a_q, _ = scaled_fp8_quant(x, layer.w13_input_scale) - - correction_bias = ( - None - if topk_config.correction_bias is None - else topk_config.correction_bias.to(x.dtype) - ) - - routing_method_type = getattr( - layer, "routing_method_type", RoutingMethodType.DeepSeekV3 - ) - - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - - if self.block_quant: - # FIXME: there is a bug in the trtllm_fp8_block_scale_moe. - # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325 - # so we put the whole function under the ``use_symmetric_memory`` context manager. - # If the bug is fixed, we can only put the output tensor allocation under the context manager. - return trtllm_fp8_block_scale_moe( - routing_logits=( - router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ), - routing_bias=correction_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=layer.w13_weight, - gemm1_weights_scale=layer.w13_weight_scale_inv, - gemm2_weights=layer.w2_weight, - gemm2_weights_scale=layer.w2_weight_scale_inv, - num_experts=layer.num_experts, - top_k=topk_config.top_k, - n_group=topk_config.num_expert_group, - topk_group=topk_config.topk_group, - intermediate_size=layer.w2_weight.shape[2], - local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, - local_num_experts=layer.num_local_experts, - routed_scaling_factor=( - routed_scaling_factor - if routed_scaling_factor is not None - else 1.0 - ), - tile_tokens_dim=None, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - tune_max_num_tokens=next_power_of_2(a_q.shape[0]), - ) - else: - routing_bias_cast = ( - None - if correction_bias is None - else correction_bias.to(torch.bfloat16) - ) - - return trtllm_fp8_per_tensor_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), - routing_bias=routing_bias_cast, - hidden_states=a_q, - gemm1_weights=layer.w13_weight, - output1_scales_scalar=layer.output1_scales_scalar, - output1_scales_gate_scalar=layer.output1_scales_gate_scalar, - gemm2_weights=layer.w2_weight, - output2_scales_scalar=layer.output2_scales_scalar, - num_experts=layer.num_experts, - top_k=topk_config.top_k, - n_group=topk_config.num_expert_group, - topk_group=topk_config.topk_group, - intermediate_size=layer.w2_weight.shape[2], - local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, - local_num_experts=layer.num_local_experts, - routed_scaling_factor=( - routed_scaling_factor - if routed_scaling_factor is not None - else 1.0 - ), - use_routing_scales_on_input=False, - routing_method_type=routing_method_type, - tune_max_num_tokens=next_power_of_2(a_q.shape[0]), - ) - def maybe_apply_hip_fused_experts( self, layer: torch.nn.Module,