| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the stride-aware FP8 quantization kernel with head_dim padding.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 9 | get_fp8_min_max, |
| 10 | ) |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.triton_utils import HAS_TRITON |
| 13 | |
| 14 | if HAS_TRITON: |
| 15 | from vllm.kernels.triton.qkv_padded_fp8_quant import ( |
| 16 | quantize_fp8_pad_head_dim_triton, |
| 17 | ) |
| 18 | |
| 19 | HEAD_DIMS = [72, 80, 128] |
| 20 | SEQ_LENS = [64, 256] |
| 21 | NUM_HEADS = [16] |
| 22 | SCALES = [0.01, 0.1, 1.0] |
| 23 | |
| 24 | |
| 25 | def _naive_fp8_quantize( |
| 26 | tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool |
| 27 | ) -> torch.Tensor: |
| 28 | """Reference FP8 quantization in PyTorch.""" |
| 29 | fp8_dtype = current_platform.fp8_dtype() |
| 30 | fp8_min, fp8_max = get_fp8_min_max() |
| 31 | |
| 32 | x = tensor.float() |
| 33 | if not skip_scale: |
| 34 | x = x / scale.item() |
| 35 | x = x.clamp(fp8_min, fp8_max) |
| 36 | return x.to(fp8_dtype) |
| 37 | |
| 38 | |
| 39 | @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") |
| 40 | @pytest.mark.parametrize("head_dim", HEAD_DIMS) |
| 41 | @pytest.mark.parametrize("seq_len", SEQ_LENS) |
| 42 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 43 | @pytest.mark.parametrize("scale_val", SCALES) |
| 44 | def test_quantize_contiguous( |
| 45 | head_dim: int, seq_len: int, num_heads: int, scale_val: float |
| 46 | ) -> None: |
| 47 | """Test quantization of contiguous 3D tensors.""" |
| 48 | torch.manual_seed(42) |
| 49 | tensor = torch.randn( |
| 50 | seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16 |
| 51 | ) |
| 52 | scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view( |
| 53 | 1, 1, 1, 1 |
| 54 | ) |
| 55 | |
| 56 | result = quantize_fp8_pad_head_dim_triton(tensor, scale) |
| 57 | |
| 58 | padded_dim = (head_dim + 15) // 16 * 16 |
| 59 | assert result.shape == (seq_len, num_heads, padded_dim) |
| 60 | assert result.is_contiguous() |
| 61 | assert result.dtype == current_platform.fp8_dtype() |
| 62 | |
| 63 | # Compare unpadded portion against reference |
| 64 | ref = _naive_fp8_quantize(tensor, scale, skip_scale=False) |
| 65 | torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float()) |
| 66 | |
| 67 | # Padded region should be zero |
| 68 | if padded_dim > head_dim: |
| 69 | assert (result[:, :, head_dim:].float() == 0).all() |
| 70 | |
| 71 | |
| 72 | @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") |
| 73 | @pytest.mark.parametrize("head_dim", [72, 80]) |
| 74 | def test_quantize_non_contiguous(head_dim: int) -> None: |
| 75 | """Test quantization from non-contiguous QKV views (interleaved buffer).""" |
| 76 | seq_len, num_heads = 64, 16 |
| 77 | # Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim) |
| 78 | qkv = torch.randn( |
| 79 | seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16 |
| 80 | ) |
| 81 | # Q is every 3rd head slice - non-contiguous view |
| 82 | q = qkv[:, 0::3, :] |
| 83 | assert not q.is_contiguous() |
| 84 | |
| 85 | scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) |
| 86 | result = quantize_fp8_pad_head_dim_triton(q, scale) |
| 87 | |
| 88 | padded_dim = (head_dim + 15) // 16 * 16 |
| 89 | assert result.shape == (seq_len, num_heads, padded_dim) |
| 90 | assert result.is_contiguous() |
| 91 | |
| 92 | # Compare against contiguous reference |
| 93 | ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False) |
| 94 | torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float()) |
| 95 | |
| 96 | |
| 97 | @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") |
| 98 | def test_skip_scale() -> None: |
| 99 | """Test skip_scale=True produces cast-only output (no division).""" |
| 100 | seq_len, num_heads, head_dim = 32, 8, 80 |
| 101 | tensor = torch.randn( |
| 102 | seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16 |
| 103 | ) |
| 104 | scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) |
| 105 | |
| 106 | result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True) |
| 107 | result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False) |
| 108 | |
| 109 | # skip_scale should just cast, not divide |
| 110 | ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True) |
| 111 | torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float()) |
| 112 | |
| 113 | # With scale != 1.0, skip and no-skip should differ |
| 114 | assert not torch.equal(result_skip.float(), result_noskip.float()) |
| 115 | |
| 116 | |
| 117 | @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") |
| 118 | def test_4d_input() -> None: |
| 119 | """Test that 4D input (B, S, H, D) is handled correctly.""" |
| 120 | B, S, H, D = 2, 32, 8, 72 |
| 121 | tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) |
| 122 | scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) |
| 123 | |
| 124 | result = quantize_fp8_pad_head_dim_triton(tensor, scale) |
| 125 | padded_dim = (D + 15) // 16 * 16 |
| 126 | assert result.shape == (B, S, H, padded_dim) |