obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-07 22:38:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support ModelOpt MXFP8 non-gated MoE (#42958)

Browse Source 
Signed-off-by: tbarnatan <[email protected]>
Co-authored-by: Michael Goin <[email protected]>
This commit is contained in:
TomerBN-Nvidia
2026-06-02 13:56:03 +00:00
committed by GitHub
co-authored by Michael Goin
parent 2fd0e52252
commit 0cbc48c4f9
1 changed files with 14 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+14 -5
View File
@@ -317,9 +317,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
assert not apply_router_weight_on_input
assert activation == MoEActivation.SILU
assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
activation_type = activation_to_flashinfer_int(activation)
assert self.topk <= global_num_experts
assert self.topk <= 10
assert global_num_experts % 4 == 0
assert self.quant_config.block_shape in [[128, 128], [1, 32]]
# Kernel expects #experts <= #threads 512
@@ -333,13 +333,19 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
use_shuffled_weight = True
weight_layout = WeightLayout.MajorK
hidden_states_scale = a1q_scale
# FlashInfer expects None for non-grouped MXFP8 routing configs.
n_group = num_expert_group or None
selected_topk_group = topk_group or None
else:
assert self.topk <= 10
fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
use_shuffled_weight = True
weight_layout = WeightLayout.BlockMajorK
hidden_states_scale = a1q_scale.t().contiguous()
n_group = num_expert_group or 0
selected_topk_group = topk_group or 0
return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(
kwargs = dict(
routing_logits=router_logits,
routing_bias=e_score_correction_bias,
hidden_states=hidden_states,
@@ -350,8 +356,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
gemm2_weights_scale=self.quant_config.w2_scale,
num_experts=global_num_experts,
top_k=self.topk,
n_group=(num_expert_group or 0),
topk_group=(topk_group or 0),
n_group=n_group,
topk_group=selected_topk_group,
intermediate_size=self.intermediate_size_per_partition,
local_expert_offset=self.ep_rank * self.local_num_experts,
local_num_experts=self.local_num_experts,
@@ -361,6 +367,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
weight_layout=weight_layout,
fp8_quantization_type=fp8_quant_type,
)
if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
kwargs["activation_type"] = activation_type
return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
def _apply_per_tensor(
self,