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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,313 lines 80.1 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import functools
5	from types import SimpleNamespace
6
7	import pytest
8	import torch
9
10	from vllm.platforms import current_platform
11	from vllm.utils.torch_utils import set_random_seed
12
13	pytestmark = pytest.mark.skipif(
14	not current_platform.is_rocm(), reason="Only used by ROCm"
15	)
16
17
18	def _on_split_decode_arch() -> bool:
19	if not current_platform.is_rocm():
20	return False
21	try:
22	from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950
23
24	return bool(_ON_GFX942 or _ON_GFX950)
25	except Exception:
26	return False
27
28
29	def _on_gfx950() -> bool:
30	if not current_platform.is_rocm():
31	return False
32	try:
33	from vllm.platforms.rocm import _ON_GFX950
34
35	return _ON_GFX950
36	except ImportError:
37	return False
38
39
40	# The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other
41	# architectures take the fallback decode kernel, so its tests are skipped there.
42	requires_split_decode_arch = pytest.mark.skipif(
43	not _on_split_decode_arch(),
44	reason="split-K decode kernel is only tuned for AMD gfx942/gfx950",
45	)
46	requires_gfx950 = pytest.mark.skipif(
47	not _on_gfx950(),
48	reason="optimized sparse decode partial is gfx950-only",
49	)
50
51	NOPE_HEAD_DIM = 448
52	ROPE_HEAD_DIM = 64
53	HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
54
55
56	@pytest.fixture
57	def enable_aiter_mqa(monkeypatch):
58	from vllm._aiter_ops import rocm_aiter_ops
59
60	with monkeypatch.context() as context:
61	context.setenv("VLLM_ROCM_USE_AITER", "1")
62	rocm_aiter_ops.refresh_env_variables()
63	yield
64	rocm_aiter_ops.refresh_env_variables()
65
66
67	@requires_gfx950
68	@pytest.mark.parametrize("column_scales", [False, True], ids=["vector", "column"])
69	@torch.inference_mode()
70	def test_unpaged_mqa_logits_preserves_intervals_and_scale_layout(
71	column_scales, monkeypatch, enable_aiter_mqa
72	) -> None:
73	"""The vLLM AITER dispatch preserves ragged masking and per-token scales."""
74	from vllm._aiter_ops import rocm_aiter_ops
75	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
76
77	assert rocm_aiter_ops.is_enabled(), "AITER is required on gfx950"
78	assert mod.mqa_logits_module() is not None, "AITER MQA logits are required"
79	torch.manual_seed(42)
80	device = torch.device("cuda")
81	dtype = current_platform.fp8_dtype()
82	q = torch.randn(7, 32, 128, device=device).to(dtype)
83	k = torch.randn(259, 128, device=device).to(dtype)
84	scales = torch.linspace(0.25, 1.25, k.shape[0], device=device)
85	if column_scales:
86	scales = scales[:, None]
87	weights = torch.rand(q.shape[:2], device=device) / q.shape[1]
88	starts = torch.tensor([0, 3, 128, 5, 0, 257, 259], device=device, dtype=torch.int32)
89	ends = torch.tensor(
90	[1, 129, 259, 5, 259, 259, 259], device=device, dtype=torch.int32
91	)
92
93	module = mod.mqa_logits_module()
94	implementation = module.fp8_mqa_logits
95	calls = 0
96
97	def traced_implementation(args, *kwargs):
98	nonlocal calls
99	calls += 1
100	return implementation(args, *kwargs)
101
102	monkeypatch.setattr(module, "fp8_mqa_logits", traced_implementation)
103	actual = mod.rocm_fp8_mqa_logits(q, (k, scales), weights, starts, ends)
104	assert calls == 1
105	score = torch.einsum("mhd,nd->mhn", q.float(), k.float()) * scales.reshape(-1)
106	expected = (score.relu() * weights[..., None]).sum(dim=1)
107	positions = torch.arange(k.shape[0], device=device)
108	valid = (positions >= starts[:, None]) & (positions < ends[:, None])
109	assert actual.shape == expected.shape
110	assert actual.dtype == torch.float32
111	assert torch.equal(torch.isneginf(actual), ~valid)
112	torch.testing.assert_close(actual[valid], expected[valid], rtol=1e-3, atol=1e-3)
113
114
115	def _ref_global_topk_ragged(
116	topk_indices: torch.Tensor,
117	token_to_req_indices: torch.Tensor,
118	block_table: torch.Tensor,
119	block_size: int,
120	is_valid_token: torch.Tensor,
121	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
122	topk = topk_indices.reshape(topk_indices.shape[0], -1)
123	valid = (topk >= 0) & is_valid_token[:, None]
124	lens = valid.sum(dim=1, dtype=torch.int32)
125	indptr = torch.zeros(lens.shape[0] + 1, dtype=torch.int32, device=topk.device)
126	torch.cumsum(lens, dim=0, out=indptr[1:])
127
128	safe_topk = torch.clamp(topk, min=0)
129	block_indices = safe_topk // block_size
130	block_offsets = safe_topk % block_size
131	req_indices = token_to_req_indices[:, None].expand_as(topk)
132	slot_ids = block_table[req_indices, block_indices] * block_size + block_offsets
133
134	offsets = torch.arange(topk.shape[1], dtype=torch.int32, device=topk.device)
135	positions = indptr[:-1, None] + offsets[None, :]
136	return slot_ids[valid], positions[valid].to(torch.long), indptr, lens
137
138
139	def _ref_sparse_prefill_ragged(
140	q: torch.Tensor,
141	kv: torch.Tensor,
142	rows: list[list[int]],
143	scale: float,
144	attn_sink: torch.Tensor \| None,
145	) -> torch.Tensor:
146	q_f32 = q.float()
147	kv_f32 = kv.float()
148	out = torch.empty_like(q_f32)
149
150	for query_idx in range(q.shape[0]):
151	row_indices = rows[query_idx]
152	for head_idx in range(q.shape[1]):
153	if row_indices:
154	selected_kv = kv_f32[row_indices]
155	scores = torch.mv(selected_kv, q_f32[query_idx, head_idx]) * scale
156	if attn_sink is not None:
157	scores_with_sink = torch.cat(
158	[scores, attn_sink[head_idx].float().reshape(1)]
159	)
160	probs = torch.softmax(scores_with_sink, dim=0)[:-1]
161	else:
162	probs = torch.softmax(scores, dim=0)
163	out[query_idx, head_idx] = torch.sum(
164	probs[:, None] * selected_kv, dim=0
165	)
166	else:
167	out[query_idx, head_idx] = 0
168	return out.to(torch.bfloat16)
169
170
171	def _pack_fp8_ds_mla_cache(
172	kv: torch.Tensor, block_size: int, use_fnuz: bool
173	) -> torch.Tensor:
174	assert kv.shape[-1] == HEAD_DIM
175	from vllm.models.deepseek_v4.common.ops.cache_utils import (
176	quantize_and_insert_k_cache,
177	)
178
179	num_tokens = kv.shape[0]
180	num_blocks = (num_tokens + block_size - 1) // block_size
181	cache = torch.zeros(
182	(num_blocks, block_size, 584),
183	dtype=torch.uint8,
184	device=kv.device,
185	)
186	slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=kv.device)
187	quantize_and_insert_k_cache(
188	kv,
189	cache,
190	slot_mapping,
191	block_size=block_size,
192	use_fnuz=use_fnuz,
193	)
194	return cache
195
196
197	def _poison_fp8_ds_mla_cache_row(
198	cache: torch.Tensor, block_size: int, slot: int = 0
199	) -> None:
200	flat = cache.flatten()
201	block_idx = slot // block_size
202	pos = slot % block_size
203	block_base = block_idx * cache.stride(0)
204	token_base = block_base + pos * 576
205	scale_base = block_base + block_size * 576 + pos * 8
206	flat[token_base] = 0x7F
207	flat[scale_base : scale_base + 7] = 255
208	flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = float(
209	"nan"
210	)
211
212
213	def _read_fp8_ds_mla_cache_rows(
214	cache: torch.Tensor,
215	slots: torch.Tensor,
216	block_size: int,
217	use_fnuz: bool,
218	) -> torch.Tensor:
219	cache_flat = cache.view(torch.uint8).flatten()
220	block_idx = slots // block_size
221	pos = slots % block_size
222	block_base = block_idx * cache.stride(0)
223	token_base = block_base + pos * 576
224	scale_base = block_base + block_size * 576 + pos * 8
225
226	fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn
227	nope_offsets = torch.arange(NOPE_HEAD_DIM, device=cache.device)
228	nope_u8 = cache_flat[token_base[:, None] + nope_offsets]
229	nope = nope_u8.view(fp8_dtype).to(torch.float32)
230	scale_offsets = torch.arange(7, device=cache.device)
231	scales = torch.exp2(
232	cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0
233	)
234	nope = nope * scales.repeat_interleave(64, dim=1)
235	rope_offsets = torch.arange(ROPE_HEAD_DIM * 2, device=cache.device)
236	rope_u8 = cache_flat[token_base[:, None] + NOPE_HEAD_DIM + rope_offsets]
237	rope = rope_u8.contiguous().view(torch.bfloat16).to(torch.float32)
238	return torch.cat([nope, rope], dim=1)
239
240
241	def _ref_sparse_decode_ragged(
242	q: torch.Tensor,
243	main_cache: torch.Tensor,
244	main_rows: list[list[int]],
245	scale: float,
246	attn_sink: torch.Tensor \| None,
247	block_size: int,
248	extra_cache: torch.Tensor \| None = None,
249	extra_rows: list[list[int]] \| None = None,
250	main_use_fnuz: bool = False,
251	extra_use_fnuz: bool = False,
252	extra_block_size: int \| None = None,
253	) -> torch.Tensor:
254	q_f32 = q.float()
255	out = torch.empty_like(q_f32)
256
257	for query_idx in range(q.shape[0]):
258	row_kv = []
259	if main_rows[query_idx]:
260	main_slots = torch.tensor(
261	main_rows[query_idx], dtype=torch.int64, device=q.device
262	)
263	row_kv.append(
264	_read_fp8_ds_mla_cache_rows(
265	main_cache, main_slots, block_size, main_use_fnuz
266	)
267	)
268	if extra_cache is not None and extra_rows is not None and extra_rows[query_idx]:
269	extra_slots = torch.tensor(
270	extra_rows[query_idx], dtype=torch.int64, device=q.device
271	)
272	row_kv.append(
273	_read_fp8_ds_mla_cache_rows(
274	extra_cache,
275	extra_slots,
276	extra_block_size or block_size,
277	extra_use_fnuz,
278	)
279	)
280
281	if not row_kv:
282	out[query_idx] = 0
283	continue
284	kv = torch.cat(row_kv)
285	for head_idx in range(q.shape[1]):
286	scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
287	if attn_sink is not None:
288	scores_with_sink = torch.cat(
289	[scores, attn_sink[head_idx].float().reshape(1)]
290	)
291	probs = torch.softmax(scores_with_sink, dim=0)[:-1]
292	else:
293	probs = torch.softmax(scores, dim=0)
294	out[query_idx, head_idx] = torch.sum(probs[:, None] * kv, dim=0)
295	return out.to(torch.bfloat16)
296
297
298	def _ragged_from_rows(
299	rows: list[list[int]], device: torch.device
300	) -> tuple[torch.Tensor, torch.Tensor]:
301	"""Flatten per-query slot lists into ragged (indices, indptr) tensors."""
302	flat = [slot for row in rows for slot in row]
303	indptr = [0]
304	for row in rows:
305	indptr.append(indptr[-1] + len(row))
306	return (
307	torch.tensor(flat, dtype=torch.int32, device=device),
308	torch.tensor(indptr, dtype=torch.int32, device=device),
309	)
310
311
312	def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
313	ends = indptr.cpu().tolist()
314	values = indices[: ends[-1]].cpu().tolist()
315	return [values[start:end] for start, end in zip(ends, ends[1:])]
316
317
318	def _launch_sparse_decode_reduce(
319	part_m: torch.Tensor,
320	part_l: torch.Tensor,
321	part_acc: torch.Tensor,
322	adaptive_splits: bool,
323	positions: torch.Tensor \| None = None,
324	cos_sin_cache: torch.Tensor \| None = None,
325	out_dtype: torch.dtype = torch.bfloat16,
326	out_scale: torch.Tensor \| None = None,
327	) -> torch.Tensor:
328	"""Run the reduce kernel, with the inverse-RoPE epilogue if given positions.
329
330	Passing ``out_scale`` selects the MXFP8 epilogue (``out_dtype`` e4m3).
331	"""
332	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
333
334	num_queries, num_splits, num_heads = part_m.shape
335	out = torch.empty(
336	(num_queries, num_heads, HEAD_DIM),
337	dtype=out_dtype,
338	device=part_m.device,
339	)
340	attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device)
341	mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](
342	part_m,
343	part_l,
344	part_acc,
345	attn_sink,
346	out,
347	positions,
348	cos_sin_cache,
349	out_scale,
350	out.stride(0),
351	out.stride(1),
352	out_scale.stride(0) if out_scale is not None else 0,
353	HEAD_DIM // 32,
354	part_m.stride(0),
355	part_m.stride(1),
356	part_acc.stride(0),
357	part_acc.stride(1),
358	part_acc.stride(2),
359	cos_sin_cache.stride(0) if cos_sin_cache is not None else 0,
360	num_heads,
361	HAS_ATTN_SINK=False,
362	ADAPTIVE_SPLITS=adaptive_splits,
363	COMB_DIM=HEAD_DIM,
364	BLOCK_H=1,
365	NUM_SPLITS=num_splits,
366	SPLITS_PAD=1 << (num_splits - 1).bit_length(),
367	FUSE_INV_ROPE=positions is not None,
368	NOPE=NOPE_HEAD_DIM,
369	HALF=ROPE_HEAD_DIM // 2,
370	QUANT_OUT=out_scale is not None,
371	num_warps=4,
372	)
373	return out
374
375
376	@torch.inference_mode()
377	def test_compute_global_topk_ragged_indices_and_indptr() -> None:
378	from vllm.models.deepseek_v4.amd.rocm import (
379	compute_global_topk_ragged_indices_and_indptr,
380	)
381
382	device = torch.device("cuda")
383	block_size = 4
384	topk_indices = torch.tensor(
385	[
386	[0, 3, 4, -1],
387	[5, 8, -1, -1],
388	[2, 7, 9, -1],
389	],
390	dtype=torch.int32,
391	device=device,
392	)
393	token_to_req_indices = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
394	block_table = torch.tensor(
395	[
396	[10, 11, 12],
397	[20, 21, 22],
398	],
399	dtype=torch.int32,
400	device=device,
401	)
402	is_valid_token = torch.tensor([True, False, True], dtype=torch.bool, device=device)
403
404	actual_ragged, actual_indptr, actual_lens = (
405	compute_global_topk_ragged_indices_and_indptr(
406	topk_indices,
407	token_to_req_indices,
408	block_table,
409	block_size,
410	is_valid_token,
411	)
412	)
413	expected_values, expected_positions, expected_indptr, expected_lens = (
414	_ref_global_topk_ragged(
415	topk_indices,
416	token_to_req_indices,
417	block_table,
418	block_size,
419	is_valid_token,
420	)
421	)
422
423	torch.testing.assert_close(actual_ragged[expected_positions], expected_values)
424	torch.testing.assert_close(actual_indptr, expected_indptr)
425	torch.testing.assert_close(actual_lens, expected_lens)
426
427
428	@pytest.mark.parametrize("width", [16, 640, 2176])
429	@pytest.mark.parametrize("num_queries", [1, 70, 1000])
430	@torch.inference_mode()
431	def test_build_ragged_indices_from_dense_drops_invalid_entries(
432	width: int, num_queries: int
433	) -> None:
434	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
435	build_ragged_indices_from_dense,
436	)
437
438	device = torch.device("cuda")
439	set_random_seed(width + num_queries)
440	num_rows = 5000
441	indices = torch.randint(0, num_rows + 100, (num_queries, width))
442	indices[torch.rand(num_queries, width) < 0.3] = -1
443	lengths = torch.randint(0, width + 8, (num_queries,))
444
445	flat, indptr = build_ragged_indices_from_dense(
446	indices.to(device, torch.int32),
447	lengths.to(device, torch.int32),
448	num_rows=num_rows,
449	)
450
451	# Each row keeps its in-range entries within its length, in order.
452	expected = [
453	[x for x in row[:n] if 0 <= x < num_rows]
454	for row, n in zip(indices.tolist(), lengths.tolist())
455	]
456	assert _rows_from_ragged(flat, indptr) == expected
457
458
459	@torch.inference_mode()
460	def test_combine_topk_swa_indices_adds_image_visibility() -> None:
461	from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
462
463	device = torch.device("cuda")
464	num_tokens = 8
465	topk_indices = torch.full((num_tokens, 1), -1, dtype=torch.int32, device=device)
466	query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
467	seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
468	gather_lens = seq_lens.clone()
469	left_visible = torch.tensor(
470	[0, 0, 0, 1, 2, 3, 4, 0], dtype=torch.int32, device=device
471	)
472	right_visible = torch.tensor(
473	[0, 0, 4, 3, 2, 1, 0, 0], dtype=torch.int32, device=device
474	)
475
476	indices, lens = combine_topk_swa_indices(
477	topk_indices,
478	query_start_loc,
479	seq_lens,
480	gather_lens,
481	window_size=4,
482	compress_ratio=1,
483	topk=0,
484	M=16,
485	N=0,
486	max_image_tokens=5,
487	left_visible=left_visible,
488	right_visible=right_visible,
489	)
490
491	expected_rows = [
492	[0],
493	[0, 1],
494	[0, 1, 2, 3, 4, 5, 6],
495	[0, 1, 2, 3, 4, 5, 6],
496	[1, 2, 3, 4, 5, 6],
497	[2, 3, 4, 5, 6],
498	[2, 3, 4, 5, 6],
499	[4, 5, 6, 7],
500	]
501	for token_idx, expected in enumerate(expected_rows):
502	actual = indices[token_idx, : lens[token_idx]].cpu().tolist()
503	assert actual == expected
504
505
506	@torch.inference_mode()
507	def test_combine_topk_swa_indices_apc_hit_inside_image() -> None:
508	from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
509
510	device = torch.device("cuda")
511	indices, lens = combine_topk_swa_indices(
512	torch.full((2, 1), -1, dtype=torch.int32, device=device),
513	torch.tensor([0, 2], dtype=torch.int32, device=device),
514	torch.tensor([10], dtype=torch.int32, device=device),
515	# Only positions [5, 10) exist in the gathered SWA workspace.
516	torch.tensor([5], dtype=torch.int32, device=device),
517	window_size=4,
518	compress_ratio=1,
519	topk=0,
520	M=10,
521	N=0,
522	max_image_tokens=10,
523	left_visible=torch.tensor([8, 9], dtype=torch.int32, device=device),
524	# Deliberately extends beyond seq_len to exercise the upper clamp too.
525	right_visible=torch.tensor([5, 5], dtype=torch.int32, device=device),
526	)
527
528	assert lens.cpu().tolist() == [5, 5]
529	assert indices[:, :5].cpu().tolist() == [list(range(5)), list(range(5))]
530
531
532	@torch.inference_mode()
533	def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None:
534	from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
535
536	device = torch.device("cuda")
537	indices, lens = combine_topk_swa_indices(
538	torch.full((1, 1), -1, dtype=torch.int32, device=device),
539	torch.tensor([0, 1], dtype=torch.int32, device=device),
540	torch.tensor([1], dtype=torch.int32, device=device),
541	torch.tensor([1], dtype=torch.int32, device=device),
542	window_size=120,
543	compress_ratio=1,
544	topk=0,
545	M=256,
546	N=0,
547	max_image_tokens=16,
548	)
549
550	assert indices.shape == (1, 256)
551	assert lens.item() == 1
552
553
554	def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None:
555	from vllm.models.deepseek_v4.amd import rocm as mod
556
557	monkeypatch.setattr(mod, "_ON_GFX950", True)
558	assert mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True)
559	assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", True, True)
560	assert not mod._trust_dsv4_extra_cache_nan_free("bfloat16", False, True)
561	assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, False)
562
563	monkeypatch.setattr(mod, "_ON_GFX950", False)
564	assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True)
565
566
567	@torch.inference_mode()
568	def test_sparse_attn_prefill_ragged_kernel() -> None:
569	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
570	_rocm_sparse_attn_prefill_ragged_triton,
571	)
572
573	device = torch.device("cuda")
574	torch.manual_seed(0)
575	q = torch.randn(3, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
576	kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
577	indices = torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device)
578	indptr = torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device)
579	attn_sink = torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device)
580	scale = HEAD_DIM**-0.5
581
582	actual = _rocm_sparse_attn_prefill_ragged_triton(
583	q=q,
584	kv=kv,
585	indices=indices,
586	indptr=indptr,
587	scale=scale,
588	attn_sink=attn_sink,
589	nope_head_dim=NOPE_HEAD_DIM,
590	rope_head_dim=ROPE_HEAD_DIM,
591	)
592	expected = _ref_sparse_prefill_ragged(
593	q, kv, [[0, 2], [1, 3, 4], []], scale, attn_sink
594	)
595
596	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
597
598
599	def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict:
600	device = torch.device("cuda")
601	set_random_seed(7)
602	head_dim = nope_dim + rope_dim
603	return dict(
604	q=torch.randn(3, 3, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
605	kv=torch.randn(5, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
606	indices=torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device),
607	indptr=torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device),
608	scale=head_dim**-0.5,
609	attn_sink=torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device),
610	nope_head_dim=nope_dim,
611	rope_head_dim=rope_dim,
612	)
613
614
615	# 448+64 is DeepSeek V4/V4.1, 512+64 is V3.2, and 256+0 stands in for the
616	# NoPE-only layouts where the destination spans the whole head dim.
617	SPARSE_PREFILL_DIMS = [(NOPE_HEAD_DIM, ROPE_HEAD_DIM), (512, 64), (256, 0)]
618
619
620	@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
621	@torch.inference_mode()
622	def test_sparse_attn_prefill_ragged_out_matches_allocated(
623	nope_dim: int, rope_dim: int
624	) -> None:
625	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
626	_rocm_sparse_attn_prefill_ragged_triton,
627	)
628
629	kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
630	allocated = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)
631
632	q = kwargs["q"]
633	dest = torch.empty(q.shape[0], q.shape[1], nope_dim, dtype=q.dtype, device=q.device)
634	returned = _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
635
636	assert returned is dest
637	torch.testing.assert_close(dest, allocated[..., :nope_dim], atol=0.0, rtol=0.0)
638
639
640	@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
641	@torch.inference_mode()
642	def test_sparse_attn_prefill_ragged_out_view_leaves_neighbors_untouched(
643	nope_dim: int, rope_dim: int
644	) -> None:
645	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
646	_rocm_sparse_attn_prefill_ragged_triton,
647	)
648
649	kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
650	q = kwargs["q"]
651	num_queries, num_heads = q.shape[0], q.shape[1]
652
653	buffer = torch.full(
654	(num_queries + 4, num_heads + 2, nope_dim),
655	-7.0,
656	dtype=q.dtype,
657	device=q.device,
658	)
659	dest = buffer[2 : 2 + num_queries, :num_heads]
660	assert not dest.is_contiguous()
661	_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
662
663	expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :nope_dim]
664	torch.testing.assert_close(dest, expected, atol=0.0, rtol=0.0)
665
666	untouched = torch.ones_like(buffer, dtype=torch.bool)
667	untouched[2 : 2 + num_queries, :num_heads] = False
668	assert torch.all(buffer[untouched] == -7.0)
669
670
671	@pytest.mark.parametrize("out_dtype", [torch.float32, torch.float16])
672	@torch.inference_mode()
673	def test_sparse_attn_prefill_ragged_out_casts_to_dest_dtype(
674	out_dtype: torch.dtype,
675	) -> None:
676	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
677	_rocm_sparse_attn_prefill_ragged_triton,
678	)
679
680	kwargs = _sparse_prefill_ragged_inputs(512, 64)
681	q = kwargs["q"]
682	dest = torch.empty(q.shape[0], q.shape[1], 512, dtype=out_dtype, device=q.device)
683	_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
684
685	expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :512]
686	assert dest.dtype == out_dtype
687	torch.testing.assert_close(dest.float(), expected.float(), atol=2e-2, rtol=2e-2)
688
689
690	@torch.inference_mode()
691	def test_sparse_attn_prefill_head_dim_wide_out_keeps_rope_rows() -> None:
692	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
693	_rocm_sparse_attn_prefill_ragged_triton,
694	)
695
696	kwargs = _sparse_prefill_ragged_inputs(NOPE_HEAD_DIM, ROPE_HEAD_DIM)
697	q = kwargs["q"]
698	dest = torch.empty_like(q)
699	_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
700
701	expected = _ref_sparse_prefill_ragged(
702	q, kwargs["kv"], [[0, 2], [1, 3, 4], []], kwargs["scale"], kwargs["attn_sink"]
703	)
704	torch.testing.assert_close(dest, expected, atol=2e-2, rtol=2e-2)
705
706
707	@pytest.mark.parametrize(
708	("num_queries", "on_gfx950", "expected"),
709	[(1023, True, False), (1024, True, True), (1024, False, False)],
710	)
711	def test_aiter_sparse_prefill_opus_selection(
712	num_queries: int, on_gfx950: bool, expected: bool
713	) -> None:
714	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
715	_can_use_aiter_sparse_prefill_opus,
716	)
717
718	device = torch.device("cuda")
719	q = torch.empty(num_queries, 16, HEAD_DIM, dtype=torch.bfloat16, device=device)
720	kv = torch.empty(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device)
721	attn_sink = torch.empty(16, dtype=torch.float32, device=device)
722	output = torch.empty_like(q)
723
724	assert (
725	_can_use_aiter_sparse_prefill_opus(
726	q, kv, attn_sink, output, on_gfx950=on_gfx950
727	)
728	is expected
729	)
730
731
732	def test_aiter_sparse_prefill_opus_selection_rejects_incompatible_inputs() -> None:
733	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
734	_can_use_aiter_sparse_prefill_opus,
735	)
736
737	device = torch.device("cuda")
738	q = torch.empty(1024, 16, HEAD_DIM, dtype=torch.bfloat16, device=device)
739	kv = torch.empty(1024, HEAD_DIM, dtype=torch.bfloat16, device=device)
740	attn_sink = torch.empty(16, dtype=torch.float32, device=device)
741	output = torch.empty_like(q)
742
743	assert _can_use_aiter_sparse_prefill_opus(q, kv, attn_sink, output, on_gfx950=True)
744	assert not _can_use_aiter_sparse_prefill_opus(
745	q,
746	kv,
747	attn_sink,
748	torch.empty(16, 1024, HEAD_DIM, dtype=q.dtype, device=device).transpose(0, 1),
749	on_gfx950=True,
750	)
751	assert not _can_use_aiter_sparse_prefill_opus(
752	q, kv, attn_sink[:-1], output, on_gfx950=True
753	)
754	assert not _can_use_aiter_sparse_prefill_opus(
755	q.cpu(), kv.cpu(), attn_sink.cpu(), output.cpu(), on_gfx950=True
756	)
757
758
759	def test_sparse_attn_prefill_aiter_opus_routing(monkeypatch) -> None:
760	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
761
762	q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
763	kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
764	indices = torch.tensor([[0], [1]], dtype=torch.int32)
765	topk_length = torch.ones(2, dtype=torch.int32)
766	attn_sink = torch.empty(1, dtype=torch.float32)
767	output = torch.empty_like(q)
768	opus_calls = 0
769
770	def fake_opus(*args, out):
771	nonlocal opus_calls
772	opus_calls += 1
773	assert args[2].dtype == torch.int32
774	assert args[3].dtype == torch.int32
775	assert args[5].numel() == 0
776	assert torch.count_nonzero(args[6]) == 0
777	out.zero_()
778	return out
779
780	monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
781	monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: fake_opus)
782	monkeypatch.setattr(
783	mod,
784	"build_ragged_indices_from_dense",
785	lambda args, *kwargs: (
786	torch.tensor([0, 1], dtype=torch.int32),
787	torch.tensor([0, 1, 2], dtype=torch.int32),
788	),
789	)
790	monkeypatch.setattr(
791	mod,
792	"_rocm_sparse_attn_prefill_triton",
793	lambda args, *kwargs: pytest.fail("unexpected dense Triton fallback"),
794	)
795	monkeypatch.setattr(
796	mod,
797	"_rocm_sparse_attn_prefill_ragged_triton",
798	lambda args, *kwargs: pytest.fail("unexpected ragged Triton fallback"),
799	)
800
801	mod.rocm_sparse_attn_prefill(
802	q=q,
803	kv=kv,
804	indices=indices,
805	topk_length=topk_length,
806	scale=HEAD_DIM**-0.5,
807	head_dim=HEAD_DIM,
808	nope_head_dim=NOPE_HEAD_DIM,
809	rope_head_dim=ROPE_HEAD_DIM,
810	attn_sink=attn_sink,
811	output=output,
812	)
813
814	assert opus_calls == 1
815	assert torch.count_nonzero(output) == 0
816
817
818	def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> None:
819	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
820
821	q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
822	kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
823	indices = torch.tensor([[0], [1]], dtype=torch.int32)
824	topk_length = torch.ones(2, dtype=torch.int32)
825	attn_sink = torch.empty(1, dtype=torch.float32)
826	output = torch.empty_like(q)
827	dense_fallback_calls = 0
828
829	def fake_dense_fallback(args, out, *kwargs):
830	nonlocal dense_fallback_calls
831	dense_fallback_calls += 1
832	out.zero_()
833	return out
834
835	monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
836	monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
837	monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_triton", fake_dense_fallback)
838	monkeypatch.setattr(
839	mod,
840	"_rocm_sparse_attn_prefill_ragged_triton",
841	lambda args, *kwargs: pytest.fail("unexpected ragged Triton fallback"),
842	)
843
844	mod.rocm_sparse_attn_prefill(
845	q=q,
846	kv=kv,
847	indices=indices,
848	topk_length=topk_length,
849	scale=HEAD_DIM**-0.5,
850	head_dim=HEAD_DIM,
851	nope_head_dim=NOPE_HEAD_DIM,
852	rope_head_dim=ROPE_HEAD_DIM,
853	attn_sink=attn_sink,
854	output=output,
855	)
856
857	assert dense_fallback_calls == 1
858	assert torch.count_nonzero(output) == 0
859
860
861	@requires_gfx950
862	@torch.inference_mode()
863	def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None:
864	opus_mod = pytest.importorskip("aiter.ops.pa_sparse_prefill_opus")
865	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
866
867	monkeypatch.setattr(
868	mod, "_get_aiter_sparse_prefill_opus", lambda: opus_mod.pa_sparse_prefill_opus
869	)
870
871	device = torch.device("cuda")
872	torch.manual_seed(4)
873	num_queries = 8
874	num_heads = 16
875	q = (
876	torch.randn(
877	num_queries,
878	num_heads,
879	HEAD_DIM,
880	dtype=torch.bfloat16,
881	device=device,
882	)
883	* 0.125
884	)
885	kv = torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
886	rows = [list(range(query_idx + 1)) for query_idx in range(num_queries)]
887	indices = torch.tensor(
888	[index for row in rows for index in row], dtype=torch.int32, device=device
889	)
890	indptr = torch.tensor(
891	[0]
892	+ [
893	sum(len(row) for row in rows[: query_idx + 1])
894	for query_idx in range(num_queries)
895	],
896	dtype=torch.int32,
897	device=device,
898	)
899	attn_sink = torch.linspace(
900	-0.25, 0.25, num_heads, dtype=torch.float32, device=device
901	)
902	output = torch.empty_like(q)
903	scale = HEAD_DIM**-0.5
904
905	assert mod._rocm_sparse_attn_prefill_ragged_aiter_opus(
906	q=q,
907	kv=kv,
908	indices=indices,
909	indptr=indptr,
910	scale=scale,
911	attn_sink=attn_sink,
912	output=output,
913	)
914	expected = _ref_sparse_prefill_ragged(q, kv, rows, scale, attn_sink)
915
916	torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
917
918
919	@torch.inference_mode()
920	def test_sparse_attn_decode_ragged_kernel() -> None:
921	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
922	_rocm_sparse_attn_decode_ragged_triton,
923	)
924
925	device = torch.device("cuda")
926	torch.manual_seed(1)
927	block_size = 4
928	main_use_fnuz = current_platform.is_fp8_fnuz()
929	q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
930	main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
931	extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
932	main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz)
933	extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False)
934	main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
935	main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
936	extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
937	extra_indptr = torch.tensor([0, 1, 3], dtype=torch.int32, device=device)
938	attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
939	scale = HEAD_DIM**-0.5
940
941	out = torch.empty_like(q)
942	actual = _rocm_sparse_attn_decode_ragged_triton(
943	q=q,
944	main_cache=main_cache,
945	main_indices=main_indices,
946	main_indptr=main_indptr,
947	scale=scale,
948	attn_sink=attn_sink,
949	nope_head_dim=NOPE_HEAD_DIM,
950	rope_head_dim=ROPE_HEAD_DIM,
951	extra_cache=extra_cache,
952	extra_indices=extra_indices,
953	extra_indptr=extra_indptr,
954	out=out,
955	)
956	expected = _ref_sparse_decode_ragged(
957	q=q,
958	main_cache=main_cache,
959	main_rows=[[0, 2], [4, 1]],
960	scale=scale,
961	attn_sink=attn_sink,
962	block_size=block_size,
963	extra_cache=extra_cache,
964	extra_rows=[[1], [3, 0]],
965	main_use_fnuz=main_use_fnuz,
966	)
967
968	assert actual.data_ptr() == out.data_ptr()
969	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
970
971
972	@requires_gfx950
973	@torch.inference_mode()
974	def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None:
975	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
976	_rocm_sparse_attn_decode_ragged_triton,
977	)
978
979	device = torch.device("cuda")
980	block_size = 4
981	main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device)
982	extra_cache = torch.zeros_like(main_cache)
983	_poison_fp8_ds_mla_cache_row(main_cache, block_size)
984	_poison_fp8_ds_mla_cache_row(extra_cache, block_size)
985	indices = torch.zeros(1, dtype=torch.int32, device=device)
986	indptr = torch.tensor([0, 1], dtype=torch.int32, device=device)
987
988	actual = _rocm_sparse_attn_decode_ragged_triton(
989	q=torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device),
990	main_cache=main_cache,
991	main_indices=indices,
992	main_indptr=indptr,
993	scale=HEAD_DIM**-0.5,
994	attn_sink=None,
995	nope_head_dim=NOPE_HEAD_DIM,
996	rope_head_dim=ROPE_HEAD_DIM,
997	extra_cache=extra_cache,
998	extra_indices=indices,
999	extra_indptr=indptr,
1000	)
1001
1002	assert not torch.isnan(actual).any()
1003	assert torch.equal(actual, torch.zeros_like(actual))
1004
1005
1006	@pytest.mark.parametrize("on_gfx950", [False, True])
1007	@torch.inference_mode()
1008	def test_rocm_ragged_graph_buffer_view_tracks_source_width(
1009	monkeypatch, on_gfx950: bool
1010	) -> None:
1011	from vllm.models.deepseek_v4.amd import rocm as rocm_mod
1012
1013	monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950)
1014
1015	indices_buffer = torch.full((16,), -1, dtype=torch.int32)
1016	indptr_buffer = torch.full((3,), -1, dtype=torch.int32)
1017	first_indices = torch.tensor([3, 5, 7], dtype=torch.int32)
1018	first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32)
1019	first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
1020	first_indices,
1021	first_indptr,
1022	indices_buffer,
1023	indptr_buffer,
1024	num_rows=2,
1025	max_entries_per_row=8,
1026	)
1027
1028	second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)
1029	second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32)
1030	second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
1031	second_indices,
1032	second_indptr,
1033	indices_buffer,
1034	indptr_buffer,
1035	num_rows=2,
1036	max_entries_per_row=8,
1037	)
1038
1039	expected_first_entries = (
1040	first_indices.numel() if on_gfx950 else indices_buffer.numel()
1041	)
1042	expected_second_entries = (
1043	second_indices.numel() if on_gfx950 else indices_buffer.numel()
1044	)
1045	assert first_view.numel() == expected_first_entries
1046	assert second_view.numel() == expected_second_entries
1047	assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr()
1048	assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr()
1049	assert torch.equal(second_view[: second_indices.numel()], second_indices)
1050	assert torch.equal(second_indptr_view, second_indptr)
1051
1052
1053	def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None:
1054	from vllm.models.deepseek_v4.amd import rocm as rocm_mod
1055	from vllm.models.deepseek_v4.sparse_mla import (
1056	DeepseekV4SparseMLAMetadataBuilder,
1057	)
1058
1059	metadata = SimpleNamespace(for_cudagraph_capture=False)
1060	monkeypatch.setattr(
1061	DeepseekV4SparseMLAMetadataBuilder,
1062	"build_for_cudagraph_capture",
1063	lambda *_: metadata,
1064	)
1065	builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder)
1066
1067	actual = builder.build_for_cudagraph_capture(SimpleNamespace())
1068
1069	assert actual is metadata
1070	assert actual.for_cudagraph_capture is _on_gfx950()
1071
1072
1073	@requires_split_decode_arch
1074	@torch.inference_mode()
1075	def test_decode_num_splits_heuristic(monkeypatch) -> None:
1076	"""Split-count heuristic added with the flash-decode split-K decode path."""
1077	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1078
1079	# Pin the CU count so the heuristic is deterministic off-device.
1080	monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
1081
1082	# A batch that already fills the device should not be split.
1083	assert mod._decode_num_splits(256, 1, avg_main_len=128.0, avg_extra_len=0.0) == 1
1084	# A tiny batch on a large device should split to add parallelism.
1085	assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1
1086
1087	# The shared gfx942 selector retains its original 16-split ceiling.
1088	assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16
1089
1090	# The chosen count always stays within the searched [1, 16] range, and a
1091	# zero-length workload never splits (no work to parallelize).
1092	for num_queries in (1, 4, 24, 224, 1024):
1093	splits = mod._decode_num_splits(
1094	num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0
1095	)
1096	assert 1 <= splits <= 16
1097	assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
1098
1099
1100	@torch.inference_mode()
1101	def test_decode_num_splits_gfx950(monkeypatch) -> None:
1102	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1103
1104	monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
1105	assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32
1106	assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4
1107	assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1
1108
1109
1110	@requires_split_decode_arch
1111	@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
1112	@pytest.mark.parametrize("with_extra", [True, False])
1113	@pytest.mark.parametrize("with_sink", [True, False])
1114	@torch.inference_mode()
1115	def test_sparse_attn_decode_split_k_kernel(
1116	monkeypatch, num_splits: int, with_extra: bool, with_sink: bool
1117	) -> None:
1118	"""Flash-decode split-K decode path (partial + reduce kernels).
1119
1120	This path is the gfx942/gfx950 production path, so the test only runs on
1121	those architectures. The split count is pinned so the partial/reduce kernels are
1122	exercised across split counts. ``num_splits=8`` drives splits past the
1123	shortest segment length, covering the empty-split edge case handled by the
1124	reduce kernel.
1125	"""
1126	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1127
1128	device = torch.device("cuda")
1129	torch.manual_seed(7)
1130	block_size = 4
1131	num_heads = 3
1132	main_use_fnuz = current_platform.is_fp8_fnuz()
1133
1134	main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]]
1135	num_queries = len(main_rows)
1136	q = (
1137	torch.randn(
1138	num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
1139	)
1140	* 0.125
1141	)
1142	main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
1143	main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz)
1144	main_indices, main_indptr = _ragged_from_rows(main_rows, device)
1145
1146	extra_rows: list[list[int]] \| None = None
1147	extra_cache: torch.Tensor \| None = None
1148	extra_indices: torch.Tensor \| None = None
1149	extra_indptr: torch.Tensor \| None = None
1150	if with_extra:
1151	rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
1152	extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
1153	extra_rows = rows
1154	extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False)
1155	extra_indices, extra_indptr = _ragged_from_rows(rows, device)
1156
1157	attn_sink = (
1158	torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
1159	if with_sink
1160	else None
1161	)
1162	scale = HEAD_DIM**-0.5
1163
1164	# Pin the split count so each parametrized value is exercised deterministically.
1165	split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits"
1166	other_split_fn = (
1167	"_decode_num_splits" if _on_gfx950() else "_decode_gfx950_num_splits"
1168	)
1169	monkeypatch.setattr(mod, split_fn, lambda args, *kwargs: num_splits)
1170	monkeypatch.setattr(
1171	mod, other_split_fn, lambda args, *kwargs: pytest.fail("wrong selector")
1172	)
1173
1174	actual = mod._rocm_sparse_attn_decode_ragged_triton(
1175	q=q,
1176	main_cache=main_cache,
1177	main_indices=main_indices,
1178	main_indptr=main_indptr,
1179	scale=scale,
1180	attn_sink=attn_sink,
1181	nope_head_dim=NOPE_HEAD_DIM,
1182	rope_head_dim=ROPE_HEAD_DIM,
1183	extra_cache=extra_cache,
1184	extra_indices=extra_indices,
1185	extra_indptr=extra_indptr,
1186	)
1187	expected = _ref_sparse_decode_ragged(
1188	q=q,
1189	main_cache=main_cache,
1190	main_rows=main_rows,
1191	scale=scale,
1192	attn_sink=attn_sink,
1193	block_size=block_size,
1194	extra_cache=extra_cache,
1195	extra_rows=extra_rows,
1196	main_use_fnuz=main_use_fnuz,
1197	)
1198
1199	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
1200
1201
1202	@requires_gfx950
1203	@torch.inference_mode()
1204	def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None:
1205	device = torch.device("cuda")
1206	part_m = torch.full(
1207	(1, 8, 1),
1208	torch.finfo(torch.float32).min,
1209	dtype=torch.float32,
1210	device=device,
1211	)
1212	part_l = torch.zeros_like(part_m)
1213	part_acc = torch.full(
1214	(1, 8, 1, HEAD_DIM),
1215	float("nan"),
1216	dtype=torch.float32,
1217	device=device,
1218	)
1219	part_m[:, :2] = 0
1220	part_l[:, :2] = 1
1221	part_acc[:, 0] = 1
1222	part_acc[:, 1] = 3
1223
1224	actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True)
1225
1226	assert torch.isfinite(actual).all()
1227	assert torch.equal(actual, torch.full_like(actual, 2))
1228
1229
1230	@requires_gfx950
1231	@pytest.mark.parametrize("adaptive_splits", [False, True])
1232	@torch.inference_mode()
1233	def test_sparse_attn_decode_reduce_inverse_rope_epilogue(adaptive_splits: bool) -> None:
1234	"""FUSE_INV_ROPE must land where the standalone rotation would.
1235
1236	The epilogue rotates a whole [H, nope + rope] row with one expression by
1237	feeding cos=1/sin=0 to the NoPE lanes, which only holds if the pair split
1238	lands on ``nope_head_dim // 2``; the NoPE lanes therefore have to come
1239	back untouched, not merely close.
1240	"""
1241	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
1242
1243	device = torch.device("cuda")
1244	torch.manual_seed(7)
1245	num_queries, num_splits, num_heads, max_pos = 5, 4, 8, 64
1246	shape = (num_queries, num_splits, num_heads)
1247	part_m = torch.randn(shape, dtype=torch.float32, device=device)
1248	part_l = torch.rand(shape, dtype=torch.float32, device=device) + 0.5
1249	part_acc = torch.randn((*shape, HEAD_DIM), dtype=torch.float32, device=device)
1250	positions = torch.randint(
1251	0, max_pos, (num_queries,), dtype=torch.int64, device=device
1252	)
1253	angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
1254	cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
1255
1256	unfused = _launch_sparse_decode_reduce(part_m, part_l, part_acc, adaptive_splits)
1257	expected = _fused_inverse_rope_gptj(
1258	unfused, positions, cos_sin_cache, ROPE_HEAD_DIM
1259	)
1260	actual = _launch_sparse_decode_reduce(
1261	part_m, part_l, part_acc, adaptive_splits, positions, cos_sin_cache
1262	)
1263
1264	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
1265	assert torch.equal(actual[..., :NOPE_HEAD_DIM], unfused[..., :NOPE_HEAD_DIM])
1266	# An epilogue that quietly did nothing would satisfy everything above.
1267	assert not torch.equal(actual[..., NOPE_HEAD_DIM:], unfused[..., NOPE_HEAD_DIM:])
1268
1269
1270	def _mxfp8_dequant(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
1271	blocks = data.float().view(*data.shape[:-1], -1, 32)
1272	return (blocks * torch.exp2(scale.float() - 127.0)[..., None]).view(data.shape)
1273
1274
1275	def _random_cos_sin_cache(max_pos: int, device: torch.device) -> torch.Tensor:
1276	angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
1277	return torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
1278
1279
1280	@requires_gfx950
1281	@torch.inference_mode()
1282	def test_sparse_attn_decode_reduce_mxfp8_epilogue() -> None:
1283	"""QUANT_OUT must equal MXFP8-quantizing the reduce's fp32 output."""
1284	from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1285	_mxfp8_e4m3_quantize_torch,
1286	)
1287
1288	device = torch.device("cuda")
1289	torch.manual_seed(7)
1290	num_queries, num_splits, num_heads = 5, 4, 8
1291	shape = (num_queries, num_splits, num_heads)
1292	part_m = torch.randn(shape, device=device)
1293	part_l = torch.rand(shape, device=device) + 0.5
1294	part_acc = torch.randn((*shape, HEAD_DIM), device=device)
1295	positions = torch.randint(0, 64, (num_queries,), device=device)
1296	cos_sin_cache = _random_cos_sin_cache(64, device)
1297	args = (part_m, part_l, part_acc, False, positions, cos_sin_cache)
1298
1299	rotated = _launch_sparse_decode_reduce(*args, out_dtype=torch.float32)
1300	expected_data, expected_scale = _mxfp8_e4m3_quantize_torch(
1301	rotated.view(num_queries, -1)
1302	)
1303	scale = torch.empty_like(expected_scale)
1304	data = _launch_sparse_decode_reduce(
1305	*args, out_dtype=torch.float8_e4m3fn, out_scale=scale
1306	)
1307
1308	assert torch.equal(scale, expected_scale)
1309	assert torch.equal(
1310	data.view(num_queries, -1).view(torch.uint8), expected_data.view(torch.uint8)
1311	)
1312
1313
1314	@requires_gfx950
1315	@torch.inference_mode()
1316	def test_inverse_rope_mxfp8_rows() -> None:
1317	"""Prefill rows get the same inverse RoPE + MXFP8 as the decode epilogue."""
1318	from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1319	_mxfp8_e4m3_quantize_torch,
1320	)
1321	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
1322	rocm_inverse_rope_mxfp8_rows,
1323	)
1324
1325	device = torch.device("cuda")
1326	torch.manual_seed(5)
1327	num_tokens, num_heads = 7, 4
1328	o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).bfloat16()
1329	positions = torch.randint(0, 64, (num_tokens,), device=device)
1330	cos_sin_cache = _random_cos_sin_cache(64, device)
1331
1332	ref = o.float()
1333	cos, sin = cos_sin_cache[positions, None].chunk(2, dim=-1)
1334	even, odd = ref[..., NOPE_HEAD_DIM::2].clone(), ref[..., NOPE_HEAD_DIM + 1 :: 2]
1335	ref[..., NOPE_HEAD_DIM::2] = even * cos + odd * sin
1336	ref[..., NOPE_HEAD_DIM + 1 :: 2] = odd * cos - even * sin
1337	ref = ref.view(num_tokens, -1)
1338	_, expected_scale = _mxfp8_e4m3_quantize_torch(ref)
1339
1340	data = torch.empty_like(ref, dtype=torch.float8_e4m3fn)
1341	scale = torch.empty_like(expected_scale)
1342	rocm_inverse_rope_mxfp8_rows(
1343	o, positions, cos_sin_cache, ROPE_HEAD_DIM, data, scale
1344	)
1345
1346	assert torch.equal(scale, expected_scale)
1347	# Within half an e4m3 step at the top of each block's range.
1348	half_step = 16.0 * torch.exp2(scale.float() - 127.0).repeat_interleave(32, -1)
1349	assert bool(((_mxfp8_dequant(data, scale) - ref).abs() <= half_step).all())
1350
1351
1352	@requires_gfx950
1353	@torch.inference_mode()
1354	def test_sparse_attn_decode_mxfp8_output() -> None:
1355	"""``out_mxfp8`` routes the real decode through the MXFP8 epilogue."""
1356	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
1357	_rocm_sparse_attn_decode_ragged_triton,
1358	)
1359
1360	device = torch.device("cuda")
1361	torch.manual_seed(1)
1362	num_heads = 16
1363	q = torch.randn(2, num_heads, HEAD_DIM, device=device).bfloat16() * 0.125
1364	kv = torch.randn(6, HEAD_DIM, device=device).bfloat16() * 0.125
1365	decode = functools.partial(
1366	_rocm_sparse_attn_decode_ragged_triton,
1367	q=q,
1368	main_cache=_pack_fp8_ds_mla_cache(kv, 4, current_platform.is_fp8_fnuz()),
1369	main_indices=torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device),
1370	main_indptr=torch.tensor([0, 2, 4], dtype=torch.int32, device=device),
1371	scale=HEAD_DIM**-0.5,
1372	attn_sink=None,
1373	nope_head_dim=NOPE_HEAD_DIM,
1374	rope_head_dim=ROPE_HEAD_DIM,
1375	inv_rope_positions=torch.tensor([3, 9], device=device),
1376	inv_rope_cos_sin_cache=_random_cos_sin_cache(16, device),
1377	)
1378
1379	expected = decode().float().view(2, -1)
1380	data = torch.empty(
1381	2, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
1382	)
1383	scale = torch.empty(2, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device)
1384	decode(out_mxfp8=(data, scale))
1385
1386	torch.testing.assert_close(
1387	_mxfp8_dequant(data, scale), expected, atol=1e-3, rtol=2**-4
1388	)
1389
1390
1391	@requires_gfx950
1392	@pytest.mark.parametrize(
1393	"num_tokens, n_groups",
1394	# One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles.
1395	[
1396	(1, 4),
1397	(20, 4),
1398	(48, 4),
1399	(77, 4),
1400	(130, 4),
1401	(300, 4),
1402	(700, 4),
1403	(1000, 4),
1404	(1100, 8),
1405	],
1406	)
1407	@pytest.mark.parametrize("block_scales", [False, True])
1408	@torch.inference_mode()
1409	def test_rocm_mxfp8_wo_a_bmm(
1410	num_tokens: int, n_groups: int, block_scales: bool
1411	) -> None:
1412	from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
1413	_as_block32_scale,
1414	)
1415	from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1416	_mxfp8_e4m3_quantize_torch,
1417	)
1418	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_mxfp8_wo_a_bmm
1419
1420	device = torch.device("cuda")
1421	torch.manual_seed(3)
1422	o_lora_rank, group_dim = 128, 2048
1423	a, a_scale = _mxfp8_e4m3_quantize_torch(
1424	torch.randn(num_tokens, n_groups * group_dim, device=device)
1425	)
1426	w, w_scale = _mxfp8_e4m3_quantize_torch(
1427	torch.randn(n_groups * o_lora_rank, group_dim, device=device)
1428	)
1429	if block_scales:
1430	# A 32x32 checkpoint loads as per-row scales that repeat every 32 rows;
1431	# RocmDotScaledMxfp8LinearKernel compacts them before wo_a sees them.
1432	w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0))
1433	assert w_scale is not None
1434	wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
1435
1436	out = rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
1437
1438	expected = torch.einsum(
1439	"tgd,grd->tgr",
1440	_mxfp8_dequant(a, a_scale).view(num_tokens, n_groups, group_dim),
1441	_mxfp8_dequant(
1442	w, w_scale.repeat_interleave(w.shape[0] // w_scale.shape[0], 0)
1443	).view(n_groups, o_lora_rank, group_dim),
1444	)
1445	torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
1446
1447
1448	@requires_gfx950
1449	@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
1450	@torch.inference_mode()
1451	def test_sparse_attn_decode_gfx950_outer64_boundaries(
1452	monkeypatch, extra_len: int
1453	) -> None:
1454	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1455
1456	device = torch.device("cuda")
1457	torch.manual_seed(13)
1458	block_size = 4
1459	num_heads = 16
1460	num_extra_rows = 80
1461	q = torch.randn(2, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device)
1462	q *= 0.125
1463	main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device)
1464	main_indices = torch.empty(0, dtype=torch.int32, device=device)
1465	main_indptr = torch.zeros(3, dtype=torch.int32, device=device)
1466	extra_cache = _pack_fp8_ds_mla_cache(
1467	torch.randn(num_extra_rows, HEAD_DIM, dtype=torch.bfloat16, device=device)
1468	* 0.125,
1469	block_size,
1470	use_fnuz=False,
1471	)
1472	_poison_fp8_ds_mla_cache_row(extra_cache, block_size)
1473
1474	raw_row = list(range(1, extra_len + 1))
1475	if extra_len > 3:
1476	raw_row[3] = -1
1477	if extra_len > 40:
1478	raw_row[40] = num_extra_rows
1479	if extra_len > 64:
1480	raw_row[64] = num_extra_rows + 1024
1481	extra_indices, extra_indptr = _ragged_from_rows([raw_row, []], device)
1482	valid_row = [slot for slot in raw_row if 0 <= slot < num_extra_rows]
1483	attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
1484
1485	monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
1486	actual = mod._rocm_sparse_attn_decode_ragged_triton(
1487	q=q,
1488	main_cache=main_cache,
1489	main_indices=main_indices,
1490	main_indptr=main_indptr,
1491	scale=HEAD_DIM**-0.5,
1492	attn_sink=attn_sink,
1493	nope_head_dim=NOPE_HEAD_DIM,
1494	rope_head_dim=ROPE_HEAD_DIM,
1495	extra_cache=extra_cache,
1496	extra_indices=extra_indices,
1497	extra_indptr=extra_indptr,
1498	)
1499	expected = _ref_sparse_decode_ragged(
1500	q=q,
1501	main_cache=main_cache,
1502	main_rows=[[], []],
1503	scale=HEAD_DIM**-0.5,
1504	attn_sink=attn_sink,
1505	block_size=block_size,
1506	extra_cache=extra_cache,
1507	extra_rows=[valid_row, []],
1508	)
1509
1510	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
1511	assert torch.equal(actual[1], torch.zeros_like(actual[1]))
1512
1513
1514	@requires_gfx950
1515	@torch.inference_mode()
1516	def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None:
1517	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1518
1519	device = torch.device("cuda")
1520	torch.manual_seed(17)
1521	block_size = 64
1522	num_queries = 16
1523	num_heads = 16
1524	num_splits = 8
1525	extra_per_query = 65 * num_splits
1526	max_extra_per_query = 8192
1527	q = (
1528	torch.randn(
1529	num_queries,
1530	num_heads,
1531	HEAD_DIM,
1532	dtype=torch.bfloat16,
1533	device=device,
1534	)
1535	* 0.125
1536	)
1537	main_cache = _pack_fp8_ds_mla_cache(
1538	torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125,
1539	block_size,
1540	use_fnuz=False,
1541	)
1542	extra_cache = _pack_fp8_ds_mla_cache(
1543	torch.randn(
1544	num_queries * extra_per_query,
1545	HEAD_DIM,
1546	dtype=torch.bfloat16,
1547	device=device,
1548	)
1549	* 0.125,
1550	block_size,
1551	use_fnuz=False,
1552	)
1553	main_rows = [[query_idx] for query_idx in range(num_queries)]
1554	extra_rows = [
1555	list(range(query_idx * extra_per_query, (query_idx + 1) * extra_per_query))
1556	for query_idx in range(num_queries)
1557	]
1558	main_indices, main_indptr = _ragged_from_rows(main_rows, device)
1559	short_extra_rows = [row[:64] for row in extra_rows]
1560	long_indices, long_indptr = _ragged_from_rows(extra_rows, device)
1561	short_indices, short_indptr = _ragged_from_rows(short_extra_rows, device)
1562	extra_indices = torch.full(
1563	(num_queries * max_extra_per_query,),
1564	-1,
1565	dtype=torch.int32,
1566	device=device,
1567	)
1568	extra_indices[: long_indices.numel()].copy_(long_indices)
1569	extra_indptr = long_indptr.clone()
1570	extra_indices_ptr = extra_indices.data_ptr()
1571	attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
1572	out = torch.empty_like(q)
1573
1574	monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: num_splits)
1575
1576	def run_decode() -> torch.Tensor:
1577	return mod._rocm_sparse_attn_decode_ragged_triton(
1578	q=q,
1579	main_cache=main_cache,
1580	main_indices=main_indices,
1581	main_indptr=main_indptr,
1582	scale=HEAD_DIM**-0.5,
1583	attn_sink=attn_sink,
1584	nope_head_dim=NOPE_HEAD_DIM,
1585	rope_head_dim=ROPE_HEAD_DIM,
1586	extra_cache=extra_cache,
1587	extra_indices=extra_indices,
1588	extra_indptr=extra_indptr,
1589	out=out,
1590	extra_cache_nan_free=True,
1591	adaptive_splits=True,
1592	)
1593
1594	run_decode()
1595	torch.accelerator.synchronize()
1596	graph = torch.cuda.CUDAGraph()
1597	with torch.cuda.graph(graph):
1598	captured_out = run_decode()
1599	torch.accelerator.synchronize()
1600	captured_long = out.clone()
1601	expected_long = _ref_sparse_decode_ragged(
1602	q=q,
1603	main_cache=main_cache,
1604	main_rows=main_rows,
1605	scale=HEAD_DIM**-0.5,
1606	attn_sink=attn_sink,
1607	block_size=block_size,
1608	extra_cache=extra_cache,
1609	extra_rows=extra_rows,
1610	)
1611	torch.testing.assert_close(captured_long, expected_long, atol=2e-2, rtol=2e-2)
1612
1613	extra_indices[: short_indices.numel()].copy_(short_indices)
1614	extra_indptr.copy_(short_indptr)
1615	graph.replay()
1616	torch.accelerator.synchronize()
1617	short_out = out.clone()
1618	expected_short = _ref_sparse_decode_ragged(
1619	q=q,
1620	main_cache=main_cache,
1621	main_rows=main_rows,
1622	scale=HEAD_DIM**-0.5,
1623	attn_sink=attn_sink,
1624	block_size=block_size,
1625	extra_cache=extra_cache,
1626	extra_rows=short_extra_rows,
1627	)
1628
1629	assert captured_out.data_ptr() == out.data_ptr()
1630	assert extra_indices.data_ptr() == extra_indices_ptr
1631	assert extra_indices.numel() == num_queries * max_extra_per_query
1632	assert not torch.equal(short_out, captured_long)
1633	torch.testing.assert_close(short_out, expected_short, atol=2e-2, rtol=2e-2)
1634
1635	extra_indices[: long_indices.numel()].copy_(long_indices)
1636	extra_indptr.copy_(long_indptr)
1637	graph.replay()
1638	torch.accelerator.synchronize()
1639	torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2)
1640
1641
1642	@requires_gfx950
1643	@torch.inference_mode()
1644	def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
1645	"""Replay the production ROCm decode path after device-only reallocation."""
1646	from tests.v1.attention.utils import create_vllm_config
1647	from vllm.models.deepseek_v4.amd.rocm import (
1648	DeepseekV4ROCMAiterMLASparseMetadataBuilder,
1649	DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
1650	)
1651	from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
1652	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1653	from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec
1654
1655	device = torch.device("cuda")
1656	num_reqs = 3
1657	upper_query_len = 8
1658	graph_tokens = num_reqs * upper_query_len
1659	block_size = 256
1660	compressed_block_size = block_size // 128
1661	window_size = 32
1662
1663	vllm_config = create_vllm_config(
1664	model_name="facebook/opt-125m",
1665	max_model_len=1024,
1666	block_size=block_size,
1667	max_num_seqs=num_reqs,
1668	max_num_batched_tokens=graph_tokens,
1669	hf_config_override={
1670	"compress_ratios": [128],
1671	"index_topk": 2048,
1672	"sliding_window": window_size,
1673	},
1674	)
1675	vllm_config.speculative_config = SimpleNamespace(
1676	num_speculative_tokens=upper_query_len - 1,
1677	parallel_drafting=False,
1678	enable_adaptive_verification=True,
1679	use_dspark=lambda: True,
1680	)
1681	mla_spec = MLAAttentionSpec(
1682	block_size=block_size,
1683	num_kv_heads=1,
1684	head_size=HEAD_DIM,
1685	dtype=torch.bfloat16,
1686	tokens_per_state=128,
1687	cache_dtype_str="fp8_ds_mla",
1688	model_version="deepseek_v4",
1689	)
1690	swa_spec = SlidingWindowMLASpec(
1691	block_size=block_size,
1692	num_kv_heads=1,
1693	head_size=HEAD_DIM,
1694	dtype=torch.bfloat16,
1695	sliding_window=window_size,
1696	cache_dtype_str="fp8_ds_mla",
1697	model_version="deepseek_v4",
1698	)
1699	mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder(
1700	mla_spec, ["c128a"], vllm_config, device
1701	)
1702	swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
1703	swa_spec, ["c128a"], vllm_config, device
1704	)
1705	assert (
1706	mla_builder.get_cudagraph_support(vllm_config, mla_spec)
1707	== AttentionCGSupport.ALWAYS
1708	)
1709	assert (
1710	swa_builder.get_cudagraph_support(vllm_config, swa_spec)
1711	== AttentionCGSupport.ALWAYS
1712	)
1713
1714	seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device)
1715	seq_lens_cpu = seq_lens.cpu()
1716	cpu_query_start_loc = torch.arange(
1717	0,
1718	graph_tokens + 1,
1719	upper_query_len,
1720	dtype=torch.int32,
1721	)
1722	max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size
1723	block_table = torch.arange(
1724	num_reqs * max_blocks, dtype=torch.int32, device=device
1725	).view(num_reqs, max_blocks)
1726
1727	def build_metadata(query_lens: list[int]):
1728	query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device)
1729	query_start_loc = torch.cat(
1730	[
1731	torch.zeros(1, dtype=torch.int32, device=device),
1732	query_lens_tensor.cumsum(0),
1733	]
1734	)
1735	active_tokens = sum(query_lens)
1736	positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device)
1737	position_rows = [
1738	torch.arange(
1739	seq_len - query_len,
1740	seq_len,
1741	dtype=torch.int64,
1742	device=device,
1743	)
1744	for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens)
1745	]
1746	positions[:active_tokens] = torch.cat(position_rows)
1747	slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device)
1748	slot_mapping[:active_tokens] = torch.arange(
1749	active_tokens, dtype=torch.int64, device=device
1750	)
1751	common = CommonAttentionMetadata(
1752	query_start_loc=query_start_loc,
1753	query_start_loc_cpu=cpu_query_start_loc,
1754	seq_lens=seq_lens,
1755	seq_lens_cpu_upper_bound=seq_lens_cpu,
1756	num_reqs=num_reqs,
1757	num_actual_tokens=graph_tokens,
1758	max_query_len=upper_query_len,
1759	max_seq_len=int(seq_lens_cpu.max()),
1760	block_table_tensor=block_table,
1761	slot_mapping=slot_mapping,
1762	positions=positions,
1763	causal=True,
1764	)
1765	return (
1766	mla_builder.build_for_cudagraph_capture(common),
1767	swa_builder.build_for_cudagraph_capture(common),
1768	)
1769
1770	mla_metadata, swa_metadata = build_metadata([8, 8, 8])
1771	assert mla_metadata.for_cudagraph_capture
1772	metadata_ptrs = (
1773	mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(),
1774	mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(),
1775	swa_metadata.decode_swa_ragged_indices.data_ptr(),
1776	swa_metadata.decode_swa_ragged_indptr.data_ptr(),
1777	)
1778
1779	torch.manual_seed(19)
1780	num_heads = 16
1781	q = (
1782	torch.randn(
1783	graph_tokens,
1784	num_heads,
1785	HEAD_DIM,
1786	dtype=torch.bfloat16,
1787	device=device,
1788	)
1789	* 0.125
1790	)
1791	num_cache_blocks = num_reqs * max_blocks
1792	swa_cache = _pack_fp8_ds_mla_cache(
1793	torch.randn(
1794	num_cache_blocks * block_size,
1795	HEAD_DIM,
1796	dtype=torch.bfloat16,
1797	device=device,
1798	)
1799	* 0.125,
1800	block_size,
1801	use_fnuz=False,
1802	)
1803	compressed_cache = _pack_fp8_ds_mla_cache(
1804	torch.randn(
1805	num_cache_blocks * compressed_block_size,
1806	HEAD_DIM,
1807	dtype=torch.bfloat16,
1808	device=device,
1809	)
1810	* 0.125,
1811	compressed_block_size,
1812	use_fnuz=False,
1813	)
1814	attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
1815	out = torch.empty_like(q)
1816	monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
1817
1818	def run_decode(mla_md, swa_md) -> None:
1819	mod.rocm_sparse_attn_decode(
1820	q=q,
1821	kv_cache=compressed_cache,
1822	swa_k_cache=swa_cache,
1823	swa_only=False,
1824	topk_indices=mla_md.c128a_global_decode_topk_indices,
1825	topk_lens=mla_md.c128a_decode_topk_lens,
1826	swa_indices=swa_md.decode_swa_indices,
1827	swa_lens=swa_md.decode_swa_lens,
1828	swa_ragged_indices=swa_md.decode_swa_ragged_indices,
1829	swa_ragged_indptr=swa_md.decode_swa_ragged_indptr,
1830	topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices,
1831	topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr,
1832	attn_sink=attn_sink,
1833	scale=HEAD_DIM**-0.5,
1834	head_dim=HEAD_DIM,
1835	nope_head_dim=NOPE_HEAD_DIM,
1836	rope_head_dim=ROPE_HEAD_DIM,
1837	output=out,
1838	extra_cache_nan_free=True,
1839	adaptive_splits=True,
1840	)
1841
1842	def reference(mla_md, swa_md) -> torch.Tensor:
1843	return _ref_sparse_decode_ragged(
1844	q=q,
1845	main_cache=swa_cache,
1846	main_rows=_rows_from_ragged(
1847	swa_md.decode_swa_ragged_indices,
1848	swa_md.decode_swa_ragged_indptr,
1849	),
1850	scale=HEAD_DIM**-0.5,
1851	attn_sink=attn_sink,
1852	block_size=block_size,
1853	extra_cache=compressed_cache,
1854	extra_rows=_rows_from_ragged(
1855	mla_md.c128a_decode_topk_ragged_indices,
1856	mla_md.c128a_decode_topk_ragged_indptr,
1857	),
1858	extra_block_size=compressed_block_size,
1859	)
1860
1861	run_decode(mla_metadata, swa_metadata)
1862	torch.accelerator.synchronize()
1863	expected_full = reference(mla_metadata, swa_metadata)
1864	torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2)
1865
1866	graph = torch.cuda.CUDAGraph()
1867	with torch.cuda.graph(graph):
1868	run_decode(mla_metadata, swa_metadata)
1869	torch.accelerator.synchronize()
1870	captured_full = out.clone()
1871
1872	reallocated_mla, reallocated_swa = build_metadata([3, 8, 5])
1873	assert metadata_ptrs == (
1874	reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(),
1875	reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(),
1876	reallocated_swa.decode_swa_ragged_indices.data_ptr(),
1877	reallocated_swa.decode_swa_ragged_indptr.data_ptr(),
1878	)
1879	expected_reallocated = reference(reallocated_mla, reallocated_swa)
1880	run_decode(reallocated_mla, reallocated_swa)
1881	torch.accelerator.synchronize()
1882	torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
1883	assert not torch.equal(captured_full, expected_reallocated)
1884
1885	graph.replay()
1886	torch.accelerator.synchronize()
1887	torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
1888
1889
1890	# ---------------------------------------------------------------------------
1891	# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
1892	# ---------------------------------------------------------------------------
1893
1894
1895	# Cache rows = max_position_embeddings * scaling_factor.
1896	_ROTARY_MAX_POS = 1024
1897	_ROTARY_SCALING_FACTOR = 4.0
1898	_ROTARY_CACHE_LEN = int(_ROTARY_MAX_POS * _ROTARY_SCALING_FACTOR)
1899
1900
1901	def _make_dsv4_rotary(device: torch.device):
1902	"""The official DSv4 rotary embedding, sized down for unit tests."""
1903	from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
1904	DeepseekV4ScalingRotaryEmbedding,
1905	)
1906
1907	# The model loader constructs layers under a default-device context;
1908	# mirror that so the fp32 cos_sin_cache lands on the GPU.
1909	with torch.device(device):
1910	rotary_emb = DeepseekV4ScalingRotaryEmbedding(
1911	head_size=ROPE_HEAD_DIM,
1912	rotary_dim=ROPE_HEAD_DIM,
1913	max_position_embeddings=_ROTARY_MAX_POS,
1914	base=10000,
1915	is_neox_style=False,
1916	scaling_factor=_ROTARY_SCALING_FACTOR,
1917	dtype=torch.bfloat16,
1918	mscale=1.0,
1919	mscale_all_dim=1.0,
1920	)
1921	rotary_emb = rotary_emb.to(device)
1922	assert rotary_emb.cos_sin_cache.shape == (_ROTARY_CACHE_LEN, ROPE_HEAD_DIM)
1923	return rotary_emb
1924
1925
1926	def _inv_rope_via_rotary_native(
1927	rotary_emb: torch.nn.Module,
1928	o: torch.Tensor,
1929	positions: torch.Tensor,
1930	) -> torch.Tensor:
1931	"""Reference: the official ``forward_native(inverse=True)`` path."""
1932	expected, _ = rotary_emb.forward_native(positions, o.clone(), None, inverse=True)
1933	return expected.to(torch.bfloat16)
1934
1935
1936	class _FakeWoA(torch.nn.Module):
1937	"""Stand-in for the wo_a linear layer holding the (optionally fp8) weight."""
1938
1939	def __init__(
1940	self, weight: torch.Tensor, weight_scale_inv: torch.Tensor \| None = None
1941	) -> None:
1942	super().__init__()
1943	self.weight = weight
1944	if weight_scale_inv is not None:
1945	self.weight_scale_inv = weight_scale_inv
1946
1947
1948	@pytest.mark.parametrize("num_tokens", [1, 7, 64])
1949	@pytest.mark.parametrize("num_heads", [1, 8])
1950	@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
1951	@torch.inference_mode()
1952	def test_fused_inverse_rope_gptj_matches_rotary_native(
1953	num_tokens: int, num_heads: int, pos_dtype: torch.dtype, default_vllm_config
1954	) -> None:
1955	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
1956
1957	device = torch.device("cuda")
1958	torch.manual_seed(0)
1959	rotary_emb = _make_dsv4_rotary(device)
1960	o = torch.randn(
1961	num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
1962	)
1963	positions = torch.randint(
1964	0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=pos_dtype, device=device
1965	)
1966
1967	actual = _fused_inverse_rope_gptj(
1968	o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
1969	)
1970	expected = _inv_rope_via_rotary_native(rotary_emb, o, positions)
1971
1972	assert actual.dtype == torch.bfloat16
1973	assert actual.shape == o.shape
1974	# NoPE lanes are a pure bf16 passthrough -> must be bit-exact.
1975	assert torch.equal(actual[..., :NOPE_HEAD_DIM], expected[..., :NOPE_HEAD_DIM])
1976	# RoPE lanes: tolerate at most ~1 bf16 ulp from fp32 fma ordering.
1977	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
1978
1979
1980	@torch.inference_mode()
1981	def test_fused_inverse_rope_gptj_empty(default_vllm_config) -> None:
1982	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
1983
1984	device = torch.device("cuda")
1985	rotary_emb = _make_dsv4_rotary(device)
1986	o = torch.empty(0, 8, HEAD_DIM, dtype=torch.bfloat16, device=device)
1987	positions = torch.empty(0, dtype=torch.int32, device=device)
1988
1989	out = _fused_inverse_rope_gptj(
1990	o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
1991	)
1992	assert out.shape == (0, 8, HEAD_DIM)
1993	assert out.dtype == torch.bfloat16
1994
1995
1996	@torch.inference_mode()
1997	def test_rocm_inv_rope_einsum_matches_rotary_native(default_vllm_config) -> None:
1998	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
1999
2000	device = torch.device("cuda")
2001	torch.manual_seed(2)
2002	num_tokens, num_heads = 5, 8
2003	n_local_groups = num_heads
2004	o_lora_rank = 16
2005	hidden_dim = num_heads * HEAD_DIM // n_local_groups # 512
2006
2007	rotary_emb = _make_dsv4_rotary(device)
2008	o = (
2009	torch.randn(
2010	num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
2011	)
2012	* 0.125
2013	)
2014	positions = torch.randint(
2015	0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=torch.int32, device=device
2016	)
2017	weight = (
2018	torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125
2019	).to(torch.bfloat16)
2020	wo_a = _FakeWoA(weight)
2021
2022	actual = rocm_inv_rope_einsum(
2023	rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a
2024	)
2025
2026	o_ref = _inv_rope_via_rotary_native(rotary_emb, o, positions)
2027	o_ref = o_ref.view(num_tokens, n_local_groups, -1)
2028	wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
2029	expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref)
2030
2031	assert actual.shape == (num_tokens, n_local_groups, o_lora_rank)
2032	torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
2033
2034
2035	@torch.inference_mode()
2036	def test_get_cached_wo_a_bf16_plain_caches() -> None:
2037	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
2038
2039	device = torch.device("cuda")
2040	torch.manual_seed(4)
2041	n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
2042	weight = torch.randn(
2043	n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
2044	)
2045	wo_a = _FakeWoA(weight)
2046
2047	out1 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
2048	expected = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
2049	assert out1.shape == (n_local_groups, o_lora_rank, hidden_dim)
2050	torch.testing.assert_close(out1, expected, atol=0, rtol=0)
2051	assert hasattr(wo_a, "_dsv4_wo_a_bf16")
2052
2053	# Mutate the source weight: the cached tensor must be returned unchanged
2054	# (proving the dequant is not recomputed per call).
2055	wo_a.weight.zero_()
2056	out2 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
2057	assert out2 is out1
2058	torch.testing.assert_close(out2, expected, atol=0, rtol=0)
2059
2060
2061	@pytest.mark.parametrize("scale_attr", ["weight_scale", "weight_scale_inv"])
2062	@torch.inference_mode()
2063	def test_get_cached_wo_a_bf16_dequantized_ignores_retained_scale(
2064	scale_attr: str,
2065	) -> None:
2066	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
2067
2068	device = torch.device("cuda")
2069	torch.manual_seed(6)
2070	n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
2071	weight = torch.randn(
2072	n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
2073	)
2074	retained_scale = torch.full(
2075	(n_local_groups, 1), 0.25, dtype=torch.float32, device=device
2076	)
2077	wo_a = _FakeWoA(weight)
2078	setattr(wo_a, scale_attr, retained_scale)
2079
2080	out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
2081
2082	expected = weight.view(n_local_groups, o_lora_rank, hidden_dim)
2083	torch.testing.assert_close(out, expected, atol=0, rtol=0)
2084
2085
2086	@torch.inference_mode()
2087	def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
2088	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
2089
2090	device = torch.device("cuda")
2091	torch.manual_seed(5)
2092	n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
2093	row_block, col_block = 2, 2
2094	row_blocks = o_lora_rank // row_block
2095	col_blocks = hidden_dim // col_block
2096
2097	fp8_dtype = current_platform.fp8_dtype()
2098	weight_f32 = (
2099	torch.randn(
2100	n_local_groups, o_lora_rank, hidden_dim, dtype=torch.float32, device=device
2101	)
2102	* 0.1
2103	)
2104	weight_fp8 = weight_f32.to(fp8_dtype)
2105	scale = (
2106	torch.rand(
2107	n_local_groups, row_blocks, col_blocks, dtype=torch.float32, device=device
2108	)
2109	* 0.5
2110	+ 0.5
2111	)
2112	wo_a = _FakeWoA(
2113	weight_fp8.reshape(n_local_groups * o_lora_rank, hidden_dim),
2114	weight_scale_inv=scale.reshape(n_local_groups * row_blocks, col_blocks),
2115	)
2116
2117	out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
2118
2119	scale_full = scale.repeat_interleave(row_block, dim=-2).repeat_interleave(
2120	col_block, dim=-1
2121	)
2122	expected = (weight_fp8.to(torch.float32) * scale_full).to(torch.bfloat16)
2123	assert out.shape == (n_local_groups, o_lora_rank, hidden_dim)
2124	torch.testing.assert_close(out, expected, atol=0, rtol=0)
2125
2126	# Second call returns the same cached object.
2127	assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out
2128
2129
2130	@requires_split_decode_arch
2131	@torch.inference_mode()
2132	def test_fp8_paged_mqa_logits_triton_matches_torch_ref() -> None:
2133	"""Block-flat Triton decode logits vs the eager torch reference.
2134
2135	Compare only valid key positions: the kernel skips the -inf pad that
2136	the torch ref writes past seq_len.
2137	"""
2138	from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
2139	from vllm.v1.worker.workspace import init_workspace_manager
2140
2141	device = torch.device("cuda")
2142	init_workspace_manager(device)
2143	torch.manual_seed(0)
2144
2145	fp8_dtype = current_platform.fp8_dtype()
2146	batch_size = 2
2147	next_n = 1
2148	num_heads = 4
2149	head_size = 128
2150	block_size = 64
2151	max_model_len = 256
2152	seq_lens = [180, 64]
2153	num_pages = [(seq_len + block_size - 1) // block_size for seq_len in seq_lens]
2154	num_blocks = sum(num_pages)
2155
2156	q = torch.randn(
2157	batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
2158	).to(fp8_dtype)
2159	weights = torch.rand(
2160	batch_size * next_n, num_heads, device=device, dtype=torch.float32
2161	)
2162	context_lens = torch.tensor(seq_lens, device=device, dtype=torch.int32)
2163	block_tables = torch.full(
2164	(batch_size, max(num_pages)), 0, device=device, dtype=torch.int32
2165	)
2166	page = 0
2167	for i, n_pages in enumerate(num_pages):
2168	block_tables[i, :n_pages] = torch.arange(
2169	page, page + n_pages, device=device, dtype=torch.int32
2170	)
2171	page += n_pages
2172
2173	kv_cache = torch.empty(
2174	num_blocks, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
2175	)
2176	values = torch.randn(
2177	num_blocks, block_size, head_size, device=device, dtype=torch.bfloat16
2178	).to(fp8_dtype)
2179	scales = torch.rand(
2180	num_blocks, block_size, device=device, dtype=torch.float32
2181	).clamp(min=1e-3)
2182	kv_flat = kv_cache.view(num_blocks, -1)
2183	scale_off = block_size * head_size
2184	kv_flat[:, :scale_off] = values.reshape(num_blocks, -1).view(torch.uint8)
2185	kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_blocks, -1)
2186
2187	got = mod.rocm_fp8_paged_mqa_logits_triton(
2188	q, kv_cache, weights, context_lens, block_tables, max_model_len
2189	)
2190	ref = mod.fp8_paged_mqa_logits_torch(
2191	q, kv_cache, weights, context_lens, block_tables, max_model_len
2192	)
2193	for i, seq_len in enumerate(seq_lens):
2194	torch.testing.assert_close(
2195	got[i, :seq_len], ref[i, :seq_len], atol=2e-3, rtol=2e-3
2196	)
2197
2198
2199	def test_indexer_k_is_c4a_block_flat_gate() -> None:
2200	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
2201	_indexer_k_is_c4a_block_flat,
2202	)
2203
2204	assert not _indexer_k_is_c4a_block_flat(1)
2205	assert not _indexer_k_is_c4a_block_flat(2)
2206	assert _indexer_k_is_c4a_block_flat(4)
2207	assert not _indexer_k_is_c4a_block_flat(128)
2208
2209
2210	@requires_split_decode_arch
2211	@torch.inference_mode()
2212	def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None:
2213	"""Ratio 1/2 stay on AITER; ratio 4 takes Triton."""
2214	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
2215	fp8_paged_mqa_logits_torch,
2216	indexer_k_quant_and_cache_triton,
2217	rocm_fp8_paged_mqa_logits,
2218	rocm_fp8_paged_mqa_logits_triton,
2219	)
2220	from vllm.v1.worker.workspace import init_workspace_manager
2221
2222	device = torch.device("cuda")
2223	init_workspace_manager(device)
2224	torch.manual_seed(1)
2225
2226	fp8_dtype = current_platform.fp8_dtype()
2227	batch_size = 1
2228	next_n = 1
2229	num_heads = 32
2230	head_size = 128
2231	block_size = 64
2232	seq_len = 180
2233	max_model_len = 256
2234	num_pages = (seq_len + block_size - 1) // block_size
2235	num_tokens = num_pages * block_size
2236
2237	q = torch.randn(
2238	batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
2239	).to(fp8_dtype)
2240	weights = torch.rand(
2241	batch_size * next_n, num_heads, device=device, dtype=torch.float32
2242	)
2243	context_lens = torch.tensor([seq_len], device=device, dtype=torch.int32)
2244	block_tables = torch.arange(num_pages, device=device, dtype=torch.int32).view(1, -1)
2245	dummy_sched = torch.zeros(8, 2, device=device, dtype=torch.int32)
2246
2247	k_bf16 = torch.randn(num_tokens, head_size, device=device, dtype=torch.bfloat16)
2248	slot = torch.arange(num_tokens, device=device, dtype=torch.int32)
2249	cache_shuffle = torch.zeros(
2250	num_pages, block_size, head_size + 4, device=device, dtype=torch.uint8
2251	)
2252	indexer_k_quant_and_cache_triton(k_bf16, cache_shuffle, slot, 128, "ue8m0")
2253	cache_shuffle = cache_shuffle.unsqueeze(2)
2254
2255	aiter = rocm_fp8_paged_mqa_logits(
2256	q,
2257	cache_shuffle,
2258	weights,
2259	context_lens,
2260	block_tables,
2261	dummy_sched,
2262	max_model_len,
2263	compress_ratio=1,
2264	)
2265	for ratio in (1, 2):
2266	got = rocm_fp8_paged_mqa_logits(
2267	q,
2268	cache_shuffle,
2269	weights,
2270	context_lens,
2271	block_tables,
2272	dummy_sched,
2273	max_model_len,
2274	compress_ratio=ratio,
2275	)
2276	torch.testing.assert_close(got[:, :seq_len], aiter[:, :seq_len], atol=0, rtol=0)
2277
2278	cache_flat = torch.empty(
2279	num_pages, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
2280	)
2281	values = torch.randn(
2282	num_pages, block_size, head_size, device=device, dtype=torch.bfloat16
2283	).to(fp8_dtype)
2284	scales = torch.rand(
2285	num_pages, block_size, device=device, dtype=torch.float32
2286	).clamp(min=1e-3)
2287	kv_flat = cache_flat.view(num_pages, -1)
2288	scale_off = block_size * head_size
2289	kv_flat[:, :scale_off] = values.reshape(num_pages, -1).view(torch.uint8)
2290	kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_pages, -1)
2291
2292	gated_c4a = rocm_fp8_paged_mqa_logits(
2293	q,
2294	cache_flat,
2295	weights,
2296	context_lens,
2297	block_tables,
2298	dummy_sched,
2299	max_model_len,
2300	compress_ratio=4,
2301	)
2302	triton = rocm_fp8_paged_mqa_logits_triton(
2303	q, cache_flat, weights, context_lens, block_tables, max_model_len
2304	)
2305	ref = fp8_paged_mqa_logits_torch(
2306	q, cache_flat, weights, context_lens, block_tables, max_model_len
2307	)
2308	torch.testing.assert_close(
2309	gated_c4a[:, :seq_len], triton[:, :seq_len], atol=0, rtol=0
2310	)
2311	torch.testing.assert_close(
2312	gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3
2313	)