A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 04:45 UTC
Support ModelOpt MXFP8 non-gated MoE (#42958)
Signed-off-by: tbarnatan <tbarnatan@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
TomerBN-Nvidia committed
Jun 2, 2026 at 16:56 UTC
0cbc48c4f98873acb88ab220db4540ffe99cbc1e
1 file changed
+14
-5
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+14
-5
@@ -317,9 +317,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
317
from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
318
319
assert not apply_router_weight_on_input
320
- assert activation == MoEActivation.SILU
320
+ assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
321
+ activation_type = activation_to_flashinfer_int(activation)
322
assert self.topk <= global_num_experts
322
- assert self.topk <= 10
323
assert global_num_experts % 4 == 0
324
assert self.quant_config.block_shape in [[128, 128], [1, 32]]
325
# Kernel expects #experts <= #threads 512
@@ -333,13 +333,19 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
333
use_shuffled_weight = True
334
weight_layout = WeightLayout.MajorK
335
hidden_states_scale = a1q_scale
336
+ # FlashInfer expects None for non-grouped MXFP8 routing configs.
337
+ n_group = num_expert_group or None
338
+ selected_topk_group = topk_group or None
339
else:
340
+ assert self.topk <= 10
341
fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
342
use_shuffled_weight = True
343
weight_layout = WeightLayout.BlockMajorK
344
hidden_states_scale = a1q_scale.t().contiguous()
345
+ n_group = num_expert_group or 0
346
+ selected_topk_group = topk_group or 0
347
342
- return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(
348
+ kwargs = dict(
349
routing_logits=router_logits,
350
routing_bias=e_score_correction_bias,
351
hidden_states=hidden_states,
@@ -350,8 +356,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
356
gemm2_weights_scale=self.quant_config.w2_scale,
357
num_experts=global_num_experts,
358
top_k=self.topk,
353
- n_group=(num_expert_group or 0),
354
- topk_group=(topk_group or 0),
359
+ n_group=n_group,
360
+ topk_group=selected_topk_group,
361
intermediate_size=self.intermediate_size_per_partition,
362
local_expert_offset=self.ep_rank * self.local_num_experts,
363
local_num_experts=self.local_num_experts,
@@ -361,6 +367,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
367
weight_layout=weight_layout,
368
fp8_quantization_type=fp8_quant_type,
369
)
370
+ if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
371
+ kwargs["activation_type"] = activation_type
372
+ return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
373
374
def _apply_per_tensor(
375
self,