@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0103a9b96f

[Quantization] Add per-token NVFP4 CuTe-DSL MoE backend (#50030)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>

Matej Sirovatka committed Sep 30, 2026 at 03:42 UTC 0103a9b96fde27d26a464bf2c30b4d977b2a134f

9 files changed +186 -66

tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py

+52 -19

@@ -29,14 +29,20 @@ from vllm.model_executor.layers.fused_moe.config import (

29 FusedMoEConfig,

30 FusedMoEParallelConfig,

31 RoutingMethodType,

32 - nvfp4_moe_quant_config,

32 )

33 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (

34 FlashInferCuteDSLExperts,

35 )

36 +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

37 + NvFp4MoeBackend,

38 + make_nvfp4_moe_quant_config,

39 +)

40 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

41 prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,

42 )

43 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

44 + quantize_nvfp4_per_token_input,

45 +)

46 from vllm.platforms import current_platform

47 from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4

48 from vllm.utils.math_utils import next_power_of_2

@@ -159,8 +165,12 @@ def _dequantize_nvfp4_linear(

165 return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype)

166

167

162 -@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)])

168 +@pytest.mark.parametrize(

169 + "m,n,k,e,topk",

170 + [(16, 128, 512, 4, 2), (16, 192, 512, 4, 2), (16, 192, 2688, 4, 2)],

171 +)

172 @pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES)

173 +@pytest.mark.parametrize("per_token_activation", [False, True])

174 @pytest.mark.parametrize("dtype", [torch.bfloat16])

175 @torch.inference_mode()

