From c296851a7d173fa89d2eefbca0243be42ae9b5e0 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:18:14 +0800 Subject: [PATCH] [MoE] Refine FlashInfer one-sided All2All integration (#51924) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/moe_kernel_features.md | 2 +- tests/distributed/test_mnnvl_alltoall.py | 30 ++++---- tests/kernels/moe/test_moe_layer.py | 2 +- vllm/config/parallel.py | 1 + .../device_communicators/all2all.py | 12 +-- .../layers/fused_moe/all2all_utils.py | 74 +++++++++++++------ .../fused_moe/experts/trtllm_fp8_moe.py | 28 ++++++- .../flashinfer_nvlink_one_sided.py | 48 ++++++------ 8 files changed, 129 insertions(+), 68 deletions(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index bdf4075102..0b1a9151ec 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -36,7 +36,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | -| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index f9b76ea1b0..636f7cb525 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -342,6 +342,8 @@ def _one_sided_lifecycle_worker(rank, world_size): top_k=2, num_experts=world_size * 8, hidden_size=4096, + x_bytes_per_token=4096 * 2, + x_sf_bytes_per_token=0, ) # Initialize @@ -409,12 +411,12 @@ def _one_sided_workspace_grow_worker(rank, world_size): hidden_size=4096, ) nvfp4_kwargs = dict( - dispatch_dtype_bytes_per_elem=0, - dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16, + x_bytes_per_token=base_kwargs["hidden_size"] // 2, + x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16, ) bf16_kwargs = dict( - dispatch_dtype_bytes_per_elem=2, - dispatch_scale_bytes_per_token=0, + x_bytes_per_token=base_kwargs["hidden_size"] * 2, + x_sf_bytes_per_token=0, ) # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16). @@ -780,11 +782,12 @@ def _one_sided_data_worker(rank, world_size): top_k=experts_per_token, num_experts=num_experts, hidden_size=hidden_size, - # Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes + x_bytes_per_token=hidden_size // 2, + # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes # per token) that is dispatched alongside the nvfp4 hidden states. # Without this the dispatch region is under-reserved and the combine # payload overflows the per-rank workspace. - dispatch_scale_bytes_per_token=hidden_size // 16, + x_sf_bytes_per_token=hidden_size // 16, ) assert manager.initialized assert manager.moe_alltoall is not None @@ -798,17 +801,17 @@ def _one_sided_data_worker(rank, world_size): # Create test data with raw tensors matching the nvfp4 payload # sizes the workspace was allocated for: - # a1q: (tokens, hidden_size // 2) — nvfp4 hidden states - # a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors + # x: (tokens, hidden_size // 2) — nvfp4 hidden states + # x_sf: (tokens, hidden_size // 16) — fp8 scaling factors torch.manual_seed(rank + 42) - a1q = torch.randint( + x = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 2), device=device, dtype=torch.uint8, ) - a1q_scale = torch.randint( + x_sf = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 16), @@ -830,15 +833,16 @@ def _one_sided_data_worker(rank, world_size): ) # --- One-sided dispatch --- - payloads = [a1q, a1q_scale, topk_ids, topk_weights] + payloads = [x, x_sf, topk_ids, topk_weights] recv_payloads = manager.moe_alltoall.dispatch( token_selected_experts=topk_ids, input_payloads=payloads, runtime_max_tokens_per_rank=runtime_max_tokens, ) assert len(recv_payloads) == 4 - recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads - assert recv_a1q.numel() > 0 + recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads + assert recv_x.numel() > 0 + assert recv_x_sf.numel() > 0 assert recv_ids.numel() > 0 # --- Round-trip exact verification --- diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 1942ce96bc..74554944a3 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -130,7 +130,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { "mori_high_throughput": {None, "fp8", "modelopt_fp8"}, "mori_low_latency": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 - "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501 + "flashinfer_nvlink_one_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 "nixl_ep": {None, "fp8_blocked", "modelopt_fp8"}, diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 83c4e6fe53..eb346ed4a3 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -677,6 +677,7 @@ class ParallelConfig: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", + "flashinfer_nvlink_one_sided", "mori_high_throughput", "mori_low_latency", "nixl_ep", diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5abe7568a2..657ad9fd0e 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -734,17 +734,13 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): top_k: int, num_experts: int, hidden_size: int, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, ): """Initialize (or grow) the MoeAlltoAll workspace.""" - if dispatch_dtype_bytes_per_elem == 0: - hidden_bytes = hidden_size // 2 - else: - hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem total_dispatch_payload_size_per_token = ( - hidden_bytes - + dispatch_scale_bytes_per_token + x_bytes_per_token + + x_sf_bytes_per_token + top_k * 4 # int32 topks ids + top_k * 4 # float32 topk weights ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 29905927c1..7861345453 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass from typing import Any import torch @@ -37,6 +38,52 @@ from vllm.utils.import_utils import ( has_nixl_ep, ) + +@dataclass(frozen=True) +class FlashInferOneSidedDispatchLayout: + x_bytes_per_token: int + x_sf_bytes_per_token: int + + +def flashinfer_one_sided_dispatch_layout( + hidden_dim: int, quant_config: FusedMoEQuantConfig +) -> FlashInferOneSidedDispatchLayout: + """Return the one-sided activation payload layout.""" + if quant_config.quant_dtype is None: + return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0) + if quant_config.quant_dtype == "nvfp4": + scale_elems = hidden_dim // 16 + return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems) + if quant_config.quant_dtype == "mxfp8": + align = quant_config.mx_alignment + padded_k = ( + ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim + ) + scale_elems = padded_k // 32 + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems) + if ( + quant_config.use_fp8_w8a8 + and quant_config.quant_dtype == current_platform.fp8_dtype() + and quant_config.block_shape == [128, 128] + ): + if hidden_dim % 128 != 0: + raise NotImplementedError( + "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch " + f"requires hidden_dim divisible by 128; got {hidden_dim}" + ) + scale_elems = hidden_dim // 128 + scale_bytes = scale_elems * torch.float32.itemsize + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes) + raise NotImplementedError( + "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " + "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 " + "(quant_dtype=None) today; got " + f"quant_dtype={quant_config.quant_dtype!r}, " + f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, " + f"block_shape={quant_config.block_shape!r}" + ) + + logger = init_logger(__name__) if current_platform.is_cuda_alike(): @@ -282,34 +329,17 @@ def maybe_make_prepare_finalize( max_num_tokens = ( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) - if quant_config.quant_dtype is None: - dispatch_dtype_bytes_per_elem = 2 - dispatch_scale_bytes_per_token = 0 - elif quant_config.quant_dtype == "nvfp4": - dispatch_dtype_bytes_per_elem = 0 - dispatch_scale_bytes_per_token = moe.hidden_dim // 16 - elif quant_config.quant_dtype == "mxfp8": - dispatch_dtype_bytes_per_elem = 1 - align = quant_config.mx_alignment - if align > 0: - padded_k = ((moe.hidden_dim + align - 1) // align) * align - else: - padded_k = moe.hidden_dim - dispatch_scale_bytes_per_token = padded_k // 32 - else: - raise NotImplementedError( - "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " - "and bf16 (quant_dtype=None) today; got " - f"quant_dtype={quant_config.quant_dtype!r}" - ) + dispatch_layout = flashinfer_one_sided_dispatch_layout( + moe.hidden_dim, quant_config + ) prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize( max_num_tokens=max_num_tokens, top_k=moe.experts_per_token, num_experts=moe.num_experts, hidden_size=moe.hidden_dim, num_dispatchers=all2all_manager.world_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=dispatch_layout.x_bytes_per_token, + x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token, ) elif moe.use_ag_rs_all2all_kernels and allow_new_interface: diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index d32877eb9a..00ab06fe95 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -36,6 +36,28 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe logger = init_logger(__name__) +def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tensor: + """Validate DeepSeek Blockwise FP8 tensors and return TRTLLM layout.""" + if x.dtype != current_platform.fp8_dtype(): + raise ValueError( + f"DeepSeekFp8 activations must use the platform E4M3 dtype; got {x.dtype}" + ) + if x.ndim != 2 or x.shape[1] % 128 != 0: + raise ValueError( + "DeepSeekFp8 activations must be [M,K] with K divisible by 128; " + f"got {tuple(x.shape)}" + ) + expected_shape = (x.shape[0], x.shape[1] // 128) + if x_sf.dtype != torch.float32 or tuple(x_sf.shape) != expected_shape: + raise ValueError( + "DeepSeekFp8 activation scales must be FP32 [M,K/128]; " + f"expected {expected_shape}, got dtype={x_sf.dtype}, " + f"shape={tuple(x_sf.shape)}" + ) + # FlashInfer TRTLLM-gen consumes [K/128,M] for DeepSeekFp8/BlockMajorK. + return x_sf.t().contiguous() + + class TrtLlmFp8ExpertsBase: """ Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic @@ -239,6 +261,10 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): # Pack topk ids and weights into format expected by the kernel. packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + if a1q_scale is None: + raise RuntimeError( + "TRT-LLM FP8 experts require precomputed activation scales" + ) assert a1q_scale is not None is_mxfp8 = self.quant_config.block_shape == [1, 32] @@ -251,7 +277,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 use_shuffled_weight = True weight_layout = WeightLayout.BlockMajorK - hidden_states_scale = a1q_scale.t().contiguous() + hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=packed_topk_ids, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index 74341e7681..0b57afb606 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -30,9 +30,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo top_k: int, num_experts: int, hidden_size: int, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, num_dispatchers: int = 1, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, ): super().__init__() self.max_num_tokens = max_num_tokens @@ -40,7 +40,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo self.num_experts = num_experts self.hidden_size = hidden_size self.num_dispatchers_ = num_dispatchers - self.scale_elems_per_token = dispatch_scale_bytes_per_token device_communicator = get_ep_group().device_communicator assert device_communicator is not None @@ -52,8 +51,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo top_k=self.top_k, num_experts=self.num_experts, hidden_size=self.hidden_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=x_bytes_per_token, + x_sf_bytes_per_token=x_sf_bytes_per_token, ) @property @@ -98,9 +97,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo ) if defer_input_quant: - a1q, a1q_scale = a1, None + dispatch_x, dispatch_x_sf = a1, None else: - a1q, a1q_scale = moe_kernel_quantize_input( + dispatch_x, dispatch_x_sf = moe_kernel_quantize_input( a1, quant_config.a1_gscale, quant_config.quant_dtype, @@ -110,10 +109,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo mx_alignment=quant_config.mx_alignment, ) - payloads = [] - payloads.append(a1q) - if a1q_scale is not None: - payloads.append(a1q_scale) + payloads = [dispatch_x] + if dispatch_x_sf is not None: + payloads.append(dispatch_x_sf) topk_ids_payload_index = len(payloads) payloads.append(topk_ids) payloads.append(topk_weights) @@ -126,23 +124,29 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo invalid_token_expert_id=-1, # Follow TRTLLM Pattern expert_id_payload_index=topk_ids_payload_index, ) - if a1q_scale is not None: - a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads + if dispatch_x_sf is not None: + recv_x, recv_x_sf, topk_ids_recv, topk_weights_recv = recv_payloads + x_sf_width = recv_x_sf.shape[-1] # Apply scale interleaving only for CUTLASS (not TRT-LLM) if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled: - a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1]) - a1q_scale_recv = a1q_scale_recv.view(torch.uint8) - a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv) - assert self.scale_elems_per_token > 0 - a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) + recv_x_sf = recv_x_sf.view(torch.uint8) + recv_x_sf = nvfp4_block_scale_interleave(recv_x_sf) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) else: - a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads - a1q_scale_recv = None - a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1]) + recv_x, topk_ids_recv, topk_weights_recv = recv_payloads + recv_x_sf = None + recv_x = recv_x.view(-1, recv_x.shape[-1]) topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1]) topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1]) - return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv + return ( + recv_x, + recv_x_sf, + None, + topk_ids_recv, + topk_weights_recv, + ) def finalize( self,