[Quantization] Select per-token NVFP4 MoE backends explicitly (#57176)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.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>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/quantization/test_online.py

@@ -36,6 +36,7 @@ from vllm.model_executor.kernels.linear.mxfp8.marlin import (
 )
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     LinearBase,
@@ -503,6 +504,76 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:
         ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")
 
 
+@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(
+    monkeypatch, per_token_activation: bool, input_dtype: torch.dtype
+) -> None:
+    from tests.kernels.moe.utils import make_dummy_moe_config
+    from vllm.model_executor.layers.fused_moe.all2all_utils import (
+        flashinfer_one_sided_dispatch_layout,
+    )
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+        TrtLlmNvFp4ExpertsModular,
+    )
+
+    config = make_dummy_moe_config(hidden_dim=6144, in_dtype=input_dtype)
+    config.moe_parallel_config.dp_size = 2
+    config.moe_parallel_config.ep_size = 2
+    config.moe_parallel_config.use_ep = True
+    config.moe_parallel_config.all2all_backend = "flashinfer_nvlink_one_sided"
+    expert_quant_config = FusedMoEQuantConfig.make("nvfp4")
+
+    def check_dispatch_layout(*, quant_config, input_dtype, **kwargs):
+        assert quant_config is expert_quant_config
+        layout = flashinfer_one_sided_dispatch_layout(
+            config.hidden_dim, quant_config, input_dtype=input_dtype
+        )
+        if per_token_activation:
+            assert input_dtype == config.in_dtype
+            assert (
+                layout.x_bytes_per_token == config.hidden_dim * config.in_dtype.itemsize
+            )
+            assert layout.x_sf_bytes_per_token == 0
+        else:
+            assert input_dtype is None
+            assert layout.x_bytes_per_token == config.hidden_dim // 2
+            assert layout.x_sf_bytes_per_token == config.hidden_dim // 16
+        assert expert_quant_config.quant_dtype == "nvfp4"
+        raise RuntimeError("dispatch layout verified before allocation")
+
+    monkeypatch.setattr(
+        nvfp4_oracle, "maybe_make_prepare_finalize", check_dispatch_layout
+    )
+    with pytest.raises(
+        RuntimeError, match="dispatch layout verified before allocation"
+    ):
+        nvfp4_oracle.make_nvfp4_moe_kernel(
+            moe_quant_config=expert_quant_config,
+            moe_config=config,
+            experts_cls=TrtLlmNvFp4ExpertsModular,
+            backend=nvfp4_oracle.NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            per_token_activation=per_token_activation,
+        )
+
+
+@pytest.mark.parametrize(
+    "input_dtype",
+    [torch.float32, torch.float8_e4m3fn, torch.float8_e5m2, torch.int8],
+)
+def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtype):
+    from vllm.model_executor.layers.fused_moe.all2all_utils import (
+        flashinfer_one_sided_dispatch_layout,
+    )
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+
+    with pytest.raises(ValueError, match="unpacked inputs must be float16 or bfloat16"):
+        flashinfer_one_sided_dispatch_layout(
+            6144, FusedMoEQuantConfig.make("nvfp4"), input_dtype=input_dtype
+        )
+
+
 @pytest.mark.skipif(
     not is_quant_method_supported("fp8"),
     reason="FP8 is not supported on this GPU type.",

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -48,9 +48,18 @@ class FlashInferOneSidedDispatchLayout:
 
 
 def flashinfer_one_sided_dispatch_layout(
-    hidden_dim: int, quant_config: FusedMoEQuantConfig
+    hidden_dim: int,
+    quant_config: FusedMoEQuantConfig,
+    input_dtype: torch.dtype | None = None,
 ) -> FlashInferOneSidedDispatchLayout:
     """Return the one-sided activation payload layout."""
+    if input_dtype is not None:
+        if input_dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                "flashinfer_nvlink_one_sided unpacked inputs must be float16 "
+                f"or bfloat16, got {input_dtype}"
+            )
+        return FlashInferOneSidedDispatchLayout(hidden_dim * input_dtype.itemsize, 0)
     if quant_config.quant_dtype is None:
         return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)
     if quant_config.quant_dtype == "nvfp4":
