| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for FP8 per-channel online quantization. |
| 4 | |
| 5 | Per-output-channel weight scale + dynamic per-token activation scale. |
| 6 | bf16/fp16 checkpoints are quantized at load time with one fp32 scale per |
| 7 | output channel for weights and one fp32 scale per token for activations |
| 8 | (computed dynamically inside the kernel). Run via |
| 9 | `pytest tests/quantization/test_fp8_per_channel.py --forked`. |
| 10 | """ |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from tests.quantization.utils import is_quant_method_supported |
| 16 | from vllm import _custom_ops as ops |
| 17 | from vllm.config.quantization import ( |
| 18 | _ONLINE_SHORTHANDS, |
| 19 | QUANT_KEY_NAMES, |
| 20 | QuantizationConfigArgs, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization.online.base import ( |
| 23 | _ONLINE_LINEAR_METHODS, |
| 24 | _ONLINE_MOE_METHODS, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.quantization.online.fp8 import ( |
| 27 | Fp8PtpcOnlineLinearMethod, |
| 28 | Fp8PtpcOnlineMoEMethod, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 31 | kFp8StaticChannelSym, |
| 32 | ) |
| 33 | from vllm.platforms import current_platform |
| 34 | |
| 35 | |
| 36 | def test_fp8_per_channel_shorthand_registered() -> None: |
| 37 | """The `fp8_per_channel` CLI shorthand must resolve to a config that |
| 38 | dispatches the per-channel methods. Guards against regressions in |
| 39 | `_ONLINE_SHORTHANDS` / `_ONLINE_LINEAR_METHODS` / `_ONLINE_MOE_METHODS` |
| 40 | drifting out of sync. |
| 41 | """ |
| 42 | args = _ONLINE_SHORTHANDS["fp8_per_channel"] |
| 43 | assert isinstance(args, QuantizationConfigArgs) |
| 44 | assert args.linear is not None |
| 45 | assert args.moe is not None |
| 46 | assert args.linear.weight is kFp8StaticChannelSym |
| 47 | assert args.moe.weight is kFp8StaticChannelSym |
| 48 | |
| 49 | assert _ONLINE_LINEAR_METHODS[kFp8StaticChannelSym] is Fp8PtpcOnlineLinearMethod |
| 50 | assert _ONLINE_MOE_METHODS[kFp8StaticChannelSym] is Fp8PtpcOnlineMoEMethod |
| 51 | |
| 52 | assert QUANT_KEY_NAMES["fp8_per_channel_static"] is kFp8StaticChannelSym |
| 53 | |
| 54 | |
| 55 | @pytest.mark.skipif( |
| 56 | not is_quant_method_supported("fp8"), |
| 57 | reason="FP8 is not supported on this GPU type.", |
| 58 | ) |
| 59 | def test_scaled_fp8_quant_per_channel_shape() -> None: |
| 60 | """Verify the kernel call per-channel quant depends on: passing a 2D |
| 61 | weight to `ops.scaled_fp8_quant` with `use_per_token_if_dynamic=True` |
| 62 | yields one scale per output row -- a [N, 1] fp32 tensor. |
| 63 | """ |
| 64 | x = (torch.randn(size=(96, 256), device="cuda") * 13).to(torch.bfloat16) |
| 65 | y, s = ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True) |
| 66 | assert y.shape == (96, 256) |
| 67 | assert y.dtype == current_platform.fp8_dtype() |
| 68 | assert s.shape == (96, 1) |
| 69 | assert s.dtype == torch.float32 |