| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | from torch.testing import assert_close |
| 6 | |
| 7 | from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton |
| 8 | |
| 9 | |
| 10 | def test_pack_seq_basic_fp8(): |
| 11 | """Test basic functionality of pack_seq_triton with fp8 and 3D tensors.""" |
| 12 | device = "cuda" |
| 13 | dtype = torch.float8_e4m3fn |
| 14 | |
| 15 | # Test cases with 3D tensors (N, H, D) |
| 16 | test_cases = [ |
| 17 | (6, 8, 4, 2, [3, 3]), # (6, 8, 4) -> (2, 3, 8, 4) |
| 18 | (10, 4, 8, 3, [2, 4, 4]), # (10, 4, 8) -> (3, 4, 4, 8) |
| 19 | (20, 16, 32, 4, [5, 5, 5, 5]), # (20, 16, 32) -> (4, 5, 16, 32) |
| 20 | ] |
| 21 | |
| 22 | for N, H, D, B, lengths_list in test_cases: |
| 23 | # Create input tensor with small values for fp8 |
| 24 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 25 | x = x.to(dtype=dtype) |
| 26 | lengths = torch.tensor(lengths_list, device=device) |
| 27 | |
| 28 | # Pack the data |
| 29 | packed = pack_seq_triton(x, lengths) |
| 30 | |
| 31 | # Check output shape and properties |
| 32 | expected_shape = (B, max(lengths_list), H, D) |
| 33 | assert packed.shape == expected_shape |
| 34 | assert packed.dtype == dtype |
| 35 | assert packed.device == x.device |
| 36 | |
| 37 | # Check that valid data is preserved (within fp8 precision) |
| 38 | for b in range(B): |
| 39 | start_idx = sum(lengths_list[:b]) |
| 40 | seq_len = lengths_list[b] |
| 41 | |
| 42 | expected_data = x[start_idx : start_idx + seq_len].to(torch.float32) |
| 43 | actual_data = packed[b, :seq_len].to(torch.float32) |
| 44 | |
| 45 | assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2) |
| 46 | |
| 47 | |
| 48 | def test_pack_seq_custom_padding_fp8(): |
| 49 | """Test pack_seq_triton with custom padding values for fp8.""" |
| 50 | device = "cuda" |
| 51 | dtype = torch.float8_e4m3fn |
| 52 | N, H, D, B = 20, 8, 16, 2 |
| 53 | lengths = torch.tensor([10, 10], device=device) |
| 54 | |
| 55 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 56 | x = x.to(dtype=dtype) |
| 57 | |
| 58 | # Test with different padding values |
| 59 | for pad_value in [-100.0, -10.0, 0.0, 10.0, 100.0]: |
| 60 | result = pack_seq_triton(x, lengths, pad_value=pad_value) |
| 61 | |
| 62 | # Check valid data |
| 63 | for b in range(B): |
| 64 | start_idx = b * 10 |
| 65 | expected_data = x[start_idx : start_idx + 10].to(torch.float32) |
| 66 | actual_data = result[b, :10].to(torch.float32) |
| 67 | assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2) |
| 68 | |
| 69 | # Check padding (fp8 has limited range, so check for large values) |
| 70 | padded_data = result[:, 10:].to(torch.float32) |
| 71 | if pad_value < 0: |
| 72 | assert torch.all(padded_data < -50) # Large negative values |
| 73 | elif pad_value > 0: |
| 74 | assert torch.all(padded_data > 50) # Large positive values |
| 75 | else: |
| 76 | assert torch.allclose(padded_data, torch.zeros_like(padded_data), atol=1e-2) |
| 77 | |
| 78 | |
| 79 | def test_pack_seq_default_negative_inf_padding_fp8(): |
| 80 | """Test that pack_seq_triton uses -inf padding by default for fp8.""" |
| 81 | device = "cuda" |
| 82 | dtype = torch.float8_e4m3fn |
| 83 | # B = 2 |
| 84 | N, H, D = 20, 8, 16 |
| 85 | lengths = torch.tensor([10, 10], device=device) |
| 86 | |
| 87 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 88 | x = x.to(dtype=dtype) |
| 89 | result = pack_seq_triton(x, lengths) |
| 90 | |
| 91 | # Check that padding is large negative values (fp8 representation of -inf) |
| 92 | padded_data = result[:, 10:].to(torch.float32) |
| 93 | assert torch.all( |
| 94 | padded_data < -100 |
| 95 | ) # fp8 -inf is represented as large negative number |
| 96 | |
| 97 | |
| 98 | def test_pack_seq_edge_cases_fp8(): |
| 99 | """Test pack_seq_triton with edge cases for fp8.""" |
| 100 | device = "cuda" |
| 101 | dtype = torch.float8_e4m3fn |
| 102 | |
| 103 | # Test with single batch element |
| 104 | x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1 |
| 105 | x = x.to(dtype=dtype) |
| 106 | lengths = torch.tensor([10], device=device) |
| 107 | result = pack_seq_triton(x, lengths) |
| 108 | assert result.shape == (1, 10, 8, 16) |
| 109 | |
| 110 | # Test with very short sequences |
| 111 | x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1 |
| 112 | x = x.to(dtype=dtype) |
| 113 | lengths = torch.tensor([1, 1, 1], device=device) |
| 114 | result = pack_seq_triton(x, lengths) |
| 115 | assert result.shape == (3, 1, 4, 8) |
| 116 | |
| 117 | # Test with different sequence lengths |
| 118 | x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1 |
| 119 | x = x.to(dtype=dtype) |
| 120 | lengths = torch.tensor([5, 7, 3], device=device) |
| 121 | result = pack_seq_triton(x, lengths) |
| 122 | assert result.shape == (3, 7, 8, 16) |
| 123 | |
| 124 | |
| 125 | def test_pack_seq_different_block_sizes_fp8(): |
| 126 | """Test pack_seq_triton with different block sizes for fp8.""" |
| 127 | device = "cuda" |
| 128 | dtype = torch.float8_e4m3fn |
| 129 | N, H, D, B = 100, 16, 32, 4 |
| 130 | lengths = torch.tensor([25, 25, 25, 25], device=device) |
| 131 | |
| 132 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 133 | x = x.to(dtype=dtype) |
| 134 | |
| 135 | # Test different block sizes |
| 136 | for block_t, block_d in [(32, 32), (64, 64), (128, 128)]: |
| 137 | result = pack_seq_triton(x, lengths, block_t=block_t, block_d=block_d) |
| 138 | |
| 139 | assert result.shape == (B, 25, H, D) |
| 140 | |
| 141 | # Check that valid data is preserved (within fp8 precision) |
| 142 | for b in range(B): |
| 143 | start_idx = b * 25 |
| 144 | expected_data = x[start_idx : start_idx + 25].to(torch.float32) |
| 145 | actual_data = result[b, :25].to(torch.float32) |
| 146 | assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2) |
| 147 | |
| 148 | |
| 149 | def test_pack_seq_shape_consistency(): |
| 150 | """Test that pack_seq_triton maintains shape consistency.""" |
| 151 | device = "cuda" |
| 152 | dtype = torch.float8_e4m3fn |
| 153 | N, H, D, B = 20, 8, 16, 2 |
| 154 | lengths = torch.tensor([10, 10], device=device) |
| 155 | |
| 156 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 157 | x = x.to(dtype=dtype) |
| 158 | |
| 159 | result = pack_seq_triton(x, lengths) |
| 160 | |
| 161 | # Check shape consistency |
| 162 | assert result.shape[0] == B # Batch dimension |
| 163 | assert result.shape[1] == lengths.max().item() # Max sequence length |
| 164 | assert result.shape[2:] == x.shape[1:] # Feature dimensions preserved |
| 165 | |
| 166 | |
| 167 | def test_pack_unpack_roundtrip_fp8(): |
| 168 | """Test that pack -> unpack gives us back the original data for fp8.""" |
| 169 | device = "cuda" |
| 170 | dtype = torch.float8_e4m3fn |
| 171 | |
| 172 | # Test cases with 3D tensors |
| 173 | test_cases = [ |
| 174 | (6, 8, 4, 2, [3, 3]), |
| 175 | (10, 4, 8, 3, [2, 4, 4]), |
| 176 | (20, 16, 32, 4, [5, 5, 5, 5]), |
| 177 | (15, 8, 16, 3, [7, 5, 3]), |
| 178 | ] |
| 179 | |
| 180 | for N, H, D, B, lengths_list in test_cases: |
| 181 | # Create input tensor with small values for fp8 |
| 182 | x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1 |
| 183 | x = x.to(dtype=dtype) |
| 184 | lengths = torch.tensor(lengths_list, device=device) |
| 185 | |
| 186 | # Pack the data |
| 187 | packed = pack_seq_triton(x, lengths) |
| 188 | |
| 189 | # Unpack the data |
| 190 | unpacked = unpack_seq_triton(packed, lengths) |
| 191 | |
| 192 | # Check that we get back the original data (within fp8 precision) |
| 193 | assert unpacked.shape == x.shape |
| 194 | x_f32 = x.to(torch.float32) |
| 195 | unpacked_f32 = unpacked.to(torch.float32) |
| 196 | assert_close(x_f32, unpacked_f32, rtol=1e-3, atol=1e-3) |
| 197 | |
| 198 | # Unpack without explicit start locations (computed in kernel) |
| 199 | unpacked_with_loc = unpack_seq_triton(packed, lengths) |
| 200 | assert_close(x_f32, unpacked_with_loc.to(torch.float32), rtol=1e-3, atol=1e-2) |
| 201 | |
| 202 | |
| 203 | def test_unpack_seq_triton_edge_cases_fp8(): |
| 204 | """Test unpack function with edge cases for fp8.""" |
| 205 | device = "cuda" |
| 206 | dtype = torch.float8_e4m3fn |
| 207 | |
| 208 | # Test with single batch element |
| 209 | x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1 |
| 210 | x = x.to(dtype=dtype) |
| 211 | lengths = torch.tensor([10], device=device) |
| 212 | packed = pack_seq_triton(x, lengths) |
| 213 | unpacked = unpack_seq_triton(packed, lengths) |
| 214 | assert unpacked.shape == x.shape |
| 215 | assert_close(x.to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2) |
| 216 | |
| 217 | # Test with very short sequences |
| 218 | x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1 |
| 219 | x = x.to(dtype=dtype) |
| 220 | lengths = torch.tensor([1, 1, 1], device=device) |
| 221 | packed = pack_seq_triton(x, lengths) |
| 222 | unpacked = unpack_seq_triton(packed, lengths) |
| 223 | # Only compare the first 3 elements that were actually packed |
| 224 | assert_close( |
| 225 | x[:3].to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2 |
| 226 | ) |
| 227 | |
| 228 | x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1 |
| 229 | x = x.to(dtype=dtype) |
| 230 | lengths = torch.tensor([5, 7, 3], device=device) |
| 231 | packed = pack_seq_triton(x, lengths) |
| 232 | unpacked = unpack_seq_triton(packed, lengths) |
| 233 | assert unpacked.shape == x.shape |
| 234 | assert_close(x.to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2) |