@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b4806c8ee1

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 11:15 UTC

[DSV4] Add BF16 and MXFP8 A2A support for flashinfer a2a one sided (#40960)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: Zijing Liu <liuzijing2014@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Yongye Zhu committed Apr 30, 2026 at 18:33 UTC b4806c8ee12d5c5bbfebd6070b389e2f4daad1fd

12 files changed +107 -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 | 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 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |

41 !!! info "Table key"

42 1. All types: mxfp4, nvfp4, int4, int8, fp8

vllm/distributed/device_communicators/all2all.py

+8 -2

@@ -577,6 +577,8 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

577 top_k: int,

578 num_experts: int,

579 hidden_size: int,

580 + dispatch_dtype_bytes_per_elem: int = 0,

581 + dispatch_scale_bytes_per_token: int = 0,

582 ):

583 """Initialize the MoeAlltoAll workspace."""

584 if self.initialized:

@@ -607,9 +609,13 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

609 ep_config = MnnvlConfig(

610 comm_backend=CustomCommunicator(self.cpu_group),

611 )

612 + if dispatch_dtype_bytes_per_elem == 0:

613 + hidden_bytes = hidden_size // 2

614 + else:

615 + hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem

616 total_dispatch_payload_size_per_token = (

611 - hidden_size // 2 # nvfp4 hidden states

612 - + hidden_size // 16 # fp8 scaling factors

617 + hidden_bytes

618 + + dispatch_scale_bytes_per_token

619 + top_k * 4 # int32 topks ids

620 + top_k * 4 # float32 topk weights

621 )

vllm/model_executor/layers/fused_moe/all2all_utils.py

+22 -8

