| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """NVFP4 detection for ModelOpt MIXED_PRECISION checkpoints. |
| 4 | |
| 5 | A checkpoint that quantizes some layers to NVFP4 and others to FP8 declares |
| 6 | ``quant_algo: MIXED_PRECISION`` and resolves to ``modelopt_mixed``, so NVFP4 |
| 7 | has to be detected from the per-layer algorithms. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | |
| 12 | from vllm.config.model import _modelopt_mixed_has_nvfp4 |
| 13 | |
| 14 | |
| 15 | def _quantized_layers(*algos: str) -> dict: |
| 16 | return { |
| 17 | "quantized_layers": { |
| 18 | f"model.layers.{i}.proj": {"quant_algo": algo} |
| 19 | for i, algo in enumerate(algos) |
| 20 | } |
| 21 | } |
| 22 | |
| 23 | |
| 24 | def test_detects_nvfp4_alongside_another_algorithm(): |
| 25 | """The shipped shape: FP8 attention projections, NVFP4 MLP.""" |
| 26 | assert _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "NVFP4", "FP8")) |
| 27 | |
| 28 | |
| 29 | def test_ignores_a_checkpoint_with_no_nvfp4_layer(): |
| 30 | assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "MXFP8")) |
| 31 | |
| 32 | |
| 33 | def test_weight_only_nvfp4_does_not_count(): |
| 34 | """W4A16_NVFP4 leaves activations in high precision, so nothing to fuse.""" |
| 35 | assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "W4A16_NVFP4")) |
| 36 | |
| 37 | |
| 38 | @pytest.mark.parametrize( |
| 39 | "quant_config", |
| 40 | [None, {}, {"quantized_layers": None}, {"quantized_layers": ["NVFP4"]}], |
| 41 | ) |
| 42 | def test_absent_or_malformed_config_is_not_nvfp4(quant_config): |
| 43 | assert not _modelopt_mixed_has_nvfp4(quant_config) |