chenchenghao/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 17:15:04 -06: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

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

@@ -1226,12 +1226,13 @@ class FlashInferFP4MoE(FusedMoE):
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),
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_scale=self.gemm1_scales_fp4_shuffled.data.view(
torch.float8_e4m3fn

4
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
View File

@@ -442,7 +442,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
routing_logits=router_logits,
routing_bias=correction_bias,
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_scale=quant_info.gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn

4
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py
View File

@@ -331,7 +331,9 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme):
False, # is_sf_swizzled_layout
)
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 = (
None