[Quant] Add `nvfp4_per_token` online MoE quantization (#48538)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 16, 2026 at 17:25 UTC
c95c6630499f96ddee97a87ba1ca4e6b85afe0fe
7 files changed
+285
-5
tests/quantization/test_online.py
+36
@@ -16,7 +16,11 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
16
Fp8PerTensorOnlineLinearMethod,
17
Fp8PerTensorOnlineMoEMethod,
18
)
19
+from vllm.model_executor.layers.quantization.online.nvfp4 import (
20
+ Nvfp4OnlineMoEMethod,
21
+)
22
from vllm.platforms import current_platform
23
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
24
25
26
@pytest.mark.skipif(
@@ -145,6 +149,38 @@ def test_online_quantization(
149
print(outputs[0][1])
150
151
152
+@pytest.mark.skipif(
153
+ not (
154
+ current_platform.is_cuda()
155
+ and current_platform.is_device_capability_family(100)
156
+ and has_flashinfer_trtllm_fused_moe()
157
+ ),
158
+ reason="nvfp4_per_token needs a Blackwell (SM100) GPU + FlashInfer TRTLLM MoE.",
159
+)
160
+def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None:
161
+ """Online NVFP4 quantizes the MoE and leaves dense layers unquantized."""
162
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
163
+
164
+ with vllm_runner(
165
+ "ibm-granite/granite-3.0-1b-a400m-base",
166
+ quantization="nvfp4_per_token",
167
+ enforce_eager=True,
168
+ ) as llm:
169
+
170
+ def check_model(model):
171
+ layer = model.model.layers[0]
172
+ assert isinstance(
173
+ layer.block_sparse_moe.experts._quant_method, Nvfp4OnlineMoEMethod
174
+ )
175
+ assert isinstance(
176
+ layer.self_attn.o_proj.quant_method, UnquantizedLinearMethod
177
+ )
178
+
179
+ llm.apply_model(check_model)
180
+ outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
181
+ print(outputs[0][1])
182
+
183
+
184
@pytest.mark.skipif(
185
not is_quant_method_supported("fp8"),
186
reason="FP8 is not supported on this GPU type.",
vllm/config/quantization.py
+6
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
kInt8StaticChannelSym,
19
kMxfp4Dynamic,
20
kMxfp8Dynamic,
21
+ kNvfp4Static,
22
)
23
24
# User-facing names addressable from quantization_config.
@@ -134,6 +135,11 @@ _ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = {
135
"int8_per_channel_weight_only": QuantizationConfigArgs(
136
moe=QuantSpec(weight=kInt8StaticChannelSym),
137
),
138
+ # Online NVFP4 on MoE with per-token dynamic activation scales (Blackwell +
139
+ # FlashInfer TRTLLM only); linear stays unquantized (no `linear` field).
140
+ "nvfp4_per_token": QuantizationConfigArgs(
141
+ moe=QuantSpec(weight=kNvfp4Static),
142
+ ),
143
}
144
145
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+63
-5
@@ -33,6 +33,10 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
33
34
logger = init_logger(__name__)
35
36
+# Base scale for per-token NVFP4 activation quant; the kernel folds the
37
+# per-token global scale (from the activation amax) on top of it.
38
+_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0)
39
+
40
41
class TrtLlmNvFp4ExpertsBase:
42
"""
@@ -43,9 +47,13 @@ class TrtLlmNvFp4ExpertsBase:
47
self,
48
moe_config: FusedMoEConfig,
49
quant_config: FusedMoEQuantConfig,
50
+ per_token_activation: bool = False,
51
):
52
self.moe_config = moe_config
53
self.quant_config = quant_config
54
+ # Quantize the input here (deferred from prepare) to capture a per-token
55
+ # global scale, instead of a static one.
56
+ self.per_token_activation = per_token_activation
57
58
self.routing_method_type = self.moe_config.routing_method
59
self.topk = moe_config.experts_per_token
@@ -211,6 +219,27 @@ class TrtLlmNvFp4ExpertsBase:
219
def activation_format() -> mk.FusedMoEActivationFormat:
220
return mk.FusedMoEActivationFormat.Standard
221
222
+ @property
223
+ def expects_unquantized_inputs(self) -> bool:
224
+ return self.per_token_activation
225
+
226
+ def _quantize_per_token_input(
227
+ self, hidden_states: torch.Tensor
228
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
229
+ """NVFP4-quantize activations with a per-token global scale.
230
+
231
+ Returns ``(packed_fp4, block_scale, per_token_scale)``.
232
+ """
233
+ from flashinfer import SfLayout, nvfp4_quantize
234
+
235
+ hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize(
236
+ hidden_states,
237
+ _PER_TOKEN_BASE_GLOBAL_SCALE,
238
+ sfLayout=SfLayout.layout_linear,
239
+ per_token_activation=True,
240
+ )
241
+ return hs_fp4, hs_block_scale, per_token_scale
242
+
243
def _get_chunk_size(self) -> int:
244
MAX_GRID_Y = 65535
245
MAX_TILE_TOKENS_DIM = 128
@@ -255,6 +284,15 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
284
expert_tokens_meta: mk.ExpertTokensMetadata | None,
285
activation: MoEActivation,
286
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
287
+ if self.per_token_activation:
288
+ # Deferred input quant leaves K unpacked here, breaking the
289
+ # workspace assumptions below. Per-token NVFP4 is only supported on
290
+ # the monolithic (non-EP) path for now.
291
+ raise NotImplementedError(
292
+ "NVFP4 per-token activation is only supported on the monolithic "
293
+ "(non-EP) FlashInfer TRTLLM MoE path."
294
+ )
295
+
296
# The workspaces for this implementation are managed by flashinfer.
297
workspace1 = (0,)
298
workspace2 = (0,)
@@ -286,6 +324,15 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
324
assert self.quant_config.w1_scale is not None
325
assert self.quant_config.w2_scale is not None
326
327
+ # Per-token: input is unquantized, quantize it here. Otherwise it was
328
+ # already quantized in prepare() with the static global scale.
329
+ if self.per_token_activation:
330
+ hidden_states, block_scale, per_token_scale = (
331
+ self._quantize_per_token_input(hidden_states)
332
+ )
333
+ else:
334
+ block_scale, per_token_scale = a1q_scale, None
335
+
336
# Pack topk ids and weights into format expected by the kernel.
337
packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
338
output1_scale_gate_scalar = self.quant_config.g1_alphas
@@ -295,7 +342,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
342
topk_ids=packed_tensor,
343
routing_bias=None,
344
hidden_states=hidden_states,
298
- hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(
345
+ hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(
346
*hidden_states.shape[:-1], -1
347
),
348
gemm1_weights=w1,
@@ -321,6 +368,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
368
routing_method_type=1, # not used
369
do_finalize=True,
370
activation_type=activation_to_flashinfer_int(activation),
371
+ per_token_scale=per_token_scale,
372
output=output,
373
tune_max_num_tokens=min(
374
fi_moe_largest_bucket(self.moe_config), self._get_chunk_size()
@@ -346,7 +394,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
394
apply_router_weight_on_input: bool,
395
):
396
assert self._supports_activation(activation)
349
- assert a1q_scale is not None
397
+ # Per-token defers input quant to _invoke_kernel, so a1q_scale is None.
398
+ assert a1q_scale is not None or self.per_token_activation
399
400
M = hidden_states.shape[0]
401
chunk_size = self._get_chunk_size()
@@ -375,7 +424,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
424
topk_ids[start:end],
425
activation,
426
global_num_experts,
378
- a1q_scale[start:end],
427
+ None if a1q_scale is None else a1q_scale[start:end],
428
)
429
430
@@ -439,7 +488,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
488
import flashinfer
489
490
assert self._supports_activation(activation)
442
- assert a1q_scale is not None
491
+ assert a1q_scale is not None or self.per_token_activation
492
assert self.quant_config.w1_scale is not None
493
assert self.quant_config.w2_scale is not None
494
assert (
@@ -450,6 +499,14 @@ class TrtLlmNvFp4ExpertsMonolithic(
499
and self.routing_method_type != RoutingMethodType.Llama4
500
)
501
502
+ # Per-token: input is unquantized, quantize it here (see modular apply).
503
+ if self.per_token_activation:
504
+ hidden_states, block_scale, per_token_scale = (
505
+ self._quantize_per_token_input(hidden_states)
506
+ )
507
+ else:
508
+ block_scale, per_token_scale = a1q_scale, None
509
+
510
output1_scale_gate_scalar = self.quant_config.g1_alphas
511
512
# Invoke kernel.
@@ -459,7 +516,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
516
routing_logits=router_logits,
517
routing_bias=e_score_correction_bias,
518
hidden_states=hidden_states,
462
- hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(
519
+ hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(
520
*hidden_states.shape[:-1], -1
521
),
522
gemm1_weights=w1,
@@ -485,5 +542,6 @@ class TrtLlmNvFp4ExpertsMonolithic(
542
routing_method_type=self.routing_method_type,
543
do_finalize=True,
544
activation_type=activation_to_flashinfer_int(activation),
545
+ per_token_scale=per_token_scale,
546
tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
547
)[0]
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+3
@@ -529,6 +529,7 @@ def make_nvfp4_moe_kernel(
529
backend: NvFp4MoeBackend,
530
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
531
layer: torch.nn.Module | None = None,
532
+ per_token_activation: bool = False,
533
) -> mk.FusedMoEKernel:
534
# Create Prepare/Finalize.
535
prepare_finalize = maybe_make_prepare_finalize(
@@ -546,6 +547,8 @@ def make_nvfp4_moe_kernel(
547
if backend == NvFp4MoeBackend.HUMMING:
548
assert layer is not None
549
extra_kwargs = {"layer": layer}
550
+ if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and per_token_activation:
551
+ extra_kwargs["per_token_activation"] = True
552
553
# Create Experts.
554
if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
vllm/model_executor/layers/quantization/__init__.py
+1
@@ -42,6 +42,7 @@ QuantizationMethods = Literal[
42
"fp8_per_block",
43
"fp8_per_channel",
44
"int8_per_channel_weight_only",
45
+ "nvfp4_per_token",
46
"mxfp8",
47
]
48
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
vllm/model_executor/layers/quantization/online/base.py
+5
@@ -40,6 +40,9 @@ from vllm.model_executor.layers.quantization.online.mxfp8 import (
40
Mxfp8OnlineLinearMethod,
41
Mxfp8OnlineMoEMethod,
42
)
43
+from vllm.model_executor.layers.quantization.online.nvfp4 import (
44
+ Nvfp4OnlineMoEMethod,
45
+)
46
from vllm.model_executor.layers.quantization.utils.quant_utils import (
47
QuantKey,
48
kFp8Static128BlockSym,
@@ -47,6 +50,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
50
kFp8StaticTensorSym,
51
kInt8StaticChannelSym,
52
kMxfp8Dynamic,
53
+ kNvfp4Static,
54
)
55
56
logger = init_logger(__name__)
@@ -68,6 +72,7 @@ _ONLINE_MOE_METHODS: dict[QuantKey, type] = {
72
kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod,
73
kMxfp8Dynamic: Mxfp8OnlineMoEMethod,
74
kInt8StaticChannelSym: Int8OnlineMoEMethod,
75
+ kNvfp4Static: Nvfp4OnlineMoEMethod,
76
}
77
78
vllm/model_executor/layers/quantization/online/nvfp4.py
new
+171
@@ -0,0 +1,171 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch.nn import Module
6
+
7
+from vllm._custom_ops import scaled_fp4_quant
8
+from vllm.model_executor.layers.fused_moe import RoutedExperts
9
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
11
+ convert_to_nvfp4_moe_kernel_format,
12
+ make_nvfp4_moe_kernel,
13
+ make_nvfp4_moe_quant_config,
14
+ select_nvfp4_moe_backend,
15
+)
16
+from vllm.model_executor.layers.quantization.online.moe_base import (
17
+ OnlineMoEMethodBase,
18
+)
19
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
20
+ FLOAT4_E2M1_MAX,
21
+)
22
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
+ kNvfp4Dynamic,
24
+ kNvfp4Static,
25
+)
26
+from vllm.model_executor.utils import replace_parameter
27
+from vllm.platforms import current_platform
28
+
29
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
30
+
31
+
32
+def _quantize_moe_weight_to_nvfp4(
33
+ weight: torch.Tensor,
34
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
35
+ """Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4.
36
+
37
+ One FP32 global scale per expert plus per-block (group-16) FP8 scales,
38
+ matching the ModelOpt NVFP4 checkpoint layout. Returns packed FP4 weights
39
+ ``(E, N, K // 2)``, block scales ``(E, N, K // 16)``, and the per-expert
40
+ global scale ``(E,)`` stored as ``amax / (fp4_max * fp8_max)``.
41
+ """
42
+ assert weight.dim() == 3, f"expected 3D expert weights, got {weight.shape}"
43
+ num_experts, n, k = weight.shape
44
+ assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}"
45
+
46
+ amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8)
47
+ global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax
48
+ weight_scale_2 = (1.0 / global_scale).to(torch.float32)
49
+
50
+ # scaled_fp4_quant(w, g) == scaled_fp4_quant(w * g, 1), so fold each
51
+ # expert's scale in and quantize all experts in one call (fp32 to keep the
52
+ # large scale precise), rather than looping per expert.
53
+ scaled = (weight.float() * global_scale[:, None, None]).to(weight.dtype)
54
+ scaled = scaled.reshape(-1, k)
55
+ one = torch.ones((), device=weight.device, dtype=torch.float32)
56
+ qweight, block_scale = scaled_fp4_quant(scaled, one, is_sf_swizzled_layout=False)
57
+ return (
58
+ qweight.reshape(num_experts, n, k // 2),
59
+ block_scale.reshape(num_experts, n, k // 16),
60
+ weight_scale_2,
61
+ )
62
+
63
+
64
+class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
65
+ """Online NVFP4 MoE quantization with per-token activation scales.
66
+
67
+ Quantizes fp16/bf16 expert weights to NVFP4 at load time; the FlashInfer
68
+ TRTLLM kernel computes per-token activation scales at runtime. Blackwell
69
+ (SM100) only.
70
+ """
71
+
72
+ def __init__(
73
+ self,
74
+ *,
75
+ layer: torch.nn.Module,
76
+ ):
77
+ if not current_platform.is_device_capability_family(100):
78
+ raise ValueError(
79
+ "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU."
80
+ )
81
+ super().__init__(layer.moe_config)
82
+ self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
83
+ config=self.moe,
84
+ weight_key=kNvfp4Static,
85
+ activation_key=kNvfp4Dynamic,
86
+ )
87
+
88
+ def process_weights_after_loading(self, layer: Module) -> None:
89
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
90
+ return
91
+
92
+ self._quantize_weights(layer)
93
+ self._setup_kernel(layer)
94
+
95
+ layer._already_called_process_weights_after_loading = True
96
+
97
+ def _quantize_weights(self, layer: Module) -> None:
98
+ w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w13_weight)
99
+ w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w2_weight)
100
+
101
+ replace_parameter(layer, "w13_weight", w13)
102
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
103
+ replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
104
+ replace_parameter(layer, "w2_weight", w2)
105
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
106
+ replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
107
+
108
+ # Neutral (1.0) activation global scales: the kernel derives per-token
109
+ # scales at runtime, so the output scalars reduce to the weight scales.
110
+ ones = torch.ones(layer.num_experts, device=w13.device, dtype=torch.float32)
111
+ replace_parameter(layer, "w13_input_scale", ones)
112
+ replace_parameter(layer, "w2_input_scale", ones.clone())
113
+
114
+ def _setup_kernel(self, layer: RoutedExperts) -> None:
115
+ (
116
+ w13,
117
+ w13_scale,
118
+ w13_scale_2,
119
+ a13_scale,
120
+ w2,
121
+ w2_scale,
122
+ w2_scale_2,
123
+ a2_scale,
124
+ ) = convert_to_nvfp4_moe_kernel_format(
125
+ nvfp4_backend=self.nvfp4_backend,
126
+ layer=layer,
127
+ w13=layer.w13_weight,
128
+ w13_scale=layer.w13_weight_scale,
129
+ w13_scale_2=layer.w13_weight_scale_2,
130
+ a13_scale=layer.w13_input_scale,
131
+ w2=layer.w2_weight,
132
+ w2_scale=layer.w2_weight_scale,
133
+ w2_scale_2=layer.w2_weight_scale_2,
134
+ a2_scale=layer.w2_input_scale,
135
+ is_act_and_mul=self.moe.is_act_and_mul,
136
+ )
137
+
138
+ replace_parameter(layer, "w13_weight", w13)
139
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
140
+ replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
141
+ replace_parameter(layer, "w13_input_scale", a13_scale)
142
+ replace_parameter(layer, "w2_weight", w2)
143
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
144
+ replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
145
+ replace_parameter(layer, "w2_input_scale", a2_scale)
146
+
147
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
148
+ assert self.experts_cls is not None
149
+ self.moe_kernel = make_nvfp4_moe_kernel(
150
+ moe_quant_config=self.moe_quant_config,
151
+ moe_config=self.moe,
152
+ experts_cls=self.experts_cls,
153
+ backend=self.nvfp4_backend,
154
+ routing_tables=layer._expert_routing_tables(),
155
+ layer=layer,
156
+ per_token_activation=True,
157
+ )
158
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
159
+
160
+ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
161
+ return make_nvfp4_moe_quant_config(
162
+ backend=self.nvfp4_backend,
163
+ w13_scale=layer.w13_weight_scale,
164
+ w2_scale=layer.w2_weight_scale,
165
+ w13_scale_2=layer.w13_weight_scale_2,
166
+ w2_scale_2=layer.w2_weight_scale_2,
167
+ a13_scale=layer.w13_input_scale,
168
+ a2_scale=layer.w2_input_scale,
169
+ swiglu_limit=getattr(layer, "swiglu_limit", None),
170
+ layer=layer,
171
+ )