test_prefix_prefill.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,022 lines 34.7 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import math
5	import random
6	import time
7	from collections.abc import Callable
8
9	import pytest
10	import torch
11	import torch.nn.functional as F
12
13	from vllm.platforms import current_platform
14	from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
15	from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
16	chunked_prefill_paged_decode,
17	)
18	from vllm.v1.attention.ops.prefix_prefill import context_attention_fwd
19
20	pytestmark = pytest.mark.skip_global_cleanup
21
22	NUM_HEADS = [64]
23	NUM_QUERIES_PER_KV = [1, 64]
24	HEAD_SIZES = [24, 128]
25	DTYPES = [torch.float16]
26	CUDA_DEVICES = [
27	f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
28	]
29	SLIDING_WINDOW = [0, 16, 2048]
30	KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
31
32	OPS = [chunked_prefill_paged_decode, context_attention_fwd]
33
34
35	def create_causal_attention_mask_for_sdpa(
36	query_lens: list[int],
37	seq_lens: list[int],
38	sliding_window: int = 0,
39	device: torch.device = None,
40	dtype: torch.dtype = None,
41	) -> torch.Tensor:
42	total_queries = sum(query_lens)
43	total_keys = sum(seq_lens)
44
45	# Create a mask filled with -inf
46	mask = torch.full(
47	(total_queries, total_keys), float("-inf"), device=device, dtype=dtype
48	)
49
50	query_start = 0
51	key_start = 0
52
53	for query_len, seq_len in zip(query_lens, seq_lens):
54	query_end = query_start + query_len
55	key_end = key_start + seq_len
56	q_indices = torch.arange(query_len, device=device)
57	k_indices = torch.arange(seq_len, device=device)
58	q_pos_in_seq = seq_len - query_len + q_indices
59
60	valid_mask = k_indices[None, :] <= q_pos_in_seq[:, None]
61
62	if sliding_window > 0:
63	valid_mask &= k_indices[None, :] >= (
64	q_pos_in_seq[:, None] - sliding_window + 1
65	)
66
67	mask[query_start:query_end, key_start:key_end][valid_mask] = 0.0
68
69	query_start = query_end
70	key_start = key_end
71
72	return mask
73
74
75	def create_alibi_causal_mask(
76	query_len: int,
77	seq_len: int,
78	alibi_slopes: torch.Tensor,
79	device: torch.device,
80	dtype: torch.dtype,
81	) -> torch.Tensor:
82	query_pos = torch.arange(
83	seq_len - query_len, seq_len, device=device, dtype=torch.float32
84	)
85	key_pos = torch.arange(seq_len, device=device, dtype=torch.float32)
86
87	rel_pos = key_pos[None, :] - query_pos[:, None]
88
89	# Apply ALiBi slopes: [num_heads, query_len, seq_len]
90	alibi_bias = alibi_slopes[:, None, None] * rel_pos[None, :, :]
91	alibi_bias = alibi_bias.to(dtype)
92
93	# Apply causal mask: prevent attending to future positions
94	# causal_mask[i, j] = True if key_pos[j] <= query_pos[i]
95	causal_mask = key_pos[None, :] <= query_pos[:, None]
96	alibi_bias = alibi_bias.masked_fill(~causal_mask[None, :, :], float("-inf"))
97
98	# Add batch dimension: [1, num_heads, query_len, seq_len]
99	# SDPA expects batch dimension even for single sequences
100	return alibi_bias.unsqueeze(0)
101
102
103	@pytest.mark.parametrize("num_heads", NUM_HEADS)
104	@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
105	@pytest.mark.parametrize("head_size", HEAD_SIZES)
106	@pytest.mark.parametrize("dtype", DTYPES)
107	@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
108	@pytest.mark.parametrize("device", CUDA_DEVICES)
109	@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
110	@pytest.mark.parametrize("op", OPS)
111	@torch.inference_mode()
112	def test_contexted_kv_attention(
113	num_heads: int,
114	num_queries_per_kv: int,
115	head_size: int,
116	sliding_window: int,
117	dtype: torch.dtype,
118	kv_cache_dtype: str,
119	device: str,
120	op: Callable,
121	block_size: int = 32,
122	) -> None:
123	if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
124	pytest.skip(
125	"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
126	)
127
128	if (
129	current_platform.is_rocm()
130	and op is chunked_prefill_paged_decode
131	and kv_cache_dtype == "fp8_e5m2"
132	):
133	pytest.skip("ROCm custom paged attention does not support fp8_e5m2 KV cache")
134
135	set_random_seed(0)
136	torch.set_default_device(device)
137
138	# Need this, otherwise when we capture the graph the process
139	# for GPU 1 would run on both GPU0 and GPU1 and things would hang
140	#
141	# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
142	torch.accelerator.set_device_index(device)
143
144	MAX_SEQ_LEN = 1024
145	MAX_CTX_LEN = 1024
146	BS = 10
147	cache_size = 640
148	max_block_per_request = 64
149	query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
150	# ensure one sequence in batch is a decode
151	query_lens[-1] = 1
152
153	ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
154	seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
155	num_kv_heads = num_heads // num_queries_per_kv
156
157	num_tokens = sum(query_lens)
158	query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
159	query.uniform_(-1e-3, 1e-3)
160	output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
161
162	kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
163	kv.uniform_(-1e-3, 1e-3)
164	key, value = kv.unbind(dim=1)
165
166	if kv_cache_dtype == "auto":
167	cache_dtype = dtype
168	else:
169	cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
170	k_cache = torch.zeros(
171	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
172	)
173	v_cache = torch.zeros(
174	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
175	)
176	k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
177	v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
178	values = torch.arange(0, cache_size, dtype=torch.int32)
179	values = values[torch.randperm(cache_size)]
180	block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
181	b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
182	b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
183	b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
184	max_input_len = MAX_SEQ_LEN
185	# copy kv to cache
186	b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
187	torch.int32
188	)
189	for i in range(BS):
190	for j in range(query_lens[i]):
191	k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] + j])
192	v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] + b_ctx_len[i] + j])
193	cur_ctx = 0
194	block_id = 0
195	while cur_ctx < b_ctx_len[i]:
196	start_loc = b_seq_start_loc[i] + cur_ctx
197	if cur_ctx + block_size > b_ctx_len[i]:
198	end_loc = b_seq_start_loc[i] + b_ctx_len[i]
199	else:
200	end_loc = start_loc + block_size
201	start_slot = block_table[i, block_id] * block_size
202	end_slot = start_slot + end_loc - start_loc
203	k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
204	key[start_loc:end_loc]
205	)
206	v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
207	value[start_loc:end_loc]
208	)
209	cur_ctx += block_size
210	block_id += 1
211	# transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
212	# to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
213	k_cache = (
214	k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
215	.permute(0, 2, 3, 1, 4)
216	.contiguous()
217	)
218	# transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
219	# to V_cache[num_blocks, num_kv_heads, head_size, block_size]
220	v_cache = (
221	v_cache.view(-1, block_size, num_kv_heads, head_size)
222	.permute(0, 2, 3, 1)
223	.contiguous()
224	)
225	k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
226
227	# Warm up the Triton kernel by calling it once before actually measuring
228	# generation time
229	op(
230	query,
231	k,
232	v,
233	output,
234	kv_cache_dtype,
235	k_cache,
236	v_cache,
237	block_table,
238	b_start_loc,
239	b_seq_len,
240	MAX_CTX_LEN,
241	max_input_len,
242	k_scale,
243	v_scale,
244	sliding_window=sliding_window,
245	)
246	torch.accelerator.synchronize()
247	start_time = time.time()
248	op(
249	query,
250	k,
251	v,
252	output,
253	kv_cache_dtype,
254	k_cache,
255	v_cache,
256	block_table,
257	b_start_loc,
258	b_seq_len,
259	MAX_CTX_LEN,
260	max_input_len,
261	k_scale,
262	v_scale,
263	sliding_window=sliding_window,
264	)
265	torch.accelerator.synchronize()
266	end_time = time.time()
267	print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")
268
269	scale = float(1.0 / (head_size**0.5))
270
271	# Reshape for SDPA: (seq_len, num_heads, head_size) ->
272	# (1, num_heads, seq_len, head_size)
273	query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size)
274	query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
275	1, num_heads, num_tokens, head_size
276	)
277
278	# Expand key and value for GQA/MQA to match query heads
279	key_sdpa = key[:, :, None, :].expand(
280	key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
281	)
282	key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(
283	1, num_heads, sum(seq_lens), head_size
284	)
285
286	value_sdpa = value[:, :, None, :].expand(
287	value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
288	)
289	value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
290	1, num_heads, sum(seq_lens), head_size
291	)
292
293	attn_mask = create_causal_attention_mask_for_sdpa(
294	query_lens, seq_lens, sliding_window, device=device, dtype=dtype
295	)
296
297	output_ref = F.scaled_dot_product_attention(
298	query_sdpa,
299	key_sdpa,
300	value_sdpa,
301	attn_mask=attn_mask,
302	dropout_p=0.0,
303	scale=scale,
304	)
305	torch.accelerator.synchronize()
306	start_time = time.time()
307	output_ref = F.scaled_dot_product_attention(
308	query_sdpa,
309	key_sdpa,
310	value_sdpa,
311	attn_mask=attn_mask,
312	dropout_p=0.0,
313	scale=scale,
314	)
315	torch.accelerator.synchronize()
316	end_time = time.time()
317	print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")
318
319	# Reshape output back to (num_tokens, num_heads, head_size)
320	output_ref = output_ref.view(num_heads, num_tokens, head_size)
321	output_ref = output_ref.permute(1, 0, 2).contiguous()
322	atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4
323	torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
324
325
326	@pytest.mark.parametrize("num_heads", NUM_HEADS)
327	@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
328	@pytest.mark.parametrize("head_size", HEAD_SIZES)
329	@pytest.mark.parametrize("dtype", DTYPES)
330	@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
331	@pytest.mark.parametrize("device", CUDA_DEVICES)
332	@torch.inference_mode()
333	def test_contexted_kv_attention_cached_kv(
334	num_heads: int,
335	num_queries_per_kv: int,
336	head_size: int,
337	dtype: torch.dtype,
338	kv_cache_dtype: str,
339	device: str,
340	block_size: int = 32,
341	) -> None:
342	# Exercises the KV_FROM_CACHE path of context_attention_fwd: the current
343	# chunk K/V are not passed as dense tensors (k=v=None); they are read back
344	# from the paged KV cache, as done by layers that re-attend an already
345	# cached sequence with the query only (e.g. IQuest LoopCoder's
346	# attn(q, None, None)). The whole sequence therefore lives in the cache and
347	# the result must still match a dense causal SDPA reference.
348	if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
349	pytest.skip(
350	"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
351	)
352
353	set_random_seed(0)
354	torch.set_default_device(device)
355	torch.accelerator.set_device_index(device)
356
357	MAX_SEQ_LEN = 1024
358	MAX_CTX_LEN = 1024
359	BS = 10
360	cache_size = 640
361	max_block_per_request = 64
362	query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
363	# ensure one sequence in batch is a decode
364	query_lens[-1] = 1
365	ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
366	seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
367	num_kv_heads = num_heads // num_queries_per_kv
368
369	num_tokens = sum(query_lens)
370	query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
371	query.uniform_(-1e-3, 1e-3)
372	output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
373
374	kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
375	kv.uniform_(-1e-3, 1e-3)
376	key, value = kv.unbind(dim=1)
377
378	if kv_cache_dtype == "auto":
379	cache_dtype = dtype
380	else:
381	cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
382	k_cache = torch.zeros(
383	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
384	)
385	v_cache = torch.zeros(
386	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
387	)
388	values = torch.arange(0, cache_size, dtype=torch.int32)
389	values = values[torch.randperm(cache_size)]
390	block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
391	b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
392	b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
393	max_input_len = MAX_SEQ_LEN
394	b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
395	torch.int32
396	)
397	# Unlike the dense test, write the WHOLE sequence (context + current chunk)
398	# into the paged cache, since the current chunk is read back from the cache.
399	for i in range(BS):
400	cur = 0
401	block_id = 0
402	while cur < seq_lens[i]:
403	start_loc = b_seq_start_loc[i] + cur
404	if cur + block_size > seq_lens[i]:
405	end_loc = b_seq_start_loc[i] + seq_lens[i]
406	else:
407	end_loc = start_loc + block_size
408	start_slot = block_table[i, block_id] * block_size
409	end_slot = start_slot + end_loc - start_loc
410	k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
411	key[start_loc:end_loc]
412	)
413	v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
414	value[start_loc:end_loc]
415	)
416	cur += block_size
417	block_id += 1
418	# transpose to the paged cache layouts the kernel expects:
419	# K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
420	# V_cache[num_blocks, num_kv_heads, head_size, block_size]
421	k_cache = (
422	k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
423	.permute(0, 2, 3, 1, 4)
424	.contiguous()
425	)
426	v_cache = (
427	v_cache.view(-1, block_size, num_kv_heads, head_size)
428	.permute(0, 2, 3, 1)
429	.contiguous()
430	)
431	k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
432
433	# Cached-K/V path: current-chunk k and v are None.
434	context_attention_fwd(
435	query,
436	None,
437	None,
438	output,
439	kv_cache_dtype,
440	k_cache,
441	v_cache,
442	block_table,
443	b_start_loc,
444	b_seq_len,
445	MAX_CTX_LEN,
446	max_input_len,
447	k_scale,
448	v_scale,
449	sliding_window=0,
450	)
451	torch.accelerator.synchronize()
452
453	scale = float(1.0 / (head_size**0.5))
454
455	query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size)
456	query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
457	1, num_heads, num_tokens, head_size
458	)
459	key_sdpa = key[:, :, None, :].expand(
460	key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
461	)
462	key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(
463	1, num_heads, sum(seq_lens), head_size
464	)
465	value_sdpa = value[:, :, None, :].expand(
466	value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
467	)
468	value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
469	1, num_heads, sum(seq_lens), head_size
470	)
471
472	attn_mask = create_causal_attention_mask_for_sdpa(
473	query_lens, seq_lens, 0, device=device, dtype=dtype
474	)
475	output_ref = F.scaled_dot_product_attention(
476	query_sdpa,
477	key_sdpa,
478	value_sdpa,
479	attn_mask=attn_mask,
480	dropout_p=0.0,
481	scale=scale,
482	)
483	output_ref = output_ref.view(num_heads, num_tokens, head_size)
484	output_ref = output_ref.permute(1, 0, 2).contiguous()
485	atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4
486	torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
487
488
489	@pytest.mark.parametrize("device", CUDA_DEVICES)
490	@torch.inference_mode()
491	def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None:
492	# Boundary guard for the KV_FROM_CACHE block-table load. With an
493	# exact-sized block table (row length == number of blocks for the
494	# sequence) and a query that ends the sequence, the last K/V tile has
495	# padded lanes whose absolute positions step past the sequence end. Those
496	# lanes must not read a block-table entry past this batch's row.
497	#
498	# seq_len is a whole number of blocks, so bn_logical for a padded lane
499	# lands exactly on num_blocks (one past the last valid row entry), and
500	# query_len is not tile-aligned so the overshoot is actually exercised.
501	# The result must still match a dense causal SDPA reference.
502	set_random_seed(0)
503	torch.set_default_device(device)
504	torch.accelerator.set_device_index(device)
505
506	dtype = torch.float16
507	kv_cache_dtype = "auto"
508	num_heads = 4
509	num_queries_per_kv = 1
510	num_kv_heads = num_heads // num_queries_per_kv
511	head_size = 32
512	block_size = 32
513	num_blocks = 4
514	seq_len = num_blocks * block_size # 128 == exact multiple of block_size
515	query_len = 99 # < seq_len and not a multiple of the kernel tile
516
517	query = torch.empty(query_len, num_heads, head_size, dtype=dtype)
518	query.uniform_(-1e-3, 1e-3)
519	output = torch.empty(query_len, num_heads, head_size, dtype=dtype)
520
521	kv = torch.empty(seq_len, 2, num_kv_heads, head_size, dtype=dtype)
522	kv.uniform_(-1e-3, 1e-3)
523	key, value = kv.unbind(dim=1)
524
525	k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)
526	v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)
527	# Exact-sized block table: row length == num_blocks (identity mapping).
528	block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)
529	b_seq_len = torch.tensor([seq_len], dtype=torch.int32)
530	b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)
531
532	# Write the whole sequence (context + current chunk) into the paged cache.
533	for cur in range(0, seq_len, block_size):
534	block_id = cur // block_size
535	end = min(cur + block_size, seq_len)
536	start_slot = block_table[0, block_id] * block_size
537	end_slot = start_slot + (end - cur)
538	k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
539	key[cur:end]
540	)
541	v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
542	value[cur:end]
543	)
544
545	k_cache = (
546	k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
547	.permute(0, 2, 3, 1, 4)
548	.contiguous()
549	)
550	v_cache = (
551	v_cache.view(-1, block_size, num_kv_heads, head_size)
552	.permute(0, 2, 3, 1)
553	.contiguous()
554	)
555	k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
556
557	# Cached-K/V path: current-chunk k and v are None.
558	context_attention_fwd(
559	query,
560	None,
561	None,
562	output,
563	kv_cache_dtype,
564	k_cache,
565	v_cache,
566	block_table,
567	b_start_loc,
568	b_seq_len,
569	seq_len,
570	query_len,
571	k_scale,
572	v_scale,
573	sliding_window=0,
574	)
575	torch.accelerator.synchronize()
576
577	scale = float(1.0 / (head_size**0.5))
578	query_sdpa = query.view(query_len, num_kv_heads, num_queries_per_kv, head_size)
579	query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
580	1, num_heads, query_len, head_size
581	)
582	key_sdpa = key[:, :, None, :].expand(
583	seq_len, num_kv_heads, num_queries_per_kv, head_size
584	)
585	key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size)
586	value_sdpa = value[:, :, None, :].expand(
587	seq_len, num_kv_heads, num_queries_per_kv, head_size
588	)
589	value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
590	1, num_heads, seq_len, head_size
591	)
592
593	attn_mask = create_causal_attention_mask_for_sdpa(
594	[query_len], [seq_len], 0, device=device, dtype=dtype
595	)
596	output_ref = F.scaled_dot_product_attention(
597	query_sdpa,
598	key_sdpa,
599	value_sdpa,
600	attn_mask=attn_mask,
601	dropout_p=0.0,
602	scale=scale,
603	)
604	output_ref = output_ref.view(num_heads, query_len, head_size)
605	output_ref = output_ref.permute(1, 0, 2).contiguous()
606	torch.testing.assert_close(output, output_ref, atol=1e-4, rtol=0)
607
608
609	@pytest.mark.parametrize("device", CUDA_DEVICES)
610	@torch.inference_mode()
611	def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None:
612	# The cached-K/V (k=None) path is not supported together with ALiBi; the
613	# entry point must reject it up-front with a clear NotImplementedError
614	# rather than launching the kernel with an unsupported combination.
615	set_random_seed(0)
616	torch.set_default_device(device)
617	torch.accelerator.set_device_index(device)
618
619	num_heads = 4
620	num_kv_heads = 4
621	head_size = 16
622	x = 8
623	block_size = 16
624	num_blocks = 4
625	query_len = 8
626
627	query = torch.empty(query_len, num_heads, head_size, dtype=torch.float16)
628	query.uniform_(-1e-3, 1e-3)
629	output = torch.empty_like(query)
630	k_cache = torch.zeros(
631	num_blocks, num_kv_heads, head_size // x, block_size, x, dtype=torch.float16
632	)
633	v_cache = torch.zeros(
634	num_blocks, num_kv_heads, head_size, block_size, dtype=torch.float16
635	)
636	block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)
637	b_seq_len = torch.tensor([query_len], dtype=torch.int32)
638	b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)
639	k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
640	alibi_slopes = torch.ones(num_heads, dtype=torch.float32, device=device)
641
642	with pytest.raises(NotImplementedError):
643	context_attention_fwd(
644	query,
645	None,
646	None,
647	output,
648	"auto",
649	k_cache,
650	v_cache,
651	block_table,
652	b_start_loc,
653	b_seq_len,
654	query_len,
655	query_len,
656	k_scale,
657	v_scale,
658	alibi_slopes=alibi_slopes,
659	)
660
661
662	@pytest.mark.parametrize("num_heads", NUM_HEADS)
663	@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
664	@pytest.mark.parametrize("head_size", HEAD_SIZES)
665	@pytest.mark.parametrize("dtype", DTYPES)
666	@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
667	@pytest.mark.parametrize("device", CUDA_DEVICES)
668	@pytest.mark.parametrize("op", OPS)
669	@torch.inference_mode()
670	def test_contexted_kv_attention_alibi(
671	num_heads: int,
672	num_queries_per_kv: int,
673	head_size: int,
674	dtype: torch.dtype,
675	kv_cache_dtype: str,
676	device: str,
677	op: Callable,
678	block_size: int = 32,
679	) -> None:
680	if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
681	pytest.skip(
682	"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
683	)
684
685	if (
686	current_platform.is_rocm()
687	and op is chunked_prefill_paged_decode
688	and kv_cache_dtype == "fp8_e5m2"
689	):
690	pytest.skip("ROCm custom paged attention does not support fp8_e5m2 KV cache")
691
692	set_random_seed(0)
693	torch.set_default_device(device)
694
695	# Need this, otherwise when we capture the graph the process
696	# for GPU 1 would run on both GPU0 and GPU1 and things would hang
697	#
698	# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
699	torch.accelerator.set_device_index(device)
700
701	def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
702	# Fork from: vllm/vllm/model_executor/models/bloom.py#L44
703	closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
704	base = torch.tensor(
705	2 (-(2 -(math.log2(closest_power_of_2) - 3))),
706	dtype=torch.float32,
707	)
708	powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
709	slopes = torch.pow(base, powers)
710
711	if closest_power_of_2 != total_num_heads:
712	extra_base = torch.tensor(
713	2 (-(2 -(math.log2(2 * closest_power_of_2) - 3))),
714	dtype=torch.float32,
715	)
716	num_remaining_heads = min(
717	closest_power_of_2, total_num_heads - closest_power_of_2
718	)
719	extra_powers = torch.arange(
720	start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32
721	)
722	slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
723	return slopes
724
725	alibi_slopes = _get_alibi_slopes(num_heads).to(device)
726
727	MAX_SEQ_LEN = 1024
728	MAX_CTX_LEN = 1024
729	BS = 10
730	cache_size = 640
731	max_block_per_request = 64
732	query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
733	ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
734	seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
735	num_kv_heads = num_heads // num_queries_per_kv
736
737	num_tokens = sum(query_lens)
738	query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
739	query.uniform_(-1e-3, 1e-3)
740	output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
741
742	kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
743	kv.uniform_(-1e-3, 1e-3)
744	key, value = kv.unbind(dim=1)
745	if kv_cache_dtype == "auto":
746	cache_dtype = dtype
747	else:
748	cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
749	k_cache = torch.zeros(
750	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
751	)
752	v_cache = torch.zeros(
753	cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
754	)
755	k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
756	v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
757	values = torch.arange(0, cache_size, dtype=torch.int32)
758	values = values[torch.randperm(cache_size)]
759	block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
760	b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
761	b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
762	b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
763	max_input_len = MAX_SEQ_LEN
764	# copy kv to cache
765	b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
766	torch.int32
767	)
768	for i in range(BS):
769	for j in range(query_lens[i]):
770	k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] + j])
771	v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] + b_ctx_len[i] + j])
772	cur_ctx = 0
773	block_id = 0
774	while cur_ctx < b_ctx_len[i]:
775	start_loc = b_seq_start_loc[i] + cur_ctx
776	if cur_ctx + block_size > b_ctx_len[i]:
777	end_loc = b_seq_start_loc[i] + b_ctx_len[i]
778	else:
779	end_loc = start_loc + block_size
780	start_slot = block_table[i, block_id] * block_size
781	end_slot = start_slot + end_loc - start_loc
782	k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
783	key[start_loc:end_loc]
784	)
785	v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
786	value[start_loc:end_loc]
787	)
788	cur_ctx += block_size
789	block_id += 1
790	# transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
791	# to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
792	k_cache = (
793	k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
794	.permute(0, 2, 3, 1, 4)
795	.contiguous()
796	)
797	# transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
798	# to V_cache[num_blocks, num_kv_heads, head_size, block_size]
799	v_cache = (
800	v_cache.view(-1, block_size, num_kv_heads, head_size)
801	.permute(0, 2, 3, 1)
802	.contiguous()
803	)
804	k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
805
806	# Warm up the Triton kernel by calling it once before actually measuring
807	# generation time
808	op(
809	query,
810	k,
811	v,
812	output,
813	kv_cache_dtype,
814	k_cache,
815	v_cache,
816	block_table,
817	b_start_loc,
818	b_seq_len,
819	MAX_CTX_LEN,
820	max_input_len,
821	k_scale,
822	v_scale,
823	alibi_slopes=alibi_slopes,
824	)
825	torch.accelerator.synchronize()
826	start_time = time.time()
827	op(
828	query,
829	k,
830	v,
831	output,
832	kv_cache_dtype,
833	k_cache,
834	v_cache,
835	block_table,
836	b_start_loc,
837	b_seq_len,
838	MAX_CTX_LEN,
839	max_input_len,
840	k_scale,
841	v_scale,
842	alibi_slopes=alibi_slopes,
843	)
844	torch.accelerator.synchronize()
845	end_time = time.time()
846	print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")
847	scale = float(1.0 / (head_size**0.5))
848
849	# Prepare query, key, value for SDPA
850	# Expand key and value for GQA/MQA to match query heads
851	key_expanded = key[:, :, None, :].expand(
852	key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
853	)
854	value_expanded = value[:, :, None, :].expand(
855	value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
856	)
857
858	output_ref = torch.empty_like(output)
859
860	torch.accelerator.synchronize()
861	start_time = time.time()
862
863	query_start = 0
864	key_start = 0
865	for i, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):
866	query_end = query_start + query_len
867	key_end = key_start + seq_len
868
869	# Get query, key, value for this sequence
870	q = query[query_start:query_end] # [query_len, num_heads, head_size]
871	k = key_expanded[
872	key_start:key_end
873	] # [seq_len, num_kv_heads, num_queries_per_kv, head_size]
874	v = value_expanded[
875	key_start:key_end
876	] # [seq_len, num_kv_heads, num_queries_per_kv, head_size]
877
878	# Reshape for SDPA: (batch=1, num_heads, seq_len, head_size)
879	q_sdpa = q.view(query_len, num_kv_heads, num_queries_per_kv, head_size)
880	q_sdpa = (
881	q_sdpa.permute(1, 2, 0, 3)
882	.reshape(1, num_heads, query_len, head_size)
883	.contiguous()
884	)
885
886	k_sdpa = (
887	k.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size).contiguous()
888	)
889	v_sdpa = (
890	v.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size).contiguous()
891	)
892
893	# Create ALiBi causal mask for this sequence using utility function
894	alibi_mask = create_alibi_causal_mask(
895	query_len, seq_len, alibi_slopes, device, dtype
896	)
897
898	# Compute attention
899	out = F.scaled_dot_product_attention(
900	q_sdpa,
901	k_sdpa,
902	v_sdpa,
903	attn_mask=alibi_mask,
904	dropout_p=0.0,
905	scale=scale,
906	)
907
908	# Reshape output back to [query_len, num_heads, head_size]
909	out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
910	output_ref[query_start:query_end].copy_(out)
911
912	query_start = query_end
913	key_start = key_end
914
915	torch.accelerator.synchronize()
916	end_time = time.time()
917	print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")
918	atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-6
919	torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
920
921
922	# These tests are optional to only run when explicitly invoked
923	#
924	# pytest -v -s --optional \
925	# tests/kernels/test_prefix_prefill.py::test_contexted_kv_attention_f32
926	#
927	# These tests are useful to test model dtype float32 on Turing devices.
928	# We skip them to not increase the time when running tests on CI
929	@pytest.mark.optional
930	@pytest.mark.parametrize("num_heads", NUM_HEADS)
931	@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
932	@pytest.mark.parametrize("head_size", HEAD_SIZES)
933	@pytest.mark.parametrize("dtype", [torch.float32])
934	@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
935	@pytest.mark.parametrize("device", CUDA_DEVICES)
936	@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
937	@pytest.mark.parametrize("op", OPS)
938	@torch.inference_mode()
939	def test_contexted_kv_attention_f32(
940	num_heads: int,
941	num_queries_per_kv: int,
942	head_size: int,
943	sliding_window: int,
944	dtype: torch.dtype,
945	kv_cache_dtype: str,
946	device: str,
947	op: Callable,
948	) -> None:
949	test_contexted_kv_attention(
950	num_heads,
951	num_queries_per_kv,
952	head_size,
953	sliding_window,
954	dtype,
955	kv_cache_dtype,
956	device,
957	op,
958	)
959
960
961	@pytest.mark.optional
962	@pytest.mark.parametrize("num_heads", NUM_HEADS)
963	@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
964	@pytest.mark.parametrize("head_size", HEAD_SIZES)
965	@pytest.mark.parametrize("dtype", [torch.float32])
966	@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
967	@pytest.mark.parametrize("device", CUDA_DEVICES)
968	@pytest.mark.parametrize("op", OPS)
969	@torch.inference_mode()
970	def test_contexted_kv_attention_alibi_f32(
971	num_heads: int,
972	num_queries_per_kv: int,
973	head_size: int,
974	dtype: torch.dtype,
975	kv_cache_dtype: str,
976	device: str,
977	op: Callable,
978	) -> None:
979	test_contexted_kv_attention_alibi(
980	num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op
981	)
982
983
984	# Hybrid mamba + full-attention models get a non-power-of-2 attention page.
985	NONSTANDARD_BLOCK_SIZE_SHAPES = [
986	(64, 1, 128, 544),
987	(8, 4, 256, 1040),
988	(8, 4, 256, 1056),
989	]
990
991
992	@pytest.mark.parametrize(
993	"num_heads,num_queries_per_kv,head_size,block_size", NONSTANDARD_BLOCK_SIZE_SHAPES
994	)
995	@pytest.mark.parametrize("dtype", DTYPES)
996	@pytest.mark.parametrize("device", CUDA_DEVICES)
997	@pytest.mark.parametrize("op", OPS)
998	@torch.inference_mode()
999	def test_qwen3_nonstandard_block_size(
1000	num_heads: int,
1001	num_queries_per_kv: int,
1002	head_size: int,
1003	block_size: int,
1004	dtype: torch.dtype,
1005	device: str,
1006	op: Callable,
1007	) -> None:
1008	"""Non-power-of-2 pages must match, even when a tile straddles a page."""
1009	if not current_platform.is_rocm():
1010	pytest.skip("Non-power-of-2 block sizes are only exercised on ROCm CI.")
1011
1012	test_contexted_kv_attention(
1013	num_heads=num_heads,
1014	num_queries_per_kv=num_queries_per_kv,
1015	head_size=head_size,
1016	block_size=block_size,
1017	sliding_window=0,
1018	dtype=dtype,
1019	kv_cache_dtype="auto",
1020	device=device,
1021	op=op,
1022	)