From ab8b0fe338d02df87b0844ead99b0a0f2cfb638c Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Tue, 16 Jun 2026 01:42:05 +0300 Subject: [PATCH] nixl_ep: Skip post-receive quantization for NVFP4 (#45606) Signed-off-by: Itay Alroy --- .../fused_moe/experts/flashinfer_cutedsl_batched_moe.py | 9 ++++++--- .../layers/fused_moe/prepare_finalize/nixl_ep.py | 9 ++------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py index 253d1dae71..d269c6f109 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py @@ -49,6 +49,9 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular): "Only nvfp4 quantization are currently supported." ) self.out_dtype = moe_config.in_dtype + self.use_deep_ep_ll_nvfp4_dispatch = ( + envs.VLLM_DEEPEPLL_NVFP4_DISPATCH and moe_config.use_deepep_ll_kernels + ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -123,7 +126,7 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular): # We use global_num_experts due to how moe_align_block_size handles # expert_maps. - K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K + K_dim = K * 2 if self.use_deep_ep_ll_nvfp4_dispatch else K output_shape = (local_num_experts, M, K_dim) workspace2 = (local_num_experts, M, N) workspace1 = output_shape @@ -161,11 +164,11 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular): assert self.w2_scale.ndim == 3 input_global_scale = ( - None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale + None if self.use_deep_ep_ll_nvfp4_dispatch else self.a1_gscale ) flashinfer_hidden_states = ( (hidden_states, a1q_scale) - if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH + if self.use_deep_ep_ll_nvfp4_dispatch else hidden_states ) flashinfer_cutedsl_moe_masked( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index 850f54df4b..ce44cd6a3f 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -6,7 +6,6 @@ import nixl_ep import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.config import get_current_vllm_config from vllm.distributed import get_ep_group from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager from vllm.logger import init_logger @@ -192,13 +191,9 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): x = x.view((-1, hidden_dim)) q_dtype = quant_config.quant_dtype - moe_backend = get_current_vllm_config().kernel_config.moe_backend - if moe_backend == "flashinfer_cutedsl": - logger.info_once( - "Skip quantization when using FlashInfer CUTEDSL " - "(--moe-backend flashinfer_cutedsl) for ModelOptNvFp4FusedMoE." - ) + if q_dtype == "nvfp4": q_dtype = None + logger.debug_once("Using NIXL EP bfloat16 dispatch for NVFP4 MoE.") x, x_scales = moe_kernel_quantize_input( x,