[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>

tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py

@@ -29,14 +29,20 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
     RoutingMethodType,
-    nvfp4_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
     FlashInferCuteDSLExperts,
 )
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    NvFp4MoeBackend,
+    make_nvfp4_moe_quant_config,
+)
 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
     prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
 )
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    quantize_nvfp4_per_token_input,
+)
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4
 from vllm.utils.math_utils import next_power_of_2
@@ -159,8 +165,12 @@ def _dequantize_nvfp4_linear(
     return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype)
 
 
-@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)])
+@pytest.mark.parametrize(
+    "m,n,k,e,topk",
+    [(16, 128, 512, 4, 2), (16, 192, 512, 4, 2), (16, 192, 2688, 4, 2)],
+)
 @pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES)
+@pytest.mark.parametrize("per_token_activation", [False, True])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @torch.inference_mode()
 def test_flashinfer_cutedsl_fp4_moe(
@@ -173,6 +183,7 @@ def test_flashinfer_cutedsl_fp4_moe(
     alpha: float | None,
     beta: float | None,
     limit: float | None,
+    per_token_activation: bool,
     dtype: torch.dtype,
     workspace_init,
 ):
@@ -181,6 +192,10 @@ def test_flashinfer_cutedsl_fp4_moe(
         VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
     ):
         hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+        # Keep GEMM1 activation variance stable as K changes, so the fixed
+        # tolerance against the BF16 GEMM2 reference measures comparable FP4
+        # quantization error at each hidden width.
+        hidden_states *= (512 / k) ** 0.5
 
         w1_rows = 2 * n if activation.is_gated else n
         w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15
@@ -193,7 +208,6 @@ def test_flashinfer_cutedsl_fp4_moe(
             hidden_states, score, topk, renormalize=False
         )
 
-        activation = MoEActivation.RELU2_NO_MUL
         moe_config = FusedMoEConfig(
             num_experts=e,
             experts_per_token=topk,
@@ -231,19 +245,30 @@ def test_flashinfer_cutedsl_fp4_moe(
             w2_scale_2=(1.0 / w2_global_scale),
             a2_scale=a2_scale,
         )
-        quant_config = nvfp4_moe_quant_config(
-            g1_alphas=w1_alpha,
-            g2_alphas=w2_alpha,
-            a1_gscale=(1.0 / a1_scale),
-            a2_gscale=(1.0 / a2_scale),
-            w1_scale=w1_scale_cutedsl,
+        # I=192 must remain unpadded for gated experts, but become 256 for
+        # non-gated experts. The reference below uses the checkpoint shapes.
+        expected_n = n if activation.is_gated else (n + 127) // 128 * 128
+        assert moe_config.intermediate_size_per_partition == expected_n
+        assert w1_cutedsl.shape[1] == expected_n * (2 if activation.is_gated else 1)
+        assert w2_cutedsl.shape[2] * 2 == expected_n
+        expected_k = (k + 255) // 256 * 256
+        assert moe_config.hidden_dim == expected_k
+        assert moe_config.hidden_dim_unpadded == k
+        assert w1_cutedsl.shape[2] * 2 == expected_k
+        assert w2_cutedsl.shape[1] == expected_k
+        quant_config = make_nvfp4_moe_quant_config(
+            backend=NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+            w13_scale_2=w1_alpha,
+            w2_scale_2=w2_alpha,
+            a13_scale=a1_scale,
+            a2_scale=a2_scale,
+            w13_scale=w1_scale_cutedsl,
             w2_scale=w2_scale_cutedsl,
-            is_scale_swizzled=False,
             # Unset params must be omitted rather than forwarded as None into
             # the kernel's float-typed SwiGLU arguments.
-            gemm1_alpha=alpha,
-            gemm1_beta=beta,
-            gemm1_clamp_limit=limit,
+            swiglu_alpha=alpha,
+            swiglu_beta=beta,
+            swiglu_limit=limit,
         )
 
         cutedsl_experts = mk.FusedMoEKernel(
@@ -256,11 +281,13 @@ def test_flashinfer_cutedsl_fp4_moe(
             FlashInferCuteDSLExperts(
                 moe_config=moe_config,
                 quant_config=quant_config,
+                per_token_activation=per_token_activation,
             ),
         )
 
         cutedsl_output = cutedsl_experts.apply(
-            hidden_states=hidden_states,
+            # The model runner pads inputs and trims outputs around this kernel.
+            hidden_states=torch.nn.functional.pad(hidden_states, (0, expected_k - k)),
             w1=w1_cutedsl,
             w2=w2_cutedsl,
             topk_weights=topk_weights,
@@ -270,19 +297,25 @@ def test_flashinfer_cutedsl_fp4_moe(
             expert_map=None,
             apply_router_weight_on_input=False,
         )
+        cutedsl_output = cutedsl_output[:, :k]
 
         a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32)
-        a_q, a_scale = ops.scaled_fp4_quant(
-            hidden_states,
-            a_global_scale,
-            is_sf_swizzled_layout=False,
-        )
+        if per_token_activation:
+            a_q, a_scale, a_row_scale = quantize_nvfp4_per_token_input(hidden_states)
+        else:
+            a_q, a_scale = ops.scaled_fp4_quant(
+                hidden_states,
+                a_global_scale,
+                is_sf_swizzled_layout=False,
+            )
         a_in_dtype = _dequantize_nvfp4_linear(
             a_q,
             a_scale,
             a_global_scale,
             dtype=dtype,
         )
+        if per_token_activation:
+            a_in_dtype = (a_in_dtype.float() * a_row_scale.view(-1, 1)).to(dtype)
 
         w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype)
         w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)

tests/quantization/test_online.py

@@ -605,6 +605,30 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:
         ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")
 
 
+def test_nvfp4_per_token_backend_contract() -> None:
+    from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
+        FlashInferCuteDSLExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
+        FlashInferExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+        TrtLlmNvFp4ExpertsModular,
+        TrtLlmNvFp4ExpertsMonolithic,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        kNvfp4DynamicToken,
+        kNvfp4Static,
+    )
+
+    scheme = (kNvfp4Static, kNvfp4DynamicToken)
+    assert TrtLlmNvFp4ExpertsMonolithic._supports_quant_scheme(*scheme)
+    assert TrtLlmNvFp4ExpertsModular._supports_quant_scheme(*scheme)
+    assert FlashInferCuteDSLExperts._supports_quant_scheme(*scheme)
+    assert FlashInferCuteDSLExperts._supports_no_act_and_mul()
+    assert not FlashInferExperts._supports_quant_scheme(*scheme)
+
+
 @pytest.mark.parametrize("per_token_activation", [False, True])
 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
 def test_nvfp4_one_sided_sizes_dispatched_activations(

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

@@ -11,8 +11,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
     prepare_nvfp4_moe_layer_for_fi_or_cutlass,
 )
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    align_fp4_moe_hidden_dim_for_fi,
     align_fp4_moe_weights_for_fi,
-    align_trtllm_fp4_moe_hidden_dim_for_fi,
 )
 
 