@@ -228,23 +228,37 @@ def maybe_make_prepare_finalize(

228

229 elif moe.use_fi_nvl_one_sided_kernels:

230 assert quant_config is not None

231 - if quant_config.quant_dtype != "nvfp4":

232 - raise ValueError(

233 - "The 'flashinfer_nvlink_one_sided' all2all backend only "

234 - "supports nvfp4 activation quantization, but got "

235 - f"quant_dtype={quant_config.quant_dtype!r}. Use a different "

236 - "all2all backend (e.g. 'flashinfer_nvlink_two_sided' or "

237 - "'allgather_reducescatter') for non-nvfp4 models."

238 - )

231 max_num_tokens = (

232 get_current_vllm_config().scheduler_config.max_num_batched_tokens

233 )

234 + if quant_config.quant_dtype is None:

235 + dispatch_dtype_bytes_per_elem = 2

236 + dispatch_scale_bytes_per_token = 0

237 + elif quant_config.quant_dtype == "nvfp4":

238 + dispatch_dtype_bytes_per_elem = 0

239 + dispatch_scale_bytes_per_token = moe.hidden_dim // 16

240 + elif quant_config.quant_dtype == "mxfp8":

241 + dispatch_dtype_bytes_per_elem = 1

242 + align = quant_config.mx_alignment

243 + if align > 0:

244 + padded_k = ((moe.hidden_dim + align - 1) // align) * align

245 + else:

246 + padded_k = moe.hidden_dim

247 + dispatch_scale_bytes_per_token = padded_k // 32

248 + else:

249 + raise NotImplementedError(

250 + "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "

251 + "and bf16 (quant_dtype=None) today; got "

252 + f"quant_dtype={quant_config.quant_dtype!r}"

253 + )

254 prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(

255 max_num_tokens=max_num_tokens,

256 top_k=moe.experts_per_token,

257 num_experts=moe.num_experts,

258 hidden_size=moe.hidden_dim,

259 num_dispatchers=all2all_manager.world_size,

260 + dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,

261 + dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,

262 )

263

264 elif moe.use_ag_rs_all2all_kernels and allow_new_interface:

vllm/model_executor/layers/fused_moe/config.py

@@ -254,6 +254,8 @@ class FusedMoEQuantConfig:

254 gemm1_beta: float | None = None

255 gemm1_clamp_limit: float | None = None

256

257 + mx_alignment: int = 0

258 +

259 def __post_init__(self):

260 assert not self.per_act_token_quant or self.block_shape is None, (

261 "illegal quantization"

@@ -712,6 +714,7 @@ def mxfp4_mxfp8_moe_quant_config(

714 gemm1_alpha: float | None = None,

715 gemm1_beta: float | None = None,

716 gemm1_clamp_limit: float | None = None,

717 + mx_alignment: int = 0,

718 ) -> FusedMoEQuantConfig:

719 """

720 Construct a quant config for mxfp4 activations and mxfp4 weights.

@@ -724,6 +727,7 @@ def mxfp4_mxfp8_moe_quant_config(

727 gemm1_alpha=gemm1_alpha,

728 gemm1_beta=gemm1_beta,

729 gemm1_clamp_limit=gemm1_clamp_limit,

730 + mx_alignment=mx_alignment,

731 )

732

733

vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py

+17 -36

@@ -44,6 +44,9 @@ class TrtLlmMxfp4ExpertsBase:

44 moe_config.intermediate_size_per_partition

45 )

46 self.hidden_dim = moe_config.hidden_dim

47 + self.hidden_dim_unpadded = (

48 + moe_config.hidden_dim_unpadded or moe_config.hidden_dim

49 + )

50 self.local_num_experts = moe_config.num_local_experts

51 self.ep_rank = moe_config.moe_parallel_config.ep_rank

@@ -82,9 +85,6 @@ class TrtLlmMxfp4ExpertsBase:

85 get_current_vllm_config().compilation_config.max_cudagraph_capture_size

86 )

85 - # P1-5 fix: use public quant_dtype property instead of private _a1

86 - self.use_mxfp8_input = quant_config.quant_dtype == "mxfp8"

87 -

88 @staticmethod

89 def _supports_current_device() -> bool:

90 p = current_platform

@@ -121,8 +121,7 @@ class TrtLlmMxfp4ExpertsBase:

121

122 @property

123 def expects_unquantized_inputs(self) -> bool:

124 - # Expert handles MXFP8 quantization internally if needed

125 - return True

124 + return False

125

126

127 class TrtLlmMxfp4ExpertsMonolithic(

@@ -181,24 +180,19 @@ class TrtLlmMxfp4ExpertsMonolithic(

180 ) -> torch.Tensor:

181 from flashinfer import trtllm_fp4_block_scale_moe

182

184 - # Handle input quantization

185 - if self.use_mxfp8_input:

186 - from flashinfer import mxfp8_quantize

187 -

188 - x_quant, x_scale = mxfp8_quantize(

189 - hidden_states,

190 - is_sf_swizzled_layout=False,

191 - alignment=256,

192 - )

193 - x_scale = x_scale.view(torch.float8_e4m3fn).reshape(

194 - *hidden_states.shape[:-1], -1

195 - )

183 + if a1q_scale is not None:

184 + x_quant = hidden_states

185 + x_scale = a1q_scale.view(torch.float8_e4m3fn)

186 else:

187 assert hidden_states.dtype == torch.bfloat16

188 x_quant = hidden_states

189 x_scale = None

200 -

201 - output = torch.empty_like(hidden_states)

190 + output = torch.empty(

191 + *hidden_states.shape[:-1],

192 + self.hidden_dim_unpadded,

193 + dtype=torch.bfloat16,

194 + device=hidden_states.device,

195 + )

196

197 from vllm.utils.flashinfer import _is_fi_autotuning, autotune

198

@@ -244,10 +238,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula

238 Moved from trtllm_moe.py.

239 """

240

247 - @property

248 - def expects_unquantized_inputs(self) -> bool:

249 - return True

250 -

241 @staticmethod

242 def _supports_parallel_config(

243 moe_parallel_config: FusedMoEParallelConfig,

@@ -284,7 +274,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula

274 # The workspaces for this implementation are managed by flashinfer.

275 workspace1 = (0,)

276 workspace2 = (0,)

287 - output = (M, K)

277 + output = (M, self.hidden_dim_unpadded)

278 return (workspace1, workspace2, output)

279

280 def apply(

@@ -310,18 +300,9 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula

300 intermediate_size = self.intermediate_size_per_partition

301 local_expert_offset = self.moe_config.ep_rank * local_num_experts

302

313 - # Handle input quantization

314 - if self.use_mxfp8_input:

315 - from flashinfer import mxfp8_quantize

316 -

317 - x_quant, x_scale = mxfp8_quantize(

318 - hidden_states,

319 - is_sf_swizzled_layout=False,

320 - alignment=256,

321 - )

322 - x_scale = x_scale.view(torch.float8_e4m3fn).reshape(

323 - *hidden_states.shape[:-1], -1

324 - )

303 + if a1q_scale is not None:

304 + x_quant = hidden_states

305 + x_scale = a1q_scale.view(torch.float8_e4m3fn)

306 else:

307 assert hidden_states.dtype == torch.bfloat16

308 x_quant = hidden_states

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+12 -5

@@ -1195,10 +1195,18 @@ def make_mxfp4_moe_quant_config(

1195 gemm1_beta=gemm1_beta,

1196 gemm1_clamp_limit=swiglu_limit,

1197 )

1198 - elif mxfp4_backend in (

1199 - Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,

1200 - Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,

1201 - ):

1198 + elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:

1199 + return mxfp4_mxfp8_moe_quant_config(

1200 + w1_bias=w1_bias,

1201 + w2_bias=w2_bias,

1202 + w1_scale=w1_scale,

1203 + w2_scale=w2_scale,

1204 + gemm1_alpha=gemm1_alpha,

1205 + gemm1_beta=gemm1_beta,

1206 + gemm1_clamp_limit=swiglu_limit,

1207 + mx_alignment=256,

1208 + )

1209 + elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:

1210 return mxfp4_mxfp8_moe_quant_config(

1211 w1_bias=w1_bias,

1212 w2_bias=w2_bias,

@@ -1250,7 +1258,6 @@ def make_mxfp4_moe_kernel(

1258 """Create a FusedMoEKernel for the given MXFP4 backend."""

1259 is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)

1260

1253 - # Create Prepare/Finalize.

1261 prepare_finalize = maybe_make_prepare_finalize(

1262 moe=moe_config,

1263 quant_config=moe_quant_config,

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

+22 -9

@@ -31,6 +31,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo

31 num_experts: int,

32 hidden_size: int,

33 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

@@ -38,6 +40,7 @@ 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

45 device_communicator = get_ep_group().device_communicator

46 assert device_communicator is not None

@@ -49,6 +52,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo

52 top_k=self.top_k,

53 num_experts=self.num_experts,

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

57 )

59 @property

@@ -92,19 +97,24 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo

97 else a1.shape[0]

98 )

95 - a1q, a1q_scale = moe_kernel_quantize_input(

96 - a1,

97 - quant_config.a1_gscale,

98 - quant_config.quant_dtype,

99 - quant_config.per_act_token_quant,

100 - quant_config.block_shape,

101 - is_fp4_scale_swizzled=False, # delay swizzle to after comm

102 - )

100 + if defer_input_quant:

101 + a1q, a1q_scale = a1, None

102 + else:

103 + a1q, a1q_scale = moe_kernel_quantize_input(

104 + a1,

105 + quant_config.a1_gscale,

106 + quant_config.quant_dtype,

107 + quant_config.per_act_token_quant,

108 + quant_config.block_shape,

109 + is_fp4_scale_swizzled=False, # delay swizzle to after comm

110 + mx_alignment=quant_config.mx_alignment,

111 + )

112

113 payloads = []

114 payloads.append(a1q)

115 if a1q_scale is not None:

116 payloads.append(a1q_scale)

117 + topk_ids_payload_index = len(payloads)

118 payloads.append(topk_ids)

119 payloads.append(topk_weights)

120

@@ -113,6 +123,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo

123 token_selected_experts=topk_ids,

124 input_payloads=payloads,

125 runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,

126 + invalid_token_expert_id=-1, # Follow TRTLLM Pattern

127 + expert_id_payload_index=topk_ids_payload_index,

128 )

129 if a1q_scale is not None:

130 a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads

@@ -124,7 +136,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo

136 a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])

137 a1q_scale_recv = a1q_scale_recv.view(torch.uint8)

138 a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)

127 - a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)

