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

[Perf] Eliminate the slice op for Flashinfer trtllm_fp4_block_scale_moe (#15731)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
This commit is contained in:
elvischenv
2026-02-26 22:23:40 -08:00
committed by GitHub
co-authored by Brayden Zhong
parent d566816d83
commit af3eccc9ab
1 changed files with 3 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/mxfp4.py
+3 -7
View File
@@ -751,13 +751,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
# When bf16 mode is enabled, we don't need to quantize the input, # When bf16 mode is enabled, we don't need to quantize the input,
# TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
# which can theoretically improve performance # which can theoretically improve performance
origin_hidden_states_dim = x.shape[-1]
if self.flashinfer_mxfp4_moe_precision == "bf16": if self.flashinfer_mxfp4_moe_precision == "bf16":
assert x.dtype == torch.bfloat16 assert x.dtype == torch.bfloat16
x_quant = x x_quant = x
x_scale = None x_scale = None
# May be fused later if this code branch is frequently needed # May be fused later if this code branch is frequently needed
origin_hidden_states_dim = x_quant.shape[-1]
if self.hidden_size != origin_hidden_states_dim: if self.hidden_size != origin_hidden_states_dim:
x_quant = torch.nn.functional.pad( x_quant = torch.nn.functional.pad(
x_quant, x_quant,
@@ -781,11 +781,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
get_tp_group(), disabled=not is_allocation_symmetric() get_tp_group(), disabled=not is_allocation_symmetric()
): ):
num_tokens = x_quant.shape[0] num_tokens = x_quant.shape[0]
hidden_size = ( hidden_size = origin_hidden_states_dim
x_quant.shape[-1] * 2
if x_quant.dtype == torch.uint8
else x_quant.shape[-1]
)
symm_output = torch.empty( symm_output = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device
) )
@@ -813,7 +809,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.intermediate_size_per_partition, # padded to multiple of 256 self.intermediate_size_per_partition, # padded to multiple of 256
layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset
layer.num_local_experts, # local num experts layer.num_local_experts, # local num experts
None, None, # routed_scaling_factor
1, # routing_method_type, renormalize 1, # routing_method_type, renormalize
True, # do finalize True, # do finalize
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),