chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

MoE Refactor: Refactor fp8.py -> flashinfer_trllm.py (#15151)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
b8zhong
2026-01-07 15:35:00 -08:00
committed by GitHub
parent 3a4767daa3
commit 24b30f7757
4 changed files with 307 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -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.

238
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py Normal file
View File

@@ -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)

9
python/sglang/srt/layers/moe/moe_runner/runner.py
View File

@@ -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(

258
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -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,