139 + assert self.scale_elems_per_token > 0

140 + a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token)

141 else:

142 a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads

143 a1q_scale_recv = None

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

@@ -174,6 +174,7 @@ def flashinfer_alltoall_dispatch(

174 # the hidden states, breaking the A2A kernel. So, we

175 # delay the swizzling until after the A2A.

176 is_fp4_scale_swizzled=False,

177 + mx_alignment=quant_config.mx_alignment,

178 )

179

180 x = MnnvlMoe.mnnvl_moe_alltoallv(

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

@@ -40,6 +40,7 @@ def _quantize_and_setup_dispatch(

40 per_act_token_quant=quant_config.per_act_token_quant,

41 block_shape=quant_config.block_shape,

42 is_fp4_scale_swizzled=False,

43 + mx_alignment=quant_config.mx_alignment,

44 )

46 # Skip gathering scales if we have static quantization

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

@@ -31,6 +31,7 @@ def _quantize_input(

31 per_act_token_quant=quant_config.per_act_token_quant,

32 block_shape=quant_config.block_shape,

33 is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,

34 + mx_alignment=quant_config.mx_alignment,

35 )

37 return a1q, a1q_scale

vllm/model_executor/layers/fused_moe/utils.py

+5 -2

@@ -208,11 +208,12 @@ def _mxfp8_e4m3_quantize(

208 per_act_token_quant: bool,

209 block_shape: list[int] | None = None,

210 is_sf_swizzled_layout: bool = False,

211 + mx_alignment: int = 0,

212 ) -> tuple[torch.Tensor, torch.Tensor]:

213 assert A_scale is None

214 assert not per_act_token_quant

215 assert block_shape is None or block_shape == [1, 32]

215 - return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout)