@@ -174,6 +183,7 @@ def maybe_make_prepare_finalize(
     allow_new_interface: bool = False,
     use_monolithic: bool = False,
     all2all_manager: Any | None = None,
+    input_dtype: torch.dtype | None = None,
 ) -> FusedMoEPrepareAndFinalize | None:
     if not moe.moe_parallel_config.use_all2all_kernels:
         if not allow_new_interface:
@@ -367,7 +377,7 @@ def maybe_make_prepare_finalize(
             get_current_vllm_config().scheduler_config.max_num_batched_tokens
         )
         dispatch_layout = flashinfer_one_sided_dispatch_layout(
-            moe.hidden_dim, quant_config
+            moe.hidden_dim, quant_config, input_dtype=input_dtype
         )
         prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
             max_num_tokens=max_num_tokens,

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

@@ -28,6 +28,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kNvfp4Dynamic,
+    kNvfp4DynamicToken,
     kNvfp4Static,
 )
 from vllm.model_executor.utils import is_weights_pre_processed
@@ -222,10 +223,10 @@ class TrtLlmNvFp4ExpertsBase:
         activation_key: QuantKey | None,
     ) -> bool:
         """Supports Nvfp4 quantization."""
-        SUPPORTED_W_A = [
-            (kNvfp4Static, kNvfp4Dynamic),
-        ]
-        return (weight_key, activation_key) in SUPPORTED_W_A
+        return weight_key == kNvfp4Static and activation_key in (
+            kNvfp4Dynamic,
+            kNvfp4DynamicToken,
+        )
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
@@ -316,22 +317,12 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        if self.per_token_activation:
-            # Deferred input quant leaves K unpacked here, breaking the
-            # workspace assumptions below. Per-token NVFP4 is only supported on
-            # the monolithic (non-EP) path for now.
-            raise NotImplementedError(
-                "NVFP4 per-token activation is only supported on the monolithic "
-                "(non-EP) FlashInfer TRTLLM MoE path."
-            )
-
         # The workspaces for this implementation are managed by flashinfer.
         workspace1 = (0,)
         workspace2 = (0,)
 
-        # Hidden states are Nvfp4, packed into int8 dtype, so we
-        # need to multiply K by 2 to get the output shape right.
-        assert self.hidden_dim == K * 2
+        # Per-token inputs are unpacked; otherwise each byte holds two FP4 values.
+        assert self.hidden_dim == (K if self.expects_unquantized_inputs else K * 2)
         output = (M, self.hidden_dim)
 
         return (workspace1, workspace2, output)
@@ -349,7 +340,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         topk_ids: torch.Tensor,
         activation: MoEActivation,
         global_num_experts: int,
-        a1q_scale: torch.Tensor,
+        a1q_scale: torch.Tensor | None,
     ):
         import flashinfer
 

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

@@ -605,6 +605,7 @@ def make_nvfp4_moe_kernel(
         routing_tables=routing_tables,
         allow_new_interface=True,
         use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),
+        input_dtype=moe_config.in_dtype if per_token_activation else None,
     )
     assert prepare_finalize is not None
 

vllm/model_executor/layers/quantization/online/nvfp4.py

@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     amax_for_moe_weight_quant,
-    kNvfp4Dynamic,
+    kNvfp4DynamicToken,
     kNvfp4Static,
     weight_amax,
 )
@@ -100,7 +100,7 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
         self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
             config=self.moe,
             weight_key=kNvfp4Static,
-            activation_key=kNvfp4Dynamic,
+            activation_key=kNvfp4DynamicToken,
         )
 
     def process_weights_after_loading(self, layer: Module) -> None:

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

@@ -208,6 +208,9 @@ kNvfp4DynamicGroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))
 kNvfp4Dynamic = QuantKey(
     FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kStaticTensorScale
 )
+kNvfp4DynamicToken = QuantKey(
+    FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale
+)
 
 kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16))
 kNvfp4Static = QuantKey(