176 def test_flashinfer_cutedsl_fp4_moe(

@@ -173,6 +183,7 @@ def test_flashinfer_cutedsl_fp4_moe(

183 alpha: float | None,

184 beta: float | None,

185 limit: float | None,

186 + per_token_activation: bool,

187 dtype: torch.dtype,

188 workspace_init,

189 ):

@@ -181,6 +192,10 @@ def test_flashinfer_cutedsl_fp4_moe(

192 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

193 ):

194 hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10

195 + # Keep GEMM1 activation variance stable as K changes, so the fixed

196 + # tolerance against the BF16 GEMM2 reference measures comparable FP4

197 + # quantization error at each hidden width.

198 + hidden_states *= (512 / k) ** 0.5

199

200 w1_rows = 2 * n if activation.is_gated else n

201 w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15

@@ -193,7 +208,6 @@ def test_flashinfer_cutedsl_fp4_moe(

208 hidden_states, score, topk, renormalize=False

209 )

210

196 - activation = MoEActivation.RELU2_NO_MUL

211 moe_config = FusedMoEConfig(

212 num_experts=e,

213 experts_per_token=topk,

@@ -231,19 +245,30 @@ def test_flashinfer_cutedsl_fp4_moe(

245 w2_scale_2=(1.0 / w2_global_scale),

246 a2_scale=a2_scale,

247 )

234 - quant_config = nvfp4_moe_quant_config(

235 - g1_alphas=w1_alpha,

236 - g2_alphas=w2_alpha,

237 - a1_gscale=(1.0 / a1_scale),

238 - a2_gscale=(1.0 / a2_scale),

239 - w1_scale=w1_scale_cutedsl,

248 + # I=192 must remain unpadded for gated experts, but become 256 for

249 + # non-gated experts. The reference below uses the checkpoint shapes.

250 + expected_n = n if activation.is_gated else (n + 127) // 128 * 128

251 + assert moe_config.intermediate_size_per_partition == expected_n

252 + assert w1_cutedsl.shape[1] == expected_n * (2 if activation.is_gated else 1)

253 + assert w2_cutedsl.shape[2] * 2 == expected_n

254 + expected_k = (k + 255) // 256 * 256

255 + assert moe_config.hidden_dim == expected_k

256 + assert moe_config.hidden_dim_unpadded == k

257 + assert w1_cutedsl.shape[2] * 2 == expected_k

258 + assert w2_cutedsl.shape[1] == expected_k

259 + quant_config = make_nvfp4_moe_quant_config(

260 + backend=NvFp4MoeBackend.FLASHINFER_CUTEDSL,

261 + w13_scale_2=w1_alpha,

262 + w2_scale_2=w2_alpha,

263 + a13_scale=a1_scale,

264 + a2_scale=a2_scale,

265 + w13_scale=w1_scale_cutedsl,

266 w2_scale=w2_scale_cutedsl,

241 - is_scale_swizzled=False,

267 # Unset params must be omitted rather than forwarded as None into

268 # the kernel's float-typed SwiGLU arguments.

244 - gemm1_alpha=alpha,

245 - gemm1_beta=beta,

246 - gemm1_clamp_limit=limit,

269 + swiglu_alpha=alpha,

270 + swiglu_beta=beta,

271 + swiglu_limit=limit,

272 )

273

274 cutedsl_experts = mk.FusedMoEKernel(

@@ -256,11 +281,13 @@ def test_flashinfer_cutedsl_fp4_moe(

281 FlashInferCuteDSLExperts(

282 moe_config=moe_config,

283 quant_config=quant_config,

284 + per_token_activation=per_token_activation,

285 ),

286 )

287

288 cutedsl_output = cutedsl_experts.apply(

263 - hidden_states=hidden_states,

289 + # The model runner pads inputs and trims outputs around this kernel.

290 + hidden_states=torch.nn.functional.pad(hidden_states, (0, expected_k - k)),

291 w1=w1_cutedsl,

292 w2=w2_cutedsl,

293 topk_weights=topk_weights,

@@ -270,19 +297,25 @@ def test_flashinfer_cutedsl_fp4_moe(

297 expert_map=None,

298 apply_router_weight_on_input=False,

299 )

300 + cutedsl_output = cutedsl_output[:, :k]

301

302 a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32)

275 - a_q, a_scale = ops.scaled_fp4_quant(

276 - hidden_states,

277 - a_global_scale,

278 - is_sf_swizzled_layout=False,

279 - )

303 + if per_token_activation:

304 + a_q, a_scale, a_row_scale = quantize_nvfp4_per_token_input(hidden_states)

305 + else:

306 + a_q, a_scale = ops.scaled_fp4_quant(

307 + hidden_states,

308 + a_global_scale,

309 + is_sf_swizzled_layout=False,

310 + )

311 a_in_dtype = _dequantize_nvfp4_linear(

312 a_q,

313 a_scale,

314 a_global_scale,

315 dtype=dtype,

316 )

317 + if per_token_activation:

318 + a_in_dtype = (a_in_dtype.float() * a_row_scale.view(-1, 1)).to(dtype)

319

320 w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype)

321 w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)

tests/quantization/test_online.py

+24

@@ -605,6 +605,30 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:

605 ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")

606

607

608 +def test_nvfp4_per_token_backend_contract() -> None:

609 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (

610 + FlashInferCuteDSLExperts,

611 + )

612 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (

613 + FlashInferExperts,

614 + )

615 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (

616 + TrtLlmNvFp4ExpertsModular,

617 + TrtLlmNvFp4ExpertsMonolithic,

618 + )

619 + from vllm.model_executor.layers.quantization.utils.quant_utils import (

620 + kNvfp4DynamicToken,

621 + kNvfp4Static,

622 + )

623 +

624 + scheme = (kNvfp4Static, kNvfp4DynamicToken)

625 + assert TrtLlmNvFp4ExpertsMonolithic._supports_quant_scheme(*scheme)

626 + assert TrtLlmNvFp4ExpertsModular._supports_quant_scheme(*scheme)

627 + assert FlashInferCuteDSLExperts._supports_quant_scheme(*scheme)

628 + assert FlashInferCuteDSLExperts._supports_no_act_and_mul()

629 + assert not FlashInferExperts._supports_quant_scheme(*scheme)

630 +

631 +

632 @pytest.mark.parametrize("per_token_activation", [False, True])

633 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])

634 def test_nvfp4_one_sided_sizes_dispatched_activations(

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

+3 -3

@@ -11,8 +11,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

11 prepare_nvfp4_moe_layer_for_fi_or_cutlass,

12 )

13 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

14 + align_fp4_moe_hidden_dim_for_fi,

15 align_fp4_moe_weights_for_fi,

15 - align_trtllm_fp4_moe_hidden_dim_for_fi,

16 )

@@ -63,7 +63,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_noop():

63 w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)

65 out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (

66 - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)

66 + align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)

67 )

69 assert padded_hidden == 512

@@ -90,7 +90,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():

90 )

92 out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (

93 - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)

93 + align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)

94 )

96 assert out_hidden_dim == padded_hidden_dim

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

+39 -5

@@ -14,11 +14,14 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

14 TopKWeightAndReduceNoOP,

15 )

16 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

17 + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,