216 + return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout, mx_alignment)

217

218

219 def _mxfp6_e3m2_quantize(

@@ -258,6 +259,7 @@ def moe_kernel_quantize_input(

259 is_fp4_scale_swizzled: bool = True,

260 ocp_mx_scheme: str | None = None,

261 quantization_emulation: bool = False,

262 + mx_alignment: int = 0,

263 ) -> tuple[torch.Tensor, torch.Tensor | None]:

264 # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation

265 if ocp_mx_scheme is not None:

@@ -320,7 +322,8 @@ def moe_kernel_quantize_input(

322 A_scale,

323 per_act_token_quant,

324 block_shape,

323 - is_sf_swizzled_layout=is_fp4_scale_swizzled,

325 + is_sf_swizzled_layout=False,

326 + mx_alignment=mx_alignment,

327 )

328 elif quant_dtype == "mxfp6_e3m2":

329 if not quantization_emulation:

vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

+13 -5

@@ -85,7 +85,9 @@ def _mxfp8_e4m3_quantize_torch(

87 def _mxfp8_e4m3_quantize_impl(

88 - x: torch.Tensor, is_sf_swizzled_layout: bool = False

88 + x: torch.Tensor,

89 + is_sf_swizzled_layout: bool = False,

90 + alignment: int = 0,

91 ) -> tuple[torch.Tensor, torch.Tensor]:

92 from vllm.platforms import current_platform

@@ -93,7 +95,9 @@ def _mxfp8_e4m3_quantize_impl(

95 from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize

97 x_q, x_scales = flashinfer_mxfp8_quantize(

96 - x, is_sf_swizzled_layout=is_sf_swizzled_layout

98 + x,

99 + is_sf_swizzled_layout=is_sf_swizzled_layout,

100 + alignment=alignment if alignment > 0 else 32,

101 )

102 if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:

103 x_scales = x_scales.view(x.size(0), -1)

@@ -103,9 +107,11 @@ def _mxfp8_e4m3_quantize_impl(

107

108

109 def mxfp8_e4m3_quantize(

106 - x: torch.Tensor, is_sf_swizzled_layout: bool = False

110 + x: torch.Tensor,

111 + is_sf_swizzled_layout: bool = False,

112 + alignment: int = 0,

113 ) -> tuple[torch.Tensor, torch.Tensor]:

108 - return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout)

114 + return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout, alignment)

115

116

117 def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:

@@ -125,7 +131,9 @@ def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor

131

132

133 def mxfp8_e4m3_quantize_fake(

128 - x: torch.Tensor, is_sf_swizzled_layout: bool = False

134 + x: torch.Tensor,

135 + is_sf_swizzled_layout: bool = False,

136 + alignment: int = 0,

137 ) -> tuple[torch.Tensor, torch.Tensor]:

138 """Fake implementation for torch.compile tracing."""

139 fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE)