main
py 940 lines 34.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Test ModelOpt quantization method setup and weight loading.
4
5 Run `pytest tests/quantization/test_modelopt.py`.
6 """
7
8 import os
9 from typing import Any, NoReturn
10 from unittest.mock import MagicMock, Mock, patch
11
12 import pytest
13 import torch
14
15 from tests.quantization.utils import (
16 is_quant_method_supported,
17 load_model_without_vllm_runner,
18 )
19 from vllm.config import VllmConfig, set_current_vllm_config
20 from vllm.config.model import ModelConfig
21 from vllm.forward_context import set_forward_context
22 from vllm.model_executor.kernels.linear import (
23 FlashInferCuteDslNvFp4W4A16LinearKernel,
24 HummingNvFp4LinearKernel,
25 MarlinNvFp4LinearKernel,
26 )
27 from vllm.model_executor.layers.attention import Attention
28 from vllm.model_executor.layers.fused_moe import RoutedExperts
29 from vllm.model_executor.layers.linear import UnquantizedLinearMethod
30 from vllm.model_executor.layers.quantization.modelopt import (
31 LINEAR_ALGOS,
32 ModelOptFp8Config,
33 ModelOptLinearMethod,
34 ModelOptMixedPrecisionConfig,
35 ModelOptMxFp8Config,
36 ModelOptNvFp4Config,
37 )
38 from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 kFp8StaticTensorSym,
40 kMxfp8Dynamic,
41 kMxfp8Static,
42 kNvfp4Dynamic,
43 kNvfp4Static,
44 )
45 from vllm.model_executor.layers.vocab_parallel_embedding import (
46 ParallelLMHead,
47 VocabParallelEmbedding,
48 )
49 from vllm.platforms import current_platform
50
51
52 @pytest.fixture(scope="function", autouse=True)
53 def enable_pickle(monkeypatch):
54 """`LLM.apply_model` requires pickling a function."""
55 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
56
57
58 def _skip(msg: str) -> NoReturn:
59 pytest.skip(msg)
60 raise RuntimeError(msg)
61
62
63 def _snapshot_download_or_skip(model_id: str) -> str:
64 try:
65 from vllm.transformers_utils.repo_utils import hf_api
66 except Exception as e: # pragma: no cover
67 _skip(f"huggingface_hub is required to download {model_id}: {e}")
68
69 try:
70 return hf_api().snapshot_download(
71 repo_id=model_id,
72 repo_type="model",
73 # These checkpoints are already small; download full repo for simplicity.
74 allow_patterns=["*"],
75 )
76 except Exception as e:
77 _skip(f"Failed to download {model_id} from the HF Hub: {e}")
78
79
80 def _mock_lm_head() -> Mock:
81 lm_head = Mock(spec=ParallelLMHead)
82 lm_head.__class__ = ParallelLMHead
83 return lm_head
84
85
86 def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionConfig:
87 return ModelOptMixedPrecisionConfig(
88 kv_cache_quant_method=None,
89 exclude_modules=[],
90 quantized_layers=quantized_layers,
91 fp8_config=ModelOptFp8Config(
92 quant_method="FP8",
93 is_checkpoint_fp8_serialized=True,
94 kv_cache_quant_method=None,
95 exclude_modules=[],
96 ),
97 nvfp4_config=ModelOptNvFp4Config(
98 is_checkpoint_nvfp4_serialized=True,
99 kv_cache_quant_algo=None,
100 exclude_modules=[],
101 ),
102 w4a16_nvfp4_config=ModelOptNvFp4Config(
103 quant_method="W4A16_NVFP4",
104 is_checkpoint_nvfp4_serialized=True,
105 kv_cache_quant_algo=None,
106 exclude_modules=[],
107 ),
108 mxfp8_config=ModelOptMxFp8Config(
109 is_checkpoint_mxfp8_serialized=True,
110 kv_cache_quant_algo=None,
111 exclude_modules=[],
112 ),
113 )
114
115
116 def test_modelopt_nvfp4_quantizes_parallel_lm_head():
117 config = ModelOptNvFp4Config(
118 is_checkpoint_nvfp4_serialized=True,
119 kv_cache_quant_algo=None,
120 exclude_modules=[],
121 )
122
123 method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
124
125 assert isinstance(method, ModelOptLinearMethod)
126 assert method.spec.weight is kNvfp4Static
127 assert method.spec.activation is kNvfp4Dynamic
128
129
130 def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch):
131 """Standard MXFP8 checkpoints already have one scale row per weight row."""
132 from vllm.model_executor.layers.linear import ReplicatedLinear
133
134 kernel = Mock()
135 kernel.input_quant_key.return_value = None
136 monkeypatch.setattr(
137 "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel",
138 lambda **kwargs: kernel,
139 )
140 config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
141 linear = ReplicatedLinear(64, 64, bias=False, quant_config=config)
142 scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2)
143 linear.weight_scale.weight_loader(linear.weight_scale, scales)
144 assert torch.equal(linear.weight_scale, scales)
145
146
147 def test_modelopt_mxfp8_pre_processed_weights_follow_kernel(monkeypatch):
148 """The weight cache daemon can only serve MXFP8 layers whose kernel needs
149 no post-load state beyond the parameters it exports."""
150 from vllm.config.quantization import QuantSpec
151 from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearKernel
152
153 method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
154 method.spec = QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic)
155 method.kernel = None
156 assert not method.supports_pre_processed_weights
157
158 method.kernel = Mock(spec=Mxfp8LinearKernel)
159 method.kernel.supports_pre_processed_weights = False
160 assert not method.supports_pre_processed_weights
161 method.kernel.supports_pre_processed_weights = True
162 assert method.supports_pre_processed_weights
163
164 monkeypatch.setattr(
165 "vllm.model_executor.layers.quantization.modelopt.is_weights_pre_processed",
166 lambda: True,
167 )
168 layer = torch.nn.Module()
169 method.process_weights_after_loading(layer)
170 method.kernel.process_weights_after_loading.assert_not_called()
171
172 method.kernel.supports_pre_processed_weights = False
173 with pytest.raises(RuntimeError, match="pre-processed"):
174 method.process_weights_after_loading(layer)
175
176
177 def test_modelopt_fp8_updates_weight_dims_after_transpose():
178 """Humming reads weight.input_dim/output_dim. Swapping the
179 ModelWeightParameter for a plain Parameter drops them, so the per-tensor
180 FP8 scheme must restore them for the transposed [in, out] layout.
181 """
182 from vllm.config.quantization import QuantSpec
183 from vllm.model_executor.layers.quantization.modelopt import (
184 SCHEME_FOR,
185 CkptCtx,
186 FormatScheme,
187 )
188
189 layer = torch.nn.Module()
190 layer.register_parameter(
191 "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False)
192 )
193 layer.register_parameter(
194 "weight_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
195 )
196 layer.register_parameter(
197 "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
198 )
199 layer.logical_widths = [3]
200
201 method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
202 method.spec = QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym)
203 method.ctx = CkptCtx()
204 method.fmt = FormatScheme()
205 method.wkey = SCHEME_FOR[kFp8StaticTensorSym]
206 method.akey = SCHEME_FOR[kFp8StaticTensorSym]
207 method.kernel = Mock()
208 method.process_weights_after_loading(layer)
209
210 assert layer.weight.shape == (2, 3)
211 assert layer.weight.input_dim == 0
212 assert layer.weight.output_dim == 1
213 method.kernel.process_weights_after_loading.assert_called_once_with(layer)
214
215
216 def test_modelopt_linear_algos_table_matches_resolve():
217 """LINEAR_ALGOS is the single source of truth for supported linear algos.
218
219 Every entry must be dispatchable by resolve(), and every config's
220 validation list must be derived from it -- so adding a format is one row
221 here plus one row in resolve(), with nothing else to keep in sync.
222 """
223 from vllm.model_executor.layers.quantization.modelopt import (
224 QUANT_ALGOS,
225 algos_owned_by,
226 resolve,
227 )
228
229 class _Cfg:
230 group_size = 16
231
232 for algo in LINEAR_ALGOS:
233 spec, _, _ = resolve(algo, _Cfg(), "layer")
234 assert spec.weight is not None, algo
235
236 assert list(QUANT_ALGOS) == [*LINEAR_ALGOS, "MIXED_PRECISION"]
237 assert algos_owned_by("modelopt") == (
238 "FP8",
239 "FP8_PER_CHANNEL_PER_TOKEN",
240 "FP8_PB_WO",
241 )
242 assert algos_owned_by("modelopt_fp4") == ("NVFP4", "W4A16_NVFP4")
243 assert algos_owned_by("modelopt_mxfp8") == ("MXFP8",)
244
245
246 @pytest.mark.parametrize("algo", list(LINEAR_ALGOS))
247 def test_modelopt_mixed_precision_dispatches_every_linear_algo(algo):
248 """Mixed precision must route every algo in LINEAR_ALGOS through the
249 generic method. FP8_PER_CHANNEL_PER_TOKEN and FP8_PB_WO used to fall
250 through to UnquantizedLinearMethod, which loses the checkpoint's scales.
251 """
252 from vllm.model_executor.layers.linear import LinearBase
253 from vllm.model_executor.layers.quantization import modelopt as m
254
255 config = m.ModelOptMixedPrecisionConfig.from_config(
256 {
257 "quantization": {
258 "quant_algo": "MIXED_PRECISION",
259 "kv_cache_quant_algo": None,
260 "exclude_modules": [],
261 "group_size": 16,
262 "quantized_layers": {
263 "model.layers.0.mlp.down_proj": {"quant_algo": algo}
264 },
265 }
266 }
267 )
268 method = config.get_quant_method(
269 MagicMock(spec=LinearBase), "model.layers.0.mlp.down_proj"
270 )
271
272 assert isinstance(method, ModelOptLinearMethod), (algo, type(method).__name__)
273
274
275 @pytest.mark.parametrize("algo", ["FP8_PB_WO", "FP8_BLOCK_SCALES"])
276 def test_modelopt_mixed_precision_dispatches_block_fp8_moe(algo):
277 config = ModelOptMixedPrecisionConfig.from_config(
278 {
279 "quantization": {
280 "quant_algo": "MIXED_PRECISION",
281 "quantized_layers": {
282 "mtp.layers.48.mlp.experts": {
283 "quant_algo": algo,
284 "group_size": 128,
285 }
286 },
287 }
288 }
289 )
290 layer = MagicMock(spec=RoutedExperts)
291 expected = object()
292
293 with patch(
294 "vllm.model_executor.layers.quantization.modelopt.Fp8MoEMethod",
295 return_value=expected,
296 ) as method_cls:
297 method = config.get_quant_method(layer, "mtp.layers.48.mlp.experts")
298
299 assert method is expected
300 fp8_config, called_layer = method_cls.call_args.args
301 assert called_layer is layer
302 assert fp8_config.weight_block_size == [128, 128]
303 assert fp8_config.activation_scheme == "dynamic"
304
305
306 def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
307 config = ModelOptNvFp4Config(
308 is_checkpoint_nvfp4_serialized=True,
309 kv_cache_quant_algo=None,
310 exclude_modules=["lm_head"],
311 )
312
313 method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
314
315 assert isinstance(method, UnquantizedLinearMethod)
316
317
318 def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
319 config = _mixed_precision_config(
320 {"lm_head": {"quant_algo": "NVFP4", "group_size": 16}}
321 )
322
323 method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
324
325 assert isinstance(method, ModelOptLinearMethod)
326 assert method.spec.weight is kNvfp4Static
327 assert method.spec.activation is kNvfp4Dynamic
328
329
330 def test_modelopt_mixed_precision_resolves_declared_packed_projection():
331 config = _mixed_precision_config(
332 {
333 "model.layers.0.self_attn.q_proj": {"quant_algo": "MXFP8"},
334 "model.layers.0.self_attn.k_proj": {"quant_algo": "MXFP8"},
335 "model.layers.0.self_attn.v_proj": {"quant_algo": "MXFP8"},
336 }
337 )
338 config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
339
340 assert config._resolve_quant_algo("model.layers.0.self_attn.qkv_proj") == "MXFP8"
341
342
343 def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
344 config = _mixed_precision_config(
345 {
346 "model.layers.0.linear_attn.in_proj_qkv": {"quant_algo": "FP8"},
347 "model.layers.0.linear_attn.in_proj_z": {"quant_algo": "FP8"},
348 "model.layers.0.linear_attn.out_proj": {"quant_algo": "FP8"},
349 }
350 )
351 config.packed_modules_mapping = {
352 "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
353 "in_proj_ba": ["in_proj_b", "in_proj_a"],
354 }
355
356 assert (
357 config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_qkvz") == "FP8"
358 )
359 assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
360
361
362 def test_modelopt_mixed_precision_composes_gemma4_mappers():
363 from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
364 from vllm.model_executor.models.gemma4_mm import (
365 Gemma4ForConditionalGeneration,
366 )
367
368 config = _mixed_precision_config(
369 {
370 "model.language_model.layers.0.experts": {
371 "quant_algo": "NVFP4",
372 "group_size": 16,
373 },
374 "model.language_model.layers.1.experts.gate_up_proj": {
375 "quant_algo": "NVFP4",
376 "group_size": 16,
377 },
378 }
379 )
380
381 config.apply_vllm_mapper(
382 Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_rename_mapper()
383 )
384 config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper())
385
386 expected_prefix = "language_model.model.layers.0.experts"
387 assert set(config.quantized_layers) == {
388 expected_prefix,
389 "language_model.model.layers.1.experts.gate_up_proj",
390 }
391 assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
392
393
394 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
395 from vllm.model_executor.layers.linear import LinearBase
396
397 config = _mixed_precision_config(
398 {
399 "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"},
400 "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"},
401 }
402 )
403
404 fake_layer = MagicMock(spec=LinearBase)
405 method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
406
407 assert isinstance(method, ModelOptLinearMethod)
408 assert method.spec.weight is kNvfp4Static
409 assert method.spec.activation is kNvfp4Dynamic
410
411
412 @pytest.mark.parametrize(
413 ("quantized_prefix", "missing_prefix"),
414 [
415 ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"),
416 ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"),
417 ],
418 )
419 def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear(
420 quantized_prefix, missing_prefix
421 ):
422 from vllm.model_executor.layers.linear import LinearBase
423
424 config = _mixed_precision_config(
425 {
426 quantized_prefix: {"quant_algo": "NVFP4"},
427 }
428 )
429
430 fake_layer = MagicMock(spec=LinearBase)
431 method = config.get_quant_method(fake_layer, missing_prefix)
432
433 assert isinstance(method, UnquantizedLinearMethod)
434
435
436 def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
437 holder = Mock()
438 scale = torch.nn.Parameter(torch.empty(1))
439 loaded_scale = torch.tensor(2.0)
440
441 VocabParallelEmbedding.weight_loader(holder, scale, loaded_scale)
442
443 assert torch.equal(scale, loaded_scale.reshape(1))
444
445
446 @pytest.mark.skipif(
447 not is_quant_method_supported("modelopt"),
448 reason="ModelOpt FP8 is not supported on this GPU type.",
449 )
450 def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
451 """Test ModelOpt FP8 checkpoint loading and structure validation."""
452 # TODO: provide a small publicly available test checkpoint
453 model_path = (
454 "/home/scratch.omniml_data_1/zhiyu/ckpts/test_ckpts/"
455 "TinyLlama-1.1B-Chat-v1.0-fp8-0710"
456 )
457
458 # Skip test if checkpoint doesn't exist
459 if not os.path.exists(model_path):
460 pytest.skip(
461 f"Test checkpoint not found at {model_path}. "
462 "This test requires a local ModelOpt FP8 checkpoint."
463 )
464
465 # Set model config as model_config.dtype is required in ModelOptLinearMethod.
466 default_vllm_config.model_config = ModelConfig()
467 with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
468
469 def check_model(model):
470 layer = model.model.layers[0]
471
472 qkv_proj = layer.self_attn.qkv_proj
473 o_proj = layer.self_attn.o_proj
474 gate_up_proj = layer.mlp.gate_up_proj
475 down_proj = layer.mlp.down_proj
476
477 # Check that ModelOpt quantization method is properly applied
478 assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
479 assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
480 assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
481 assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
482
483 # Check weight dtype is FP8
484 assert qkv_proj.weight.dtype == torch.float8_e4m3fn
485 assert o_proj.weight.dtype == torch.float8_e4m3fn
486 assert gate_up_proj.weight.dtype == torch.float8_e4m3fn
487 assert down_proj.weight.dtype == torch.float8_e4m3fn
488
489 # Check scales are present and have correct dtype
490 assert hasattr(qkv_proj, "weight_scale")
491 assert hasattr(qkv_proj, "input_scale")
492 assert qkv_proj.weight_scale.dtype == torch.float32
493 assert qkv_proj.input_scale.dtype == torch.float32
494
495 assert hasattr(o_proj, "weight_scale")
496 assert hasattr(o_proj, "input_scale")
497 assert o_proj.weight_scale.dtype == torch.float32
498 assert o_proj.input_scale.dtype == torch.float32
499
500 assert hasattr(gate_up_proj, "weight_scale")
501 assert hasattr(gate_up_proj, "input_scale")
502 assert gate_up_proj.weight_scale.dtype == torch.float32
503 assert gate_up_proj.input_scale.dtype == torch.float32
504
505 assert hasattr(down_proj, "weight_scale")
506 assert hasattr(down_proj, "input_scale")
507 assert down_proj.weight_scale.dtype == torch.float32
508 assert down_proj.input_scale.dtype == torch.float32
509
510 llm.apply_model(check_model)
511
512 # Run a simple generation test to ensure the model works
513 output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
514 assert output
515 print(f"ModelOpt FP8 output: {output}")
516
517
518 @pytest.mark.skipif(
519 not is_quant_method_supported("modelopt"),
520 reason="ModelOpt FP8 is not supported on this GPU type.",
521 )
522 def test_modelopt_fp8_pc_pt_checkpoint_setup(monkeypatch, dist_init, workspace_init):
523 """Test ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoint setup."""
524 model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pc-pt"
525 model_path = _snapshot_download_or_skip(model_id)
526
527 model, vllm_config = load_model_without_vllm_runner(
528 model_path,
529 quantization="modelopt",
530 )
531 layer = model.model.layers[0]
532
533 qkv_proj = layer.self_attn.qkv_proj
534 o_proj = layer.self_attn.o_proj
535 gate_up_proj = layer.mlp.gate_up_proj
536 down_proj = layer.mlp.down_proj
537
538 assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
539 assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
540 assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
541 assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
542
543 fp8_dtype = current_platform.fp8_dtype()
544 assert qkv_proj.weight.dtype == fp8_dtype
545 assert o_proj.weight.dtype == fp8_dtype
546 assert gate_up_proj.weight.dtype == fp8_dtype
547 assert down_proj.weight.dtype == fp8_dtype
548
549 # Per-channel scales; activations are dynamically scaled per token.
550 for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
551 assert hasattr(projection, "weight_scale")
552 assert projection.weight_scale.dtype == torch.float32
553 assert projection.weight_scale.dim() == 1
554 assert not hasattr(projection, "input_scale")
555
556 monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
557 input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
558 positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
559 with (
560 set_current_vllm_config(vllm_config),
561 set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
562 ):
563 hidden_states = model(input_ids, positions, None)
564 logits = model.compute_logits(hidden_states)
565 assert torch.isfinite(logits).all()
566
567
568 @pytest.mark.skipif(
569 not is_quant_method_supported("modelopt"),
570 reason="ModelOpt FP8 is not supported on this GPU type.",
571 )
572 def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_init):
573 """Test ModelOpt FP8_PB_WO checkpoint setup."""
574 model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pb-wo"
575 model_path = _snapshot_download_or_skip(model_id)
576
577 model, vllm_config = load_model_without_vllm_runner(
578 model_path,
579 quantization="modelopt",
580 )
581 layer = model.model.layers[0]
582
583 qkv_proj = layer.self_attn.qkv_proj
584 o_proj = layer.self_attn.o_proj
585 gate_up_proj = layer.mlp.gate_up_proj
586 down_proj = layer.mlp.down_proj
587
588 assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod)
589 assert isinstance(o_proj.quant_method, ModelOptLinearMethod)
590 assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod)
591 assert isinstance(down_proj.quant_method, ModelOptLinearMethod)
592
593 fp8_dtype = current_platform.fp8_dtype()
594 assert qkv_proj.weight.dtype == fp8_dtype
595 assert o_proj.weight.dtype == fp8_dtype
596 assert gate_up_proj.weight.dtype == fp8_dtype
597 assert down_proj.weight.dtype == fp8_dtype
598
599 # Block scales are materialized as a 2D [out_blk, in_blk] tensor.
600 for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
601 assert hasattr(projection, "weight_scale")
602 assert projection.weight_scale.dtype == torch.float32
603 assert projection.weight_scale.dim() == 2
604
605 monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
606 input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
607 positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
608 with (
609 set_current_vllm_config(vllm_config),
610 set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
611 ):
612 hidden_states = model(input_ids, positions, None)
613 logits = model.compute_logits(hidden_states)
614 assert torch.isfinite(logits).all()
615
616
617 def test_modelopt_nvfp4_config_dispatches_w4a4_method():
618 """``quant_method="NVFP4"`` (W4A4) resolves to a
619 ``(kNvfp4Static, kNvfp4Dynamic)`` QuantSpec under the generic
620 ``ModelOptLinearMethod``."""
621 from vllm.model_executor.layers.linear import LinearBase
622
623 config = ModelOptNvFp4Config(
624 quant_method="NVFP4",
625 is_checkpoint_nvfp4_serialized=True,
626 kv_cache_quant_algo=None,
627 exclude_modules=[],
628 )
629 assert config.quant_method == "NVFP4"
630
631 method = config.get_quant_method(
632 MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
633 )
634 assert isinstance(method, ModelOptLinearMethod)
635 assert method.spec.weight is kNvfp4Static
636 assert method.spec.activation is kNvfp4Dynamic
637
638
639 def test_modelopt_nvfp4_config_dispatches_w4a16_method():
640 """``quant_method="W4A16_NVFP4"`` resolves to a weight-only QuantSpec
641 (``activation=None``) — distinct from the W4A4 sibling.
642
643 A regression here would mean a W4A16 NVFP4 checkpoint silently loaded
644 with a dynamic fp4 activation key, registering an ``input_scale`` and
645 routing to the cutlass W4A4 NVFP4 GEMM instead of FP4 Marlin.
646 """
647 from vllm.model_executor.layers.linear import LinearBase
648
649 config = ModelOptNvFp4Config(
650 quant_method="W4A16_NVFP4",
651 is_checkpoint_nvfp4_serialized=True,
652 kv_cache_quant_algo=None,
653 exclude_modules=[],
654 )
655 assert config.quant_method == "W4A16_NVFP4"
656
657 method = config.get_quant_method(
658 MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
659 )
660 assert isinstance(method, ModelOptLinearMethod)
661 assert method.spec.weight is kNvfp4Static
662 assert method.spec.activation is None
663
664
665 def test_modelopt_linear_method_builder_registry_override(monkeypatch):
666 """The bespoke-method escape hatch: a format registered in
667 ``LINEAR_METHOD_BUILDERS`` routes that algo to its own method instead of the
668 generic ``ModelOptLinearMethod``. This is how a format that cannot be a
669 ``(weight, activation)`` key pair plugs into dispatch."""
670 from vllm.model_executor.layers.linear import LinearBase
671 from vllm.model_executor.layers.quantization import modelopt as m
672
673 sentinel = object()
674 monkeypatch.setitem(m.LINEAR_METHOD_BUILDERS, "NVFP4", lambda cfg, prefix: sentinel)
675
676 config = ModelOptNvFp4Config(
677 quant_method="NVFP4",
678 is_checkpoint_nvfp4_serialized=True,
679 kv_cache_quant_algo=None,
680 exclude_modules=[],
681 )
682 method = config.get_quant_method(
683 MagicMock(spec=LinearBase), "model.layers.0.fake_proj"
684 )
685 assert method is sentinel # bespoke builder wins over the generic path
686
687
688 @pytest.mark.parametrize(
689 ("linear_backend", "kernel_cls"),
690 [
691 ("auto", MarlinNvFp4LinearKernel),
692 ("humming", HummingNvFp4LinearKernel),
693 ("flashinfer_cutedsl", FlashInferCuteDslNvFp4W4A16LinearKernel),
694 ],
695 )
696 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
697 def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):
698 """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``:
699 ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The
700 generic method routes this through ``select_linear_kernel``."""
701 from vllm.config.quantization import QuantSpec
702 from vllm.model_executor.layers.quantization.modelopt import (
703 RuntimeDtypes,
704 select_linear_kernel,
705 )
706
707 if linear_backend != "auto":
708 is_supported, reason = kernel_cls.is_supported()
709 if not is_supported:
710 pytest.skip(reason)
711
712 vllm_config = VllmConfig()
713 vllm_config.kernel_config.linear_backend = linear_backend
714 spec = QuantSpec(weight=kNvfp4Static, activation=None)
715 rt = RuntimeDtypes(torch.bfloat16, torch.bfloat16)
716 with set_current_vllm_config(vllm_config):
717 kernel = select_linear_kernel(spec, MagicMock(), rt)
718 assert isinstance(kernel, kernel_cls)
719
720
721 @pytest.mark.parametrize(
722 "quant_method, expected_use_a16, act_key_is_none",
723 [
724 ("NVFP4", False, False), # W4A4 default
725 ("W4A16_NVFP4", True, True), # native W4A16 ckpt
726 ],
727 )
728 def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16(
729 quant_method, expected_use_a16, act_key_is_none
730 ):
731 """``ModelOptNvFp4FusedMoE``: when the ckpt's ``quant_method`` is
732 ``W4A16_NVFP4``, the MoE class must pass ``activation_key=None`` to
733 ``select_nvfp4_moe_backend``. That filters out every W4A4 backend
734 (their ``_supports_quant_scheme`` requires
735 ``(kNvfp4Static, kNvfp4Dynamic)`` exactly); Marlin survives because
736 it only checks ``weight_key``. A regression here would mean a W4A16
737 ckpt silently went to the cutlass W4A4 path.
738 """
739 from vllm.model_executor.layers.quantization.modelopt import (
740 ModelOptNvFp4Config,
741 ModelOptNvFp4FusedMoE,
742 )
743 from vllm.model_executor.layers.quantization.utils.quant_utils import (
744 kNvfp4Dynamic,
745 kNvfp4Static,
746 )
747
748 config = ModelOptNvFp4Config(
749 quant_method=quant_method,
750 is_checkpoint_nvfp4_serialized=True,
751 kv_cache_quant_algo=None,
752 exclude_modules=[],
753 group_size=16,
754 )
755
756 mock_select = MagicMock(return_value=(MagicMock(), MagicMock()))
757 with (
758 patch(
759 "vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend",
760 mock_select,
761 ),
762 patch(
763 "vllm.model_executor.layers.quantization.modelopt."
764 "is_global_sf_supported_for_nvfp4_backend",
765 return_value=False,
766 ),
767 ):
768 moe = ModelOptNvFp4FusedMoE(config, MagicMock())
769
770 assert moe.use_a16 is expected_use_a16
771 _, kwargs = mock_select.call_args
772 assert kwargs["weight_key"] is kNvfp4Static
773 if act_key_is_none:
774 assert kwargs["activation_key"] is None
775 else:
776 assert kwargs["activation_key"] is kNvfp4Dynamic
777
778
779 @pytest.mark.parametrize(
780 "per_layer_algo, expected_weight, expected_activation",
781 [
782 ("NVFP4", kNvfp4Static, kNvfp4Dynamic),
783 ("W4A16_NVFP4", kNvfp4Static, None),
784 ("FP8", kFp8StaticTensorSym, kFp8StaticTensorSym),
785 ("MXFP8", kMxfp8Static, kMxfp8Dynamic),
786 ],
787 )
788 def test_modelopt_mixed_precision_dispatches_linear_layer(
789 per_layer_algo, expected_weight, expected_activation
790 ):
791 """``ModelOptMixedPrecisionConfig.get_quant_method`` routes a Linear layer
792 to the generic ``ModelOptLinearMethod`` with the ``QuantSpec`` resolved
793 from its per-layer ``quant_algo`` entry in ``quantized_layers``. A
794 regression here would mean a layer got the wrong ``(weight, activation)``
795 key pair or fell through to ``UnquantizedLinearMethod`` — e.g. a W4A16
796 layer picking up a dynamic fp4 activation key (cutlass W4A4 path) instead
797 of the weight-only Marlin path.
798 """
799 from vllm.model_executor.layers.linear import LinearBase
800 from vllm.model_executor.layers.quantization import modelopt as m
801
802 hf_quant_config: dict[str, Any] = {
803 "quantization": {
804 "quant_algo": "MIXED_PRECISION",
805 "kv_cache_quant_algo": None,
806 "exclude_modules": [],
807 "group_size": 16,
808 "quantized_layers": {
809 "model.layers.0.fake_proj": {"quant_algo": per_layer_algo},
810 },
811 }
812 }
813 config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)
814
815 fake_layer = MagicMock(spec=LinearBase)
816 method = config.get_quant_method(fake_layer, "model.layers.0.fake_proj")
817
818 assert isinstance(method, m.ModelOptLinearMethod)
819 assert method.spec.weight is expected_weight
820 assert method.spec.activation is expected_activation
821
822
823 def test_modelopt_mixed_precision_builds_w4a16_sibling_config():
824 """Sanity: ``ModelOptMixedPrecisionConfig._from_config`` builds **two**
825 NVFP4 sub-configs — one for W4A4 (default) and one tagged
826 ``quant_method='W4A16_NVFP4'`` — so per-layer dispatch can hand
827 Marlin-bound layers the right config without re-instantiating it on
828 every call.
829 """
830 from vllm.model_executor.layers.quantization import modelopt as m
831
832 hf_quant_config: dict[str, Any] = {
833 "quantization": {
834 "quant_algo": "MIXED_PRECISION",
835 "kv_cache_quant_algo": None,
836 "exclude_modules": [],
837 "group_size": 16,
838 "quantized_layers": {
839 "model.layers.0.a": {"quant_algo": "NVFP4"},
840 "model.layers.0.b": {"quant_algo": "W4A16_NVFP4"},
841 },
842 }
843 }
844 config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)
845
846 assert isinstance(config, m.ModelOptMixedPrecisionConfig)
847 assert config.nvfp4_config.quant_method == "NVFP4"
848 assert config.w4a16_nvfp4_config.quant_method == "W4A16_NVFP4"
849
850
851 def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch):
852 """FP8_PB_WO output width that is not a multiple of 128 (a partial trailing
853 block) is padded up to a block boundary before the kernel post-load, the
854 GEMM runs on the padded weight, and the output is trimmed back to the
855 logical width with bias added after. Faithful port of wei-zhao #53132's
856 test_modelopt_fp8_pb_wo_hides_output_padding for the generic method +
857 _Fp8PbWoPartialBlock FormatScheme.
858
859 Width is the real motivating case -- GLM's fused qkv_a_proj, q_a 2048 +
860 kv_a 576 = 2624 (replicated, so no TP degree makes it a 128-multiple),
861 padded to 2688 = 21 * 128.
862 """
863 from vllm.config.quantization import QuantSpec
864 from vllm.model_executor.layers.quantization import modelopt as mo
865 from vllm.model_executor.layers.quantization.utils.quant_utils import (
866 kFp8Dynamic128Sym,
867 kFp8Static128BlockSym,
868 )
869
870 kernel = Mock()
871 init_fp8_linear_kernel = Mock(return_value=kernel)
872 monkeypatch.setattr(mo, "init_fp8_linear_kernel", init_fp8_linear_kernel)
873 monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, k: None)
874
875 method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
876 method.spec = QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym)
877 method.ctx = mo.CkptCtx()
878 method.fmt = mo._PB_WO_PARTIAL_BLOCK
879 method.wkey = mo.SCHEME_FOR[kFp8Static128BlockSym]
880 method.akey = mo.SCHEME_FOR[kFp8Dynamic128Sym]
881 method.input_dtype = method.out_dtype = torch.bfloat16
882 method.marlin_input_dtype = None
883
884 layer = torch.nn.Module()
885 with (
886 patch(
887 "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
888 return_value=0,
889 ),
890 patch(
891 "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
892 return_value=1,
893 ),
894 ):
895 # output 2624 = 2048 + 576, input 128
896 method.create_weights(
897 layer, 128, [2048, 576], 128, 2624, torch.bfloat16, weight_loader=Mock()
898 )
899
900 # loaded at logical size; scale is cdiv(2624, 128) = 21 block rows
901 assert layer.weight.shape == (2624, 128)
902 assert layer.weight_scale.shape == (21, 1, 1, 1)
903 assert init_fp8_linear_kernel.call_args.kwargs["weight_shape"] == (2688, 128)
904
905 layer.weight.data.fill_(1)
906 method.process_weights_after_loading(layer)
907
908 # weight padded to the block boundary; the pad rows are zero
909 assert layer.weight.shape == (2688, 128)
910 assert torch.count_nonzero(layer.weight[2624:].float()) == 0
911 kernel.process_weights_after_loading.assert_called_once_with(layer)
912
913 # apply: GEMM on padded weight (bias=None), output trimmed + bias added
914 physical_output = torch.randn(4, 2688, dtype=torch.bfloat16)
915 kernel.apply_weights.return_value = physical_output
916 bias = torch.randn(2624, dtype=torch.bfloat16)
917 output = method.apply(layer, torch.randn(4, 128), bias)
918
919 torch.testing.assert_close(output, physical_output[:, :2624] + bias)
920 assert output.shape == (4, 2624)
921 assert output.is_contiguous()
922 kernel.apply_weights.assert_called_once()
923 assert kernel.apply_weights.call_args.kwargs["bias"] is None
924
925
926 def test_modelopt_fp8_pb_wo_rejects_non_128_input():
927 """Input width must still be a multiple of 128 (same as #53132, which only
928 pads the output). A partial input block is refused loudly rather than
929 silently loading wrong scales."""
930 from vllm.model_executor.layers.quantization import modelopt as mo
931 from vllm.model_executor.layers.quantization.utils.quant_utils import (
932 kFp8Static128BlockSym,
933 )
934
935 scheme = mo.SCHEME_FOR[kFp8Static128BlockSym]
936 shapes = mo.Shapes([128], 100, torch.bfloat16) # input 100 not divisible by 128
937 with pytest.raises(ValueError, match="in divisible by 128"):
938 scheme.create_weights(
939 torch.nn.Module(), mo.WEIGHT, mo.CkptCtx(), shapes, Mock()
940 )