From 9127170b880dc86e032c4acaf649c225fdeff3b1 Mon Sep 17 00:00:00 2001 From: Matej Sirovatka <54212263+S1ro1@users.noreply.github.com> Date: Wed, 23 Sep 2026 00:30:58 -0700 Subject: [PATCH] [Quantization] Select per-token NVFP4 MoE backends explicitly (#57176) Signed-off-by: aoshen02 Signed-off-by: S1ro1 Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: OpenAI Codex Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/quantization/test_online.py | 71 +++++++++++++++++++ .../layers/fused_moe/all2all_utils.py | 14 +++- .../fused_moe/experts/trtllm_nvfp4_moe.py | 25 +++---- .../layers/fused_moe/oracle/nvfp4.py | 1 + .../layers/quantization/online/nvfp4.py | 4 +- .../layers/quantization/utils/quant_utils.py | 3 + 6 files changed, 97 insertions(+), 21 deletions(-) diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 846b981478..4fc908f484 100644 --- a/tests/quantization/test_online.py +++ b/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.", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index eb7edaef53..3186b5d3c4 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index cea985830d..eec846b5a0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 0d794584d7..0fecb65e4b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/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 diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index 7abefce257..2faa8f7838 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/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: diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index f3f72adc78..4c0a8a3ff7 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/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(