18 activation_to_flashinfer_int,

19 + quantize_nvfp4_per_token_input,

20 )

21 from vllm.model_executor.layers.quantization.utils.quant_utils import (

22 QuantKey,

23 kNvfp4Dynamic,

24 + kNvfp4DynamicToken,

25 kNvfp4Static,

26 )

27 from vllm.platforms import current_platform

@@ -40,6 +43,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

43 self,

44 moe_config: FusedMoEConfig,

45 quant_config: FusedMoEQuantConfig,

46 + per_token_activation: bool = False,

47 ):

48 super().__init__(

49 moe_config=moe_config,

@@ -63,6 +67,14 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

67 self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit

68 self.situ_beta = moe_config.activation_situ_beta

69 self.situ_linear_beta = moe_config.activation_situ_linear_beta

70 + self.per_token_activation = per_token_activation

71 + self.per_token_global_scale = None

72 + if per_token_activation:

73 + assert quant_config.a2_gscale is not None

74 + self.per_token_global_scale = quant_config.a2_gscale.new_full(

75 + (1,),

76 + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,

77 + )

79 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

80 layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)

@@ -92,6 +104,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

104 ) -> bool:

105 SUPPORTED_W_A = [

106 (kNvfp4Static, kNvfp4Dynamic),

107 + (kNvfp4Static, kNvfp4DynamicToken),

108 ]

109 return (weight_key, activation_key) in SUPPORTED_W_A

110

@@ -114,6 +127,10 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

127 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

128 return TopKWeightAndReduceNoOP()

129

130 + @property

131 + def expects_unquantized_inputs(self) -> bool:

132 + return self.per_token_activation

133 +

134 def workspace_shapes(

135 self,

136 M: int,

@@ -127,8 +144,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

144 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

145 workspace1 = (0,)

146 workspace2 = (0,)

130 - # K is packed (K//2 for uint8), so output uses hidden_dim.

131 - assert self.hidden_dim == K * 2

147 + expected_hidden_dim = K if self.expects_unquantized_inputs else K * 2

148 + assert self.hidden_dim == expected_hidden_dim

149 output = (M, self.hidden_dim)

150 return (workspace1, workspace2, output)

151

@@ -151,13 +168,26 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

168 apply_router_weight_on_input: bool | None,

169 ):

170 assert self.quant_dtype == "nvfp4"

154 - assert a1q_scale is not None

171 assert self.w1_scale is not None

172 assert self.w2_scale is not None

173

174 + if self.expects_unquantized_inputs:

175 + hidden_states, block_scale, per_token_scale = (

176 + quantize_nvfp4_per_token_input(hidden_states)

177 + )

178 + fc2_input_scale = self.per_token_global_scale

179 + else:

180 + assert a1q_scale is not None

181 + block_scale = a1q_scale

182 + per_token_scale = None

183 + fc2_input_scale = self.a2_gscale

184 +

185 + assert block_scale is not None

186 + assert fc2_input_scale is not None

187 +

188 # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize.

189 # The functional API expects x_sf with trailing dim: (M, K//16, 1).

160 - x_sf = a1q_scale.unsqueeze(-1)

190 + x_sf = block_scale.unsqueeze(-1)

191

192 # The kernel defaults swiglu_{alpha,beta,limit} to the plain-SwiGLU

193 # values, so only forward the ones the model actually sets.

@@ -188,6 +218,9 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

218 }

219 swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None}

220

221 + per_token_kwargs = (

222 + {"per_token_scale": per_token_scale} if self.per_token_activation else {}

223 + )

224 flashinfer_cute_dsl_fused_moe_nvfp4(

225 x=hidden_states,

226 x_sf=x_sf,

@@ -196,7 +229,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

229 w1_weight=w1,

230 w1_weight_sf=self.w1_scale,

231 w1_alpha=self.g1_alphas,

199 - fc2_input_scale=self.a2_gscale,

232 + fc2_input_scale=fc2_input_scale,

233 w2_weight=w2,

234 w2_weight_sf=self.w2_scale,

235 w2_alpha=self.g2_alphas,

@@ -209,4 +242,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

242 MoEActivation.SILU if activation == MoEActivation.SITU else activation

243 ),

244 **swiglu_kwargs,

245 + **per_token_kwargs,

246 )

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

+7 -27

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket

24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

25 activation_to_flashinfer_int,

26 has_flashinfer_situ_activation,

