[Bugfix][Quantization] Refresh online NVFP4 scales before reload post-processing (#57954)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
Matej Sirovatka committed
Sep 24, 2026 at 18:36 UTC
3652f35e7c75b8a750c746ecb4d8752c2a96ae2b
4 files changed
+104
-20
tests/kernels/moe/test_trtllm_nvfp4_moe.py
+54
@@ -277,6 +277,60 @@ def test_trtllm_fp4_moe_no_graph(
277
torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1)
278
279
280
+@torch.inference_mode()
281
+def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params():
282
+ """Weight reloads rerun post-processing on the same experts object; the
283
+ g1_alphas fold of clamp/beta must start from the unfolded values."""
284
+ e, n, k = 8, 256, 256
285
+ clamp, beta = 7.0, 1.0
286
+ with set_current_vllm_config(
287
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
288
+ ):
289
+ _, _, quant_config = make_test_quant_config(
290
+ e,
291
+ n,
292
+ k,
293
+ in_dtype=torch.bfloat16,
294
+ quant_dtype="nvfp4",
295
+ block_shape=None,
296
+ per_act_token_quant=False,
297
+ is_scale_swizzled=False,
298
+ )
299
+ quant_config.gemm1_clamp_limit = clamp
300
+ quant_config.gemm1_beta = beta
301
+ moe_config = FusedMoEConfig(
302
+ num_experts=e,
303
+ experts_per_token=2,
304
+ hidden_dim=k,
305
+ intermediate_size=n,
306
+ num_local_experts=e,
307
+ num_logical_experts=e,
308
+ activation=MoEActivation.SILU,
309
+ device="cuda",
310
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
311
+ in_dtype=torch.bfloat16,
312
+ routing_method=RoutingMethodType.TopK,
313
+ max_num_tokens=16,
314
+ )
315
+ experts = TrtLlmNvFp4ExpertsModular(
316
+ moe_config=moe_config, quant_config=quant_config
317
+ )
318
+ g1_alphas = quant_config.g1_alphas
319
+ assert g1_alphas is not None
320
+ layer = torch.nn.Module()
321
+ layer.w13_weight_scale_2 = g1_alphas
322
+ layer.w2_weight_scale_2 = quant_config.g2_alphas
323
+ layer.w13_input_scale = torch.ones_like(g1_alphas)
324
+ layer.w2_input_scale = torch.ones_like(g1_alphas)
325
+
326
+ for new_alphas in (0.5, 0.25):
327
+ g1_alphas.fill_(new_alphas)
328
+ experts.process_weights_after_loading(layer)
329
+ expected = torch.full_like(g1_alphas, 1.0 / new_alphas)
330
+ torch.testing.assert_close(experts.gemm1_clamp_limit, clamp * expected)
331
+ torch.testing.assert_close(experts.gemm1_beta, beta * expected)
332
+
333
+
334
if __name__ == "__main__":
335
test_trtllm_fp4_moe_no_graph(
336
64,
tests/quantization/test_online.py
+29
-8
@@ -216,16 +216,29 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
216
method.moe_kernel = None
217
218
layer = Mock()
219
- converted_weights = tuple(object() for _ in range(8))
220
- convert_weights = Mock(return_value=converted_weights)
221
- process_weights = Mock()
219
+ converted_weights = [
220
+ tuple(torch.full((2,), value + index) for index in range(8))
221
+ for value in (1.0, 2.0, 0.5)
222
+ ]
223
+ convert_weights = Mock(side_effect=converted_weights)
224
+ quant_config = SimpleNamespace(
225
+ g1_alphas=converted_weights[0][2].clone(),
226
+ g2_alphas=converted_weights[0][6].clone(),
227
+ )
228
+ alpha_ptrs = (quant_config.g1_alphas.data_ptr(), quant_config.g2_alphas.data_ptr())
229
+ observed_scales = []
230
+ process_weights = Mock(
231
+ side_effect=lambda _: observed_scales.append(
232
+ (quant_config.g1_alphas.clone(), quant_config.g2_alphas.clone())
233
+ )
234
+ )
235
kernel = SimpleNamespace(
236
fused_experts=SimpleNamespace(
237
process_weights_after_loading=process_weights,
238
)
239
)
240
make_kernel = Mock(return_value=kernel)
228
- get_quant_config = Mock(return_value=object())
241
+ get_quant_config = Mock(return_value=quant_config)
242
method.get_fused_moe_quant_config = get_quant_config
243
244
monkeypatch.setattr(
@@ -242,14 +255,22 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
255
make_kernel,
256
)
257
245
- method._setup_kernel(layer)
246
- method._setup_kernel(layer)
258
+ for weights in converted_weights:
259
+ method._setup_kernel(layer)
260
+ # Expert post-processing must observe fresh scales, with stable storage.
261
+ torch.testing.assert_close(observed_scales[-1][0], weights[2])
262
+ torch.testing.assert_close(observed_scales[-1][1], weights[6])
263
+ assert method.moe_quant_config is quant_config
264
+ assert (
265
+ quant_config.g1_alphas.data_ptr(),
266
+ quant_config.g2_alphas.data_ptr(),
267
+ ) == alpha_ptrs
268
269
assert method.moe_kernel is kernel
249
- assert convert_weights.call_count == 2
270
+ assert convert_weights.call_count == 3
271
make_kernel.assert_called_once()
272
get_quant_config.assert_called_once()
252
- assert process_weights.call_count == 2
273
+ assert process_weights.call_count == 3
274
275
276
def _fully_quantized_quark_config() -> QuarkConfig:
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+13
-12
@@ -125,6 +125,10 @@ class TrtLlmNvFp4ExpertsBase:
125
self.gemm1_beta = _per_expert(situ_linear_beta)
126
self.gemm1_clamp_limit = None
127
128
+ # Unfolded values: process_weights_after_loading reruns on weight reload.
129
+ self._gemm1_clamp_limit_unfolded = self.gemm1_clamp_limit
130
+ self._gemm1_beta_unfolded = self.gemm1_beta
131
+
132
logger.debug_once(
133
"activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
134
moe_config.activation,
@@ -164,15 +168,15 @@ class TrtLlmNvFp4ExpertsBase:
168
169
# Pre-fold the per-expert g1_alphas (= output1_scale_gate_scalar)
170
# division so the TRTLLM kernel receives the raw-GEMM-space clamp
167
- # directly. g1_alphas is set once here in process_weights_after_loading
168
- # (via the in-place mul above) and never changes again, so this is a
169
- # static, per-expert constant. Register on the layer so EPLB
170
- # rearranges it alongside the other expert tensors.
171
+ # directly. Fold from the unfolded value, since g1_alphas changes when
172
+ # weights are reloaded. Register on the layer so EPLB rearranges it
173
+ # alongside the other expert tensors.
174
# SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the
175
# raw GEMM1 accumulator, so they are registered as-is without the
176
# g1_alphas fold used by the SwiGLU-OAI clamp/beta below.
174
- if self.gemm1_clamp_limit is not None and not self.is_situ:
175
- gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas
177
+ clamp_limit = self._gemm1_clamp_limit_unfolded
178
+ if clamp_limit is not None and not self.is_situ:
179
+ gemm1_clamp_limit = clamp_limit / self.quant_config.g1_alphas
180
layer.register_parameter(
181
"gemm1_clamp_limit",
182
torch.nn.Parameter(gemm1_clamp_limit, requires_grad=False),
@@ -183,12 +187,9 @@ class TrtLlmNvFp4ExpertsBase:
187
# clamp limit. alpha is applied to the dequantized gate, so it stays
188
# raw. Register both on the layer so EPLB rearranges them with the
189
# other per-expert tensors.
186
- if self.gemm1_beta is not None:
187
- gemm1_beta = (
188
- self.gemm1_beta
189
- if self.is_situ
190
- else self.gemm1_beta / self.quant_config.g1_alphas
191
- )
190
+ beta = self._gemm1_beta_unfolded
191
+ if beta is not None:
192
+ gemm1_beta = beta if self.is_situ else beta / self.quant_config.g1_alphas
193
layer.register_parameter(
194
"gemm1_beta",
195
torch.nn.Parameter(gemm1_beta, requires_grad=False),
vllm/model_executor/layers/quantization/online/nvfp4.py
+8
@@ -178,6 +178,14 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
178
routing_tables=layer._expert_routing_tables(),
179
per_token_activation=True,
180
)
181
+ else:
182
+ # Reload creates new scale tensors; derived kernel scales must use
183
+ # their new values before layerwise reload restores captured storage.
184
+ assert self.moe_quant_config is not None
185
+ assert self.moe_quant_config.g1_alphas is not None
186
+ assert self.moe_quant_config.g2_alphas is not None
187
+ self.moe_quant_config.g1_alphas.copy_(w13_scale_2)
188
+ self.moe_quant_config.g2_alphas.copy_(w2_scale_2)
189
190
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
191