@@ -63,7 +63,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_noop():
     w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)
 
     out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (
-        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+        align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
     )
 
     assert padded_hidden == 512
@@ -90,7 +90,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
     )
 
     out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (
-        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+        align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
     )
 
     assert out_hidden_dim == padded_hidden_dim

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

@@ -14,11 +14,14 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,
     activation_to_flashinfer_int,
+    quantize_nvfp4_per_token_input,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kNvfp4Dynamic,
+    kNvfp4DynamicToken,
     kNvfp4Static,
 )
 from vllm.platforms import current_platform
@@ -40,6 +43,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         self,
         moe_config: FusedMoEConfig,
         quant_config: FusedMoEQuantConfig,
+        per_token_activation: bool = False,
     ):
         super().__init__(
             moe_config=moe_config,
@@ -63,6 +67,14 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
         self.situ_beta = moe_config.activation_situ_beta
         self.situ_linear_beta = moe_config.activation_situ_linear_beta
+        self.per_token_activation = per_token_activation
+        self.per_token_global_scale = None
+        if per_token_activation:
+            assert quant_config.a2_gscale is not None
+            self.per_token_global_scale = quant_config.a2_gscale.new_full(
+                (1,),
+                NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,
+            )
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -92,6 +104,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
     ) -> bool:
         SUPPORTED_W_A = [
             (kNvfp4Static, kNvfp4Dynamic),
+            (kNvfp4Static, kNvfp4DynamicToken),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
@@ -114,6 +127,10 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return self.per_token_activation
+
     def workspace_shapes(
         self,
         M: int,
@@ -127,8 +144,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
         workspace1 = (0,)
         workspace2 = (0,)
-        # K is packed (K//2 for uint8), so output uses hidden_dim.
-        assert self.hidden_dim == K * 2
+        expected_hidden_dim = K if self.expects_unquantized_inputs else K * 2
+        assert self.hidden_dim == expected_hidden_dim
         output = (M, self.hidden_dim)
         return (workspace1, workspace2, output)
 
@@ -151,13 +168,26 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         apply_router_weight_on_input: bool | None,
     ):
         assert self.quant_dtype == "nvfp4"
-        assert a1q_scale is not None
         assert self.w1_scale is not None
         assert self.w2_scale is not None
 
+        if self.expects_unquantized_inputs:
+            hidden_states, block_scale, per_token_scale = (
+                quantize_nvfp4_per_token_input(hidden_states)
+            )
+            fc2_input_scale = self.per_token_global_scale
+        else:
+            assert a1q_scale is not None
+            block_scale = a1q_scale
+            per_token_scale = None
+            fc2_input_scale = self.a2_gscale
+
+        assert block_scale is not None
+        assert fc2_input_scale is not None
+
         # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize.
         # The functional API expects x_sf with trailing dim: (M, K//16, 1).
-        x_sf = a1q_scale.unsqueeze(-1)
+        x_sf = block_scale.unsqueeze(-1)
 
         # The kernel defaults swiglu_{alpha,beta,limit} to the plain-SwiGLU
         # values, so only forward the ones the model actually sets.
@@ -188,6 +218,9 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
             }
         swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None}
 