27 + quantize_nvfp4_per_token_input,

28 )

29 from vllm.model_executor.layers.quantization.utils.quant_utils import (

30 QuantKey,

@@ -37,10 +38,6 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

39 logger = init_logger(__name__)

40 -# Base scale for per-token NVFP4 activation quant; the kernel folds the

41 -# per-token global scale (from the activation amax) on top of it.

42 -_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0)

43 -

42 class TrtLlmNvFp4ExpertsBase:

43 """NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface."""

@@ -259,23 +256,6 @@ class TrtLlmNvFp4ExpertsBase:

256 def expects_unquantized_inputs(self) -> bool:

257 return self.per_token_activation

258

262 - def _quantize_per_token_input(

263 - self, hidden_states: torch.Tensor

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

265 - """NVFP4-quantize activations with a per-token global scale.

266 -

267 - Returns ``(packed_fp4, block_scale, per_token_scale)``.

268 - """

269 - from flashinfer import SfLayout, nvfp4_quantize

270 -

271 - hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize(

272 - hidden_states,

273 - _PER_TOKEN_BASE_GLOBAL_SCALE,

274 - sfLayout=SfLayout.layout_linear,

275 - per_token_activation=True,

276 - )

277 - return hs_fp4, hs_block_scale, per_token_scale

278 -

259 def _get_chunk_size(self) -> int:

260 MAX_GRID_Y = 65535

261 MAX_TILE_TOKENS_DIM = 128

@@ -361,9 +341,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

341

342 # Per-token: input is unquantized, quantize it here. Otherwise it was

343 # already quantized in prepare() with the static global scale.

364 - if self.per_token_activation:

344 + if self.expects_unquantized_inputs:

345 hidden_states, block_scale, per_token_scale = (

366 - self._quantize_per_token_input(hidden_states)

346 + quantize_nvfp4_per_token_input(hidden_states)

347 )

348 else:

349 block_scale, per_token_scale = a1q_scale, None

@@ -437,7 +417,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

417 ) -> UnfinalizedMoEOutput | None:

418 assert self._supports_activation(activation)

419 # Per-token defers input quant to _invoke_kernel, so a1q_scale is None.

440 - assert a1q_scale is not None or self.per_token_activation

420 + assert a1q_scale is not None or self.expects_unquantized_inputs

421

422 # DeepEP produces int64 indexes.

423 topk_ids = topk_ids.to(dtype=torch.int32)

@@ -539,7 +519,7 @@ class TrtLlmNvFp4ExpertsMonolithic(

519 import flashinfer

520

521 assert self._supports_activation(activation)

542 - assert a1q_scale is not None or self.per_token_activation

522 + assert a1q_scale is not None or self.expects_unquantized_inputs

523 assert self.quant_config.w1_scale is not None

524 assert self.quant_config.w2_scale is not None

525 assert (

@@ -551,9 +531,9 @@ class TrtLlmNvFp4ExpertsMonolithic(

531 )

532

533 # Per-token: input is unquantized, quantize it here (see modular apply).

554 - if self.per_token_activation:

534 + if self.expects_unquantized_inputs:

535 hidden_states, block_scale, per_token_scale = (

556 - self._quantize_per_token_input(hidden_states)

536 + quantize_nvfp4_per_token_input(hidden_states)

537 )

538 else:

539 block_scale, per_token_scale = a1q_scale, None

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

+8 -1

@@ -612,7 +612,14 @@ def make_nvfp4_moe_kernel(

612 logger.info_once("Using %s", prepare_finalize.__class__.__name__)

613

614 extra_kwargs = {}

615 - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and per_token_activation:

615 + if (

616 + backend

617 + in (

618 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

619 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,

620 + )

621 + and per_token_activation

622 + ):

623 extra_kwargs["per_token_activation"] = True

624

625 # Create Experts.

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

+1 -1

@@ -720,7 +720,7 @@ class MoERunner(MoERunnerInterface):

720 )

721

722 # Record before `_maybe_pad_hidden_states` pads activations to match

723 - # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`

723 + # `moe_config.hidden_dim`, e.g. after `align_fp4_moe_hidden_dim_for_fi`

724 # so routed output can be trimmed before

725 # shared+routed add / latent up proj if needed.

726

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

+34 -8

@@ -9,8 +9,8 @@ import torch

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

12 + align_fp4_moe_hidden_dim_for_fi,

13 align_fp4_moe_weights_for_fi,

13 - align_trtllm_fp4_moe_hidden_dim_for_fi,

14 )

15 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

16 swizzle_blockscale,

@@ -128,9 +128,9 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(

128 ]:

129 """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.

130

131 - Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper,

132 - and interleaves w13 gate/linear rows for gated activations. Non-gated

133 - activations use a single w13 projection and keep its row order unchanged.

131 + Pads the runtime expert tensors to the kernel's GEMM alignment, converts

132 + weight scale factors to the MMA layout expected by CuteDslMoEWrapper, and

133 + interleaves w13 gate/linear rows for gated activations.

134 """

135 # Global scaling factors (same as other FlashInfer backends).

136 num_experts = w13.shape[0]

@@ -140,12 +140,38 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(

140 )

141 a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts)

