A high-throughput and memory-efficient inference and serving engine for LLMs
nixl_ep: Skip post-receive quantization for NVFP4 (#45606)
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Itay Alroy committed
Jun 16, 2026 at 01:42 UTC
ab8b0fe338d02df87b0844ead99b0a0f2cfb638c
2 files changed
+8
-10
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
+6
-3
@@ -49,6 +49,9 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
49
"Only nvfp4 quantization are currently supported."
50
)
51
self.out_dtype = moe_config.in_dtype
52
+ self.use_deep_ep_ll_nvfp4_dispatch = (
53
+ envs.VLLM_DEEPEPLL_NVFP4_DISPATCH and moe_config.use_deepep_ll_kernels
54
+ )
55
56
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
57
layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -123,7 +126,7 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
126
127
# We use global_num_experts due to how moe_align_block_size handles
128
# expert_maps.
126
- K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
129
+ K_dim = K * 2 if self.use_deep_ep_ll_nvfp4_dispatch else K
130
output_shape = (local_num_experts, M, K_dim)
131
workspace2 = (local_num_experts, M, N)
132
workspace1 = output_shape
@@ -161,11 +164,11 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
164
assert self.w2_scale.ndim == 3
165
166
input_global_scale = (
164
- None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
167
+ None if self.use_deep_ep_ll_nvfp4_dispatch else self.a1_gscale
168
)
169
flashinfer_hidden_states = (
170
(hidden_states, a1q_scale)
168
- if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
171
+ if self.use_deep_ep_ll_nvfp4_dispatch
172
else hidden_states
173
)
174
flashinfer_cutedsl_moe_masked(
vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
+2
-7
@@ -6,7 +6,6 @@ import nixl_ep
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
-from vllm.config import get_current_vllm_config
9
from vllm.distributed import get_ep_group
10
from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager
11
from vllm.logger import init_logger
@@ -192,13 +191,9 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
191
x = x.view((-1, hidden_dim))
192
q_dtype = quant_config.quant_dtype
193
195
- moe_backend = get_current_vllm_config().kernel_config.moe_backend
196
- if moe_backend == "flashinfer_cutedsl":
197
- logger.info_once(
198
- "Skip quantization when using FlashInfer CUTEDSL "
199
- "(--moe-backend flashinfer_cutedsl) for ModelOptNvFp4FusedMoE."
200
- )
194
+ if q_dtype == "nvfp4":
195
q_dtype = None
196
+ logger.debug_once("Using NIXL EP bfloat16 dispatch for NVFP4 MoE.")
197
198
x, x_scales = moe_kernel_quantize_input(
199
x,