From 3652f35e7c75b8a750c746ecb4d8752c2a96ae2b Mon Sep 17 00:00:00 2001 From: Matej Sirovatka <54212263+S1ro1@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:36:08 -0700 Subject: [PATCH] [Bugfix][Quantization] Refresh online NVFP4 scales before reload post-processing (#57954) Signed-off-by: S1ro1 Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5 Co-authored-by: aoshen02 --- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 54 +++++++++++++++++++ tests/quantization/test_online.py | 37 ++++++++++--- .../fused_moe/experts/trtllm_nvfp4_moe.py | 25 ++++----- .../layers/quantization/online/nvfp4.py | 8 +++ 4 files changed, 104 insertions(+), 20 deletions(-) diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index b9684396b2..bc0711200b 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -277,6 +277,60 @@ def test_trtllm_fp4_moe_no_graph( torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1) +@torch.inference_mode() +def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params(): + """Weight reloads rerun post-processing on the same experts object; the + g1_alphas fold of clamp/beta must start from the unfolded values.""" + e, n, k = 8, 256, 256 + clamp, beta = 7.0, 1.0 + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + _, _, quant_config = make_test_quant_config( + e, + n, + k, + in_dtype=torch.bfloat16, + quant_dtype="nvfp4", + block_shape=None, + per_act_token_quant=False, + is_scale_swizzled=False, + ) + quant_config.gemm1_clamp_limit = clamp + quant_config.gemm1_beta = beta + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=2, + hidden_dim=k, + intermediate_size=n, + num_local_experts=e, + num_logical_experts=e, + activation=MoEActivation.SILU, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + routing_method=RoutingMethodType.TopK, + max_num_tokens=16, + ) + experts = TrtLlmNvFp4ExpertsModular( + moe_config=moe_config, quant_config=quant_config + ) + g1_alphas = quant_config.g1_alphas + assert g1_alphas is not None + layer = torch.nn.Module() + layer.w13_weight_scale_2 = g1_alphas + layer.w2_weight_scale_2 = quant_config.g2_alphas + layer.w13_input_scale = torch.ones_like(g1_alphas) + layer.w2_input_scale = torch.ones_like(g1_alphas) + + for new_alphas in (0.5, 0.25): + g1_alphas.fill_(new_alphas) + experts.process_weights_after_loading(layer) + expected = torch.full_like(g1_alphas, 1.0 / new_alphas) + torch.testing.assert_close(experts.gemm1_clamp_limit, clamp * expected) + torch.testing.assert_close(experts.gemm1_beta, beta * expected) + + if __name__ == "__main__": test_trtllm_fp4_moe_no_graph( 64, diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 743a10ecd1..645ad67062 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -216,16 +216,29 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( method.moe_kernel = None layer = Mock() - converted_weights = tuple(object() for _ in range(8)) - convert_weights = Mock(return_value=converted_weights) - process_weights = Mock() + converted_weights = [ + tuple(torch.full((2,), value + index) for index in range(8)) + for value in (1.0, 2.0, 0.5) + ] + convert_weights = Mock(side_effect=converted_weights) + quant_config = SimpleNamespace( + g1_alphas=converted_weights[0][2].clone(), + g2_alphas=converted_weights[0][6].clone(), + ) + alpha_ptrs = (quant_config.g1_alphas.data_ptr(), quant_config.g2_alphas.data_ptr()) + observed_scales = [] + process_weights = Mock( + side_effect=lambda _: observed_scales.append( + (quant_config.g1_alphas.clone(), quant_config.g2_alphas.clone()) + ) + ) kernel = SimpleNamespace( fused_experts=SimpleNamespace( process_weights_after_loading=process_weights, ) ) make_kernel = Mock(return_value=kernel) - get_quant_config = Mock(return_value=object()) + get_quant_config = Mock(return_value=quant_config) method.get_fused_moe_quant_config = get_quant_config monkeypatch.setattr( @@ -242,14 +255,22 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( make_kernel, ) - method._setup_kernel(layer) - method._setup_kernel(layer) + for weights in converted_weights: + method._setup_kernel(layer) + # Expert post-processing must observe fresh scales, with stable storage. + torch.testing.assert_close(observed_scales[-1][0], weights[2]) + torch.testing.assert_close(observed_scales[-1][1], weights[6]) + assert method.moe_quant_config is quant_config + assert ( + quant_config.g1_alphas.data_ptr(), + quant_config.g2_alphas.data_ptr(), + ) == alpha_ptrs assert method.moe_kernel is kernel - assert convert_weights.call_count == 2 + assert convert_weights.call_count == 3 make_kernel.assert_called_once() get_quant_config.assert_called_once() - assert process_weights.call_count == 2 + assert process_weights.call_count == 3 def _fully_quantized_quark_config() -> QuarkConfig: 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 eec846b5a0..646b4c8ad6 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 @@ -125,6 +125,10 @@ class TrtLlmNvFp4ExpertsBase: self.gemm1_beta = _per_expert(situ_linear_beta) self.gemm1_clamp_limit = None + # Unfolded values: process_weights_after_loading reruns on weight reload. + self._gemm1_clamp_limit_unfolded = self.gemm1_clamp_limit + self._gemm1_beta_unfolded = self.gemm1_beta + logger.debug_once( "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", moe_config.activation, @@ -164,15 +168,15 @@ class TrtLlmNvFp4ExpertsBase: # Pre-fold the per-expert g1_alphas (= output1_scale_gate_scalar) # division so the TRTLLM kernel receives the raw-GEMM-space clamp - # directly. g1_alphas is set once here in process_weights_after_loading - # (via the in-place mul above) and never changes again, so this is a - # static, per-expert constant. Register on the layer so EPLB - # rearranges it alongside the other expert tensors. + # directly. Fold from the unfolded value, since g1_alphas changes when + # weights are reloaded. Register on the layer so EPLB rearranges it + # alongside the other expert tensors. # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the # raw GEMM1 accumulator, so they are registered as-is without the # g1_alphas fold used by the SwiGLU-OAI clamp/beta below. - if self.gemm1_clamp_limit is not None and not self.is_situ: - gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas + clamp_limit = self._gemm1_clamp_limit_unfolded + if clamp_limit is not None and not self.is_situ: + gemm1_clamp_limit = clamp_limit / self.quant_config.g1_alphas layer.register_parameter( "gemm1_clamp_limit", torch.nn.Parameter(gemm1_clamp_limit, requires_grad=False), @@ -183,12 +187,9 @@ class TrtLlmNvFp4ExpertsBase: # clamp limit. alpha is applied to the dequantized gate, so it stays # raw. Register both on the layer so EPLB rearranges them with the # other per-expert tensors. - if self.gemm1_beta is not None: - gemm1_beta = ( - self.gemm1_beta - if self.is_situ - else self.gemm1_beta / self.quant_config.g1_alphas - ) + beta = self._gemm1_beta_unfolded + if beta is not None: + gemm1_beta = beta if self.is_situ else beta / self.quant_config.g1_alphas layer.register_parameter( "gemm1_beta", torch.nn.Parameter(gemm1_beta, requires_grad=False), diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index 2faa8f7838..efc97705c5 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -178,6 +178,14 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): routing_tables=layer._expert_routing_tables(), per_token_activation=True, ) + else: + # Reload creates new scale tensors; derived kernel scales must use + # their new values before layerwise reload restores captured storage. + assert self.moe_quant_config is not None + assert self.moe_quant_config.g1_alphas is not None + assert self.moe_quant_config.g2_alphas is not None + self.moe_quant_config.g1_alphas.copy_(w13_scale_2) + self.moe_quant_config.g2_alphas.copy_(w2_scale_2) self.moe_kernel.fused_experts.process_weights_after_loading(layer)