142

143 - if layer.activation.is_gated:

143 + gated = layer.activation.is_gated

144 + if gated:

145 w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(

146 layer.activation, w13, w13_scale

147 )

148

148 - # Interleave up/gate rows for w13 weights and scales.

149 + # GEMM1's output dimension must be a multiple of 128: 2I for gated

150 + # activations (also required by interleaving), but only I for non-gated.

151 + # Keep the checkpoint tensors unchanged and pad only the kernel's runtime

152 + # representation. Zero rows also make the padded GEMM2 contraction a no-op.

153 + w13, w13_scale, w2, w2_scale, padded_intermediate = align_fp4_moe_weights_for_fi(

154 + w13,

155 + w13_scale,

156 + w2,

157 + w2_scale,

158 + is_act_and_mul=gated,

159 + min_alignment=64 if gated else 128,

160 + )

161 + layer.moe_config.intermediate_size_per_partition = padded_intermediate

162 +

163 + # GEMM1 gathers full 256-element K tiles of activations and their scales

164 + # without a K-tail predicate. Pad runtime weights to match; the MoE runner

165 + # pads activations and trims the output using the original hidden size.

166 + w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(

167 + w13, w13_scale, w2, w2_scale

168 + )

169 + if layer.moe_config.hidden_dim_unpadded is None:

170 + layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim

171 + layer.moe_config.hidden_dim = padded_hidden

172 +

173 + if gated:

174 + # Interleave up/gate rows for the fused gated activation.

175 w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)

176 w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)

177

@@ -371,8 +397,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

397

398 # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.

399 if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

374 - w13, w13_scale, w2, w2_scale, padded_hidden = (

375 - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)

400 + w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(

401 + w13, w13_scale, w2, w2_scale

402 )

403 if layer.moe_config.hidden_dim_unpadded is None:

404 layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim

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

+18 -2

@@ -13,6 +13,8 @@ if TYPE_CHECKING:

14 logger = init_logger(__name__)

16 +NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0)

17 +

19 def activation_to_flashinfer_int(activation: MoEActivation) -> int:

20 return activation_to_flashinfer_type(activation).value

@@ -50,6 +52,20 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"

52 return ACTIVATION_TO_FI_ACTIVATION[activation]

55 +def quantize_nvfp4_per_token_input(

56 + hidden_states: torch.Tensor,

57 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

58 + """Quantize NVFP4 activations with one FP32 decode scale per token."""

59 + from flashinfer import SfLayout, nvfp4_quantize

60 +

61 + return nvfp4_quantize(

62 + hidden_states,

63 + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,

64 + sfLayout=SfLayout.layout_linear,

65 + per_token_activation=True,

66 + )

67 +

68 +

69 def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:

70 return (

71 x.reshape(-1, 2, x.shape[-2] // 2, x.shape[-1]).flip(dims=[1]).reshape(x.shape)

@@ -264,7 +280,7 @@ def align_fp4_moe_weights_for_fi(

280 return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate

281

282

267 -def align_trtllm_fp4_moe_hidden_dim_for_fi(

283 +def align_fp4_moe_hidden_dim_for_fi(

284 w13: torch.Tensor,

285 w13_scale: torch.Tensor,

286 w2: torch.Tensor,

@@ -279,7 +295,7 @@ def align_trtllm_fp4_moe_hidden_dim_for_fi(

295 return w13, w13_scale, w2, w2_scale, hidden_size

296

297 logger.warning_once(

282 - "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. "

298 + "Padding hidden size from %d to %d for FlashInfer NVFP4 MoE weights. "

299 "This requires activation slicing at runtime and may cause "

300 "performance degradation.",

301 hidden_size,