@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 374 lines 12.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import random
5
6 import pytest
7 import torch
8
9 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
10 from tests.kernels.utils import opcheck
11 from vllm import _custom_ops as ops
12 from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
13 from vllm.platforms import current_platform
14 from vllm.utils.mem_utils import get_max_shared_memory_bytes
15 from vllm.utils.torch_utils import set_random_seed
16
17 pytestmark = pytest.mark.skip_global_cleanup
18
19 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
20 # This will change depending on the compute capability.
21 # - 512 as a buffer
22 MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
23 # There may not be enough gpu memory due to large NUM_BLOCKS.
24 # Reduce NUM_BLOCKS when it happens.
25 NUM_BLOCKS = 4321 # Arbitrary values for testing
26 PARTITION_SIZE_ROCM = 256
27 DTYPES = [torch.bfloat16, torch.float16]
28 NUM_GEN_SEQS = [7] # Arbitrary values for testing
29 NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing
30
31 # Head sizes supported by the ROCm paged attention kernel.
32 HEAD_SIZES = [64, 128]
33
34 BLOCK_SIZES = [16, 32]
35 USE_ALIBI = [False, True]
36 KV_CACHE_DTYPE = ["auto", "fp8"]
37 SEEDS = [0]
38 CUDA_DEVICES = [
39 f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
40 ]
41
42
43 def ref_masked_attention(
44 query: torch.Tensor,
45 key: torch.Tensor,
46 value: torch.Tensor,
47 scale: float,
48 attn_mask: torch.Tensor | None = None,
49 ) -> torch.Tensor:
50 attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
51 if attn_mask is not None:
52 attn_weights = attn_weights + attn_mask.float()
53 attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
54 out = torch.einsum("hqk,khd->qhd", attn_weights, value)
55 return out
56
57
58 def ref_single_query_cached_kv_attention(
59 output: torch.Tensor,
60 query: torch.Tensor,
61 num_queries_per_kv: int,
62 key_cache: torch.Tensor,
63 value_cache: torch.Tensor,
64 block_tables: torch.Tensor,
65 seq_lens: torch.Tensor,
66 scale: float,
67 alibi_slopes: torch.Tensor | None,
68 ) -> None:
69 num_query_heads = query.shape[1]
70 num_kv_heads = value_cache.shape[1]
71 head_size = value_cache.shape[2]
72 block_size = value_cache.shape[3]
73 num_seqs = query.shape[0]
74
75 block_tables_lst = block_tables.cpu().tolist()
76 seq_lens_lst = seq_lens.cpu().tolist()
77 for i in range(num_seqs):
78 q = query[i].unsqueeze(0)
79 block_table = block_tables_lst[i]
80 seq_len = int(seq_lens_lst[i])
81
82 keys_lst: list[torch.Tensor] = []
83 values_lst: list[torch.Tensor] = []
84 for j in range(seq_len):
85 block_number = int(block_table[j // block_size])
86 block_offset = j % block_size
87
88 k = key_cache[block_number, :, :, block_offset, :]
89 k = k.reshape(num_kv_heads, head_size)
90 keys_lst.append(k)
91
92 v = value_cache[block_number, :, :, block_offset]
93 values_lst.append(v)
94 keys = torch.stack(keys_lst, dim=0)
95 values = torch.stack(values_lst, dim=0)
96 if num_queries_per_kv > 1:
97 # Handle MQA and GQA
98 keys = torch.repeat_interleave(keys, num_queries_per_kv, dim=1)
99 values = torch.repeat_interleave(values, num_queries_per_kv, dim=1)
100
101 alibi_bias = None
102 if alibi_slopes is not None:
103 # Create the ALiBi bias used in the paged attention kernel.
104 position_ids = torch.arange(seq_len).int()
105 alibi_bias = (position_ids - seq_len + 1).float()
106 alibi_bias = alibi_slopes.view(-1, 1, 1) * alibi_bias.view(1, 1, -1)
107
108 out = ref_masked_attention(q, keys, values, scale, alibi_bias)
109 out = out.view(num_query_heads, head_size)
110 output[i].copy_(out, non_blocking=True)
111
112
113 @pytest.mark.skipif(
114 not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
115 )
116 @pytest.mark.parametrize("num_seqs", NUM_GEN_SEQS)
117 @pytest.mark.parametrize("num_heads", NUM_HEADS)
118 @pytest.mark.parametrize("head_size", HEAD_SIZES)
119 @pytest.mark.parametrize("use_alibi", USE_ALIBI)
120 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
121 @pytest.mark.parametrize("dtype", DTYPES)
122 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
123 @pytest.mark.parametrize("seed", SEEDS)
124 @pytest.mark.parametrize("device", CUDA_DEVICES)
125 def test_paged_attention(
126 kv_cache_factory,
127 num_seqs: int,
128 num_heads: tuple[int, int],
129 head_size: int,
130 use_alibi: bool,
131 block_size: int,
132 dtype: torch.dtype,
133 kv_cache_dtype: str,
134 seed: int,
135 device: str,
136 ) -> None:
137 if current_platform.is_navi() and (
138 kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
139 ):
140 pytest.skip()
141
142 set_random_seed(seed)
143 torch.set_default_device(device)
144 scale = float(1.0 / (head_size**0.5))
145 num_query_heads, num_kv_heads = num_heads
146 query = torch.empty(num_seqs, num_query_heads, head_size, dtype=dtype)
147 query.uniform_(-scale, scale)
148
149 assert num_query_heads % num_kv_heads == 0
150 num_queries_per_kv = num_query_heads // num_kv_heads
151 alibi_slopes = None
152 if use_alibi:
153 alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)
154
155 seq_lens_lst = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
156 seq_lens_lst[-1] = MAX_SEQ_LEN
157 max_seq_len = max(seq_lens_lst)
158 seq_lens = torch.tensor(seq_lens_lst, dtype=torch.int)
159
160 # Create the block tables.
161 max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
162 block_tables_lst: list[list[int]] = []
163 for seq_idx, seq_len in enumerate(seq_lens_lst):
164 block_table = [
165 random.randint(num_seqs, NUM_BLOCKS - 1)
166 for _ in range(max_num_blocks_per_seq)
167 ]
168 # Keep the last block private to this sequence so its padding can be
169 # poisoned without corrupting valid tokens from another sequence.
170 block_table[(seq_len - 1) // block_size] = seq_idx
171 block_tables_lst.append(block_table)
172
173 block_tables = torch.tensor(block_tables_lst, dtype=torch.int)
174
175 # Create the KV caches.
176 key_caches, value_caches = kv_cache_factory(
177 NUM_BLOCKS,
178 block_size,
179 1,
180 num_kv_heads,
181 head_size,
182 kv_cache_dtype,
183 dtype,
184 seed,
185 device,
186 )
187 key_cache, value_cache = key_caches[0], value_caches[0]
188
189 k_scale_value = 0.5 if kv_cache_dtype == "fp8" else 1.0
190 v_scale_value = 0.25 if kv_cache_dtype == "fp8" else 1.0
191 k_scale = torch.tensor(k_scale_value, dtype=torch.float32, device=device)
192 v_scale = torch.tensor(v_scale_value, dtype=torch.float32, device=device)
193
194 if kv_cache_dtype == "fp8":
195 padding_nan = 0x80 if current_platform.is_fp8_fnuz() else 0x7F
196 else:
197 padding_nan = torch.nan
198 for seq_idx, seq_len in enumerate(seq_lens_lst):
199 padding_start = seq_len % block_size
200 if padding_start:
201 last_block_idx = block_tables_lst[seq_idx][seq_len // block_size]
202 value_cache[last_block_idx, :, :, padding_start:] = padding_nan
203
204 # Call the paged attention kernel.
205 output = torch.empty_like(query)
206 num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
207 assert PARTITION_SIZE_ROCM % block_size == 0
208 num_seqs, num_heads, head_size = output.shape
209 tmp_output = torch.empty(
210 size=(num_seqs, num_heads, num_partitions, head_size),
211 dtype=output.dtype,
212 )
213 exp_sums = torch.empty(
214 size=(num_seqs, num_heads, num_partitions),
215 dtype=torch.float32,
216 )
217 max_logits = torch.empty_like(exp_sums)
218 ops.paged_attention_rocm(
219 output,
220 exp_sums,
221 max_logits,
222 tmp_output,
223 query,
224 key_cache,
225 value_cache,
226 num_kv_heads,
227 scale,
228 block_tables,
229 seq_lens,
230 None,
231 block_size,
232 max_seq_len,
233 alibi_slopes,
234 kv_cache_dtype,
235 k_scale,
236 v_scale,
237 )
238
239 opcheck(
240 torch.ops._rocm_C.paged_attention,
241 (
242 output,
243 exp_sums,
244 max_logits,
245 tmp_output,
246 query,
247 key_cache,
248 value_cache,
249 num_kv_heads,
250 scale,
251 block_tables,
252 seq_lens,
253 None,
254 block_size,
255 max_seq_len,
256 alibi_slopes,
257 kv_cache_dtype,
258 k_scale,
259 v_scale,
260 None,
261 "f16",
262 ),
263 cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
264 )
265
266 # Run the reference implementation.
267 if kv_cache_dtype == "fp8":
268 # Convert cache data back to dtype.
269 x = 16 // torch.tensor([], dtype=dtype).element_size()
270 key_cache_shape = (NUM_BLOCKS, num_kv_heads, head_size // x, block_size, x)
271 dequantized_key_cache = torch.empty(
272 size=key_cache_shape, dtype=dtype, device=device
273 )
274 ops.convert_fp8(dequantized_key_cache, key_cache, k_scale_value)
275 key_cache = dequantized_key_cache
276
277 value_cache_shape = value_cache.shape
278 dequantized_value_cache = torch.empty(
279 size=value_cache_shape, dtype=dtype, device=device
280 )
281 ops.convert_fp8(dequantized_value_cache, value_cache, v_scale_value)
282 value_cache = dequantized_value_cache
283
284 ref_output = torch.empty_like(query)
285 ref_single_query_cached_kv_attention(
286 ref_output,
287 query,
288 num_queries_per_kv,
289 key_cache,
290 value_cache,
291 block_tables,
292 seq_lens,
293 scale,
294 alibi_slopes,
295 )
296
297 # NOTE(woosuk): Due to the kernel-level differences in the two
298 # implementations, there is a small numerical difference in the two
299 # outputs. Thus, we use a relaxed tolerance for the test.
300 atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
301 rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
302
303 # NOTE(zhaoyang): FP8 KV Cache will introduce quantization error,
304 # so we use a relaxed tolerance for the test.
305 atol, rtol = 1e-3, 1e-5
306 if kv_cache_dtype == "fp8":
307 atol, rtol = 1e-2, 1e-5
308 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
309
310
311 @pytest.mark.parametrize("attention_cls", [Attention, MMEncoderAttention])
312 def test_num_heads_not_divisible_by_num_kv_heads(attention_cls: type) -> None:
313 head_size = 64
314 scale = float(1.0 / (head_size**0.5))
315 num_heads = 16
316 num_kv_heads = 5
317 with pytest.raises(AssertionError):
318 _ = attention_cls(
319 num_heads=num_heads,
320 head_size=head_size,
321 scale=scale,
322 num_kv_heads=num_kv_heads,
323 )
324
325
326 UNSUPPORTED_HEAD_SIZES = [32, 80, 96, 160, 192, 224, 256]
327
328
329 @pytest.mark.skipif(
330 not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
331 )
332 @pytest.mark.parametrize("head_size", UNSUPPORTED_HEAD_SIZES)
333 def test_paged_attention_unsupported_head_sizes(
334 kv_cache_factory, head_size: int
335 ) -> None:
336 """Verify ROCm paged attention rejects unsupported head sizes."""
337 torch.set_default_device("cuda")
338 num_seqs, num_kv_heads, block_size, max_seq_len = 1, 8, 16, 128
339 scale = head_size**-0.5
340
341 query = torch.empty(num_seqs, num_kv_heads, head_size, dtype=torch.bfloat16)
342 key_caches, value_caches = kv_cache_factory(
343 NUM_BLOCKS, block_size, 1, num_kv_heads, head_size, "auto", torch.bfloat16, 0
344 )
345 block_tables = torch.zeros(num_seqs, 1, dtype=torch.int32)
346 seq_lens = torch.tensor([max_seq_len], dtype=torch.int32)
347
348 output = torch.empty_like(query)
349 num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
350 tmp_output = torch.empty(num_seqs, num_kv_heads, num_partitions, head_size)
351 exp_sums = torch.empty(num_seqs, num_kv_heads, num_partitions, dtype=torch.float32)
352 k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
353
354 with pytest.raises(RuntimeError, match="Unsupported head size"):
355 ops.paged_attention_rocm(
356 output,
357 exp_sums,
358 torch.empty_like(exp_sums),
359 tmp_output,
360 query,
361 key_caches[0],
362 value_caches[0],
363 num_kv_heads,
364 scale,
365 block_tables,
366 seq_lens,
367 None,
368 block_size,
369 max_seq_len,
370 None,
371 "auto",
372 k_scale,
373 v_scale,
374 )