| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | from nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX, dequantize_nvfp4_to_dtype |
| 6 | |
| 7 | from vllm import _custom_ops as ops |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.torch_utils import set_random_seed |
| 10 | |
| 11 | if not current_platform.has_device_capability(100): |
| 12 | pytest.skip( |
| 13 | reason="Nvfp4 Requires compute capability of 10 or above.", |
| 14 | allow_module_level=True, |
| 15 | ) |
| 16 | |
| 17 | DTYPES = [torch.float16, torch.bfloat16] |
| 18 | # m, n, k |
| 19 | SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)] |
| 20 | PAD_SHAPES = [(150, 128, 64), (128, 128, 96)] |
| 21 | SHAPES.extend(PAD_SHAPES) |
| 22 | |
| 23 | SEEDS = [42] |
| 24 | CUDA_DEVICES = ["cuda:0"] |
| 25 | |
| 26 | |
| 27 | def get_ref_results( |
| 28 | a_fp4, |
| 29 | b_fp4, |
| 30 | a_sf, |
| 31 | b_sf, |
| 32 | a_global_scale, |
| 33 | b_global_scale, |
| 34 | m, |
| 35 | n, |
| 36 | dtype, |
| 37 | block_size, |
| 38 | device, |
| 39 | ): |
| 40 | _, m_k = a_fp4.shape |
| 41 | _, n_k = b_fp4.shape |
| 42 | assert m_k == n_k |
| 43 | a_in_dtype = dequantize_nvfp4_to_dtype( |
| 44 | a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size |
| 45 | ) |
| 46 | b_in_dtype = dequantize_nvfp4_to_dtype( |
| 47 | b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size |
| 48 | ) |
| 49 | return torch.matmul(a_in_dtype, b_in_dtype.t()) |
| 50 | |
| 51 | |
| 52 | @pytest.mark.parametrize("dtype", DTYPES) |
| 53 | @pytest.mark.parametrize("shape", SHAPES) |
| 54 | @pytest.mark.parametrize("seed", SEEDS) |
| 55 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 56 | @torch.inference_mode() |
| 57 | def test_nvfp4_gemm( |
| 58 | dtype: torch.dtype, |
| 59 | shape: tuple[int, int, int], |
| 60 | seed: int, |
| 61 | device: str, |
| 62 | ) -> None: |
| 63 | set_random_seed(seed) |
| 64 | m, n, packed_k = shape |
| 65 | k = packed_k * 2 |
| 66 | block_size = 16 |
| 67 | a_dtype = torch.randn((m, k), dtype=dtype, device=device) |
| 68 | b_dtype = torch.randn((n, k), dtype=dtype, device=device) |
| 69 | |
| 70 | a_global_scale = ( |
| 71 | (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1) |
| 72 | ).to(torch.float32) |
| 73 | b_global_scale = ( |
| 74 | (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1) |
| 75 | ).to(torch.float32) |
| 76 | alpha = 1.0 / (a_global_scale * b_global_scale) |
| 77 | # ops.scaled_fp4_quant returns swizzled scales, while weights |
| 78 | # from checkpoints are in linear scales. |
| 79 | a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale) |
| 80 | b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale) |
| 81 | |
| 82 | # get_ref_results unswizzles the scales internally. |
| 83 | expected_out = get_ref_results( |
| 84 | a_fp4, |
| 85 | b_fp4, |
| 86 | a_scale_interleaved, |
| 87 | b_scale_interleaved, |
| 88 | a_global_scale, |
| 89 | b_global_scale, |
| 90 | m, |
| 91 | n, |
| 92 | dtype, |
| 93 | block_size, |
| 94 | device, |
| 95 | ) |
| 96 | out = ops.cutlass_scaled_fp4_mm( |
| 97 | a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype |
| 98 | ) |
| 99 | |
| 100 | torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1) |