diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index fd1c826322c3..4b54d31e2154 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -887,15 +887,19 @@ def _make_all2all_kwargs( hidden: int, num_topk: int, use_fp8_dispatch: bool, + use_nvfp4_dispatch: bool = False, ) -> dict: + assert not (use_fp8_dispatch and use_nvfp4_dispatch), \ + "Cannot use both FP8 and NVFP4 dispatch simultaneously" return dict( - group=self._device_group - if self._device_group is not None + group=self._device_group if self._device_group is not None else self.cpu_group, num_max_tokens_per_rank=num_max_tokens_per_rank, hidden=hidden, num_topk=num_topk, use_fp8_dispatch=use_fp8_dispatch, + use_nvfp4_dispatch=use_nvfp4_dispatch, + use_fp8_sf=use_nvfp4_dispatch, allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE, prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP, allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION), diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 1351e87b5b51..b136d174f71f 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -213,12 +213,17 @@ def maybe_make_prepare_finalize( and quant_config.quant_dtype == current_platform.fp8_dtype() and quant_config.is_block_quantized ) + use_nvfp4_dispatch = ( + quant_config is not None + and quant_config.quant_dtype == "nvfp4" + ) all_to_all_args = dict( num_max_tokens_per_rank=moe.max_num_tokens, hidden=moe.hidden_dim, num_topk=moe.experts_per_token, num_experts=moe.num_experts, use_fp8_dispatch=use_fp8_dispatch, + use_nvfp4_dispatch=use_nvfp4_dispatch, ) handle = all2all_manager.get_handle(all_to_all_args) vllm_config = get_current_vllm_config() @@ -232,6 +237,7 @@ def maybe_make_prepare_finalize( num_experts=moe.num_experts, num_topk=moe.experts_per_token, use_fp8_dispatch=use_fp8_dispatch, + use_nvfp4_dispatch=use_nvfp4_dispatch, use_cudagraph=use_cudagraph, ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index 6495e1203e0a..59321bac6bd2 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -63,9 +63,12 @@ def __init__( num_experts: int, num_topk: int, use_fp8_dispatch: bool = False, + use_nvfp4_dispatch: bool = False, use_cudagraph: bool = False, ): super().__init__() + assert not (use_fp8_dispatch and use_nvfp4_dispatch), \ + "Cannot use both FP8 and NVFP4 dispatch simultaneously" self.buffer = buffer self.num_dispatchers_ = num_dispatchers self.dp_size = dp_size @@ -73,6 +76,7 @@ def __init__( self.num_experts = num_experts self.num_topk = num_topk self.use_fp8_dispatch = use_fp8_dispatch + self.use_nvfp4_dispatch = use_nvfp4_dispatch self.use_cudagraph = use_cudagraph # DBO microbatching: one handle slot per micro-batch. @@ -105,6 +109,13 @@ def _do_dispatch( quant_config: FusedMoEQuantConfig, defer_input_quant: bool, ) -> Callable: + if self.use_nvfp4_dispatch and token_scales is None: + tokens, token_scales = moe_kernel_quantize_input( + tokens, a1_scale, quant_config.quant_dtype, + False, quant_config.block_shape, False, + ) + token_scales = token_scales.view(torch.int32) + has_scales = token_scales is not None token_data = tokens @@ -218,7 +229,13 @@ def _receiver( device=expert_x.device, ) - if not quant_config.is_block_quantized and not defer_input_quant: + if self.use_nvfp4_dispatch and expert_x_scale is not None: + expert_x_scale = expert_x_scale.view(torch.uint8).view( + torch.float8_e4m3fn) + + if (not quant_config.is_block_quantized + and not defer_input_quant + and not self.use_nvfp4_dispatch): expert_x_scale = None if expert_x.numel() != 0: expert_x, expert_x_scale = moe_kernel_quantize_input(