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,