[Quantization] Let ModelOpt MXFP8 layers load pre-processed weights Purpose (#57316)
[Quantization] Let ModelOpt MXFP8 layers load pre-processed weights Purpose (#57316) Signed-off-by: liusy58 <mg21330037@smail.nju.edu.cn> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com>
siyu committed
Sep 19, 2026 at 17:16 UTC
4cc15f2121b3421478f3186aac38de901b4a44ef
6 files changed
+57
-4
tests/quantization/test_modelopt.py
+30
@@ -144,6 +144,36 @@ def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypat
144
assert torch.equal(linear.weight_scale, scales)
145
146
147
+def test_modelopt_mxfp8_pre_processed_weights_follow_kernel(monkeypatch):
148
+ """The weight cache daemon can only serve MXFP8 layers whose kernel needs
149
+ no post-load state beyond the parameters it exports."""
150
+ from vllm.config.quantization import QuantSpec
151
+ from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearKernel
152
+
153
+ method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
154
+ method.spec = QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic)
155
+ method.kernel = None
156
+ assert not method.supports_pre_processed_weights
157
+
158
+ method.kernel = Mock(spec=Mxfp8LinearKernel)
159
+ method.kernel.supports_pre_processed_weights = False
160
+ assert not method.supports_pre_processed_weights
161
+ method.kernel.supports_pre_processed_weights = True
162
+ assert method.supports_pre_processed_weights
163
+
164
+ monkeypatch.setattr(
165
+ "vllm.model_executor.layers.quantization.modelopt.is_weights_pre_processed",
166
+ lambda: True,
167
+ )
168
+ layer = torch.nn.Module()
169
+ method.process_weights_after_loading(layer)
170
+ method.kernel.process_weights_after_loading.assert_not_called()
171
+
172
+ method.kernel.supports_pre_processed_weights = False
173
+ with pytest.raises(RuntimeError, match="pre-processed"):
174
+ method.process_weights_after_loading(layer)
175
+
176
+
177
def test_modelopt_fp8_updates_weight_dims_after_transpose():
178
"""Humming reads weight.input_dim/output_dim. Swapping the
179
ModelWeightParameter for a plain Parameter drops them, so the per-tensor
vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
+4
@@ -27,6 +27,10 @@ class Mxfp8LinearKernel(ABC):
27
Marlin, emulation).
28
"""
29
30
+ supports_pre_processed_weights: bool = False
31
+ """True if ``process_weights_after_loading`` only rewrites parameters, so
32
+ weights exported by the weight cache daemon can be used as-is."""
33
+
34
def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
35
assert self.can_implement(c)[0]
36
assert self.is_supported()[0]
vllm/model_executor/kernels/linear/mxfp8/emulation.py
+2
@@ -16,6 +16,8 @@ from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
16
class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
17
"""Software emulation fallback for MXFP8 (dequant to BF16)."""
18
19
+ supports_pre_processed_weights = True
20
+
21
@classmethod
22
def is_supported(
23
cls, compute_capability: int | None = None
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+4
@@ -27,6 +27,8 @@ from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
27
class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
28
"""MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
29
30
+ supports_pre_processed_weights = True
31
+
32
@classmethod
33
def is_supported(
34
cls, compute_capability: int | None = None
@@ -117,6 +119,8 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
119
class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
120
"""MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
121
122
+ supports_pre_processed_weights = True
123
+
124
@classmethod
125
def is_supported(
126
cls, compute_capability: int | None = None
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+2
@@ -212,6 +212,8 @@ _DOT_SCALED_K_ALIGN = 128
212
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
213
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
214
215
+ supports_pre_processed_weights = True
216
+
217
@classmethod
218
def is_supported(
219
cls, compute_capability: int | None = None
vllm/model_executor/layers/quantization/modelopt.py
+15
-4
@@ -2482,9 +2482,16 @@ class ModelOptLinearMethod(LinearMethodBase):
2482
2483
@property
2484
def supports_pre_processed_weights(self) -> bool: # type: ignore[override]
2485
- # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack.
2485
+ # TODO(Isotr0py): support fp8 ModelOpt kernels transpose/repack.
2486
w = self.spec.weight
2487
- return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE
2487
+ return isinstance(w, QuantKey) and (
2488
+ w.dtype == FP4_DTYPE
2489
+ or (
2490
+ w == kMxfp8Static
2491
+ and self.kernel is not None
2492
+ and self.kernel.supports_pre_processed_weights
2493
+ )
2494
+ )
2495
2496
def create_weights(
2497
self,
@@ -2524,7 +2531,13 @@ class ModelOptLinearMethod(LinearMethodBase):
2531
expose_input_quant_key(layer, self.kernel)
2532
2533
def process_weights_after_loading(self, layer) -> None:
2534
+ if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
2535
+ self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
2536
if is_weights_pre_processed():
2537
+ if not self.supports_pre_processed_weights:
2538
+ raise RuntimeError(
2539
+ f"{type(self.kernel).__name__} cannot use pre-processed weights"
2540
+ )
2541
return
2542
self.fmt.pre_process(layer)
2543
self.wkey.process(layer, WEIGHT)
@@ -2556,8 +2569,6 @@ class ModelOptLinearMethod(LinearMethodBase):
2569
persistent=False,
2570
)
2571
layer._nvfp4_group_size_for_gather = self.ctx.group_size
2559
- if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
2560
- self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
2572
self.kernel.process_weights_after_loading(layer)
2573
2574
def apply(self, layer, x, bias=None):