+        per_token_kwargs = (
+            {"per_token_scale": per_token_scale} if self.per_token_activation else {}
+        )
         flashinfer_cute_dsl_fused_moe_nvfp4(
             x=hidden_states,
             x_sf=x_sf,
@@ -196,7 +229,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
             w1_weight=w1,
             w1_weight_sf=self.w1_scale,
             w1_alpha=self.g1_alphas,
-            fc2_input_scale=self.a2_gscale,
+            fc2_input_scale=fc2_input_scale,
             w2_weight=w2,
             w2_weight_sf=self.w2_scale,
             w2_alpha=self.g2_alphas,
@@ -209,4 +242,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
                 MoEActivation.SILU if activation == MoEActivation.SITU else activation
             ),
             **swiglu_kwargs,
+            **per_token_kwargs,
         )

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

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
     has_flashinfer_situ_activation,
+    quantize_nvfp4_per_token_input,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
@@ -37,10 +38,6 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 
 logger = init_logger(__name__)
 
-# Base scale for per-token NVFP4 activation quant; the kernel folds the
-# per-token global scale (from the activation amax) on top of it.
-_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0)
-
 
 class TrtLlmNvFp4ExpertsBase:
     """NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface."""
@@ -259,23 +256,6 @@ class TrtLlmNvFp4ExpertsBase:
     def expects_unquantized_inputs(self) -> bool:
         return self.per_token_activation
 
-    def _quantize_per_token_input(
-        self, hidden_states: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """NVFP4-quantize activations with a per-token global scale.
-
-        Returns ``(packed_fp4, block_scale, per_token_scale)``.
-        """
-        from flashinfer import SfLayout, nvfp4_quantize
-
-        hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize(
-            hidden_states,
-            _PER_TOKEN_BASE_GLOBAL_SCALE,
-            sfLayout=SfLayout.layout_linear,
-            per_token_activation=True,
-        )
-        return hs_fp4, hs_block_scale, per_token_scale
-
     def _get_chunk_size(self) -> int:
         MAX_GRID_Y = 65535
         MAX_TILE_TOKENS_DIM = 128
@@ -361,9 +341,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
 
         # Per-token: input is unquantized, quantize it here. Otherwise it was
         # already quantized in prepare() with the static global scale.
-        if self.per_token_activation:
+        if self.expects_unquantized_inputs:
             hidden_states, block_scale, per_token_scale = (
-                self._quantize_per_token_input(hidden_states)
+                quantize_nvfp4_per_token_input(hidden_states)
             )
         else:
             block_scale, per_token_scale = a1q_scale, None
@@ -437,7 +417,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
     ) -> UnfinalizedMoEOutput | None:
         assert self._supports_activation(activation)
         # Per-token defers input quant to _invoke_kernel, so a1q_scale is None.
-        assert a1q_scale is not None or self.per_token_activation
+        assert a1q_scale is not None or self.expects_unquantized_inputs
 
         # DeepEP produces int64 indexes.
         topk_ids = topk_ids.to(dtype=torch.int32)
@@ -539,7 +519,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
         import flashinfer
 
         assert self._supports_activation(activation)
