Add NVFP4 dispatch support for DeepEP v2 ElasticBuffer

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),

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,
         )
 

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py

@@ -63,16 +63,20 @@
         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
         self.rank_expert_offset = rank_expert_offset
         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,10 +109,17 @@
         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
         if has_scales:
             token_data = (tokens, token_scales)
 
         # Decode: do_expand=False + do_cpu_sync=False (cudagraph-safe)
@@ -218,7 +229,13 @@
             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(