iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Quantization] Refresh online NVFP4 scales before reload post-processing (#57954)

Browse Source 
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>
This commit is contained in:
Matej Sirovatka
2026-09-24 18:36:08 -07:00
committed by GitHub
co-authored by aoshen02 Claude Opus 5 aoshen02
parent 16070ed7e9
commit 3652f35e7c
4 changed files with 104 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_trtllm_nvfp4_moe.py
+54
View File
@@ -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,
tests/quantization/test_online.py
+29 -8
View File
@@ -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:
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+13 -12
View File
@@ -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),
vllm/model_executor/layers/quantization/online/nvfp4.py
+8
View File
@@ -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)