-        assert a1q_scale is not None or self.per_token_activation
+        assert a1q_scale is not None or self.expects_unquantized_inputs
         assert self.quant_config.w1_scale is not None
         assert self.quant_config.w2_scale is not None
         assert (
@@ -551,9 +531,9 @@ class TrtLlmNvFp4ExpertsMonolithic(
         )
 
         # Per-token: input is unquantized, quantize it here (see modular apply).
-        if self.per_token_activation:
+        if self.expects_unquantized_inputs:
             hidden_states, block_scale, per_token_scale = (
-                self._quantize_per_token_input(hidden_states)
+                quantize_nvfp4_per_token_input(hidden_states)
             )
         else:
             block_scale, per_token_scale = a1q_scale, None

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

@@ -612,7 +612,14 @@ def make_nvfp4_moe_kernel(
     logger.info_once("Using %s", prepare_finalize.__class__.__name__)
 
     extra_kwargs = {}
-    if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and per_token_activation:
+    if (
+        backend
+        in (
+            NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        )
+        and per_token_activation
+    ):
         extra_kwargs["per_token_activation"] = True
 
     # Create Experts.

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

@@ -720,7 +720,7 @@ class MoERunner(MoERunnerInterface):
             )
 
         # Record before `_maybe_pad_hidden_states` pads activations to match
-        # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`
+        # `moe_config.hidden_dim`, e.g. after `align_fp4_moe_hidden_dim_for_fi`
         # so routed output can be trimmed before
         # shared+routed add / latent up proj if needed.
 

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

@@ -9,8 +9,8 @@ import torch
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    align_fp4_moe_hidden_dim_for_fi,
     align_fp4_moe_weights_for_fi,
-    align_trtllm_fp4_moe_hidden_dim_for_fi,
 )
 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
     swizzle_blockscale,
@@ -128,9 +128,9 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
 ]:
     """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.
 
-    Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper,
-    and interleaves w13 gate/linear rows for gated activations. Non-gated
-    activations use a single w13 projection and keep its row order unchanged.
+    Pads the runtime expert tensors to the kernel's GEMM alignment, converts
+    weight scale factors to the MMA layout expected by CuteDslMoEWrapper, and
+    interleaves w13 gate/linear rows for gated activations.
     """
     # Global scaling factors (same as other FlashInfer backends).
     num_experts = w13.shape[0]
@@ -140,12 +140,38 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
     )
     a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts)
 
-    if layer.activation.is_gated:
+    gated = layer.activation.is_gated
+    if gated:
         w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
             layer.activation, w13, w13_scale
         )
 
-        # Interleave up/gate rows for w13 weights and scales.
+    # GEMM1's output dimension must be a multiple of 128: 2I for gated
+    # activations (also required by interleaving), but only I for non-gated.
+    # Keep the checkpoint tensors unchanged and pad only the kernel's runtime
+    # representation. Zero rows also make the padded GEMM2 contraction a no-op.
+    w13, w13_scale, w2, w2_scale, padded_intermediate = align_fp4_moe_weights_for_fi(
+        w13,
+        w13_scale,
+        w2,
+        w2_scale,
+        is_act_and_mul=gated,
+        min_alignment=64 if gated else 128,
+    )
+    layer.moe_config.intermediate_size_per_partition = padded_intermediate
+
+    # GEMM1 gathers full 256-element K tiles of activations and their scales
+    # without a K-tail predicate. Pad runtime weights to match; the MoE runner
+    # pads activations and trims the output using the original hidden size.
+    w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(
+        w13, w13_scale, w2, w2_scale
+    )
+    if layer.moe_config.hidden_dim_unpadded is None:
+        layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim
+    layer.moe_config.hidden_dim = padded_hidden
+
+    if gated:
+        # Interleave up/gate rows for the fused gated activation.
         w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
         w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
 
@@ -371,8 +397,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
 
     # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.
     if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
-        w13, w13_scale, w2, w2_scale, padded_hidden = (
-            align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+        w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(
+            w13, w13_scale, w2, w2_scale
         )
         if layer.moe_config.hidden_dim_unpadded is None:
             layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim

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

@@ -13,6 +13,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0)
+
 
 def activation_to_flashinfer_int(activation: MoEActivation) -> int:
     return activation_to_flashinfer_type(activation).value
@@ -50,6 +52,20 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
     return ACTIVATION_TO_FI_ACTIVATION[activation]
 
 
+def quantize_nvfp4_per_token_input(
+    hidden_states: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize NVFP4 activations with one FP32 decode scale per token."""
+    from flashinfer import SfLayout, nvfp4_quantize
+
+    return nvfp4_quantize(
+        hidden_states,
+        NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,
+        sfLayout=SfLayout.layout_linear,
+        per_token_activation=True,
+    )
+
+
 def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
     return (
         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(
     return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
 
 
-def align_trtllm_fp4_moe_hidden_dim_for_fi(
+def align_fp4_moe_hidden_dim_for_fi(
     w13: torch.Tensor,
     w13_scale: torch.Tensor,
     w2: torch.Tensor,
@@ -279,7 +295,7 @@ def align_trtllm_fp4_moe_hidden_dim_for_fi(
         return w13, w13_scale, w2, w2_scale, hidden_size
 
     logger.warning_once(
-        "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. "
+        "Padding hidden size from %d to %d for FlashInfer NVFP4 MoE weights. "
         "This requires activation slicing at runtime and may cause "
         "performance degradation.",
         hidden_size,