[ModelOpt] Redesign the LinearMethod classes using the generic QuantKey-driven method (#49381)
Signed-off-by: Juhi Mittal <juhim@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Juhi Mittal committed
Sep 1, 2026 at 13:28 UTC
7a977c0699b9891e97f6b30b4e28b6864c4761fc
6 files changed
+1234
-1017
tests/quantization/test_modelopt.py
+312
-152
@@ -27,14 +27,19 @@ from vllm.model_executor.kernels.linear import (
27
from vllm.model_executor.layers.attention import Attention
28
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
29
from vllm.model_executor.layers.quantization.modelopt import (
30
+ LINEAR_ALGOS,
31
ModelOptFp8Config,
31
- ModelOptFp8LinearMethod,
32
- ModelOptFp8PbWoLinearMethod,
32
+ ModelOptLinearMethod,
33
ModelOptMixedPrecisionConfig,
34
ModelOptMxFp8Config,
35
ModelOptNvFp4Config,
36
- ModelOptNvFp4LinearMethod,
37
- ModelOptNvFp4W4A16LinearMethod,
36
+)
37
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
+ kFp8StaticTensorSym,
39
+ kMxfp8Dynamic,
40
+ kMxfp8Static,
41
+ kNvfp4Dynamic,
42
+ kNvfp4Static,
43
)
44
from vllm.model_executor.layers.vocab_parallel_embedding import (
45
ParallelLMHead,
@@ -114,15 +119,25 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
119
exclude_modules=[],
120
)
121
117
- with patch(
118
- "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
119
- ):
120
- method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
122
+ method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
123
122
- assert isinstance(method, ModelOptNvFp4LinearMethod)
124
+ assert isinstance(method, ModelOptLinearMethod)
125
+ assert method.spec.weight is kNvfp4Static
126
+ assert method.spec.activation is kNvfp4Dynamic
127
128
129
def test_modelopt_fp8_updates_weight_dims_after_transpose():
130
+ """Humming reads weight.input_dim/output_dim. Swapping the
131
+ ModelWeightParameter for a plain Parameter drops them, so the per-tensor
132
+ FP8 scheme must restore them for the transposed [in, out] layout.
133
+ """
134
+ from vllm.config.quantization import QuantSpec
135
+ from vllm.model_executor.layers.quantization.modelopt import (
136
+ SCHEME_FOR,
137
+ CkptCtx,
138
+ FormatScheme,
139
+ )
140
+
141
layer = torch.nn.Module()
142
layer.register_parameter(
143
"weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False)
@@ -133,73 +148,80 @@ def test_modelopt_fp8_updates_weight_dims_after_transpose():
148
layer.register_parameter(
149
"input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
150
)
136
-
137
- method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod)
138
- method.fp8_linear = Mock()
151
+ layer.logical_widths = [3]
152
+
153
+ method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
154
+ method.spec = QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym)
155
+ method.ctx = CkptCtx()
156
+ method.fmt = FormatScheme()
157
+ method.wkey = SCHEME_FOR[kFp8StaticTensorSym]
158
+ method.akey = SCHEME_FOR[kFp8StaticTensorSym]
159
+ method.kernel = Mock()
160
method.process_weights_after_loading(layer)
161
162
assert layer.weight.shape == (2, 3)
163
assert layer.weight.input_dim == 0
164
assert layer.weight.output_dim == 1
144
- method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer)
165
+ method.kernel.process_weights_after_loading.assert_called_once_with(layer)
166
167
147
-def test_modelopt_fp8_pb_wo_hides_output_padding(default_vllm_config):
148
- config = ModelOptFp8Config(
149
- quant_method="FP8_PB_WO",
150
- is_checkpoint_fp8_serialized=True,
151
- kv_cache_quant_method=None,
152
- exclude_modules=[],
168
+def test_modelopt_linear_algos_table_matches_resolve():
169
+ """LINEAR_ALGOS is the single source of truth for supported linear algos.
170
+
171
+ Every entry must be dispatchable by resolve(), and every config's
172
+ validation list must be derived from it -- so adding a format is one row
173
+ here plus one row in resolve(), with nothing else to keep in sync.
174
+ """
175
+ from vllm.model_executor.layers.quantization.modelopt import (
176
+ QUANT_ALGOS,
177
+ algos_owned_by,
178
+ resolve,
179
)
154
- default_vllm_config.model_config = Mock(dtype=torch.bfloat16)
155
- kernel = Mock()
180
157
- with (
158
- set_current_vllm_config(default_vllm_config),
159
- patch(
160
- "vllm.model_executor.layers.quantization.modelopt.init_fp8_linear_kernel",
161
- return_value=kernel,
162
- ),
163
- patch(
164
- "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
165
- return_value=0,
166
- ),
167
- patch(
168
- "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
169
- return_value=1,
170
- ),
171
- ):
172
- method = ModelOptFp8PbWoLinearMethod(config)
173
- layer = torch.nn.Module()
174
- method.create_weights(
175
- layer,
176
- input_size_per_partition=128,
177
- output_partition_sizes=[128, 65],
178
- input_size=128,
179
- output_size=193,
180
- params_dtype=torch.bfloat16,
181
- )
181
+ class _Cfg:
182
+ group_size = 16
183
183
- assert layer.weight.shape == (193, 128)
184
- assert layer.weight_scale.shape == (2, 1, 1, 1)
184
+ for algo in LINEAR_ALGOS:
185
+ spec, _, _ = resolve(algo, _Cfg(), "layer")
186
+ assert spec.weight is not None, algo
187
186
- layer.weight.data.fill_(1)
187
- method.process_weights_after_loading(layer)
188
+ assert list(QUANT_ALGOS) == [*LINEAR_ALGOS, "MIXED_PRECISION"]
189
+ assert algos_owned_by("modelopt") == (
190
+ "FP8",
191
+ "FP8_PER_CHANNEL_PER_TOKEN",
192
+ "FP8_PB_WO",
193
+ )
194
+ assert algos_owned_by("modelopt_fp4") == ("NVFP4", "W4A16_NVFP4")
195
+ assert algos_owned_by("modelopt_mxfp8") == ("MXFP8",)
196
189
- assert layer.weight.shape == (256, 128)
190
- assert torch.count_nonzero(layer.weight[193:].float()) == 0
191
- kernel.process_weights_after_loading.assert_called_once_with(layer)
197
193
- physical_output = torch.randn(2, 256, dtype=torch.bfloat16)
194
- kernel.apply_weights.return_value = physical_output
195
- bias = torch.randn(193, dtype=torch.bfloat16)
196
- output = method.apply(layer, torch.randn(2, 128), bias)
198
+@pytest.mark.parametrize("algo", list(LINEAR_ALGOS))
199
+def test_modelopt_mixed_precision_dispatches_every_linear_algo(algo):
200
+ """Mixed precision must route every algo in LINEAR_ALGOS through the
201
+ generic method. FP8_PER_CHANNEL_PER_TOKEN and FP8_PB_WO used to fall
202
+ through to UnquantizedLinearMethod, which loses the checkpoint's scales.
203
+ """
204
+ from vllm.model_executor.layers.linear import LinearBase
205
+ from vllm.model_executor.layers.quantization import modelopt as m
206
198
- torch.testing.assert_close(output, physical_output[:, :193] + bias)
199
- assert output.shape == (2, 193)
200
- assert output.is_contiguous()
201
- kernel.apply_weights.assert_called_once()
202
- assert kernel.apply_weights.call_args.args[2] is None
207
+ config = m.ModelOptMixedPrecisionConfig.from_config(
208
+ {
209
+ "quantization": {
210
+ "quant_algo": "MIXED_PRECISION",
211
+ "kv_cache_quant_algo": None,
212
+ "exclude_modules": [],
213
+ "group_size": 16,
214
+ "quantized_layers": {
215
+ "model.layers.0.mlp.down_proj": {"quant_algo": algo}
216
+ },
217
+ }
218
+ }
219
+ )
220
+ method = config.get_quant_method(
221
+ MagicMock(spec=LinearBase), "model.layers.0.mlp.down_proj"
222
+ )
223
+
224
+ assert isinstance(method, ModelOptLinearMethod), (algo, type(method).__name__)
225
226
227
def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
@@ -219,12 +241,11 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
241
{"lm_head": {"quant_algo": "NVFP4", "group_size": 16}}
242
)
243
222
- with patch(
223
- "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
224
- ):
225
- method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
244
+ method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
245
227
- assert isinstance(method, ModelOptNvFp4LinearMethod)
246
+ assert isinstance(method, ModelOptLinearMethod)
247
+ assert method.spec.weight is kNvfp4Static
248
+ assert method.spec.activation is kNvfp4Dynamic
249
250
251
def test_modelopt_mixed_precision_resolves_declared_packed_projection():
@@ -302,12 +323,11 @@ def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
323
)
324
325
fake_layer = MagicMock(spec=LinearBase)
305
- with patch(
306
- "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
307
- ):
308
- method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
326
+ method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
327
310
- assert isinstance(method, ModelOptNvFp4LinearMethod)
328
+ assert isinstance(method, ModelOptLinearMethod)
329
+ assert method.spec.weight is kNvfp4Static
330
+ assert method.spec.activation is kNvfp4Dynamic
331
332
333
@pytest.mark.parametrize(
@@ -363,7 +383,7 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
383
"This test requires a local ModelOpt FP8 checkpoint."
384
)
385
366
- # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
386
+ # Set model config as model_config.dtype is required in ModelOptLinearMethod.
387
default_vllm_config.model_config = ModelConfig()
388
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
389
@@ -376,14 +396,10 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
396
down_proj = layer.mlp.down_proj
397
398
# Check that ModelOpt quantization method is properly applied
379
- from vllm.model_executor.layers.quantization.modelopt import (
380
- ModelOptFp8LinearMethod,
381
- )
382
-
383
- assert isinstance(qkv_proj.quant_method, ModelOptFp8LinearMethod)
384
- assert isinstance(o_proj.quant_method, ModelOptFp8LinearMethod)
385
- assert isinstance(gate_up_proj.quant_method, ModelOptFp8LinearMethod)
386
- assert isinstance(down_proj.quant_method, ModelOptFp8LinearMethod)
399
+ assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
400
+ assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
401
+ assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
402
+ assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
403
404
# Check weight dtype is FP8
405
assert qkv_proj.weight.dtype == torch.float8_e4m3fn
@@ -440,14 +456,10 @@ def test_modelopt_fp8_pc_pt_checkpoint_setup(monkeypatch, dist_init, workspace_i
456
gate_up_proj = layer.mlp.gate_up_proj
457
down_proj = layer.mlp.down_proj
458
443
- from vllm.model_executor.layers.quantization.modelopt import (
444
- ModelOptFp8PcPtLinearMethod,
445
- )
446
-
447
- assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
448
- assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
449
- assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
450
- assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
459
+ assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
460
+ assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
461
+ assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
462
+ assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
463
464
fp8_dtype = current_platform.fp8_dtype()
465
assert qkv_proj.weight.dtype == fp8_dtype
@@ -494,14 +506,10 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i
506
gate_up_proj = layer.mlp.gate_up_proj
507
down_proj = layer.mlp.down_proj
508
497
- from vllm.model_executor.layers.quantization.modelopt import (
498
- ModelOptFp8PbWoLinearMethod,
499
- )
500
-
501
- assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
502
- assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
503
- assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
504
- assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
509
+ assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
510
+ assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
511
+ assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
512
+ assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
513
514
fp8_dtype = current_platform.fp8_dtype()
515
assert qkv_proj.weight.dtype == fp8_dtype
@@ -528,12 +536,10 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i
536
537
538
def test_modelopt_nvfp4_config_dispatches_w4a4_method():
531
- """``quant_method="NVFP4"`` (W4A4 default) routes to the existing
532
- ``ModelOptNvFp4LinearMethod``."""
533
- from vllm.model_executor.layers.quantization.modelopt import (
534
- ModelOptNvFp4Config,
535
- ModelOptNvFp4LinearMethod,
536
- )
539
+ """``quant_method="NVFP4"`` (W4A4) resolves to a
540
+ ``(kNvfp4Static, kNvfp4Dynamic)`` QuantSpec under the generic
541
+ ``ModelOptLinearMethod``."""
542
+ from vllm.model_executor.layers.linear import LinearBase
543
544
config = ModelOptNvFp4Config(
545
quant_method="NVFP4",
@@ -541,24 +547,25 @@ def test_modelopt_nvfp4_config_dispatches_w4a4_method():
547
kv_cache_quant_algo=None,
548
exclude_modules=[],
549
)
544
- assert config.LinearMethodCls is ModelOptNvFp4LinearMethod
550
assert config.quant_method == "NVFP4"
551
552
+ method = config.get_quant_method(
553
+ MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
554
+ )
555
+ assert isinstance(method, ModelOptLinearMethod)
556
+ assert method.spec.weight is kNvfp4Static
557
+ assert method.spec.activation is kNvfp4Dynamic
558
+
559
560
def test_modelopt_nvfp4_config_dispatches_w4a16_method():
549
- """``quant_method="W4A16_NVFP4"`` routes to
550
- ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling.
561
+ """``quant_method="W4A16_NVFP4"`` resolves to a weight-only QuantSpec
562
+ (``activation=None``) — distinct from the W4A4 sibling.
563
552
- Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects
553
- one of three FP8 LinearMethods on ``quant_method``); a regression
554
- here would mean a W4A16 NVFP4 checkpoint silently loaded under the
555
- W4A4 activation-quantization path.
564
+ A regression here would mean a W4A16 NVFP4 checkpoint silently loaded
565
+ with a dynamic fp4 activation key, registering an ``input_scale`` and
566
+ routing to the cutlass W4A4 NVFP4 GEMM instead of FP4 Marlin.
567
"""
557
- from vllm.model_executor.layers.quantization.modelopt import (
558
- ModelOptNvFp4Config,
559
- ModelOptNvFp4LinearMethod,
560
- ModelOptNvFp4W4A16LinearMethod,
561
- )
568
+ from vllm.model_executor.layers.linear import LinearBase
569
570
config = ModelOptNvFp4Config(
571
quant_method="W4A16_NVFP4",
@@ -566,10 +573,38 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method():
573
kv_cache_quant_algo=None,
574
exclude_modules=[],
575
)
569
- assert config.LinearMethodCls is ModelOptNvFp4W4A16LinearMethod
570
- assert config.LinearMethodCls is not ModelOptNvFp4LinearMethod
576
assert config.quant_method == "W4A16_NVFP4"
577
578
+ method = config.get_quant_method(
579
+ MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
580
+ )
581
+ assert isinstance(method, ModelOptLinearMethod)
582
+ assert method.spec.weight is kNvfp4Static
583
+ assert method.spec.activation is None
584
+
585
+
586
+def test_modelopt_linear_method_builder_registry_override(monkeypatch):
587
+ """The bespoke-method escape hatch: a format registered in
588
+ ``LINEAR_METHOD_BUILDERS`` routes that algo to its own method instead of the
589
+ generic ``ModelOptLinearMethod``. This is how a format that cannot be a
590
+ ``(weight, activation)`` key pair plugs into dispatch."""
591
+ from vllm.model_executor.layers.linear import LinearBase
592
+ from vllm.model_executor.layers.quantization import modelopt as m
593
+
594
+ sentinel = object()
595
+ monkeypatch.setitem(m.LINEAR_METHOD_BUILDERS, "NVFP4", lambda cfg, prefix: sentinel)
596
+
597
+ config = ModelOptNvFp4Config(
598
+ quant_method="NVFP4",
599
+ is_checkpoint_nvfp4_serialized=True,
600
+ kv_cache_quant_algo=None,
601
+ exclude_modules=[],
602
+ )
603
+ method = config.get_quant_method(
604
+ MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
605
+ )
606
+ assert method is sentinel # bespoke builder wins over the generic path
607
+
608
609
@pytest.mark.parametrize(
610
("linear_backend", "kernel_cls"),
@@ -581,6 +616,15 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method():
616
)
617
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
618
def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):
619
+ """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``:
620
+ ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The
621
+ generic method routes this through ``select_linear_kernel``."""
622
+ from vllm.config.quantization import QuantSpec
623
+ from vllm.model_executor.layers.quantization.modelopt import (
624
+ RuntimeDtypes,
625
+ select_linear_kernel,
626
+ )
627
+
628
if linear_backend != "auto":
629
is_supported, reason = kernel_cls.is_supported()
630
if not is_supported:
@@ -588,11 +632,52 @@ def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):
632
633
vllm_config = VllmConfig()
634
vllm_config.kernel_config.linear_backend = linear_backend
635
+ spec = QuantSpec(weight=kNvfp4Static, activation=None)
636
+ rt = RuntimeDtypes(torch.bfloat16, torch.bfloat16)
637
with set_current_vllm_config(vllm_config):
592
- method = ModelOptNvFp4W4A16LinearMethod(
593
- ModelOptNvFp4Config(quant_method="W4A16_NVFP4")
638
+ kernel = select_linear_kernel(spec, MagicMock(), rt)
639
+ assert isinstance(kernel, kernel_cls)
640
+
641
+
642
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
643
+def test_modelopt_linear_exposes_humming_layer_attrs(dist_init, monkeypatch):
644
+ """``prepare_humming_linear_layer_config`` reads ``output_partition_sizes``
645
+ and ``has_bias`` straight off the layer, so ``--linear-backend=humming``
646
+ needs create_weights to leave both there. Nothing else sets
647
+ ``output_partition_sizes``; ``LinearBase`` sets ``has_bias`` but
648
+ ``ParallelLMHead`` does not.
649
+ """
650
+ from vllm.config.quantization import QuantSpec
651
+ from vllm.model_executor.layers.quantization import modelopt as mo
652
+
653
+ monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: Mock())
654
+ monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, kernel: None)
655
+
656
+ def build(layer):
657
+ method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
658
+ method.spec = QuantSpec(weight=kNvfp4Static, activation=None)
659
+ method.ctx = mo.CkptCtx(group_size=16)
660
+ method.fmt = mo.FormatScheme()
661
+ method.wkey = mo.SCHEME_FOR[kNvfp4Static]
662
+ method.akey = None
663
+ method.input_dtype = method.out_dtype = torch.bfloat16
664
+ method.marlin_input_dtype = None
665
+ method.create_weights(
666
+ layer, 64, [32, 32], 64, 64, torch.bfloat16, weight_loader=Mock()
667
)
595
- assert isinstance(method.kernel, kernel_cls)
668
+
669
+ # ParallelLMHead-style: a bias slot but no has_bias attribute.
670
+ lm_head = torch.nn.Module()
671
+ lm_head.register_parameter("bias", None)
672
+ build(lm_head)
673
+ assert lm_head.output_partition_sizes == [32, 32]
674
+ assert lm_head.has_bias is False
675
+
676
+ # LinearBase already decided has_bias; we must not overwrite it.
677
+ linear = torch.nn.Module()
678
+ linear.has_bias = True
679
+ build(linear)
680
+ assert linear.has_bias is True
681
682
683
@pytest.mark.parametrize(
@@ -654,40 +739,28 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16(
739
740
741
@pytest.mark.parametrize(
657
- "per_layer_algo, expected_linear_cls_name",
742
+ "per_layer_algo, expected_weight, expected_activation",
743
[
659
- ("NVFP4", "ModelOptNvFp4LinearMethod"),
660
- ("W4A16_NVFP4", "ModelOptNvFp4W4A16LinearMethod"),
744
+ ("NVFP4", kNvfp4Static, kNvfp4Dynamic),
745
+ ("W4A16_NVFP4", kNvfp4Static, None),
746
+ ("FP8", kFp8StaticTensorSym, kFp8StaticTensorSym),
747
+ ("MXFP8", kMxfp8Static, kMxfp8Dynamic),
748
],
749
)
663
-def test_modelopt_mixed_precision_dispatches_w4a16_layer(
664
- per_layer_algo, expected_linear_cls_name
750
+def test_modelopt_mixed_precision_dispatches_linear_layer(
751
+ per_layer_algo, expected_weight, expected_activation
752
):
666
- """``ModelOptMixedPrecisionConfig.get_quant_method`` must route a Linear
667
- layer to the right LinearMethod based on its per-layer ``quant_algo``
668
- entry in ``quantized_layers``. Verifies the new ``W4A16_NVFP4`` branch
669
- coexists with the existing ``NVFP4`` branch without regression. A
670
- regression here would mean a W4A16 layer in a mixed-precision ckpt
671
- silently fell through to ``UnquantizedLinearMethod``.
672
-
673
- NOTE: FP8 dispatch (the third branch of get_quant_method) is not
674
- covered here because ``ModelOptFp8LinearMethod.__init__`` reads
675
- ``get_current_vllm_config().model_config.dtype``, which requires a
676
- fully constructed ``ModelConfig`` (real model path). FP8 routing in
677
- mixed-precision is exercised by the existing integration tests
678
- above that use the ``vllm_runner`` fixture (e.g.
679
- ``test_modelopt_fp8_checkpoint_setup``). Our PR doesn't change the
680
- FP8 branch, so this isn't a coverage gap.
753
+ """``ModelOptMixedPrecisionConfig.get_quant_method`` routes a Linear layer
754
+ to the generic ``ModelOptLinearMethod`` with the ``QuantSpec`` resolved
755
+ from its per-layer ``quant_algo`` entry in ``quantized_layers``. A
756
+ regression here would mean a layer got the wrong ``(weight, activation)``
757
+ key pair or fell through to ``UnquantizedLinearMethod`` — e.g. a W4A16
758
+ layer picking up a dynamic fp4 activation key (cutlass W4A4 path) instead
759
+ of the weight-only Marlin path.
760
"""
761
from vllm.model_executor.layers.linear import LinearBase
762
from vllm.model_executor.layers.quantization import modelopt as m
763
685
- if (
686
- expected_linear_cls_name == "ModelOptNvFp4W4A16LinearMethod"
687
- and current_platform.is_rocm()
688
- ):
689
- pytest.skip("ModelOptNvFp4W4A16LinearMethod is not supported with rocm")
690
-
764
hf_quant_config: dict[str, Any] = {
765
"quantization": {
766
"quant_algo": "MIXED_PRECISION",
@@ -704,10 +777,9 @@ def test_modelopt_mixed_precision_dispatches_w4a16_layer(
777
fake_layer = MagicMock(spec=LinearBase)
778
method = config.get_quant_method(fake_layer, "model.layers.0.fake_proj")
779
707
- expected_cls = getattr(m, expected_linear_cls_name)
708
- assert isinstance(method, expected_cls), (
709
- f"Expected {expected_linear_cls_name}, got {type(method).__name__}"
710
- )
780
+ assert isinstance(method, m.ModelOptLinearMethod)
781
+ assert method.spec.weight is expected_weight
782
+ assert method.spec.activation is expected_activation
783
784
785
def test_modelopt_mixed_precision_builds_w4a16_sibling_config():
@@ -734,6 +806,94 @@ def test_modelopt_mixed_precision_builds_w4a16_sibling_config():
806
config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)
807
808
assert config.nvfp4_config.quant_method == "NVFP4"
737
- assert config.nvfp4_config.LinearMethodCls is m.ModelOptNvFp4LinearMethod
809
assert config.w4a16_nvfp4_config.quant_method == "W4A16_NVFP4"
739
- assert config.w4a16_nvfp4_config.LinearMethodCls is m.ModelOptNvFp4W4A16LinearMethod
810
+
811
+
812
+def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch):
813
+ """FP8_PB_WO output width that is not a multiple of 128 (a partial trailing
814
+ block) is padded up to a block boundary before the kernel post-load, the
815
+ GEMM runs on the padded weight, and the output is trimmed back to the
816
+ logical width with bias added after. Faithful port of wei-zhao #53132's
817
+ test_modelopt_fp8_pb_wo_hides_output_padding for the generic method +
818
+ _Fp8PbWoPartialBlock FormatScheme.
819
+
820
+ Width is the real motivating case -- GLM's fused qkv_a_proj, q_a 2048 +
821
+ kv_a 576 = 2624 (replicated, so no TP degree makes it a 128-multiple),
822
+ padded to 2688 = 21 * 128.
823
+ """
824
+ from vllm.config.quantization import QuantSpec
825
+ from vllm.model_executor.layers.quantization import modelopt as mo
826
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
827
+ kFp8Dynamic128Sym,
828
+ kFp8Static128BlockSym,
829
+ )
830
+
831
+ kernel = Mock()
832
+ monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: kernel)
833
+ monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, k: None)
834
+
835
+ method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
836
+ method.spec = QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym)
837
+ method.ctx = mo.CkptCtx()
838
+ method.fmt = mo._PB_WO_PARTIAL_BLOCK
839
+ method.wkey = mo.SCHEME_FOR[kFp8Static128BlockSym]
840
+ method.akey = mo.SCHEME_FOR[kFp8Dynamic128Sym]
841
+ method.input_dtype = method.out_dtype = torch.bfloat16
842
+ method.marlin_input_dtype = None
843
+
844
+ layer = torch.nn.Module()
845
+ with (
846
+ patch(
847
+ "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
848
+ return_value=0,
849
+ ),
850
+ patch(
851
+ "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
852
+ return_value=1,
853
+ ),
854
+ ):
855
+ # output 2624 = 2048 + 576, input 128
856
+ method.create_weights(
857
+ layer, 128, [2048, 576], 128, 2624, torch.bfloat16, weight_loader=Mock()
858
+ )
859
+
860
+ # loaded at logical size; scale is cdiv(2624, 128) = 21 block rows
861
+ assert layer.weight.shape == (2624, 128)
862
+ assert layer.weight_scale.shape == (21, 1, 1, 1)
863
+
864
+ layer.weight.data.fill_(1)
865
+ method.process_weights_after_loading(layer)
866
+
867
+ # weight padded to the block boundary; the pad rows are zero
868
+ assert layer.weight.shape == (2688, 128)
869
+ assert torch.count_nonzero(layer.weight[2624:].float()) == 0
870
+ kernel.process_weights_after_loading.assert_called_once_with(layer)
871
+
872
+ # apply: GEMM on padded weight (bias=None), output trimmed + bias added
873
+ physical_output = torch.randn(4, 2688, dtype=torch.bfloat16)
874
+ kernel.apply_weights.return_value = physical_output
875
+ bias = torch.randn(2624, dtype=torch.bfloat16)
876
+ output = method.apply(layer, torch.randn(4, 128), bias)
877
+
878
+ torch.testing.assert_close(output, physical_output[:, :2624] + bias)
879
+ assert output.shape == (4, 2624)
880
+ assert output.is_contiguous()
881
+ kernel.apply_weights.assert_called_once()
882
+ assert kernel.apply_weights.call_args.kwargs["bias"] is None
883
+
884
+
885
+def test_modelopt_fp8_pb_wo_rejects_non_128_input():
886
+ """Input width must still be a multiple of 128 (same as #53132, which only
887
+ pads the output). A partial input block is refused loudly rather than
888
+ silently loading wrong scales."""
889
+ from vllm.model_executor.layers.quantization import modelopt as mo
890
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
891
+ kFp8Static128BlockSym,
892
+ )
893
+
894
+ scheme = mo.SCHEME_FOR[kFp8Static128BlockSym]
895
+ shapes = mo.Shapes([128], 100, torch.bfloat16) # input 100 not divisible by 128
896
+ with pytest.raises(ValueError, match="in divisible by 128"):
897
+ scheme.create_weights(
898
+ torch.nn.Module(), mo.WEIGHT, mo.CkptCtx(), shapes, Mock()
899
+ )
vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
+8
@@ -6,6 +6,8 @@ from dataclasses import dataclass
6
7
import torch
8
9
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
10
+
11
12
@dataclass
13
class Mxfp8LinearLayerConfig:
@@ -30,6 +32,12 @@ class Mxfp8LinearKernel(ABC):
32
assert self.is_supported()[0]
33
self.config = c
34
35
+ def input_quant_key(self) -> QuantKey | None:
36
+ """Return the input quantization key supported by this kernel. If the kernel
37
+ does not support input quantization outside of the kernel, return None.
38
+ """
39
+ return None
40
+
41
@classmethod
42
@abstractmethod
43
def is_supported(
vllm/model_executor/layers/attention/mla_attention.py
+8
-2
@@ -377,12 +377,18 @@ def _get_kv_b_proj_input_dtype(
377
return None
378
if weight_dtype == current_platform.fp8_dtype():
379
from vllm.model_executor.layers.quantization.modelopt import (
380
- ModelOptFp8PbWoLinearMethod,
380
+ ModelOptLinearMethod,
381
+ )
382
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
383
+ kFp8Static128BlockSym,
384
)
385
386
quant_method = getattr(kv_b_proj, "quant_method", None)
387
# FP8_PB_WO dynamically quantizes BF16/FP16 inputs in the linear method.
385
- if isinstance(quant_method, ModelOptFp8PbWoLinearMethod):
388
+ if (
389
+ isinstance(quant_method, ModelOptLinearMethod)
390
+ and quant_method.spec.weight is kFp8Static128BlockSym
391
+ ):
392
return quant_method.input_dtype
393
if not use_fp8_prefill:
394
return None
vllm/model_executor/layers/linear.py
+2
-5
@@ -55,13 +55,10 @@ WEIGHT_LOADER_V2_SUPPORTED = [
55
"AutoGPTQLinearMethod",
56
"Fp8LinearMethod",
57
"FBGEMMFp8LinearMethod",
58
- "ModelOptFp8LinearMethod",
59
- "ModelOptFp8PcPtLinearMethod",
60
- "ModelOptFp8PbWoLinearMethod",
58
"QuarkLinearMethod",
62
- "ModelOptNvFp4LinearMethod",
63
- "ModelOptNvFp4W4A16LinearMethod",
59
"HummingLinearMethod",
60
+ # ModelOptLinearMethod self-registers via
61
+ # register_weight_loader_v2_supported_method (see modelopt.py).
62
]
63
64
vllm/model_executor/layers/quantization/modelopt.py
+899
-858
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Callable
5
+from dataclasses import dataclass
6
+from enum import Enum
7
from fnmatch import fnmatch
8
from typing import TYPE_CHECKING, Any, cast
9
@@ -8,7 +11,8 @@ import torch
11
from torch.nn.parameter import Parameter
12
13
import vllm.envs as envs
11
-from vllm.config import get_current_vllm_config
14
+from vllm.config import get_current_vllm_config_or_none
15
+from vllm.config.quantization import QuantSpec
16
from vllm.logger import init_logger
17
from vllm.model_executor.kernels.linear import (
18
init_fp8_linear_kernel,
@@ -49,6 +53,7 @@ from vllm.model_executor.layers.linear import (
53
LinearBase,
54
LinearMethodBase,
55
UnquantizedLinearMethod,
56
+ register_weight_loader_v2_supported_method,
57
)
58
from vllm.model_executor.layers.quantization import QuantizationMethods
59
from vllm.model_executor.layers.quantization.base_config import (
@@ -57,6 +62,7 @@ from vllm.model_executor.layers.quantization.base_config import (
62
)
63
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
64
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
65
+ FP8_SCALE_SENTINEL,
66
process_fp8_input_tensor_strategy_moe,
67
process_fp8_weight_channel_strategy,
68
process_fp8_weight_tensor_strategy_moe,
@@ -70,12 +76,16 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
76
MXFP8_VALUE_DTYPE,
77
)
78
from vllm.model_executor.layers.quantization.utils.quant_utils import (
73
- GroupShape,
74
- create_fp8_quant_key,
79
+ FP4_DTYPE,
80
+ QuantKey,
81
is_layer_skipped,
82
+ kFp8Dynamic128Sym,
83
kFp8DynamicTokenSym,
84
+ kFp8Static128BlockSym,
85
kFp8StaticTensorSym,
86
kFp8StaticTokenSym,
87
+ kMxfp8Dynamic,
88
+ kMxfp8Static,
89
kNvfp4Dynamic,
90
kNvfp4Static,
91
)
@@ -86,33 +96,49 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
96
from vllm.model_executor.parameter import (
97
BlockQuantScaleParameter,
98
ChannelQuantScaleParameter,
89
- GroupQuantScaleParameter,
99
ModelWeightParameter,
100
PerTensorScaleParameter,
101
)
102
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
103
+from vllm.utils.math_utils import cdiv
104
105
if TYPE_CHECKING:
106
from vllm.model_executor.models.utils import WeightsMapper
107
108
logger = init_logger(__name__)
109
100
-QUANT_ALGOS = [
101
- # FP8 (per-tensor weight + optional static activation scale).
102
- "FP8",
110
+# Single source of truth for the ModelOpt linear algos.
111
+#
112
+# Each entry is (owning config's name, mixed-precision sub-config attribute).
113
+# ``resolve`` turns the algo into a QuantSpec; this table only records which
114
+# config object holds the checkpoint parameters resolve may read -- today just
115
+# group_size, and only for the NVFP4 family.
116
+#
117
+# To add a format: one row here plus one row in resolve(). Every consumer below
118
+# is derived, so nothing else needs editing.
119
+LINEAR_ALGOS: dict[str, tuple[str, str]] = {
120
+ # FP8 per-tensor weight + optional static activation scale.
121
+ "FP8": ("modelopt", "fp8_config"),
122
# FP8 per-channel weight scale + per-token activation scale.
104
- "FP8_PER_CHANNEL_PER_TOKEN",
105
- # FP8 per-block weight-only (ModelOpt may emit this as lowercase).
106
- "FP8_PB_WO",
107
- # NVFP4 W4A4 (4-bit float weights AND 4-bit float activations).
108
- "NVFP4",
109
- # W4A16 NVFP4 (4-bit float weights, fp16/bf16 activations).
110
- "W4A16_NVFP4",
111
- # MXFP8
112
- "MXFP8",
113
- # MIXED_PRECISION,
114
- "MIXED_PRECISION",
115
-]
123
+ "FP8_PER_CHANNEL_PER_TOKEN": ("modelopt", "fp8_config"),
124
+ # FP8 128x128 per-block weight (ModelOpt may emit this lowercase).
125
+ "FP8_PB_WO": ("modelopt", "fp8_config"),
126
+ # NVFP4 W4A4: 4-bit float weights AND 4-bit float activations.
127
+ "NVFP4": ("modelopt_fp4", "nvfp4_config"),
128
+ # W4A16 NVFP4: 4-bit float weights, fp16/bf16 activations.
129
+ "W4A16_NVFP4": ("modelopt_fp4", "w4a16_nvfp4_config"),
130
+ # MXFP8: e4m3 weights with per-32-block e8m0 scales.
131
+ "MXFP8": ("modelopt_mxfp8", "mxfp8_config"),
132
+}
133
+
134
+# MIXED_PRECISION is not a linear algo; it selects a per-layer algo from the
135
+# table above.
136
+QUANT_ALGOS = [*LINEAR_ALGOS, "MIXED_PRECISION"]
137
+
138
+
139
+def algos_owned_by(config_name: str) -> tuple[str, ...]:
140
+ """The linear algos a given config accepts, for its quant_algo validation."""
141
+ return tuple(a for a, (owner, _) in LINEAR_ALGOS.items() if owner == config_name)
142
143
144
class ModelOptKVCacheMethod(BaseKVCacheMethod):
@@ -125,7 +151,11 @@ class ModelOptKVCacheMethod(BaseKVCacheMethod):
151
152
153
class ModelOptQuantConfigBase(QuantizationConfig):
128
- LinearMethodCls: type = LinearMethodBase
154
+ # ModelOpt quant-algo string, set by each subclass. Fed to resolve() to
155
+ # build the QuantSpec for the generic ModelOptLinearMethod. The mixed
156
+ # config resolves the algo per-prefix instead and does not set this.
157
+ quant_method: str
158
+
159
FusedMoEMethodCls: type = FusedMoEMethodBase
160
KVCacheMethodCls: type = BaseKVCacheMethod
161
@@ -201,10 +231,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
231
232
# now, the layer is quantized, handle it here
233
if isinstance(layer, (LinearBase, ParallelLMHead)):
204
- quant_method = self.LinearMethodCls(self)
205
- if getattr(quant_method, "backend", "") == "marlin":
206
- quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
207
- return quant_method
234
+ return build_linear_method(self, self.quant_method, prefix)
235
elif isinstance(layer, RoutedExperts):
236
quant_method = self.FusedMoEMethodCls(
237
quant_config=self, moe_config=layer.moe_config
@@ -375,27 +402,22 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
402
self.quant_method = quant_method
403
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
404
self.kv_cache_quant_method = kv_cache_quant_method
378
- if is_checkpoint_fp8_serialized:
379
- logger.warning(
380
- "Detected ModelOpt fp8 checkpoint (quant_algo=%s). Please note "
381
- "that the format is experimental and could change.",
382
- quant_method,
383
- )
405
385
- # Select LinearMethod implementation based on quant_algo.
386
- if self.quant_method == "FP8":
387
- self.LinearMethodCls = ModelOptFp8LinearMethod
388
- elif self.quant_method == "FP8_PER_CHANNEL_PER_TOKEN":
389
- self.LinearMethodCls = ModelOptFp8PcPtLinearMethod
390
- elif self.quant_method == "FP8_PB_WO":
391
- self.LinearMethodCls = ModelOptFp8PbWoLinearMethod
392
- else:
406
+ # Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch.
407
+ supported = algos_owned_by("modelopt")
408
+ if self.quant_method not in supported:
409
raise ValueError(
410
"Unsupported ModelOpt FP8 quant_algo for vLLM: "
395
- f"{self.quant_method}. Supported: FP8 / "
396
- "FP8_PER_CHANNEL_PER_TOKEN / FP8_PB_WO."
411
+ f"{self.quant_method}. Supported: {' / '.join(supported)}."
412
)
413
414
+ def has_blocked_weights(self) -> bool:
415
+ # Gates "+quant_fp8" in VllmConfig. FP8_PB_WO runs the block-scaled
416
+ # kernels, whose DeepGEMM path needs UE8M0-packed group scales; without
417
+ # the op enabled QuantFP8 falls back to forward_native, which emits
418
+ # unpacked fp32 scales and the GEMM fails to launch.
419
+ return self.quant_method == "FP8_PB_WO"
420
+
421
def get_name(self) -> QuantizationMethods:
422
return "modelopt"
423
@@ -435,334 +457,6 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
457
)
458
459
438
-class ModelOptFp8LinearMethod(LinearMethodBase):
439
- """Linear method for Model Optimizer static quantization.
440
- Supports loading FP8 checkpoints with static weight scale and
441
- activation scale. Future support might be added for dynamic
442
- scales.
443
-
444
- Limitations:
445
- 1. Only support per-tensor quantization due to torch._scaled_mm support.
446
- 2. Only support float8_e4m3fn datatype
447
- Args: quant_config: The ModelOpt quantization config.
448
- """
449
-
450
- def __init__(self, quant_config: ModelOptFp8Config) -> None:
451
- self.quant_config = quant_config
452
- self.out_dtype = get_current_vllm_config().model_config.dtype
453
- self.input_dtype = get_current_vllm_config().model_config.dtype
454
-
455
- def create_weights(
456
- self,
457
- layer: torch.nn.Module,
458
- input_size_per_partition: int,
459
- output_partition_sizes: list[int],
460
- input_size: int,
461
- output_size: int,
462
- params_dtype: torch.dtype,
463
- **extra_weight_attrs,
464
- ):
465
- del input_size, output_size
466
- output_size_per_partition = sum(output_partition_sizes)
467
- weight_loader = extra_weight_attrs.get("weight_loader")
468
- layer.logical_widths = output_partition_sizes
469
- layer.input_size_per_partition = input_size_per_partition
470
- layer.output_size_per_partition = output_size_per_partition
471
- layer.orig_dtype = params_dtype
472
- weight_dtype = (
473
- torch.float8_e4m3fn
474
- if self.quant_config.is_checkpoint_fp8_serialized
475
- else params_dtype
476
- )
477
- weight = ModelWeightParameter(
478
- data=torch.empty(
479
- output_size_per_partition, input_size_per_partition, dtype=weight_dtype
480
- ),
481
- input_dim=1,
482
- output_dim=0,
483
- weight_loader=weight_loader,
484
- )
485
- layer.register_parameter("weight", weight)
486
-
487
- if self.quant_config.is_checkpoint_fp8_serialized:
488
- # WEIGHT SCALE
489
- weight_scale = PerTensorScaleParameter(
490
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
491
- weight_loader=weight_loader,
492
- )
493
- weight_scale[:] = torch.finfo(torch.float32).min
494
- layer.register_parameter("weight_scale", weight_scale)
495
- # INPUT SCALE
496
- scale = PerTensorScaleParameter(
497
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
498
- weight_loader=weight_loader,
499
- )
500
-
501
- scale[:] = torch.finfo(torch.float32).min
502
- layer.register_parameter("input_scale", scale)
503
-
504
- self.fp8_linear = init_fp8_linear_kernel(
505
- activation_quant_key=kFp8StaticTensorSym,
506
- weight_quant_key=kFp8StaticTensorSym,
507
- weight_shape=layer.weight.shape,
508
- input_dtype=self.input_dtype,
509
- out_dtype=self.out_dtype,
510
- module_name=self.__class__.__name__,
511
- )
512
-
513
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
514
- weight = layer.weight
515
- max_w_scale = layer.weight_scale.max()
516
- if not (layer.weight_scale == layer.weight_scale[0]).all():
517
- max_w_scale, weight = requantize_with_max_scale(
518
- layer.weight, layer.weight_scale, layer.logical_widths
519
- )
520
- layer.weight = Parameter(weight.t(), requires_grad=False)
521
- layer.weight.input_dim = 0
522
- layer.weight.output_dim = 1
523
- layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
524
- layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
525
- self.fp8_linear.process_weights_after_loading(layer)
526
-
527
- def apply(
528
- self,
529
- layer: torch.nn.Module,
530
- x: torch.Tensor,
531
- bias: torch.Tensor | None = None,
532
- ) -> torch.Tensor:
533
- return self.fp8_linear.apply_weights(layer, x, bias)
534
-
535
-
536
-class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
537
- """Linear method for ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoints.
538
-
539
- Expected checkpoint structure (per Linear):
540
- - weight: fp8-e4m3fn, shape [out, in]
541
- - weight_scale: fp32, shape [out] (per-output-channel)
542
- - no input_scale (activations are dynamically quantized per-token)
543
- """
544
-
545
- def __init__(self, quant_config: ModelOptFp8Config) -> None:
546
- self.quant_config = quant_config
547
- self.out_dtype = get_current_vllm_config().model_config.dtype
548
- self.input_dtype = get_current_vllm_config().model_config.dtype
549
-
550
- def create_weights(
551
- self,
552
- layer: torch.nn.Module,
553
- input_size_per_partition: int,
554
- output_partition_sizes: list[int],
555
- input_size: int,
556
- output_size: int,
557
- params_dtype: torch.dtype,
558
- **extra_weight_attrs,
559
- ):
560
- del input_size, output_size
561
-
562
- if not self.quant_config.is_checkpoint_fp8_serialized:
563
- raise ValueError(
564
- "FP8_PER_CHANNEL_PER_TOKEN currently only supports "
565
- "FP8-serialized checkpoints."
566
- )
567
-
568
- output_size_per_partition = sum(output_partition_sizes)
569
- weight_loader = extra_weight_attrs.get("weight_loader")
570
- layer.logical_widths = output_partition_sizes
571
- layer.input_size_per_partition = input_size_per_partition
572
- layer.output_size_per_partition = output_size_per_partition
573
-
574
- weight = ModelWeightParameter(
575
- data=torch.empty(
576
- output_size_per_partition,
577
- input_size_per_partition,
578
- dtype=torch.float8_e4m3fn,
579
- ),
580
- input_dim=1,
581
- output_dim=0,
582
- weight_loader=weight_loader,
583
- )
584
- layer.register_parameter("weight", weight)
585
-
586
- weight_scale = ChannelQuantScaleParameter(
587
- data=torch.empty(output_size_per_partition, dtype=torch.float32),
588
- output_dim=0,
589
- weight_loader=weight_loader,
590
- )
591
- weight_scale[:] = torch.finfo(torch.float32).min
592
- layer.register_parameter("weight_scale", weight_scale)
593
-
594
- self.fp8_linear = init_fp8_linear_kernel(
595
- activation_quant_key=kFp8DynamicTokenSym,
596
- weight_quant_key=kFp8StaticTokenSym,
597
- weight_shape=layer.weight.shape,
598
- input_dtype=self.input_dtype,
599
- out_dtype=self.out_dtype,
600
- module_name=self.__class__.__name__,
601
- )
602
-
603
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
604
- weight, weight_scale, _ = process_fp8_weight_channel_strategy(
605
- layer.weight, layer.weight_scale.data
606
- )
607
- layer.weight = Parameter(weight.t(), requires_grad=False)
608
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
609
- self.fp8_linear.process_weights_after_loading(layer)
610
-
611
- def apply(
612
- self,
613
- layer: torch.nn.Module,
614
- x: torch.Tensor,
615
- bias: torch.Tensor | None = None,
616
- ) -> torch.Tensor:
617
- return self.fp8_linear.apply_weights(layer, x, bias)
618
-
619
-
620
-class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
621
- """Linear method for ModelOpt FP8_PB_WO checkpoints.
622
-
623
- ModelOpt exports `weight_scale` as a 4D tensor:
624
- [out_blk, 1, in_blk, 1]
625
- where block size is typically 128 for both dims.
626
-
627
- vLLM executes it as FP8 GEMM with *dynamic per-token* activation quant.
628
- Output widths that are not block-aligned are padded and restored
629
- to their logical width before returning to the model.
630
- """
631
-
632
- _WEIGHT_BLOCK_SIZE: tuple[int, int] = (128, 128)
633
-
634
- def __init__(self, quant_config: ModelOptFp8Config) -> None:
635
- self.quant_config = quant_config
636
- block_n, block_k = self._WEIGHT_BLOCK_SIZE
637
- self.weight_block_size = list(self._WEIGHT_BLOCK_SIZE)
638
-
639
- self.activation_quant_key = create_fp8_quant_key(
640
- static=False, group_shape=GroupShape(1, block_k)
641
- )
642
- self.weight_quant_key = create_fp8_quant_key(
643
- static=True, group_shape=GroupShape(block_n, block_k)
644
- )
645
-
646
- self.out_dtype = get_current_vllm_config().model_config.dtype
647
- self.input_dtype = get_current_vllm_config().model_config.dtype
648
-
649
- def create_weights(
650
- self,
651
- layer: torch.nn.Module,
652
- input_size_per_partition: int,
653
- output_partition_sizes: list[int],
654
- input_size: int,
655
- output_size: int,
656
- params_dtype: torch.dtype,
657
- **extra_weight_attrs,
658
- ):
659
- del input_size, output_size
660
-
661
- if not self.quant_config.is_checkpoint_fp8_serialized:
662
- raise ValueError(
663
- "FP8_PB_WO currently only supports FP8-serialized checkpoints."
664
- )
665
-
666
- output_size_per_partition = sum(output_partition_sizes)
667
- weight_loader = extra_weight_attrs.get("weight_loader")
668
- layer.logical_widths = output_partition_sizes
669
- layer.input_size_per_partition = input_size_per_partition
670
- layer.output_size_per_partition = output_size_per_partition
671
-
672
- # Expose block size so the v2 weight loaders can translate offsets from
673
- # element-space -> block-space for BlockQuantScaleParameter.
674
- layer.weight_block_size = self.weight_block_size
675
-
676
- block_n, block_k = self._WEIGHT_BLOCK_SIZE
677
- remainder = output_size_per_partition % block_n
678
- self.output_padding = 0 if remainder == 0 else block_n - remainder
679
- self.logical_output_size = output_size_per_partition
680
- physical_output_size = output_size_per_partition + self.output_padding
681
-
682
- weight = ModelWeightParameter(
683
- data=torch.empty(
684
- output_size_per_partition,
685
- input_size_per_partition,
686
- dtype=torch.float8_e4m3fn,
687
- ),
688
- input_dim=1,
689
- output_dim=0,
690
- weight_loader=weight_loader,
691
- )
692
- layer.register_parameter("weight", weight)
693
-
694
- if input_size_per_partition % block_k != 0:
695
- raise ValueError(
696
- "ModelOpt FP8_PB_WO requires in_features divisible by "
697
- f"{block_k}, got {input_size_per_partition}."
698
- )
699
-
700
- out_blks = physical_output_size // block_n
701
- in_blks = input_size_per_partition // block_k
702
-
703
- # Match ModelOpt's exported shape so weight loading works without a
704
- # custom loader: [out_blk, 1, in_blk, 1]
705
- weight_scale = BlockQuantScaleParameter(
706
- data=torch.empty((out_blks, 1, in_blks, 1), dtype=torch.float32),
707
- input_dim=2,
708
- output_dim=0,
709
- weight_loader=weight_loader,
710
- )
711
- weight_scale[:] = torch.finfo(torch.float32).min
712
- layer.register_parameter("weight_scale", weight_scale)
713
-
714
- self.w8a8_block_fp8_linear = init_fp8_linear_kernel(
715
- activation_quant_key=self.activation_quant_key,
716
- weight_quant_key=self.weight_quant_key,
717
- weight_shape=(physical_output_size, input_size_per_partition),
718
- input_dtype=self.input_dtype,
719
- out_dtype=self.out_dtype,
720
- module_name=self.__class__.__name__,
721
- )
722
-
723
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
724
- # Keep weight in [out, in] layout for Fp8BlockScaledMMLinearKernel.
725
- weight = layer.weight.data
726
- if self.output_padding:
727
- padded_weight = weight.new_zeros(
728
- self.logical_output_size + self.output_padding,
729
- weight.shape[1],
730
- )
731
- padded_weight[: self.logical_output_size].copy_(weight)
732
- weight = padded_weight
733
- layer.weight = Parameter(weight, requires_grad=False)
734
-
735
- scale = layer.weight_scale
736
- if scale.dim() == 4:
737
- # [out_blk, 1, in_blk, 1] -> [out_blk, in_blk]
738
- scale = scale.squeeze(1).squeeze(-1)
739
- elif scale.dim() != 2:
740
- raise ValueError(
741
- "Unexpected ModelOpt FP8_PB_WO weight_scale shape: "
742
- f"{tuple(scale.shape)}."
743
- )
744
-
745
- layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False)
746
-
747
- self.w8a8_block_fp8_linear.process_weights_after_loading(layer)
748
-
749
- def apply(
750
- self,
751
- layer: torch.nn.Module,
752
- x: torch.Tensor,
753
- bias: torch.Tensor | None = None,
754
- ) -> torch.Tensor:
755
- kernel_bias = None if self.output_padding else bias
756
- output = self.w8a8_block_fp8_linear.apply_weights(layer, x, kernel_bias)
757
- if not self.output_padding:
758
- return output
759
-
760
- output = output[..., : self.logical_output_size].contiguous()
761
- if bias is not None:
762
- output.add_(bias)
763
- return output
764
-
765
-
460
class ModelOptFp8MoEMethod(FusedMoEMethodBase):
461
"""MoE method for ModelOpt FP8.
462
Supports loading FP8 checkpoints with static weight scale and
@@ -1004,7 +698,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
698
)
699
700
1007
-ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod
701
ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod
702
ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
703
@@ -1026,27 +719,17 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
719
self.quant_method = quant_method
720
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
721
if is_checkpoint_nvfp4_serialized:
1029
- logger.warning(
1030
- "Detected ModelOpt NVFP4 checkpoint (quant_algo=%s). Please "
1031
- "note that the format is experimental and could change in "
1032
- "future.",
1033
- quant_method,
1034
- )
1035
-
722
self.group_size = group_size
723
self.kv_cache_quant_algo = kv_cache_quant_algo
724
1039
- # Select LinearMethod implementation based on quant_algo (FP8 pattern).
1040
- # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization
1041
- # W4A16_NVFP4 -> W4A16: flashinfer cute-dsl or Marlin with 16-bit inputs
1042
- if quant_method == "NVFP4":
1043
- self.LinearMethodCls = ModelOptNvFp4LinearMethod
1044
- elif quant_method == "W4A16_NVFP4":
1045
- self.LinearMethodCls = ModelOptNvFp4W4A16LinearMethod
1046
- else:
725
+ # Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch.
726
+ # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization
727
+ # W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations
728
+ supported = algos_owned_by("modelopt_fp4")
729
+ if quant_method not in supported:
730
raise ValueError(
731
f"Unsupported ModelOpt NVFP4 quant_algo: {quant_method}. "
1049
- "Supported: NVFP4 / W4A16_NVFP4."
732
+ f"Supported: {' / '.join(supported)}."
733
)
734
735
def get_name(self) -> QuantizationMethods:
@@ -1123,363 +806,84 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
806
)
807
808
1126
-class ModelOptNvFp4LinearMethod(LinearMethodBase):
1127
- """Linear method for Model Optimizer NVFP4.
1128
- Supports loading NVFP4 checkpoints with the following structure:
1129
-
1130
- input_scale: torch.float32, scalar ,
1131
- weight: NVFP4(represented as byte) Shape: [1, X, y/2]
1132
- weight_scale: FP8-E4M3, Shape: [X, Y], aka per block scale,
1133
- weight_scale_2: torch.float32, scalar,
1134
- Args: quant_config: The ModelOpt quantization config.
809
+class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
810
+ """
811
+ MoE Method for FP4 Quantization.
812
+ Args:
813
+ quant_config: NVFP4 Quant Config
814
"""
815
1137
- def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
816
+ def __init__(
817
+ self,
818
+ quant_config: ModelOptNvFp4Config,
819
+ moe_config: FusedMoEConfig,
820
+ ) -> None:
821
+ super().__init__(moe_config)
822
self.quant_config = quant_config
1139
- self.marlin_input_dtype = None
1140
- self.kernel = init_nvfp4_linear_kernel()
823
+ # W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With
824
+ # activation_key=None every W4A4 backend's _supports_quant_scheme
825
+ # rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic)
826
+ # exactly); only Marlin survives. Marlin's MoE path drops
827
+ # activation scales in convert_to_nvfp4_moe_kernel_format, so no
828
+ # other change is needed.
829
+ self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
830
+ self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
831
+ config=self.moe,
832
+ weight_key=kNvfp4Static,
833
+ activation_key=None if self.use_a16 else kNvfp4Dynamic,
834
+ )
835
+
836
+ self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
837
+ self.nvfp4_backend
838
+ )
839
+
840
+ def uses_weight_scale_2_pattern(self) -> bool:
841
+ """
842
+ FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
843
+ """
844
+ return True
845
846
def create_weights(
847
self,
1144
- layer: torch.nn.Module,
1145
- input_size_per_partition: int,
1146
- output_partition_sizes: list[int],
1147
- input_size: int,
1148
- output_size: int,
848
+ layer: RoutedExperts,
849
+ num_experts: int,
850
+ hidden_size: int,
851
+ intermediate_size_per_partition: int,
852
params_dtype: torch.dtype,
853
**extra_weight_attrs,
854
):
1152
- del input_size, output_size
1153
- if not self.quant_config.is_checkpoint_nvfp4_serialized:
1154
- raise ValueError(
1155
- "NVFP4 quantization was selected, "
1156
- " dynamic quantization is not supported."
1157
- )
1158
- output_size_per_partition = sum(output_partition_sizes)
1159
- weight_loader = extra_weight_attrs.get("weight_loader")
1160
- layer.logical_widths = output_partition_sizes
1161
- layer.input_size_per_partition = input_size_per_partition
1162
- layer.output_size_per_partition = output_size_per_partition
855
+ assert self.quant_config.is_checkpoint_nvfp4_serialized
856
1164
- if input_size_per_partition % 16 != 0:
1165
- raise ValueError(
1166
- "Unsupported model when in features size is not multiple of 16"
1167
- )
1168
- # The nvfp4 weight is still represented as
1169
- weight_dtype = (
1170
- torch.float8_e4m3fn
1171
- if self.quant_config.is_checkpoint_nvfp4_serialized
1172
- else params_dtype
1173
- )
1174
- # Weight
1175
- weight = ModelWeightParameter(
857
+ layer.num_experts = num_experts
858
+ layer.params_dtype = params_dtype
859
+ layer.quant_config = self.quant_config
860
+ weight_dtype = torch.uint8
861
+ weight_scale_dtype = torch.float8_e4m3fn
862
+ weight_loader = extra_weight_attrs.get("weight_loader")
863
+ global_num_experts = extra_weight_attrs.get("global_num_experts")
864
+ w13_num_shards = 2 if self.moe.is_act_and_mul else 1
865
+ # GEMM 1
866
+ w13_weight = ModelWeightParameter(
867
data=torch.empty(
868
+ num_experts,
869
+ w13_num_shards * intermediate_size_per_partition,
870
# 2 fp4 items are packed in the input dimension
1178
- layer.output_size_per_partition,
1179
- layer.input_size_per_partition // 2,
1180
- dtype=torch.uint8,
871
+ hidden_size // 2,
872
+ dtype=weight_dtype,
873
),
874
input_dim=1,
1183
- output_dim=0,
1184
- weight_loader=weight_loader,
1185
- )
1186
- layer.register_parameter("weight", weight)
1187
-
1188
- # Input Global Scale
1189
- input_global_scale = PerTensorScaleParameter(
1190
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1191
- weight_loader=weight_loader,
1192
- )
1193
- layer.register_parameter("input_scale", input_global_scale)
1194
-
1195
- # Weight Global Scale
1196
- weight_global_scale = PerTensorScaleParameter(
1197
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
875
+ output_dim=2,
876
weight_loader=weight_loader,
877
)
1200
- layer.register_parameter("weight_scale_2", weight_global_scale)
878
+ layer.register_parameter("w13_weight", w13_weight)
879
1202
- # Per Block Weight Scale
1203
- weight_scale = ModelWeightParameter(
880
+ # GEMM 2
881
+ w2_weight = ModelWeightParameter(
882
data=torch.empty(
1205
- output_size_per_partition,
1206
- input_size_per_partition // self.quant_config.group_size,
1207
- dtype=weight_dtype,
1208
- ),
1209
- input_dim=1,
1210
- output_dim=0,
1211
- weight_loader=weight_loader,
1212
- )
1213
-
1214
- layer.register_parameter("weight_scale", weight_scale)
1215
-
1216
- expose_input_quant_key(layer, self.kernel)
1217
-
1218
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1219
- if (
1220
- torch.unique(layer.input_scale).numel() != 1
1221
- or torch.unique(layer.weight_scale_2).numel() != 1
1222
- ):
1223
- logger.warning_once(
1224
- "In NVFP4 linear, the global scale for input or weight are different"
1225
- " for parallel layers (e.g. q_proj, k_proj, v_proj). This "
1226
- " will likely results in reduce accuracy. Please verify the model"
1227
- " accuracy. Consider using a checkpoint with a shared global NVFP4"
1228
- " scale for parallel layers."
1229
- )
1230
-
1231
- # Rename ModelOpt checkpoint names to standardized names
1232
- input_global_scale = layer.input_scale.max().to(torch.float32)
1233
- layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
1234
- del layer.input_scale
1235
-
1236
- weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
1237
- layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
1238
- del layer.weight_scale_2
1239
-
1240
- # Pre-compute alpha and inverse for runtime quantization
1241
- layer.alpha = Parameter(
1242
- layer.input_global_scale * layer.weight_global_scale, requires_grad=False
1243
- )
1244
- layer.input_global_scale_inv = Parameter(
1245
- (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False
1246
- )
1247
-
1248
- # Convert layer to NVFP4 linear kernel format
1249
- self.kernel.process_weights_after_loading(layer)
1250
-
1251
- def apply(
1252
- self,
1253
- layer: torch.nn.Module,
1254
- x: torch.Tensor,
1255
- bias: torch.Tensor | None = None,
1256
- ) -> torch.Tensor:
1257
- return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
1258
-
1259
-
1260
-class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):
1261
- """Linear method for ModelOpt NVFP4 W4A16.
1262
-
1263
- 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names
1264
- directly (no on-disk conversion) and dispatches to a W4A16 GEMM:
1265
-
1266
- weight uint8 packed NVFP4 (2 nibbles/byte along input dim)
1267
- weight_scale fp8-e4m3 per 16-elem group along input dim
1268
- weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0)
1269
-
1270
- No activation quantization. ModelOpt stores the global scale as
1271
- amax/2688, so we rename weight_scale_2 -> weight_global_scale without
1272
- reciprocation. The selected kernel converts it to its runtime format.
1273
- The CT W4A16 path reciprocates because CT stores the inverse on disk.
1274
-
1275
- We also register a placeholder input_scale parameter so that W4A4-shaped
1276
- checkpoints (which contain *_proj.input_scale tensors) can be loaded
1277
- under this method without the per-shard loader hitting a KeyError on
1278
- the merged-name lookup. The placeholder is discarded in
1279
- process_weights_after_loading -- its value is never used.
1280
- """
1281
-
1282
- def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
1283
- self.quant_config = quant_config
1284
- self.marlin_input_dtype = None
1285
- # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds:
1286
- # 1. `use_a16=True` forces flashinfer cutedsl for sm100 and `Marlin` for others.
1287
- # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale.
1288
- # 2. Specifying e.g. `--linear-backend=humming` will override.
1289
- self.kernel = init_nvfp4_linear_kernel(use_a16=True)
1290
-
1291
- def create_weights(
1292
- self,
1293
- layer: torch.nn.Module,
1294
- input_size_per_partition: int,
1295
- output_partition_sizes: list[int],
1296
- input_size: int,
1297
- output_size: int,
1298
- params_dtype: torch.dtype,
1299
- **extra_weight_attrs,
1300
- ):
1301
- del input_size, output_size
1302
- if not self.quant_config.is_checkpoint_nvfp4_serialized:
1303
- raise ValueError(
1304
- "W4A16_NVFP4 quantization was selected; "
1305
- "dynamic quantization is not supported."
1306
- )
1307
- output_size_per_partition = sum(output_partition_sizes)
1308
- weight_loader = extra_weight_attrs.get("weight_loader")
1309
- layer.logical_widths = output_partition_sizes
1310
- layer.input_size_per_partition = input_size_per_partition
1311
- layer.output_size_per_partition = output_size_per_partition
1312
- layer.output_partition_sizes = output_partition_sizes
1313
-
1314
- if input_size_per_partition % 16 != 0:
1315
- raise ValueError(
1316
- "Unsupported model: input feature size is not a multiple of 16."
1317
- )
1318
-
1319
- # Packed NVFP4 weights: uint8, 2 nibbles per byte along the input dim.
1320
- weight = ModelWeightParameter(
1321
- data=torch.empty(
1322
- output_size_per_partition,
1323
- input_size_per_partition // 2,
1324
- dtype=torch.uint8,
1325
- ),
1326
- input_dim=1,
1327
- output_dim=0,
1328
- weight_loader=weight_loader,
1329
- )
1330
- layer.register_parameter("weight", weight)
1331
-
1332
- # Per-tensor global weight scale (fp32). ModelOpt stores
1333
- # amax / (NVFP4_max * fp8_e4m3_max) = amax / 2688. PerTensorScaleParameter
1334
- # holds one entry per fused output partition (e.g. q/k/v in a fused QKV).
1335
- weight_scale_2 = PerTensorScaleParameter(
1336
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1337
- weight_loader=weight_loader,
1338
- )
1339
- layer.register_parameter("weight_scale_2", weight_scale_2)
1340
-
1341
- # Per-group fp8 weight scale.
1342
- weight_scale = GroupQuantScaleParameter(
1343
- data=torch.empty(
1344
- output_size_per_partition,
1345
- input_size_per_partition // self.quant_config.group_size,
1346
- dtype=torch.float8_e4m3fn,
1347
- ),
1348
- input_dim=1,
1349
- output_dim=0,
1350
- weight_loader=weight_loader,
1351
- )
1352
- layer.register_parameter("weight_scale", weight_scale)
1353
-
1354
- # Placeholder input_scale param so W4A4-shaped checkpoints can be
1355
- # loaded under this method without KeyError on the merged-name
1356
- # lookup (qwen2-style stacked-loader path renames *_proj.input_scale
1357
- # to e.g. qkv_proj.input_scale and looks it up unconditionally).
1358
- # Discarded in process_weights_after_loading; never read by the kernel.
1359
- # For native W4A16 checkpoints (no input_scale on disk) the param
1360
- # stays uninitialized and is simply deleted.
1361
- input_scale = PerTensorScaleParameter(
1362
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1363
- weight_loader=weight_loader,
1364
- )
1365
- layer.register_parameter("input_scale", input_scale)
1366
-
1367
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1368
- if not hasattr(layer, "has_bias"):
1369
- layer.has_bias = getattr(layer, "bias", None) is not None
1370
-
1371
- # Discard the input_scale placeholder. Whether it carries values
1372
- # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16
1373
- # ckpt), W4A16 mode does not quantize activations, so this is unused.
1374
- if hasattr(layer, "input_scale"):
1375
- del layer.input_scale
1376
-
1377
- if torch.unique(layer.weight_scale_2).numel() != 1:
1378
- logger.warning_once(
1379
- "In W4A16_NVFP4 linear, the global weight scale "
1380
- "(weight_scale_2) differs across fused parallel layers "
1381
- "(e.g. q/k/v_proj). This will likely reduce accuracy. "
1382
- "Consider a checkpoint with a shared global scale."
1383
- )
1384
-
1385
- # Rename weight_scale_2 -> weight_global_scale. NO reciprocation:
1386
- # ModelOpt already stores amax/2688, which is exactly what Marlin
1387
- # consumes via nvfp4_marlin_process_global_scale (called inside the
1388
- # Marlin adapter's process_weights_after_loading).
1389
- layer.weight_global_scale = Parameter(
1390
- layer.weight_scale_2.max().to(torch.float32), requires_grad=False
1391
- )
1392
- del layer.weight_scale_2
1393
-
1394
- self.kernel.process_weights_after_loading(layer)
1395
-
1396
- def apply(
1397
- self,
1398
- layer: torch.nn.Module,
1399
- x: torch.Tensor,
1400
- bias: torch.Tensor | None = None,
1401
- ) -> torch.Tensor:
1402
- return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
1403
-
1404
-
1405
-class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1406
- """
1407
- MoE Method for FP4 Quantization.
1408
- Args:
1409
- quant_config: NVFP4 Quant Config
1410
- """
1411
-
1412
- def __init__(
1413
- self,
1414
- quant_config: ModelOptNvFp4Config,
1415
- moe_config: FusedMoEConfig,
1416
- ) -> None:
1417
- super().__init__(moe_config)
1418
- self.quant_config = quant_config
1419
- # W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With
1420
- # activation_key=None every W4A4 backend's _supports_quant_scheme
1421
- # rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic)
1422
- # exactly); only Marlin survives. Marlin's MoE path drops
1423
- # activation scales in convert_to_nvfp4_moe_kernel_format, so no
1424
- # other change is needed.
1425
- self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
1426
- self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
1427
- config=self.moe,
1428
- weight_key=kNvfp4Static,
1429
- activation_key=None if self.use_a16 else kNvfp4Dynamic,
1430
- )
1431
-
1432
- self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
1433
- self.nvfp4_backend
1434
- )
1435
-
1436
- def uses_weight_scale_2_pattern(self) -> bool:
1437
- """
1438
- FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
1439
- """
1440
- return True
1441
-
1442
- def create_weights(
1443
- self,
1444
- layer: RoutedExperts,
1445
- num_experts: int,
1446
- hidden_size: int,
1447
- intermediate_size_per_partition: int,
1448
- params_dtype: torch.dtype,
1449
- **extra_weight_attrs,
1450
- ):
1451
- assert self.quant_config.is_checkpoint_nvfp4_serialized
1452
-
1453
- layer.num_experts = num_experts
1454
- layer.params_dtype = params_dtype
1455
- layer.quant_config = self.quant_config
1456
- weight_dtype = torch.uint8
1457
- weight_scale_dtype = torch.float8_e4m3fn
1458
- weight_loader = extra_weight_attrs.get("weight_loader")
1459
- global_num_experts = extra_weight_attrs.get("global_num_experts")
1460
- w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1461
- # GEMM 1
1462
- w13_weight = ModelWeightParameter(
1463
- data=torch.empty(
1464
- num_experts,
1465
- w13_num_shards * intermediate_size_per_partition,
1466
- # 2 fp4 items are packed in the input dimension
1467
- hidden_size // 2,
1468
- dtype=weight_dtype,
1469
- ),
1470
- input_dim=1,
1471
- output_dim=2,
1472
- weight_loader=weight_loader,
1473
- )
1474
- layer.register_parameter("w13_weight", w13_weight)
1475
-
1476
- # GEMM 2
1477
- w2_weight = ModelWeightParameter(
1478
- data=torch.empty(
1479
- num_experts,
1480
- hidden_size,
1481
- # 2 fp4 items are packed in the input dimension
1482
- intermediate_size_per_partition // 2,
883
+ num_experts,
884
+ hidden_size,
885
+ # 2 fp4 items are packed in the input dimension
886
+ intermediate_size_per_partition // 2,
887
dtype=weight_dtype,
888
),
889
input_dim=1,
@@ -1685,7 +1089,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1089
)
1090
1091
1688
-ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod
1092
ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
1093
ModelOptNvFp4Config.KVCacheMethodCls = ModelOptKVCacheMethod
1094
@@ -1700,6 +1103,7 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1103
exclude_modules: list[str],
1104
) -> None:
1105
super().__init__(exclude_modules)
1106
+ self.quant_method = "MXFP8"
1107
self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized
1108
1109
if not is_checkpoint_mxfp8_serialized:
@@ -1708,11 +1112,6 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1112
"Dynamic quantization is not supported."
1113
)
1114
1711
- logger.warning(
1712
- "Detected ModelOpt MXFP8 checkpoint. Please note that "
1713
- "the format is experimental and could change in future."
1714
- )
1715
-
1115
self.kv_cache_quant_algo = kv_cache_quant_algo
1116
1117
def get_name(self) -> QuantizationMethods:
@@ -1783,117 +1182,6 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1182
)
1183
1184
1786
-class ModelOptMxFp8LinearMethod(LinearMethodBase):
1787
- """Linear method for ModelOpt MXFP8 quantization."""
1788
-
1789
- def __init__(self, quant_config: ModelOptMxFp8Config) -> None:
1790
- self.quant_config = quant_config
1791
-
1792
- if not self.quant_config.is_checkpoint_mxfp8_serialized:
1793
- raise ValueError(
1794
- "MXFP8 currently only supports serialized checkpoints. "
1795
- "Dynamic quantization is not supported."
1796
- )
1797
-
1798
- self.kernel = init_mxfp8_linear_kernel()
1799
-
1800
- def create_weights(
1801
- self,
1802
- layer: torch.nn.Module,
1803
- input_size_per_partition: int,
1804
- output_partition_sizes: list[int],
1805
- input_size: int,
1806
- output_size: int,
1807
- params_dtype: torch.dtype,
1808
- **extra_weight_attrs,
1809
- ):
1810
- del input_size, output_size
1811
-
1812
- if not self.quant_config.is_checkpoint_mxfp8_serialized:
1813
- raise ValueError(
1814
- "MXFP8 quantization was selected, but checkpoint is not "
1815
- "MXFP8 serialized. Dynamic quantization is not supported."
1816
- )
1817
-
1818
- output_size_per_partition = sum(output_partition_sizes)
1819
- weight_loader = extra_weight_attrs.get("weight_loader")
1820
- layer.logical_widths = output_partition_sizes
1821
- layer.input_size_per_partition = input_size_per_partition
1822
- layer.output_size_per_partition = output_size_per_partition
1823
-
1824
- if input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
1825
- raise ValueError(
1826
- f"MXFP8 requires input dimension to be divisible by "
1827
- f"{MXFP8_BLOCK_SIZE}, got {input_size_per_partition}"
1828
- )
1829
-
1830
- # Weight tensor: FP8 E4M3 format
1831
- weight = ModelWeightParameter(
1832
- data=torch.empty(
1833
- output_size_per_partition,
1834
- input_size_per_partition,
1835
- dtype=MXFP8_VALUE_DTYPE,
1836
- ),
1837
- input_dim=1,
1838
- output_dim=0,
1839
- weight_loader=weight_loader,
1840
- )
1841
- layer.register_parameter("weight", weight)
1842
-
1843
- # Weight scale tensor (E8M0 encoded as uint8), one scale per block of 32 along K
1844
- weight_scale = ModelWeightParameter(
1845
- data=torch.empty(
1846
- output_size_per_partition,
1847
- input_size_per_partition // MXFP8_BLOCK_SIZE,
1848
- dtype=MXFP8_SCALE_DTYPE,
1849
- ),
1850
- input_dim=1,
1851
- output_dim=0,
1852
- weight_loader=weight_loader,
1853
- )
1854
- layer.register_parameter("weight_scale", weight_scale)
1855
-
1856
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1857
- # Idempotent: the emulation kernel may dequant the weight to BF16 at load
1858
- # time (>=2-byte). If already converted, there is nothing left to do --
1859
- # avoid re-running the MXFP8-only validation/conversion below.
1860
- if layer.weight.element_size() >= 2:
1861
- return
1862
-
1863
- # Validate weight tensor
1864
- if layer.weight.ndim != 2:
1865
- raise ValueError(
1866
- f"MXFP8 weight must be 2D tensor [N, K], got {layer.weight.ndim}D "
1867
- f"with shape {tuple(layer.weight.shape)}"
1868
- )
1869
-
1870
- if layer.weight.dtype != MXFP8_VALUE_DTYPE:
1871
- raise ValueError(
1872
- f"MXFP8 weight must be {MXFP8_VALUE_DTYPE} (FP8 E4M3), "
1873
- f"got {layer.weight.dtype}. The checkpoint may not be properly "
1874
- f"quantized with MXFP8."
1875
- )
1876
-
1877
- # Validate weight scale tensor (should be 2D, not swizzled)
1878
- assert layer.weight_scale.ndim == 2, (
1879
- f"MXFP8 weight scale must be 2D, got {layer.weight_scale.ndim}D"
1880
- )
1881
- assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE, (
1882
- f"MXFP8 weight scale must be {MXFP8_SCALE_DTYPE},"
1883
- f" got {layer.weight_scale.dtype}"
1884
- )
1885
-
1886
- self.kernel.process_weights_after_loading(layer)
1887
-
1888
- def apply(
1889
- self,
1890
- layer: torch.nn.Module,
1891
- x: torch.Tensor,
1892
- bias: torch.Tensor | None = None,
1893
- ) -> torch.Tensor:
1894
- return self.kernel.apply_weights(layer, x, bias)
1895
-
1896
-
1185
class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1186
"""FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints."""
1187
@@ -2171,7 +1459,6 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1459
1460
1461
# Register the method classes for ModelOptMxFp8Config
2174
-ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod
1462
ModelOptMxFp8Config.FusedMoEMethodCls = ModelOptMxFp8FusedMoE
1463
ModelOptMxFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
1464
@@ -2204,6 +1491,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1491
self.w4a16_nvfp4_config = w4a16_nvfp4_config
1492
self.mxfp8_config = mxfp8_config
1493
1494
+ def has_blocked_weights(self) -> bool:
1495
+ # Same gate as ModelOptFp8Config.has_blocked_weights, resolved per
1496
+ # layer: "+quant_fp8" must be on as soon as any layer is block-scaled.
1497
+ return any(
1498
+ info.get("quant_algo", "").upper() == "FP8_PB_WO"
1499
+ for info in self.quantized_layers.values()
1500
+ )
1501
+
1502
def get_name(self) -> QuantizationMethods:
1503
return "modelopt_mixed"
1504
@@ -2281,10 +1576,10 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1576
group_size=group_size,
1577
)
1578
# Sibling config for layers that declare quant_algo: "W4A16_NVFP4".
2284
- # ModelOptNvFp4Config.__init__ keys LinearMethodCls off quant_method,
2285
- # so this instance auto-selects ModelOptNvFp4W4A16LinearMethod. The
2286
- # MoE side reads quant_config.quant_method == "W4A16_NVFP4" to set
2287
- # use_a16 → Marlin backend in ModelOptNvFp4FusedMoE.__init__.
1579
+ # get_quant_method resolves this sub-config to the (kNvfp4Static, None)
1580
+ # QuantSpec (W4A16) for the generic ModelOptLinearMethod. The MoE side
1581
+ # reads quant_config.quant_method == "W4A16_NVFP4" to set use_a16 →
1582
+ # Marlin backend in ModelOptNvFp4FusedMoE.__init__.
1583
w4a16_nvfp4_config = ModelOptNvFp4Config(
1584
quant_method="W4A16_NVFP4",
1585
is_checkpoint_nvfp4_serialized=True,
@@ -2424,18 +1719,15 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1719
quant_algo = self._resolve_quant_algo(prefix)
1720
1721
if isinstance(layer, (LinearBase, ParallelLMHead)):
2427
- if quant_algo == "FP8":
2428
- return ModelOptFp8LinearMethod(self.fp8_config)
2429
- if quant_algo == "FP8_PB_WO":
2430
- return ModelOptFp8PbWoLinearMethod(self.fp8_config)
2431
- if quant_algo == "NVFP4":
2432
- return ModelOptNvFp4LinearMethod(self.nvfp4_config)
2433
- if quant_algo == "W4A16_NVFP4":
2434
- return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config)
2435
- if quant_algo == "MXFP8":
2436
- return ModelOptMxFp8LinearMethod(self.mxfp8_config)
2437
- # Layer not in quantized_layers — leave unquantized
2438
- return UnquantizedLinearMethod()
1722
+ # Per-prefix algo -> its sub-config, then the generic linear method.
1723
+ # resolve() reads group_size off whichever sub-config it is handed.
1724
+ # Read-only — never write back, or a linear-only change leaks into
1725
+ # the MoE method that shares the sub-config.
1726
+ if quant_algo is None or quant_algo not in LINEAR_ALGOS:
1727
+ # Layer not in quantized_layers — leave unquantized
1728
+ return UnquantizedLinearMethod()
1729
+ _, subcfg_attr = LINEAR_ALGOS[quant_algo]
1730
+ return build_linear_method(getattr(self, subcfg_attr), quant_algo, prefix)
1731
1732
if isinstance(layer, RoutedExperts):
1733
if quant_algo == "FP8":
@@ -2466,3 +1758,752 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1758
super().apply_vllm_mapper(hf_to_vllm_mapper)
1759
if self.quantized_layers:
1760
self.quantized_layers = hf_to_vllm_mapper.apply_dict(self.quantized_layers)
1761
+
1762
+
1763
+# ===========================================================================
1764
+# Generic QuantKey-driven linear method
1765
+#
1766
+# ``ModelOptLinearMethod`` serves every ModelOpt linear format. It composes a
1767
+# per-QuantKey weight scheme + activation scheme (the ``QuantSpec`` pair that
1768
+# ``resolve`` produces from the checkpoint's quant_algo), runs a fixed
1769
+# create/process lifecycle, selects the kernel from the pair, and applies.
1770
+#
1771
+# Adding a format:
1772
+# * Composes as a (weight, activation) key pair -> add a QuantKeyScheme per
1773
+# new key to SCHEME_FOR, plus a resolve() row returning the QuantSpec. No
1774
+# new method class; this is how every current format is built.
1775
+# * Needs format-wide state but the same lifecycle -> also return a
1776
+# FormatScheme subclass from that resolve() row (extra_weights / pre_process
1777
+# / post_process hooks).
1778
+# * Needs a different lifecycle entirely -> write a LinearMethodBase and
1779
+# register it in LINEAR_METHOD_BUILDERS by algo.
1780
+# In all cases accept the algo in the owning config's quant_method validation,
1781
+# so get_quant_method can dispatch it.
1782
+# ===========================================================================
1783
+
1784
+
1785
+class Role(Enum):
1786
+ WEIGHT = "weight"
1787
+ ACT = "activation"
1788
+
1789
+
1790
+WEIGHT = Role.WEIGHT
1791
+ACT = Role.ACT
1792
+
1793
+
1794
+@dataclass(frozen=True)
1795
+class CkptCtx:
1796
+ """Per-checkpoint facts a QuantKey cannot carry."""
1797
+
1798
+ group_size: int | None = None
1799
+
1800
+
1801
+@dataclass(frozen=True)
1802
+class RuntimeDtypes:
1803
+ """Runtime/model dtypes kernels need — format-agnostic."""
1804
+
1805
+ input_dtype: torch.dtype
1806
+ out_dtype: torch.dtype
1807
+ marlin_input_dtype: torch.dtype | None = None
1808
+
1809
+
1810
+@dataclass(frozen=True)
1811
+class Shapes:
1812
+ """Layer geometry from create_weights args."""
1813
+
1814
+ output_partition_sizes: list[int]
1815
+ input_size_per_partition: int
1816
+ params_dtype: torch.dtype
1817
+
1818
+ @property
1819
+ def output_size_per_partition(self) -> int:
1820
+ return sum(self.output_partition_sizes)
1821
+
1822
+ @property
1823
+ def num_partitions(self) -> int:
1824
+ return len(self.output_partition_sizes)
1825
+
1826
+
1827
+class QuantKeyScheme:
1828
+ """One scheme per QuantKey. Selected by key *content*; the base supplies the
1829
+ ``role`` from the slot it fills the key into (wkey->WEIGHT, akey->ACT).
1830
+
1831
+ Every scheme branches on ``role`` explicitly and ``reject``s any role it has
1832
+ not validated — never falling through to the wrong role's registration,
1833
+ which would allocate the wrong parameters and produce garbage silently.
1834
+ """
1835
+
1836
+ key: QuantKey
1837
+
1838
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1839
+ raise NotImplementedError
1840
+
1841
+ def process(self, layer, role) -> None:
1842
+ pass
1843
+
1844
+ @staticmethod
1845
+ def reject(role) -> None:
1846
+ raise NotImplementedError(
1847
+ f"role {role!r} not validated for this QuantKey scheme"
1848
+ )
1849
+
1850
+ @staticmethod
1851
+ def register_params(layer, name, shape, dtype, cls, wl, *, init=None, **dims):
1852
+ p = cls(data=torch.empty(shape, dtype=dtype), weight_loader=wl, **dims)
1853
+ if init is not None:
1854
+ p.data.fill_(init)
1855
+ layer.register_parameter(name, p)
1856
+
1857
+
1858
+class KNvfp4Static(QuantKeyScheme):
1859
+ """NVFP4 weight scheme (W4A4 and W4A16 share it). Weight-role only today."""
1860
+
1861
+ key = kNvfp4Static
1862
+
1863
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1864
+ if role is not WEIGHT:
1865
+ self.reject(role)
1866
+ if shapes.input_size_per_partition % 16 != 0:
1867
+ raise ValueError(
1868
+ "Unsupported model when in features size is not multiple of 16"
1869
+ )
1870
+ # Packed NVFP4 weight: 2 fp4 items per byte along the input dim.
1871
+ self.register_params(
1872
+ layer,
1873
+ "weight",
1874
+ (shapes.output_size_per_partition, shapes.input_size_per_partition // 2),
1875
+ torch.uint8,
1876
+ ModelWeightParameter,
1877
+ wl,
1878
+ input_dim=1,
1879
+ output_dim=0,
1880
+ )
1881
+ # Per-tensor global weight scale.
1882
+ self.register_params(
1883
+ layer,
1884
+ "weight_scale_2",
1885
+ (shapes.num_partitions,),
1886
+ torch.float32,
1887
+ PerTensorScaleParameter,
1888
+ wl,
1889
+ )
1890
+ # Per-block (group_size) weight scale.
1891
+ self.register_params(
1892
+ layer,
1893
+ "weight_scale",
1894
+ (
1895
+ shapes.output_size_per_partition,
1896
+ shapes.input_size_per_partition // ctx.group_size,
1897
+ ),
1898
+ torch.float8_e4m3fn,
1899
+ ModelWeightParameter,
1900
+ wl,
1901
+ input_dim=1,
1902
+ output_dim=0,
1903
+ )
1904
+
1905
+ def process(self, layer, role) -> None:
1906
+ if role is not WEIGHT:
1907
+ self.reject(role)
1908
+ if torch.unique(layer.weight_scale_2).numel() != 1:
1909
+ logger.warning_once(
1910
+ "In NVFP4 linear, the global weight scale differs across "
1911
+ "parallel layers (e.g. q_proj, k_proj, v_proj). This will "
1912
+ "likely reduce accuracy. Consider a checkpoint with the same "
1913
+ "global NVFP4 scale for fused weights."
1914
+ )
1915
+ # Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688.
1916
+ weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
1917
+ layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
1918
+ del layer.weight_scale_2
1919
+
1920
+
1921
+class KNvfp4Dynamic(QuantKeyScheme):
1922
+ """NVFP4 activation scheme (W4A4). Has a static global input scale on disk;
1923
+ the per-group scale is computed at runtime inside the kernel."""
1924
+
1925
+ key = kNvfp4Dynamic
1926
+
1927
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1928
+ if role is not ACT:
1929
+ self.reject(role)
1930
+ self.register_params(
1931
+ layer,
1932
+ "input_scale",
1933
+ (shapes.num_partitions,),
1934
+ torch.float32,
1935
+ PerTensorScaleParameter,
1936
+ wl,
1937
+ )
1938
+
1939
+ def process(self, layer, role) -> None:
1940
+ if role is not ACT:
1941
+ self.reject(role)
1942
+ if torch.unique(layer.input_scale).numel() != 1:
1943
+ logger.warning_once(
1944
+ "In NVFP4 linear, the global input scale differs across "
1945
+ "parallel layers (e.g. q_proj, k_proj, v_proj). This will "
1946
+ "likely reduce accuracy. Consider a checkpoint with the same "
1947
+ "global NVFP4 scale for fused weights."
1948
+ )
1949
+ input_global_scale = layer.input_scale.max().to(torch.float32)
1950
+ layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
1951
+ layer.input_global_scale_inv = Parameter(
1952
+ (1.0 / layer.input_global_scale).to(torch.float32),
1953
+ requires_grad=False,
1954
+ )
1955
+ del layer.input_scale
1956
+
1957
+
1958
+class KFp8StaticTensor(QuantKeyScheme):
1959
+ """Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and
1960
+ the activation slot (W8A8). One key in both QuantSpec slots."""
1961
+
1962
+ key = kFp8StaticTensorSym
1963
+
1964
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1965
+ if role is WEIGHT:
1966
+ self.register_params(
1967
+ layer,
1968
+ "weight",
1969
+ (shapes.output_size_per_partition, shapes.input_size_per_partition),
1970
+ torch.float8_e4m3fn,
1971
+ ModelWeightParameter,
1972
+ wl,
1973
+ input_dim=1,
1974
+ output_dim=0,
1975
+ )
1976
+ layer.orig_dtype = shapes.params_dtype
1977
+ self.register_params(
1978
+ layer,
1979
+ "weight_scale",
1980
+ (shapes.num_partitions,),
1981
+ torch.float32,
1982
+ PerTensorScaleParameter,
1983
+ wl,
1984
+ init=FP8_SCALE_SENTINEL,
1985
+ )
1986
+ elif role is ACT:
1987
+ self.register_params(
1988
+ layer,
1989
+ "input_scale",
1990
+ (shapes.num_partitions,),
1991
+ torch.float32,
1992
+ PerTensorScaleParameter,
1993
+ wl,
1994
+ init=FP8_SCALE_SENTINEL,
1995
+ )
1996
+ else:
1997
+ self.reject(role)
1998
+
1999
+ def process(self, layer, role) -> None:
2000
+ if role is WEIGHT:
2001
+ weight = layer.weight
2002
+ max_w_scale = layer.weight_scale.max()
2003
+ if not (layer.weight_scale == layer.weight_scale[0]).all():
2004
+ max_w_scale, weight = requantize_with_max_scale(
2005
+ layer.weight, layer.weight_scale, layer.logical_widths
2006
+ )
2007
+ # Transposed here rather than in the kernel: ModelOpt maps each
2008
+ # fp8 key to exactly one kernel, so the layout is key-determined.
2009
+ layer.weight = Parameter(weight.t(), requires_grad=False)
2010
+ # The plain Parameter drops ModelWeightParameter's dim attributes;
2011
+ # restore them for the transposed layout, which Humming reads.
2012
+ layer.weight.input_dim = 0
2013
+ layer.weight.output_dim = 1
2014
+ layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
2015
+ elif role is ACT:
2016
+ if torch.unique(layer.input_scale).numel() != 1:
2017
+ logger.warning_once(
2018
+ "In FP8 linear, the static input scale differs across "
2019
+ "parallel layers (e.g. q_proj, k_proj, v_proj). Collapsing "
2020
+ "them to the max will likely reduce accuracy. Consider a "
2021
+ "checkpoint with the same input scale for fused weights."
2022
+ )
2023
+ layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
2024
+ else:
2025
+ self.reject(role)
2026
+
2027
+
2028
+class KFp8StaticChannel(QuantKeyScheme):
2029
+ """Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only —
2030
+ there is no static per-channel *activation* today."""
2031
+
2032
+ key = kFp8StaticTokenSym
2033
+
2034
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2035
+ if role is not WEIGHT:
2036
+ self.reject(role)
2037
+ self.register_params(
2038
+ layer,
2039
+ "weight",
2040
+ (shapes.output_size_per_partition, shapes.input_size_per_partition),
2041
+ torch.float8_e4m3fn,
2042
+ ModelWeightParameter,
2043
+ wl,
2044
+ input_dim=1,
2045
+ output_dim=0,
2046
+ )
2047
+ self.register_params(
2048
+ layer,
2049
+ "weight_scale",
2050
+ (shapes.output_size_per_partition,),
2051
+ torch.float32,
2052
+ ChannelQuantScaleParameter,
2053
+ wl,
2054
+ output_dim=0,
2055
+ init=FP8_SCALE_SENTINEL,
2056
+ )
2057
+
2058
+ def process(self, layer, role) -> None:
2059
+ if role is not WEIGHT:
2060
+ self.reject(role)
2061
+ weight, weight_scale, _ = process_fp8_weight_channel_strategy(
2062
+ layer.weight, layer.weight_scale.data
2063
+ )
2064
+ layer.weight = Parameter(weight.t(), requires_grad=False)
2065
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
2066
+
2067
+
2068
+class KFp8Block128(QuantKeyScheme):
2069
+ """128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt
2070
+ exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D.
2071
+ No transpose (block kernel keeps [out,in])."""
2072
+
2073
+ key = kFp8Static128BlockSym
2074
+
2075
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2076
+ if role is not WEIGHT:
2077
+ self.reject(role)
2078
+ out = shapes.output_size_per_partition
2079
+ in_ = shapes.input_size_per_partition
2080
+ # Scale sized by cdiv to match the checkpoint's block grid (the last
2081
+ # output block may be partial). Weight is registered at its logical
2082
+ # size; a non-128 output width is padded to a block boundary by the
2083
+ # _Fp8PbWoPartialBlock format scheme. Input must be 128-aligned (same as
2084
+ # #53132, which only pads the output).
2085
+ if in_ % 128 != 0:
2086
+ raise ValueError(f"FP8_PB_WO requires in divisible by 128, got {in_}")
2087
+ self.register_params(
2088
+ layer,
2089
+ "weight",
2090
+ (out, in_),
2091
+ torch.float8_e4m3fn,
2092
+ ModelWeightParameter,
2093
+ wl,
2094
+ input_dim=1,
2095
+ output_dim=0,
2096
+ )
2097
+ ob, ib = cdiv(out, 128), in_ // 128
2098
+ self.register_params(
2099
+ layer,
2100
+ "weight_scale",
2101
+ (ob, 1, ib, 1),
2102
+ torch.float32,
2103
+ BlockQuantScaleParameter,
2104
+ wl,
2105
+ input_dim=2,
2106
+ output_dim=0,
2107
+ init=FP8_SCALE_SENTINEL,
2108
+ )
2109
+ layer.weight_block_size = [128, 128]
2110
+
2111
+ def process(self, layer, role) -> None:
2112
+ if role is not WEIGHT:
2113
+ self.reject(role)
2114
+ layer.weight = Parameter(layer.weight.data, requires_grad=False)
2115
+ s = layer.weight_scale
2116
+ if s.dim() == 4:
2117
+ s = s.squeeze(1).squeeze(-1) # [ob,1,ib,1] -> [ob,ib]
2118
+ elif s.dim() != 2:
2119
+ raise ValueError(
2120
+ f"Unexpected FP8_PB_WO weight_scale shape {tuple(s.shape)}"
2121
+ )
2122
+ layer.weight_scale = Parameter(s.contiguous(), requires_grad=False)
2123
+
2124
+
2125
+class KMxfp8Static(QuantKeyScheme):
2126
+ """MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale.
2127
+ Weight-role only. process is validate-only plus an idempotency guard."""
2128
+
2129
+ key = kMxfp8Static
2130
+
2131
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2132
+ if role is not WEIGHT:
2133
+ self.reject(role)
2134
+ if shapes.input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
2135
+ raise ValueError(
2136
+ f"MXFP8 requires input size divisible by {MXFP8_BLOCK_SIZE}, "
2137
+ f"got {shapes.input_size_per_partition}"
2138
+ )
2139
+ self.register_params(
2140
+ layer,
2141
+ "weight",
2142
+ (shapes.output_size_per_partition, shapes.input_size_per_partition),
2143
+ MXFP8_VALUE_DTYPE,
2144
+ ModelWeightParameter,
2145
+ wl,
2146
+ input_dim=1,
2147
+ output_dim=0,
2148
+ )
2149
+ self.register_params(
2150
+ layer,
2151
+ "weight_scale",
2152
+ (
2153
+ shapes.output_size_per_partition,
2154
+ shapes.input_size_per_partition // MXFP8_BLOCK_SIZE,
2155
+ ),
2156
+ MXFP8_SCALE_DTYPE,
2157
+ ModelWeightParameter,
2158
+ wl,
2159
+ input_dim=1,
2160
+ output_dim=0,
2161
+ )
2162
+
2163
+ def process(self, layer, role) -> None:
2164
+ if role is not WEIGHT:
2165
+ self.reject(role)
2166
+ # Idempotency: the emulation kernel's post-load may replace the 1-byte
2167
+ # MXFP8 weight with BF16 (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD, default
2168
+ # on). On a weight reload process runs again after that swap, so skip
2169
+ # the fp8 asserts when the weight is already >=2-byte.
2170
+ if layer.weight.element_size() >= 2:
2171
+ return
2172
+ assert layer.weight.ndim == 2 and layer.weight.dtype == MXFP8_VALUE_DTYPE
2173
+ assert layer.weight_scale.ndim == 2
2174
+ assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE
2175
+
2176
+
2177
+class KDynamicNoParam(QuantKeyScheme):
2178
+ """Dynamic activation with no stored scale (W8A8): quantized at runtime in
2179
+ the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a
2180
+ non-None activation key. Activation-role only. Serves the fp8 per-token, fp8
2181
+ per-block, and mxfp8 dynamic activation keys."""
2182
+
2183
+ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2184
+ if role is not ACT:
2185
+ self.reject(role)
2186
+ # dynamic -> nothing stored
2187
+
2188
+ def process(self, layer, role) -> None:
2189
+ if role is not ACT:
2190
+ self.reject(role)
2191
+
2192
+
2193
+SCHEME_FOR: dict[QuantKey | None, QuantKeyScheme] = {
2194
+ kNvfp4Static: KNvfp4Static(),
2195
+ kNvfp4Dynamic: KNvfp4Dynamic(),
2196
+ kFp8StaticTensorSym: KFp8StaticTensor(),
2197
+ kFp8StaticTokenSym: KFp8StaticChannel(),
2198
+ kFp8Static128BlockSym: KFp8Block128(),
2199
+ kMxfp8Static: KMxfp8Static(),
2200
+ kFp8DynamicTokenSym: KDynamicNoParam(),
2201
+ kFp8Dynamic128Sym: KDynamicNoParam(),
2202
+ kMxfp8Dynamic: KDynamicNoParam(),
2203
+}
2204
+
2205
+
2206
+def maybe_fuse_global_scales(layer) -> None:
2207
+ """alpha = input_global_scale * weight_global_scale, presence-gated.
2208
+
2209
+ W4A4 has both -> computed; W4A16 has no input_global_scale -> skipped.
2210
+ """
2211
+ if hasattr(layer, "weight_global_scale") and hasattr(layer, "input_global_scale"):
2212
+ layer.alpha = Parameter(
2213
+ layer.input_global_scale * layer.weight_global_scale,
2214
+ requires_grad=False,
2215
+ )
2216
+
2217
+
2218
+def select_linear_kernel(spec: QuantSpec, layer, rt: RuntimeDtypes):
2219
+ """Thin family dispatcher on the weight key: nvfp4 / mxfp8 / fp8."""
2220
+ w = spec.weight
2221
+ assert isinstance(w, QuantKey), f"resolve() must supply a weight key, got {w!r}"
2222
+ if w.dtype == FP4_DTYPE:
2223
+ # W4A16 (activation is None) → use_a16=True defaults to Marlin *and*
2224
+ # honors --linear-backend (matches upstream ModelOptNvFp4W4A16LinearMethod
2225
+ # after #50273); W4A4 → use_a16=False.
2226
+ return init_nvfp4_linear_kernel(use_a16=spec.activation is None)
2227
+ if w.scale.dtype == MXFP8_SCALE_DTYPE:
2228
+ return init_mxfp8_linear_kernel()
2229
+ # fp8 family: init_fp8 routes block-vs-plain itself off the activation key,
2230
+ # and needs a real key -- weight-only fp8 is not a ModelOpt format.
2231
+ act = spec.activation
2232
+ assert isinstance(act, QuantKey), f"fp8 needs an activation key, got {act!r}"
2233
+ return init_fp8_linear_kernel(
2234
+ activation_quant_key=act,
2235
+ weight_quant_key=w,
2236
+ input_dtype=rt.input_dtype,
2237
+ out_dtype=rt.out_dtype,
2238
+ weight_shape=layer.weight.shape,
2239
+ module_name=type(layer).__name__,
2240
+ )
2241
+
2242
+
2243
+class FormatScheme:
2244
+ """Optional per-format hooks that compose *around* the QuantKey schemes.
2245
+
2246
+ Extension seam for residue that belongs to a format as a whole rather than a
2247
+ single QuantKey — an extra parameter the weight/activation schemes don't
2248
+ cover, or a tweak before/after their ``process``. All hooks default to
2249
+ no-ops; most formats need none (they are a pure ``(weight, activation)`` key
2250
+ pair). To add one: subclass, override the hooks you need, and return an
2251
+ instance from the format's ``resolve()`` branch — no change to
2252
+ ``ModelOptLinearMethod`` itself.
2253
+ """
2254
+
2255
+ def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None:
2256
+ """Register format-level params (after the key schemes' weights)."""
2257
+
2258
+ def pre_process(self, layer) -> None:
2259
+ """Run before the key schemes' ``process``."""
2260
+
2261
+ def post_process(self, layer) -> None:
2262
+ """Run after the key schemes' ``process``, before the kernel's."""
2263
+
2264
+ def apply(self, layer, x, bias, kernel_apply):
2265
+ """Wrap the kernel's forward. ``kernel_apply(layer, x, bias) -> Tensor``.
2266
+
2267
+ Default: delegate unchanged. A format whose residue lives at compute
2268
+ time (e.g. run the GEMM on a padded weight then slice the output back
2269
+ to its logical width, #53132-style) overrides this and calls
2270
+ ``kernel_apply`` in the middle.
2271
+ """
2272
+ return kernel_apply(layer, x, bias)
2273
+
2274
+
2275
+class _Fp8PbWoPartialBlock(FormatScheme):
2276
+ """FP8_PB_WO output width that is not a multiple of 128 (a partial trailing
2277
+ block; e.g. GLM's replicated ``fused_qkv_a_proj`` = 2048 + 576 = 2624).
2278
+
2279
+ The block kernel needs full 128-blocks, so pad the weight up to the block
2280
+ boundary with zeros before the kernel's post-load, then slice the GEMM
2281
+ output back to the logical width. This is #53132's approach, expressed as
2282
+ the format's compute-time residue on the generic method. The cdiv-sized
2283
+ scale from ``KFp8Block128`` already matches the padded block count, so only
2284
+ the weight rows and the output need adjusting.
2285
+ """
2286
+
2287
+ def post_process(self, layer) -> None:
2288
+ w = layer.weight
2289
+ out = w.shape[0]
2290
+ pad = (-out) % 128
2291
+ if not pad:
2292
+ return
2293
+ layer._pbwo_logical_out = out
2294
+ padded = w.data.new_zeros(out + pad, w.shape[1])
2295
+ padded[:out].copy_(w.data)
2296
+ layer.weight = Parameter(padded, requires_grad=False)
2297
+
2298
+ def apply(self, layer, x, bias, kernel_apply):
2299
+ logical = getattr(layer, "_pbwo_logical_out", None)
2300
+ if logical is None: # width was block-aligned: nothing to trim
2301
+ return kernel_apply(layer, x, bias)
2302
+ # Run the GEMM on the padded weight without bias, then trim to the
2303
+ # logical width and add bias (the padded rows produce zeros we drop).
2304
+ out = kernel_apply(layer, x, None)
2305
+ out = out[..., :logical].contiguous()
2306
+ return out if bias is None else out.add_(bias)
2307
+
2308
+
2309
+_PB_WO_PARTIAL_BLOCK = _Fp8PbWoPartialBlock()
2310
+
2311
+
2312
+class _DropInputScale(FormatScheme):
2313
+ """Interim: register then drop a W4A16 checkpoint's on-disk input_scale.
2314
+
2315
+ Kept for backward compat with previously-exported ModelOpt checkpoints.
2316
+ """
2317
+
2318
+ def extra_weights(self, layer, shapes, ctx, wl) -> None:
2319
+ data = torch.full((shapes.num_partitions,), torch.nan)
2320
+ layer.register_parameter(
2321
+ "input_scale", PerTensorScaleParameter(data=data, weight_loader=wl)
2322
+ )
2323
+
2324
+ def post_process(self, layer) -> None:
2325
+ scale = getattr(layer, "input_scale", None)
2326
+ if scale is None:
2327
+ return
2328
+ if not torch.isnan(scale).all():
2329
+ logger.warning_once(
2330
+ "Dropping input_scale from a weight-only W4A16 NVFP4 "
2331
+ "checkpoint (deprecated)."
2332
+ )
2333
+ del layer.input_scale
2334
+
2335
+
2336
+_DROP_INPUT_SCALE = _DropInputScale()
2337
+
2338
+
2339
+@register_weight_loader_v2_supported_method
2340
+class ModelOptLinearMethod(LinearMethodBase):
2341
+ """Generic, format-agnostic ModelOpt linear method. Holds a weight scheme +
2342
+ an activation scheme (from the QuantSpec pair) and an optional
2343
+ ``FormatScheme``, runs a fixed lifecycle, selects the kernel from the pair,
2344
+ and applies.
2345
+
2346
+ Registered for weight_loader_v2 (like every ModelOpt linear method) so the
2347
+ ``BasevLLMParameter`` params route through the v2 fused loader, not the
2348
+ legacy shape-assert path.
2349
+ """
2350
+
2351
+ def __init__(
2352
+ self,
2353
+ spec: QuantSpec,
2354
+ ctx: CkptCtx,
2355
+ format_scheme=None,
2356
+ ) -> None:
2357
+ self.spec = spec
2358
+ self.ctx = ctx
2359
+ self.fmt = format_scheme or FormatScheme()
2360
+ self.wkey = SCHEME_FOR[spec.weight]
2361
+ self.akey = None if spec.activation is None else SCHEME_FOR[spec.activation]
2362
+ # Only the fp8/mxfp8 kernels read input_dtype; nvfp4 ignores it. During
2363
+ # real serving model_config is always set; fall back defensively when
2364
+ # there is no config context (bare unit-test dispatch).
2365
+ model_config = getattr(get_current_vllm_config_or_none(), "model_config", None)
2366
+ dtype = (
2367
+ model_config.dtype
2368
+ if model_config is not None
2369
+ else torch.get_default_dtype()
2370
+ )
2371
+ self.out_dtype = dtype
2372
+ self.input_dtype = dtype
2373
+ self.marlin_input_dtype = None
2374
+ # Kernel/backend are chosen in create_weights (after get_quant_method),
2375
+ # so the front-end marlin poke stays dormant here — same as the old
2376
+ # NVFP4 methods.
2377
+ self.kernel: Any = None
2378
+
2379
+ def create_weights(
2380
+ self,
2381
+ layer,
2382
+ input_size_per_partition: int,
2383
+ output_partition_sizes: list[int],
2384
+ input_size: int,
2385
+ output_size: int,
2386
+ params_dtype: torch.dtype,
2387
+ **extra_weight_attrs,
2388
+ ) -> None:
2389
+ del input_size, output_size
2390
+ weight_loader = extra_weight_attrs.get("weight_loader")
2391
+ layer.logical_widths = output_partition_sizes
2392
+ layer.input_size_per_partition = input_size_per_partition
2393
+ layer.output_size_per_partition = sum(output_partition_sizes)
2394
+ # Humming reads both off the layer in
2395
+ # prepare_humming_linear_layer_config. LinearBase sets them itself;
2396
+ # ParallelLMHead does not, so supply them here.
2397
+ layer.output_partition_sizes = output_partition_sizes
2398
+ if not hasattr(layer, "has_bias"):
2399
+ layer.has_bias = getattr(layer, "bias", None) is not None
2400
+ shapes = Shapes(output_partition_sizes, input_size_per_partition, params_dtype)
2401
+
2402
+ self.wkey.create_weights(layer, WEIGHT, self.ctx, shapes, weight_loader)
2403
+ if self.akey:
2404
+ self.akey.create_weights(layer, ACT, self.ctx, shapes, weight_loader)
2405
+ self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader)
2406
+
2407
+ rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype)
2408
+ self.kernel = select_linear_kernel(self.spec, layer, rt)
2409
+ expose_input_quant_key(layer, self.kernel)
2410
+
2411
+ def process_weights_after_loading(self, layer) -> None:
2412
+ self.fmt.pre_process(layer)
2413
+ self.wkey.process(layer, WEIGHT)
2414
+ if self.akey:
2415
+ self.akey.process(layer, ACT)
2416
+ maybe_fuse_global_scales(layer)
2417
+ self.fmt.post_process(layer)
2418
+ self.kernel.process_weights_after_loading(layer)
2419
+
2420
+ def apply(self, layer, x, bias=None):
2421
+ return self.fmt.apply(
2422
+ layer,
2423
+ x,
2424
+ bias,
2425
+ lambda lyr, inp, b: self.kernel.apply_weights(layer=lyr, x=inp, bias=b),
2426
+ )
2427
+
2428
+
2429
+def resolve(algo: str, subcfg, prefix: str):
2430
+ """Turn a sub-config into (QuantSpec, CkptCtx, format_scheme).
2431
+
2432
+ Strictly read-only over ``subcfg``: in mixed precision the same sub-config
2433
+ is shared with the MoE method, so writing back would leak a linear-only
2434
+ change into MoE.
2435
+ """
2436
+ if algo == "FP8":
2437
+ # Plain per-tensor static FP8 (W8A8): same key in both slots (bivalent).
2438
+ ctx = CkptCtx()
2439
+ return (
2440
+ QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym),
2441
+ ctx,
2442
+ None,
2443
+ )
2444
+ if algo == "FP8_PER_CHANNEL_PER_TOKEN":
2445
+ # PcPt: per-channel static weight, dynamic per-token activation (W8A8).
2446
+ ctx = CkptCtx()
2447
+ return (
2448
+ QuantSpec(weight=kFp8StaticTokenSym, activation=kFp8DynamicTokenSym),
2449
+ ctx,
2450
+ None,
2451
+ )
2452
+ if algo == "FP8_PB_WO":
2453
+ # PbWo: 128x128 block-static weight, dynamic per-block activation (W8A8).
2454
+ # The block kernel's post-load runs here; CompressedTensors block-FP8
2455
+ # is the reference for this path. The FormatScheme pads a non-128 output
2456
+ # width to a block boundary and trims the output back (#53132) -- a no-op
2457
+ # for the common block-aligned case.
2458
+ ctx = CkptCtx()
2459
+ return (
2460
+ QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym),
2461
+ ctx,
2462
+ _PB_WO_PARTIAL_BLOCK,
2463
+ )
2464
+ if algo == "MXFP8":
2465
+ # MXFP8: block(32) e4m3 weight + e8m0 scale, dynamic activation.
2466
+ ctx = CkptCtx()
2467
+ return QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic), ctx, None
2468
+
2469
+ # NVFP4 family (W4A4 / W4A16).
2470
+ ctx = CkptCtx(group_size=subcfg.group_size)
2471
+ if algo == "NVFP4":
2472
+ # W4A4: static fp4 weight + dynamic fp4 activation (has a static global
2473
+ # input scale). alpha = weight_gs * input_gs.
2474
+ return QuantSpec(weight=kNvfp4Static, activation=kNvfp4Dynamic), ctx, None
2475
+ if algo == "W4A16_NVFP4":
2476
+ # W4A16: fp4 weight, no activation quant (-> Marlin). _DROP_INPUT_SCALE
2477
+ # is a no-op for genuine exports; keeps legacy input_scale ckpts loading.
2478
+ return QuantSpec(weight=kNvfp4Static, activation=None), ctx, _DROP_INPUT_SCALE
2479
+ raise NotImplementedError(f"resolve: unsupported ModelOpt linear algo {algo!r}")
2480
+
2481
+
2482
+# Bespoke-method escape hatch. Almost every ModelOpt linear format is a
2483
+# (weight, activation) QuantKey pair handled by the generic ModelOptLinearMethod
2484
+# (optionally + a FormatScheme), so this stays empty. A format that genuinely
2485
+# cannot be expressed that way — a fundamentally different create/process/apply
2486
+# lifecycle — registers its own LinearMethodBase here, keyed by algo:
2487
+#
2488
+# LINEAR_METHOD_BUILDERS["FOO"] = lambda cfg, prefix: ModelOptFooLinearMethod(cfg)
2489
+#
2490
+# Keep the ModelOpt* class-name prefix and decorate the class with
2491
+# @register_weight_loader_v2_supported_method if it uses BasevLLMParameter
2492
+# params. Also add the algo to the owning config's quant_method validation.
2493
+LINEAR_METHOD_BUILDERS: dict[str, Callable[..., LinearMethodBase]] = {}
2494
+
2495
+
2496
+def build_linear_method(config, algo: str, prefix: str) -> LinearMethodBase:
2497
+ """Construct the linear method for ``algo``.
2498
+
2499
+ Returns a bespoke method if one is registered in ``LINEAR_METHOD_BUILDERS``,
2500
+ else the generic ``ModelOptLinearMethod`` built from ``resolve(algo, config,
2501
+ prefix)``. Single indirection point shared by the homogeneous and
2502
+ mixed-precision dispatch, so a new format plugs in without editing either
2503
+ ``get_quant_method``.
2504
+ """
2505
+ builder = LINEAR_METHOD_BUILDERS.get(algo)
2506
+ if builder is not None:
2507
+ return builder(config, prefix)
2508
+ spec, ctx, format_scheme = resolve(algo, config, prefix)
2509
+ return ModelOptLinearMethod(spec, ctx, format_scheme)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+5
@@ -38,6 +38,11 @@ from vllm.utils.platform_utils import get_device_name_as_file_name
38
39
logger = init_logger(__name__)
40
41
+# Pre-fill value for scale parameters whose shards load independently. The
42
+# shards are combined with .max(), so an unloaded shard must never win; the
43
+# smallest representable float32 guarantees that.
44
+FP8_SCALE_SENTINEL = torch.finfo(torch.float32).min
45
+
46
47
def is_fp8(x: torch.dtype | torch.Tensor) -> bool:
48
if isinstance(x, torch.Tensor):