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

Fix fallback to default tactic (flashinfer autotuner) with trtllm_fp4_block_scale_moe (#19189)

Browse Source
This commit is contained in:
Shu Wang
2026-03-06 15:15:04 -08:00
committed by GitHub
parent e89069ee64
commit 61de303f0a
3 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+3 -2
View File
@@ -1226,12 +1226,13 @@ class FlashInferFP4MoE(FusedMoE):
symm_output = torch.empty( symm_output = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
) )
result = trtllm_fp4_block_scale_moe( result = trtllm_fp4_block_scale_moe(
routing_logits=router_logits, routing_logits=router_logits,
routing_bias=correction_bias, routing_bias=correction_bias,
hidden_states=hs_fp4, hidden_states=hs_fp4,
hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn), hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).reshape(
*hs_scale_linear.shape[:-1], -1
),
gemm1_weights=self.gemm1_weights_fp4_shuffled.data, gemm1_weights=self.gemm1_weights_fp4_shuffled.data,
gemm1_weights_scale=self.gemm1_scales_fp4_shuffled.data.view( gemm1_weights_scale=self.gemm1_scales_fp4_shuffled.data.view(
torch.float8_e4m3fn torch.float8_e4m3fn
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+3 -1
View File
@@ -442,7 +442,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
routing_logits=router_logits, routing_logits=router_logits,
routing_bias=correction_bias, routing_bias=correction_bias,
hidden_states=hs_fp4, hidden_states=hs_fp4,
hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).flatten(), hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).reshape(
*hs_scale_linear.shape[:-1], -1
),
gemm1_weights=quant_info.gemm1_weights_fp4_shuffled, gemm1_weights=quant_info.gemm1_weights_fp4_shuffled,
gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view( gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn torch.float8_e4m3fn
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py
+3 -1
View File
@@ -331,7 +331,9 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme):
False, # is_sf_swizzled_layout False, # is_sf_swizzled_layout
) )
hs_fp4 = hs_fp4_bytes.reshape(x.shape[0], x.shape[1] // 2) hs_fp4 = hs_fp4_bytes.reshape(x.shape[0], x.shape[1] // 2)
hs_scale = hs_sf_bytes.view(torch.float8_e4m3fn).reshape(-1) hs_scale = hs_sf_bytes.view(torch.float8_e4m3fn).reshape(
*hs_sf_bytes.shape[:-1], -1
)
correction_bias = ( correction_bias = (
None None