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

[NVIDIA] Enable FP4 flashinfer trtllm routed moe (#21240)

Browse Source
This commit is contained in:
Trevor Morris
2026-04-07 16:16:29 -07:00
committed by GitHub
parent 0e2a0260a1
commit 7546d04c81
2 changed files with 114 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+72 -19
View File
@@ -616,13 +616,17 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
dispatch_output: StandardDispatchOutput,
quant_info: FlashInferTrtllmFp4MoeQuantInfo,
runner_config: MoeRunnerConfig,
use_routed_topk: bool = False,
) -> StandardCombineInput:
"""FlashInfer TRTLLM FP4 MoE forward pass.
This function handles the FP4 TRTLLM MoE path that was previously in
ModelOptNvFp4FusedMoEMethod.apply.
"""
from flashinfer.fused_moe import trtllm_fp4_block_scale_moe
from flashinfer.fused_moe import (
trtllm_fp4_block_scale_moe,
trtllm_fp4_block_scale_routed_moe,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
@@ -633,16 +637,70 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
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
routing_method_type = quant_info.routing_method_type
# Quantize hidden states to FP4
hs_fp4, hs_scale_linear = quantize_hidden_states_fp4(
hidden_states, quant_info.w13_input_scale_quant
)
hs_scale = hs_scale_linear.view(torch.float8_e4m3fn).reshape(
*hs_scale_linear.shape[:-1], -1
)
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
num_tokens = hs_fp4.shape[0]
hidden_size = (
hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1]
)
symm_output = torch.empty(
num_tokens, hidden_size, dtype=hidden_states.dtype, device=hs_fp4.device
)
if use_routed_topk:
assert TopKOutputChecker.format_is_standard(topk_output)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
topk_output.topk_ids, topk_output.topk_weights
)
result = trtllm_fp4_block_scale_routed_moe(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=hs_fp4,
hidden_states_scale=hs_scale,
gemm1_weights=quant_info.gemm1_weights_fp4_shuffled,
gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn
),
gemm1_bias=None,
gemm1_alpha=None,
gemm1_beta=None,
gemm1_clamp_limit=None,
gemm2_weights=quant_info.gemm2_weights_fp4_shuffled,
gemm2_weights_scale=quant_info.gemm2_scales_fp4_shuffled.view(
torch.float8_e4m3fn
),
gemm2_bias=None,
output1_scale_scalar=quant_info.g1_scale_c,
output1_scale_gate_scalar=quant_info.g1_alphas,
output2_scale_scalar=quant_info.g2_alphas,
num_experts=quant_info.global_num_experts,
top_k=topk_output.topk_ids.shape[1],
n_group=0,
topk_group=0,
intermediate_size=quant_info.intermediate_size_per_partition,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=None,
routing_method_type=1, # Unused, but must be 1 to pass validation.
do_finalize=True,
tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]),
output=symm_output,
)[0]
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
router_logits = topk_output.router_logits
topk_config = topk_output.topk_config
routing_method_type = quant_info.routing_method_type
# DeepSeekV3 style routing requires float32 router logits
if routing_method_type == RoutingMethodType.DeepSeekV3:
@@ -653,23 +711,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
if topk_config.correction_bias is None
else topk_config.correction_bias.to(hidden_states.dtype)
)
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
num_tokens = hs_fp4.shape[0]
hidden_size = (
hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1]
)
symm_output = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
)
result = trtllm_fp4_block_scale_moe(
routing_logits=router_logits,
routing_bias=correction_bias,
hidden_states=hs_fp4,
hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).reshape(
*hs_scale_linear.shape[:-1], -1
),
hidden_states_scale=hs_scale,
gemm1_weights=quant_info.gemm1_weights_fp4_shuffled,
gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn
@@ -858,6 +904,13 @@ def fused_experts_none_to_flashinfer_trtllm_routed(
quant_info: MoeQuantInfo,
runner_config: MoeRunnerConfig,
) -> StandardCombineInput:
if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
return fused_experts_none_to_flashinfer_trtllm_fp4(
dispatch_output,
quant_info,
runner_config,
use_routed_topk=True,
)
if isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo):
return fused_experts_none_to_flashinfer_trtllm_fp8(
dispatch_output,
python/sglang/srt/layers/quantization/modelopt_quant.py
+5
View File
@@ -1534,6 +1534,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
self.enable_flashinfer_trtllm_moe = (
get_moe_runner_backend().is_flashinfer_trtllm()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
self._cache_permute_indices = {}
@@ -1904,6 +1905,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
self.runner = MoeRunner(
MoeRunnerBackend.FLASHINFER_TRTLLM, moe_runner_config
)
elif get_moe_runner_backend().is_flashinfer_trtllm_routed():
self.runner = MoeRunner(
MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, moe_runner_config
)
def apply(
self,