[Bugfix][Quantization] Reuse online NVFP4 MoE kernel across reloads (#50074)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Matej Sirovatka committed
Aug 11, 2026 at 23:06 UTC
3ac9525507b2d0de5c1b08cbca96cc94850c7c7a
2 files changed
+60
-10
tests/quantization/test_online.py
+48
@@ -4,6 +4,7 @@
4
5
from types import SimpleNamespace
6
from typing import cast
7
+from unittest.mock import Mock
8
9
import pytest
10
import torch
@@ -59,6 +60,53 @@ else:
60
DEVICE = current_platform.device_type
61
62
63
+def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
64
+ monkeypatch,
65
+) -> None:
66
+ method = object.__new__(Nvfp4OnlineMoEMethod)
67
+ method.moe = SimpleNamespace(is_act_and_mul=True)
68
+ method.nvfp4_backend = object()
69
+ method.experts_cls = object
70
+ method.moe_quant_config = None
71
+ method.moe_kernel = None
72
+
73
+ layer = Mock()
74
+ converted_weights = tuple(object() for _ in range(8))
75
+ convert_weights = Mock(return_value=converted_weights)
76
+ process_weights = Mock()
77
+ kernel = SimpleNamespace(
78
+ fused_experts=SimpleNamespace(
79
+ process_weights_after_loading=process_weights,
80
+ )
81
+ )
82
+ make_kernel = Mock(return_value=kernel)
83
+ get_quant_config = Mock(return_value=object())
84
+ method.get_fused_moe_quant_config = get_quant_config
85
+
86
+ monkeypatch.setattr(
87
+ "vllm.model_executor.layers.quantization.online.nvfp4."
88
+ "convert_to_nvfp4_moe_kernel_format",
89
+ convert_weights,
90
+ )
91
+ monkeypatch.setattr(
92
+ "vllm.model_executor.layers.quantization.online.nvfp4.replace_parameter",
93
+ Mock(),
94
+ )
95
+ monkeypatch.setattr(
96
+ "vllm.model_executor.layers.quantization.online.nvfp4.make_nvfp4_moe_kernel",
97
+ make_kernel,
98
+ )
99
+
100
+ method._setup_kernel(layer)
101
+ method._setup_kernel(layer)
102
+
103
+ assert method.moe_kernel is kernel
104
+ assert convert_weights.call_count == 2
105
+ make_kernel.assert_called_once()
106
+ get_quant_config.assert_called_once()
107
+ assert process_weights.call_count == 2
108
+
109
+
110
@pytest.mark.skipif(
111
not is_quant_method_supported("fp8"),
112
reason="FP8 is not supported on this GPU type.",
vllm/model_executor/layers/quantization/online/nvfp4.py
+12
-10
@@ -164,16 +164,18 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
164
replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
165
replace_parameter(layer, "w2_input_scale", a2_scale)
166
167
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
168
- assert self.experts_cls is not None
169
- self.moe_kernel = make_nvfp4_moe_kernel(
170
- moe_quant_config=self.moe_quant_config,
171
- moe_config=self.moe,
172
- experts_cls=self.experts_cls,
173
- backend=self.nvfp4_backend,
174
- routing_tables=layer._expert_routing_tables(),
175
- per_token_activation=True,
176
- )
167
+ if self.moe_kernel is None:
168
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
169
+ assert self.experts_cls is not None
170
+ self.moe_kernel = make_nvfp4_moe_kernel(
171
+ moe_quant_config=self.moe_quant_config,
172
+ moe_config=self.moe,
173
+ experts_cls=self.experts_cls,
174
+ backend=self.nvfp4_backend,
175
+ routing_tables=layer._expert_routing_tables(),
176
+ per_token_activation=True,
177
+ )
178
+
179
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
180
181
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: