[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>
Bo Li committed
Aug 18, 2026 at 08:18 UTC
c296851a7d173fa89d2eefbca0243be42ae9b5e0
8 files changed
+129
-68
docs/design/moe_kernel_features.md
+1
-1
@@ -36,7 +36,7 @@ th {
36
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
37
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
38
| 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] |
39
-| 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] |
39
+| 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] |
40
41
!!! info "Table key"
42
1. All types: mxfp4, nvfp4, int4, int8, fp8
tests/distributed/test_mnnvl_alltoall.py
+17
-13
@@ -342,6 +342,8 @@ def _one_sided_lifecycle_worker(rank, world_size):
342
top_k=2,
343
num_experts=world_size * 8,
344
hidden_size=4096,
345
+ x_bytes_per_token=4096 * 2,
346
+ x_sf_bytes_per_token=0,
347
)
348
349
# Initialize
@@ -409,12 +411,12 @@ def _one_sided_workspace_grow_worker(rank, world_size):
411
hidden_size=4096,
412
)
413
nvfp4_kwargs = dict(
412
- dispatch_dtype_bytes_per_elem=0,
413
- dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16,
414
+ x_bytes_per_token=base_kwargs["hidden_size"] // 2,
415
+ x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16,
416
)
417
bf16_kwargs = dict(
416
- dispatch_dtype_bytes_per_elem=2,
417
- dispatch_scale_bytes_per_token=0,
418
+ x_bytes_per_token=base_kwargs["hidden_size"] * 2,
419
+ x_sf_bytes_per_token=0,
420
)
421
422
# First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16).
@@ -780,11 +782,12 @@ def _one_sided_data_worker(rank, world_size):
782
top_k=experts_per_token,
783
num_experts=num_experts,
784
hidden_size=hidden_size,
783
- # Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes
785
+ x_bytes_per_token=hidden_size // 2,
786
+ # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes
787
# per token) that is dispatched alongside the nvfp4 hidden states.
788
# Without this the dispatch region is under-reserved and the combine
789
# payload overflows the per-rank workspace.
787
- dispatch_scale_bytes_per_token=hidden_size // 16,
790
+ x_sf_bytes_per_token=hidden_size // 16,
791
)
792
assert manager.initialized
793
assert manager.moe_alltoall is not None
@@ -798,17 +801,17 @@ def _one_sided_data_worker(rank, world_size):
801
802
# Create test data with raw tensors matching the nvfp4 payload
803
# sizes the workspace was allocated for:
801
- # a1q: (tokens, hidden_size // 2) — nvfp4 hidden states
802
- # a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors
804
+ # x: (tokens, hidden_size // 2) — nvfp4 hidden states
805
+ # x_sf: (tokens, hidden_size // 16) — fp8 scaling factors
806
torch.manual_seed(rank + 42)
804
- a1q = torch.randint(
807
+ x = torch.randint(
808
0,
809
256,
810
(tokens_per_rank, hidden_size // 2),
811
device=device,
812
dtype=torch.uint8,
813
)
811
- a1q_scale = torch.randint(
814
+ x_sf = torch.randint(
815
0,
816
256,
817
(tokens_per_rank, hidden_size // 16),
@@ -830,15 +833,16 @@ def _one_sided_data_worker(rank, world_size):
833
)
834
835
# --- One-sided dispatch ---
833
- payloads = [a1q, a1q_scale, topk_ids, topk_weights]
836
+ payloads = [x, x_sf, topk_ids, topk_weights]
837
recv_payloads = manager.moe_alltoall.dispatch(
838
token_selected_experts=topk_ids,
839
input_payloads=payloads,
840
runtime_max_tokens_per_rank=runtime_max_tokens,
841
)
842
assert len(recv_payloads) == 4
840
- recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads
841
- assert recv_a1q.numel() > 0
843
+ recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads
844
+ assert recv_x.numel() > 0
845
+ assert recv_x_sf.numel() > 0
846
assert recv_ids.numel() > 0
847
848
# --- Round-trip exact verification ---
tests/kernels/moe/test_moe_layer.py
+1
-1
@@ -130,7 +130,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
130
"mori_high_throughput": {None, "fp8", "modelopt_fp8"},
131
"mori_low_latency": {None, "fp8", "modelopt_fp8"},
132
"flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
133
- "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501
133
+ "flashinfer_nvlink_one_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
134
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
135
"deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
136
"nixl_ep": {None, "fp8_blocked", "modelopt_fp8"},
vllm/config/parallel.py
+1
@@ -677,6 +677,7 @@ class ParallelConfig:
677
"allgather_reducescatter",
678
"deepep_high_throughput",
679
"deepep_low_latency",
680
+ "flashinfer_nvlink_one_sided",
681
"mori_high_throughput",
682
"mori_low_latency",
683
"nixl_ep",
vllm/distributed/device_communicators/all2all.py
+4
-8
@@ -734,17 +734,13 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
734
top_k: int,
735
num_experts: int,
736
hidden_size: int,
737
- dispatch_dtype_bytes_per_elem: int = 0,
738
- dispatch_scale_bytes_per_token: int = 0,
737
+ x_bytes_per_token: int,
738
+ x_sf_bytes_per_token: int,
739
):
740
"""Initialize (or grow) the MoeAlltoAll workspace."""
741
- if dispatch_dtype_bytes_per_elem == 0:
742
- hidden_bytes = hidden_size // 2
743
- else:
744
- hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
741
total_dispatch_payload_size_per_token = (
746
- hidden_bytes
747
- + dispatch_scale_bytes_per_token
742
+ x_bytes_per_token
743
+ + x_sf_bytes_per_token
744
+ top_k * 4 # int32 topks ids
745
+ top_k * 4 # float32 topk weights
746
)
vllm/model_executor/layers/fused_moe/all2all_utils.py
+52
-22
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass
5
from typing import Any
6
7
import torch
@@ -37,6 +38,52 @@ from vllm.utils.import_utils import (
38
has_nixl_ep,
39
)
40
41
+
42
+@dataclass(frozen=True)
43
+class FlashInferOneSidedDispatchLayout:
44
+ x_bytes_per_token: int
45
+ x_sf_bytes_per_token: int
46
+
47
+
48
+def flashinfer_one_sided_dispatch_layout(
49
+ hidden_dim: int, quant_config: FusedMoEQuantConfig
50
+) -> FlashInferOneSidedDispatchLayout:
51
+ """Return the one-sided activation payload layout."""
52
+ if quant_config.quant_dtype is None:
53
+ return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)
54
+ if quant_config.quant_dtype == "nvfp4":
55
+ scale_elems = hidden_dim // 16
56
+ return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems)
57
+ if quant_config.quant_dtype == "mxfp8":
58
+ align = quant_config.mx_alignment
59
+ padded_k = (
60
+ ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim
61
+ )
62
+ scale_elems = padded_k // 32
63
+ return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems)
64
+ if (
65
+ quant_config.use_fp8_w8a8
66
+ and quant_config.quant_dtype == current_platform.fp8_dtype()
67
+ and quant_config.block_shape == [128, 128]
68
+ ):
69
+ if hidden_dim % 128 != 0:
70
+ raise NotImplementedError(
71
+ "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch "
72
+ f"requires hidden_dim divisible by 128; got {hidden_dim}"
73
+ )
74
+ scale_elems = hidden_dim // 128
75
+ scale_bytes = scale_elems * torch.float32.itemsize
76
+ return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes)
77
+ raise NotImplementedError(
78
+ "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
79
+ "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 "
80
+ "(quant_dtype=None) today; got "
81
+ f"quant_dtype={quant_config.quant_dtype!r}, "
82
+ f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, "
83
+ f"block_shape={quant_config.block_shape!r}"
84
+ )
85
+
86
+
87
logger = init_logger(__name__)
88
89
if current_platform.is_cuda_alike():
@@ -282,34 +329,17 @@ def maybe_make_prepare_finalize(
329
max_num_tokens = (
330
get_current_vllm_config().scheduler_config.max_num_batched_tokens
331
)
285
- if quant_config.quant_dtype is None:
286
- dispatch_dtype_bytes_per_elem = 2
287
- dispatch_scale_bytes_per_token = 0
288
- elif quant_config.quant_dtype == "nvfp4":
289
- dispatch_dtype_bytes_per_elem = 0
290
- dispatch_scale_bytes_per_token = moe.hidden_dim // 16
291
- elif quant_config.quant_dtype == "mxfp8":
292
- dispatch_dtype_bytes_per_elem = 1
293
- align = quant_config.mx_alignment
294
- if align > 0:
295
- padded_k = ((moe.hidden_dim + align - 1) // align) * align
296
- else:
297
- padded_k = moe.hidden_dim
298
- dispatch_scale_bytes_per_token = padded_k // 32
299
- else:
300
- raise NotImplementedError(
301
- "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
302
- "and bf16 (quant_dtype=None) today; got "
303
- f"quant_dtype={quant_config.quant_dtype!r}"
304
- )
332
+ dispatch_layout = flashinfer_one_sided_dispatch_layout(
333
+ moe.hidden_dim, quant_config
334
+ )
335
prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
336
max_num_tokens=max_num_tokens,
337
top_k=moe.experts_per_token,
338
num_experts=moe.num_experts,
339
hidden_size=moe.hidden_dim,
340
num_dispatchers=all2all_manager.world_size,
311
- dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
312
- dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
341
+ x_bytes_per_token=dispatch_layout.x_bytes_per_token,
342
+ x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token,
343
)
344
345
elif moe.use_ag_rs_all2all_kernels and allow_new_interface:
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+27
-1
@@ -36,6 +36,28 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
36
logger = init_logger(__name__)
37
38
39
+def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tensor:
40
+ """Validate DeepSeek Blockwise FP8 tensors and return TRTLLM layout."""
41
+ if x.dtype != current_platform.fp8_dtype():
42
+ raise ValueError(
43
+ f"DeepSeekFp8 activations must use the platform E4M3 dtype; got {x.dtype}"
44
+ )
45
+ if x.ndim != 2 or x.shape[1] % 128 != 0:
46
+ raise ValueError(
47
+ "DeepSeekFp8 activations must be [M,K] with K divisible by 128; "
48
+ f"got {tuple(x.shape)}"
49
+ )
50
+ expected_shape = (x.shape[0], x.shape[1] // 128)
51
+ if x_sf.dtype != torch.float32 or tuple(x_sf.shape) != expected_shape:
52
+ raise ValueError(
53
+ "DeepSeekFp8 activation scales must be FP32 [M,K/128]; "
54
+ f"expected {expected_shape}, got dtype={x_sf.dtype}, "
55
+ f"shape={tuple(x_sf.shape)}"
56
+ )
57
+ # FlashInfer TRTLLM-gen consumes [K/128,M] for DeepSeekFp8/BlockMajorK.
58
+ return x_sf.t().contiguous()
59
+
60
+
61
class TrtLlmFp8ExpertsBase:
62
"""
63
Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic
@@ -239,6 +261,10 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
261
# Pack topk ids and weights into format expected by the kernel.
262
packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
263
264
+ if a1q_scale is None:
265
+ raise RuntimeError(
266
+ "TRT-LLM FP8 experts require precomputed activation scales"
267
+ )
268
assert a1q_scale is not None
269
270
is_mxfp8 = self.quant_config.block_shape == [1, 32]
@@ -251,7 +277,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
277
fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
278
use_shuffled_weight = True
279
weight_layout = WeightLayout.BlockMajorK
254
- hidden_states_scale = a1q_scale.t().contiguous()
280
+ hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale)
281
282
flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
283
topk_ids=packed_topk_ids,
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+26
-22
@@ -30,9 +30,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
30
top_k: int,
31
num_experts: int,
32
hidden_size: int,
33
+ x_bytes_per_token: int,
34
+ x_sf_bytes_per_token: int,
35
num_dispatchers: int = 1,
34
- dispatch_dtype_bytes_per_elem: int = 0,
35
- dispatch_scale_bytes_per_token: int = 0,
36
):
37
super().__init__()
38
self.max_num_tokens = max_num_tokens
@@ -40,7 +40,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
40
self.num_experts = num_experts
41
self.hidden_size = hidden_size
42
self.num_dispatchers_ = num_dispatchers
43
- self.scale_elems_per_token = dispatch_scale_bytes_per_token
43
44
device_communicator = get_ep_group().device_communicator
45
assert device_communicator is not None
@@ -52,8 +51,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
51
top_k=self.top_k,
52
num_experts=self.num_experts,
53
hidden_size=self.hidden_size,
55
- dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
56
- dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
54
+ x_bytes_per_token=x_bytes_per_token,
55
+ x_sf_bytes_per_token=x_sf_bytes_per_token,
56
)
57
58
@property
@@ -98,9 +97,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
97
)
98
99
if defer_input_quant:
101
- a1q, a1q_scale = a1, None
100
+ dispatch_x, dispatch_x_sf = a1, None
101
else:
103
- a1q, a1q_scale = moe_kernel_quantize_input(
102
+ dispatch_x, dispatch_x_sf = moe_kernel_quantize_input(
103
a1,
104
quant_config.a1_gscale,
105
quant_config.quant_dtype,
@@ -110,10 +109,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
109
mx_alignment=quant_config.mx_alignment,
110
)
111
113
- payloads = []
114
- payloads.append(a1q)
115
- if a1q_scale is not None:
116
- payloads.append(a1q_scale)
112
+ payloads = [dispatch_x]
113
+ if dispatch_x_sf is not None:
114
+ payloads.append(dispatch_x_sf)
115
topk_ids_payload_index = len(payloads)
116
payloads.append(topk_ids)
117
payloads.append(topk_weights)
@@ -126,23 +124,29 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
124
invalid_token_expert_id=-1, # Follow TRTLLM Pattern
125
expert_id_payload_index=topk_ids_payload_index,
126
)
129
- if a1q_scale is not None:
130
- a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
127
+ if dispatch_x_sf is not None:
128
+ recv_x, recv_x_sf, topk_ids_recv, topk_weights_recv = recv_payloads
129
+ x_sf_width = recv_x_sf.shape[-1]
130
# Apply scale interleaving only for CUTLASS (not TRT-LLM)
131
if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:
133
- a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
134
- a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
135
- a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
136
- assert self.scale_elems_per_token > 0
137
- a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token)
132
+ recv_x_sf = recv_x_sf.view(-1, x_sf_width)
133
+ recv_x_sf = recv_x_sf.view(torch.uint8)
134
+ recv_x_sf = nvfp4_block_scale_interleave(recv_x_sf)
135
+ recv_x_sf = recv_x_sf.view(-1, x_sf_width)
136
else:
139
- a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
140
- a1q_scale_recv = None
141
- a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1])
137
+ recv_x, topk_ids_recv, topk_weights_recv = recv_payloads
138
+ recv_x_sf = None
139
+ recv_x = recv_x.view(-1, recv_x.shape[-1])
140
topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1])
141
topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1])
142
145
- return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
143
+ return (
144
+ recv_x,
145
+ recv_x_sf,
146
+ None,
147
+ topk_ids_recv,
148
+ topk_weights_recv,
149
+ )
150
151
def finalize(
152
self,