@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ef43e3101b

[ROCm][DSV4][Perf] Optimize Triton sparse-MLA decode on gfx950 (#52212)

Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Cursor Grok 4.6 <cursoragent@cursor.com>

Fangzhou Ai committed Aug 16, 2026 at 15:16 UTC ef43e3101b8fda8cd2b52de150c76b4fc177fad2
5 files changed +1550 -94
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+458 -23
@@ -24,12 +24,27 @@ def _on_split_decode_arch() -> bool:
24 return False
25
26
27 +def _on_gfx950() -> bool:
28 + if not current_platform.is_rocm():
29 + return False
30 + try:
31 + from vllm.platforms.rocm import _ON_GFX950
32 +
33 + return _ON_GFX950
34 + except ImportError:
35 + return False
36 +
37 +
38 # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other
39 # architectures take the fallback decode kernel, so its tests are skipped there.
40 requires_split_decode_arch = pytest.mark.skipif(
41 not _on_split_decode_arch(),
42 reason="split-K decode kernel is only tuned for AMD gfx942/gfx950",
43 )
44 +requires_gfx950 = pytest.mark.skipif(
45 + not _on_gfx950(),
46 + reason="optimized sparse decode partial is gfx950-only",
47 +)
48
49 NOPE_HEAD_DIM = 448
50 ROPE_HEAD_DIM = 64
@@ -118,28 +133,48 @@ def _pack_fp8_ds_mla_cache(
133 return cache
134
135
121 -def _read_fp8_ds_mla_cache(
122 - cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool
123 -) -> torch.Tensor:
124 - cache_flat = cache.view(torch.uint8).flatten()
136 +def _poison_fp8_ds_mla_cache_row(
137 + cache: torch.Tensor, block_size: int, slot: int = 0
138 +) -> None:
139 + flat = cache.flatten()
140 block_idx = slot // block_size
141 pos = slot % block_size
142 block_base = block_idx * cache.stride(0)
143 token_base = block_base + pos * 576
144 scale_base = block_base + block_size * 576 + pos * 8
145 + flat[token_base] = 0x7F
146 + flat[scale_base : scale_base + 7] = 255
147 + flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = float(
148 + "nan"
149 + )
150 +
151 +
152 +def _read_fp8_ds_mla_cache_rows(
153 + cache: torch.Tensor,
154 + slots: torch.Tensor,
155 + block_size: int,
156 + use_fnuz: bool,
157 +) -> torch.Tensor:
158 + cache_flat = cache.view(torch.uint8).flatten()
159 + block_idx = slots // block_size
160 + pos = slots % block_size
161 + block_base = block_idx * cache.stride(0)
162 + token_base = block_base + pos * 576
163 + scale_base = block_base + block_size * 576 + pos * 8
164
165 fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn
132 - nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
166 + nope_offsets = torch.arange(NOPE_HEAD_DIM, device=cache.device)
167 + nope_u8 = cache_flat[token_base[:, None] + nope_offsets]
168 nope = nope_u8.view(fp8_dtype).to(torch.float32)
169 + scale_offsets = torch.arange(7, device=cache.device)
170 scales = torch.exp2(
135 - cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0
171 + cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0
172 )
137 - nope = nope * scales.repeat_interleave(64)
138 - rope_u8 = cache_flat[
139 - token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
140 - ]
141 - rope = rope_u8.view(torch.bfloat16).to(torch.float32)
142 - return torch.cat([nope, rope])
173 + nope = nope * scales.repeat_interleave(64, dim=1)
174 + rope_offsets = torch.arange(ROPE_HEAD_DIM * 2, device=cache.device)
175 + rope_u8 = cache_flat[token_base[:, None] + NOPE_HEAD_DIM + rope_offsets]
176 + rope = rope_u8.contiguous().view(torch.bfloat16).to(torch.float32)
177 + return torch.cat([nope, rope], dim=1)
178
179
180 def _ref_sparse_decode_ragged(
@@ -158,19 +193,30 @@ def _ref_sparse_decode_ragged(
193 out = torch.empty_like(q_f32)
194
195 for query_idx in range(q.shape[0]):
161 - row_kv = [
162 - _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz)
163 - for slot in main_rows[query_idx]
164 - ]
165 - if extra_cache is not None and extra_rows is not None:
166 - row_kv.extend(
167 - _read_fp8_ds_mla_cache(
168 - extra_cache, int(slot), block_size, extra_use_fnuz
196 + row_kv = []
197 + if main_rows[query_idx]:
198 + main_slots = torch.tensor(
199 + main_rows[query_idx], dtype=torch.int64, device=q.device
200 + )
201 + row_kv.append(
202 + _read_fp8_ds_mla_cache_rows(
203 + main_cache, main_slots, block_size, main_use_fnuz
204 + )
205 + )
206 + if extra_cache is not None and extra_rows is not None and extra_rows[query_idx]:
207 + extra_slots = torch.tensor(
208 + extra_rows[query_idx], dtype=torch.int64, device=q.device
209 + )
210 + row_kv.append(
211 + _read_fp8_ds_mla_cache_rows(
212 + extra_cache, extra_slots, block_size, extra_use_fnuz
213 )
170 - for slot in extra_rows[query_idx]
214 )
215
173 - kv = torch.stack(row_kv).to(q.device)
216 + if not row_kv:
217 + out[query_idx] = 0
218 + continue
219 + kv = torch.cat(row_kv)
220 for head_idx in range(q.shape[1]):
221 scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
222 if attn_sink is not None:
@@ -198,6 +244,46 @@ def _ragged_from_rows(
244 )
245
246
247 +def _launch_sparse_decode_reduce(
248 + part_m: torch.Tensor,
249 + part_l: torch.Tensor,
250 + part_acc: torch.Tensor,
251 + adaptive_splits: bool,
252 +) -> torch.Tensor:
253 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
254 +
255 + num_queries, num_splits, num_heads = part_m.shape
256 + out = torch.empty(
257 + (num_queries, num_heads, HEAD_DIM),
258 + dtype=torch.bfloat16,
259 + device=part_m.device,
260 + )
261 + attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device)
262 + mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](
263 + part_m,
264 + part_l,
265 + part_acc,
266 + attn_sink,
267 + out,
268 + out.stride(0),
269 + out.stride(1),
270 + part_m.stride(0),
271 + part_m.stride(1),
272 + part_acc.stride(0),
273 + part_acc.stride(1),
274 + part_acc.stride(2),
275 + num_heads,
276 + HAS_ATTN_SINK=False,
277 + ADAPTIVE_SPLITS=adaptive_splits,
278 + COMB_DIM=HEAD_DIM,
279 + BLOCK_H=1,
280 + NUM_SPLITS=num_splits,
281 + SPLITS_PAD=1 << (num_splits - 1).bit_length(),
282 + num_warps=4,
283 + )
284 + return out
285 +
286 +
287 @torch.inference_mode()
288 def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None:
289 from vllm._aiter_ops import rocm_aiter_ops
@@ -312,6 +398,19 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None:
398 torch.testing.assert_close(actual_lens, expected_lens)
399
400
401 +def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None:
402 + from vllm.models.deepseek_v4.amd import rocm as mod
403 +
404 + monkeypatch.setattr(mod, "_ON_GFX950", True)
405 + assert mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True)
406 + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", True, True)
407 + assert not mod._trust_dsv4_extra_cache_nan_free("bfloat16", False, True)
408 + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, False)
409 +
410 + monkeypatch.setattr(mod, "_ON_GFX950", False)
411 + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True)
412 +
413 +
414 @torch.inference_mode()
415 def test_sparse_attn_prefill_ragged_kernel() -> None:
416 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
@@ -366,6 +465,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
465 attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
466 scale = HEAD_DIM**-0.5
467
468 + out = torch.empty_like(q)
469 actual = _rocm_sparse_attn_decode_ragged_triton(
470 q=q,
471 main_cache=main_cache,
@@ -378,6 +478,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
478 extra_cache=extra_cache,
479 extra_indices=extra_indices,
480 extra_indptr=extra_indptr,
481 + out=out,
482 )
483 expected = _ref_sparse_decode_ragged(
484 q=q,
@@ -391,9 +492,111 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
492 main_use_fnuz=main_use_fnuz,
493 )
494
495 + assert actual.data_ptr() == out.data_ptr()
496 torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
497
498
499 +@requires_gfx950
500 +@torch.inference_mode()
501 +def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None:
502 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
503 + _rocm_sparse_attn_decode_ragged_triton,
504 + )
505 +
506 + device = torch.device("cuda")
507 + block_size = 4
508 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device)
509 + extra_cache = torch.zeros_like(main_cache)
510 + _poison_fp8_ds_mla_cache_row(main_cache, block_size)
511 + _poison_fp8_ds_mla_cache_row(extra_cache, block_size)
512 + indices = torch.zeros(1, dtype=torch.int32, device=device)
513 + indptr = torch.tensor([0, 1], dtype=torch.int32, device=device)
514 +
515 + actual = _rocm_sparse_attn_decode_ragged_triton(
516 + q=torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device),
517 + main_cache=main_cache,
518 + main_indices=indices,
519 + main_indptr=indptr,
520 + scale=HEAD_DIM**-0.5,
521 + attn_sink=None,
522 + nope_head_dim=NOPE_HEAD_DIM,
523 + rope_head_dim=ROPE_HEAD_DIM,
524 + extra_cache=extra_cache,
525 + extra_indices=indices,
526 + extra_indptr=indptr,
527 + )
528 +
529 + assert not torch.isnan(actual).any()
530 + assert torch.equal(actual, torch.zeros_like(actual))
531 +
532 +
533 +@pytest.mark.parametrize("on_gfx950", [False, True])
534 +@torch.inference_mode()
535 +def test_rocm_ragged_graph_buffer_view_tracks_source_width(
536 + monkeypatch, on_gfx950: bool
537 +) -> None:
538 + from vllm.models.deepseek_v4.amd import rocm as rocm_mod
539 +
540 + monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950)
541 +
542 + indices_buffer = torch.full((16,), -1, dtype=torch.int32)
543 + indptr_buffer = torch.full((3,), -1, dtype=torch.int32)
544 + first_indices = torch.tensor([3, 5, 7], dtype=torch.int32)
545 + first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32)
546 + first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
547 + first_indices,
548 + first_indptr,
549 + indices_buffer,
550 + indptr_buffer,
551 + num_rows=2,
552 + max_entries_per_row=8,
553 + )
554 +
555 + second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)
556 + second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32)
557 + second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
558 + second_indices,
559 + second_indptr,
560 + indices_buffer,
561 + indptr_buffer,
562 + num_rows=2,
563 + max_entries_per_row=8,
564 + )
565 +
566 + expected_first_entries = (
567 + first_indices.numel() if on_gfx950 else indices_buffer.numel()
568 + )
569 + expected_second_entries = (
570 + second_indices.numel() if on_gfx950 else indices_buffer.numel()
571 + )
572 + assert first_view.numel() == expected_first_entries
573 + assert second_view.numel() == expected_second_entries
574 + assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr()
575 + assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr()
576 + assert torch.equal(second_view[: second_indices.numel()], second_indices)
577 + assert torch.equal(second_indptr_view, second_indptr)
578 +
579 +
580 +def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None:
581 + from vllm.models.deepseek_v4.amd import rocm as rocm_mod
582 + from vllm.models.deepseek_v4.sparse_mla import (
583 + DeepseekV4SparseMLAMetadataBuilder,
584 + )
585 +
586 + metadata = SimpleNamespace(for_cudagraph_capture=False)
587 + monkeypatch.setattr(
588 + DeepseekV4SparseMLAMetadataBuilder,
589 + "build_for_cudagraph_capture",
590 + lambda *_: metadata,
591 + )
592 + builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder)
593 +
594 + actual = builder.build_for_cudagraph_capture(SimpleNamespace())
595 +
596 + assert actual is metadata
597 + assert actual.for_cudagraph_capture is _on_gfx950()
598 +
599 +
600 @requires_split_decode_arch
601 @torch.inference_mode()
602 def test_decode_num_splits_heuristic(monkeypatch) -> None:
@@ -408,6 +611,9 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None:
611 # A tiny batch on a large device should split to add parallelism.
612 assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1
613
614 + # The shared gfx942 selector retains its original 16-split ceiling.
615 + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16
616 +
617 # The chosen count always stays within the searched [1, 16] range, and a
618 # zero-length workload never splits (no work to parallelize).
619 for num_queries in (1, 4, 24, 224, 1024):
@@ -418,6 +624,16 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None:
624 assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
625
626
627 +@torch.inference_mode()
628 +def test_decode_num_splits_gfx950(monkeypatch) -> None:
629 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
630 +
631 + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
632 + assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32
633 + assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4
634 + assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1
635 +
636 +
637 @requires_split_decode_arch
638 @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
639 @pytest.mark.parametrize("with_extra", [True, False])
@@ -473,7 +689,14 @@ def test_sparse_attn_decode_split_k_kernel(
689 scale = HEAD_DIM**-0.5
690
691 # Pin the split count so each parametrized value is exercised deterministically.
476 - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits)
692 + split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits"
693 + other_split_fn = (
694 + "_decode_num_splits" if _on_gfx950() else "_decode_gfx950_num_splits"
695 + )
696 + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits)
697 + monkeypatch.setattr(
698 + mod, other_split_fn, lambda *args, **kwargs: pytest.fail("wrong selector")
699 + )
700
701 actual = mod._rocm_sparse_attn_decode_ragged_triton(
702 q=q,
@@ -503,6 +726,218 @@ def test_sparse_attn_decode_split_k_kernel(
726 torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
727
728
729 +@requires_gfx950
730 +@torch.inference_mode()
731 +def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None:
732 + device = torch.device("cuda")
733 + part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device)
734 + part_l = torch.zeros_like(part_m)
735 + part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device)
736 + part_m[:, :2] = 0
737 + part_l[:, :2] = 1
738 + part_acc[:, 0] = 1
739 + part_acc[:, 1] = 3
740 +
741 + actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True)
742 +
743 + assert torch.isfinite(actual).all()
744 + assert torch.equal(actual, torch.full_like(actual, 2))
745 +
746 +
747 +@requires_gfx950
748 +@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
749 +@torch.inference_mode()
750 +def test_sparse_attn_decode_gfx950_outer64_boundaries(
751 + monkeypatch, extra_len: int
752 +) -> None:
753 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
754 +
755 + device = torch.device("cuda")
756 + torch.manual_seed(13)
757 + block_size = 4
758 + num_heads = 16
759 + num_extra_rows = 80
760 + q = torch.randn(2, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device)
761 + q *= 0.125
762 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device)
763 + main_indices = torch.empty(0, dtype=torch.int32, device=device)
764 + main_indptr = torch.zeros(3, dtype=torch.int32, device=device)
765 + extra_cache = _pack_fp8_ds_mla_cache(
766 + torch.randn(num_extra_rows, HEAD_DIM, dtype=torch.bfloat16, device=device)
767 + * 0.125,
768 + block_size,
769 + use_fnuz=False,
770 + )
771 + _poison_fp8_ds_mla_cache_row(extra_cache, block_size)
772 +
773 + raw_row = list(range(1, extra_len + 1))
774 + if extra_len > 3:
775 + raw_row[3] = -1
776 + if extra_len > 40:
777 + raw_row[40] = num_extra_rows
778 + if extra_len > 64:
779 + raw_row[64] = num_extra_rows + 1024
780 + extra_indices, extra_indptr = _ragged_from_rows([raw_row, []], device)
781 + valid_row = [slot for slot in raw_row if 0 <= slot < num_extra_rows]
782 + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
783 +
784 + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
785 + actual = mod._rocm_sparse_attn_decode_ragged_triton(
786 + q=q,
787 + main_cache=main_cache,
788 + main_indices=main_indices,
789 + main_indptr=main_indptr,
790 + scale=HEAD_DIM**-0.5,
791 + attn_sink=attn_sink,
792 + nope_head_dim=NOPE_HEAD_DIM,
793 + rope_head_dim=ROPE_HEAD_DIM,
794 + extra_cache=extra_cache,
795 + extra_indices=extra_indices,
796 + extra_indptr=extra_indptr,
797 + )
798 + expected = _ref_sparse_decode_ragged(
799 + q=q,
800 + main_cache=main_cache,
801 + main_rows=[[], []],
802 + scale=HEAD_DIM**-0.5,
803 + attn_sink=attn_sink,
804 + block_size=block_size,
805 + extra_cache=extra_cache,
806 + extra_rows=[valid_row, []],
807 + )
808 +
809 + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
810 + assert torch.equal(actual[1], torch.zeros_like(actual[1]))
811 +
812 +
813 +@requires_gfx950
814 +@torch.inference_mode()
815 +def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None:
816 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
817 +
818 + device = torch.device("cuda")
819 + torch.manual_seed(17)
820 + block_size = 64
821 + num_queries = 16
822 + num_heads = 16
823 + num_splits = 8
824 + extra_per_query = 65 * num_splits
825 + max_extra_per_query = 8192
826 + q = (
827 + torch.randn(
828 + num_queries,
829 + num_heads,
830 + HEAD_DIM,
831 + dtype=torch.bfloat16,
832 + device=device,
833 + )
834 + * 0.125
835 + )
836 + main_cache = _pack_fp8_ds_mla_cache(
837 + torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125,
838 + block_size,
839 + use_fnuz=False,
840 + )
841 + extra_cache = _pack_fp8_ds_mla_cache(
842 + torch.randn(
843 + num_queries * extra_per_query,
844 + HEAD_DIM,
845 + dtype=torch.bfloat16,
846 + device=device,
847 + )
848 + * 0.125,
849 + block_size,
850 + use_fnuz=False,
851 + )
852 + main_rows = [[query_idx] for query_idx in range(num_queries)]
853 + extra_rows = [
854 + list(range(query_idx * extra_per_query, (query_idx + 1) * extra_per_query))
855 + for query_idx in range(num_queries)
856 + ]
857 + main_indices, main_indptr = _ragged_from_rows(main_rows, device)
858 + short_extra_rows = [row[:64] for row in extra_rows]
859 + long_indices, long_indptr = _ragged_from_rows(extra_rows, device)
860 + short_indices, short_indptr = _ragged_from_rows(short_extra_rows, device)
861 + extra_indices = torch.full(
862 + (num_queries * max_extra_per_query,),
863 + -1,
864 + dtype=torch.int32,
865 + device=device,
866 + )
867 + extra_indices[: long_indices.numel()].copy_(long_indices)
868 + extra_indptr = long_indptr.clone()
869 + extra_indices_ptr = extra_indices.data_ptr()
870 + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
871 + out = torch.empty_like(q)
872 +
873 + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: num_splits)
874 +
875 + def run_decode() -> torch.Tensor:
876 + return mod._rocm_sparse_attn_decode_ragged_triton(
877 + q=q,
878 + main_cache=main_cache,
879 + main_indices=main_indices,
880 + main_indptr=main_indptr,
881 + scale=HEAD_DIM**-0.5,
882 + attn_sink=attn_sink,
883 + nope_head_dim=NOPE_HEAD_DIM,
884 + rope_head_dim=ROPE_HEAD_DIM,
885 + extra_cache=extra_cache,
886 + extra_indices=extra_indices,
887 + extra_indptr=extra_indptr,
888 + out=out,
889 + extra_cache_nan_free=True,
890 + adaptive_splits=True,
891 + )
892 +
893 + run_decode()
894 + torch.accelerator.synchronize()
895 + graph = torch.cuda.CUDAGraph()
896 + with torch.cuda.graph(graph):
897 + captured_out = run_decode()
898 + torch.accelerator.synchronize()
899 + captured_long = out.clone()
900 + expected_long = _ref_sparse_decode_ragged(
901 + q=q,
902 + main_cache=main_cache,
903 + main_rows=main_rows,
904 + scale=HEAD_DIM**-0.5,
905 + attn_sink=attn_sink,
906 + block_size=block_size,
907 + extra_cache=extra_cache,
908 + extra_rows=extra_rows,
909 + )
910 + torch.testing.assert_close(captured_long, expected_long, atol=2e-2, rtol=2e-2)
911 +
912 + extra_indices[: short_indices.numel()].copy_(short_indices)
913 + extra_indptr.copy_(short_indptr)
914 + graph.replay()
915 + torch.accelerator.synchronize()
916 + short_out = out.clone()
917 + expected_short = _ref_sparse_decode_ragged(
918 + q=q,
919 + main_cache=main_cache,
920 + main_rows=main_rows,
921 + scale=HEAD_DIM**-0.5,
922 + attn_sink=attn_sink,
923 + block_size=block_size,
924 + extra_cache=extra_cache,
925 + extra_rows=short_extra_rows,
926 + )
927 +
928 + assert captured_out.data_ptr() == out.data_ptr()
929 + assert extra_indices.data_ptr() == extra_indices_ptr
930 + assert extra_indices.numel() == num_queries * max_extra_per_query
931 + assert not torch.equal(short_out, captured_long)
932 + torch.testing.assert_close(short_out, expected_short, atol=2e-2, rtol=2e-2)
933 +
934 + extra_indices[: long_indices.numel()].copy_(long_indices)
935 + extra_indptr.copy_(long_indptr)
936 + graph.replay()
937 + torch.accelerator.synchronize()
938 + torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2)
939 +
940 +
941 # ---------------------------------------------------------------------------
942 # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
943 # ---------------------------------------------------------------------------
tests/kernels/test_compressor_kv_cache.py
+239 -4
@@ -28,16 +28,81 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
28 _fused_kv_compress_norm_rope_insert_indexer_attn,
29 _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
30 _launch_two_stage_sparse_attn_compressor,
31 + compress_norm_rope_store_triton,
32 )
33 from vllm.models.deepseek_v4.compressor import _get_c128_boundary
34 from vllm.platforms import current_platform
35 from vllm.v1.attention.backends.mla.compressor_utils import (
36 get_dspark_swa_index_width,
37 )
38 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
39 + cp_gather_indexer_k_quant_cache_triton,
40 + indexer_k_quant_and_cache_triton,
41 +)
42
43 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
44
45
46 +def _on_gfx950() -> bool:
47 + if not current_platform.is_rocm():
48 + return False
49 + try:
50 + from vllm.platforms.rocm import _ON_GFX950
51 +
52 + return _ON_GFX950
53 + except Exception:
54 + return False
55 +
56 +
57 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch")
58 +def test_cp_gather_despecialized_kernel_is_gfx950_only(monkeypatch):
59 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
60 +
61 + class FakeKernel:
62 + def __init__(self):
63 + self.calls = []
64 +
65 + def __getitem__(self, grid):
66 + def launch(*args):
67 + self.calls.append((grid, args))
68 +
69 + return launch
70 +
71 + legacy_kernel = FakeKernel()
72 + gfx950_kernel = FakeKernel()
73 + monkeypatch.setattr(mod, "_cp_gather_indexer_quant_cache_kernel", legacy_kernel)
74 + monkeypatch.setattr(
75 + mod,
76 + "_cp_gather_indexer_quant_cache_gfx950_kernel",
77 + gfx950_kernel,
78 + )
79 +
80 + k_cache = torch.zeros((4, 1, 132), dtype=torch.uint8)
81 + k_fp8 = torch.empty((5, 128), dtype=current_platform.fp8_dtype())
82 + k_scale = torch.empty((5, 4), dtype=torch.uint8)
83 + block_table = torch.zeros((2, 7), dtype=torch.int32)
84 + cu_seqlen = torch.tensor([0, 2, 5], dtype=torch.int32)
85 + token_to_seq = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32)
86 + args = (k_cache, k_fp8, k_scale, block_table, cu_seqlen, token_to_seq)
87 +
88 + monkeypatch.setattr(mod, "_ON_GFX950", True)
89 + mod.cp_gather_indexer_k_quant_cache_triton(*args)
90 + assert len(gfx950_kernel.calls) == 1
91 + assert not legacy_kernel.calls
92 + gfx950_grid, gfx950_args = gfx950_kernel.calls[0]
93 + assert gfx950_grid == (5,)
94 + assert len(gfx950_args) == 18
95 + assert gfx950_args[-3:] == (2, 7, 4)
96 +
97 + monkeypatch.setattr(mod, "_ON_GFX950", False)
98 + mod.cp_gather_indexer_k_quant_cache_triton(*args)
99 + assert len(legacy_kernel.calls) == 1
100 + legacy_grid, legacy_args = legacy_kernel.calls[0]
101 + assert legacy_grid == (5,)
102 + assert len(legacy_args) == 19
103 + assert legacy_args[-4:] == (5, 2, 7, 4)
104 +
105 +
106 @pytest.mark.parametrize(
107 ("window_size", "num_speculative_tokens", "expected"),
108 [(128, 5, 192), (512, 5, 576), (1024, 0, 1024)],
@@ -91,6 +156,129 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
156 return x_fp8, scales
157
158
159 +def _decode_dsv4_cache_row(
160 + cache: torch.Tensor, block_size: int, scrub_nan: bool
161 +) -> torch.Tensor:
162 + flat = cache.flatten()
163 + nope = flat[:448].view(torch.float8_e4m3fn).to(torch.bfloat16)
164 + encoded = flat[block_size * 576 : block_size * 576 + 7]
165 + scales = torch.exp2(encoded.to(torch.float32) - 127.0).to(torch.bfloat16)
166 + nope = nope * scales.repeat_interleave(64)
167 + rope = flat[448:576].view(torch.bfloat16)
168 + decoded = torch.cat((nope, rope))
169 + if scrub_nan:
170 + decoded = torch.where(decoded == decoded, decoded, 0.0)
171 + return decoded
172 +
173 +
174 +def _assert_nan_free_cache_matches_legacy_scrub(
175 + cache: torch.Tensor, block_size: int
176 +) -> None:
177 + flat = cache.flatten()
178 + scale_base = block_size * 576
179 + scale_codes = flat[scale_base : scale_base + 8]
180 + assert scale_codes[0].item() == 254
181 + assert scale_codes[1].item() == 247
182 + assert scale_codes[:7].max().item() <= 254
183 + nope_bytes = flat[:448]
184 + assert not ((nope_bytes == 0x7F) | (nope_bytes == 0xFF)).any()
185 +
186 + rope = flat[448:576].view(torch.bfloat16)
187 + assert not torch.isnan(rope).any()
188 + assert torch.isposinf(rope[0])
189 + assert torch.equal(rope[1:4], torch.zeros_like(rope[1:4]))
190 +
191 + legacy_cache = cache.clone()
192 + legacy_flat = legacy_cache.flatten()
193 + legacy_flat[scale_base] = 255
194 + legacy_rope = legacy_flat[448:576].view(torch.bfloat16)
195 + legacy_rope[1:4] = float("nan")
196 + canonical = _decode_dsv4_cache_row(cache, block_size, scrub_nan=False)
197 + legacy = _decode_dsv4_cache_row(legacy_cache, block_size, scrub_nan=True)
198 + torch.testing.assert_close(canonical, legacy, rtol=0, atol=0)
199 + assert torch.isinf(canonical[0])
200 + assert torch.isposinf(canonical[64])
201 +
202 +
203 +@pytest.mark.skipif(
204 + not _on_gfx950(),
205 + reason="NaN-free fp8_ds_mla compressed-cache contract is gfx950-only",
206 +)
207 +@pytest.mark.parametrize("writer", ["single_pass", "two_stage_finalizer"])
208 +def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None:
209 + head_dim = 512
210 + rope_dim = 64
211 + block_size = 4
212 + device = "cuda"
213 +
214 + positions = torch.zeros(1, dtype=torch.int64, device=device)
215 + slot_mapping = torch.zeros(1, dtype=torch.int64, device=device)
216 + rms_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
217 + rms_weight[0] = float("inf")
218 + rms_weight[64] = torch.finfo(torch.bfloat16).max
219 + rms_weight[448] = float("inf")
220 + rms_weight[450] = float("nan")
221 + cos_sin_cache = torch.zeros(1, rope_dim, dtype=torch.float32, device=device)
222 + cos_sin_cache[:, : rope_dim // 2] = 1.0
223 + cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device)
224 +
225 + state_cache = torch.zeros(1, 1, 2 * head_dim, dtype=torch.float32, device=device)
226 + state_cache[..., :head_dim] = 1.0
227 + token_to_req = torch.zeros(1, dtype=torch.int32, device=device)
228 + block_table = torch.zeros(1, 1, dtype=torch.int32, device=device)
229 +
230 + if writer == "single_pass":
231 + compress_norm_rope_store_triton(
232 + state_cache=state_cache,
233 + num_actual=1,
234 + token_to_req_indices=token_to_req,
235 + positions=positions,
236 + slot_mapping=slot_mapping,
237 + block_table=block_table,
238 + block_size=1,
239 + state_width=head_dim,
240 + cos_sin_cache=cos_sin_cache,
241 + kv_cache=cache,
242 + k_cache_metadata=SimpleNamespace(slot_mapping=slot_mapping),
243 + pdl_kwargs={},
244 + head_dim=head_dim,
245 + rope_head_dim=rope_dim,
246 + compress_ratio=1,
247 + overlap=False,
248 + use_fp4_cache=False,
249 + rms_norm_weight=rms_weight,
250 + rms_norm_eps=1e-6,
251 + quant_block=64,
252 + token_stride=576,
253 + scale_dim=8,
254 + )
255 + else:
256 + _launch_two_stage_sparse_attn_compressor(
257 + state_cache,
258 + token_to_req,
259 + positions,
260 + slot_mapping,
261 + block_table,
262 + 1,
263 + head_dim,
264 + 1,
265 + cos_sin_cache,
266 + cache,
267 + slot_mapping,
268 + rms_weight,
269 + 1e-6,
270 + 64,
271 + 576,
272 + 8,
273 + head_dim,
274 + rope_dim,
275 + 1,
276 + torch.empty(1, head_dim, dtype=torch.float32, device=device),
277 + )
278 +
279 + _assert_nan_free_cache_matches_legacy_scrub(cache, block_size)
280 +
281 +
282 @pytest.mark.parametrize(
283 ("starts", "query_start_loc", "expected"),
284 [
@@ -402,7 +590,8 @@ def test_indexer_gather_accepts_upper_bound_output():
590 valid_tokens = 9
591 upper_bound_tokens = 13
592 block_size = 16
405 - num_blocks = 2
593 + num_seqs = 3
594 + num_blocks = num_seqs
595 sentinel = 123
596 device = "cuda"
597
@@ -410,13 +599,15 @@ def test_indexer_gather_accepts_upper_bound_output():
599 kv_cache = torch.zeros(
600 num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device
601 )
413 - slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device)
602 + slot_mapping = torch.tensor(
603 + [0, 1, 2, 16, 17, 18, 32, 33, 34], dtype=torch.int64, device=device
604 + )
605 ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0")
606
607 block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
417 - 0
608 + 1
609 )
419 - cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device)
610 + cu_seq_lens = torch.tensor([0, 3, 6, 9], dtype=torch.int32, device=device)
611 dst_k = torch.full(
612 (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device
613 )
@@ -431,8 +622,52 @@ def test_indexer_gather_accepts_upper_bound_output():
622 ops.cp_gather_indexer_k_quant_cache(
623 kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
624 )
625 +
626 + if current_platform.is_rocm():
627 + triton_kv_cache = torch.zeros_like(kv_cache)
628 + indexer_k_quant_and_cache_triton(
629 + k,
630 + triton_kv_cache,
631 + slot_mapping,
632 + quant_block_size,
633 + "ue8m0",
634 + )
635 + triton_dst_k = torch.full_like(dst_k, sentinel)
636 + triton_dst_scale = torch.full_like(dst_scale, sentinel)
637 + token_to_seq = torch.cat(
638 + (
639 + torch.repeat_interleave(
640 + torch.arange(num_seqs, dtype=torch.int32, device=device), 3
641 + ),
642 + torch.full(
643 + (upper_bound_tokens - valid_tokens,),
644 + -1,
645 + dtype=torch.int32,
646 + device=device,
647 + ),
648 + )
649 + )
650 + cp_gather_indexer_k_quant_cache_triton(
651 + triton_kv_cache,
652 + triton_dst_k.view(current_platform.fp8_dtype()),
653 + triton_dst_scale,
654 + block_table,
655 + cu_seq_lens,
656 + token_to_seq,
657 + )
658 torch.accelerator.synchronize()
659
660 + if current_platform.is_rocm():
661 + triton_recovered = triton_dst_k[:valid_tokens].view(
662 + current_platform.fp8_dtype()
663 + ).float() * triton_dst_scale[:valid_tokens].view(torch.float32)
664 + triton_error = (triton_recovered - k.float()).abs().amax(dim=1)
665 + max_triton_error = (
666 + 16.0 * triton_dst_scale[:valid_tokens].view(torch.float32).flatten()
667 + )
668 + assert torch.all(triton_error <= max_triton_error)
669 + assert torch.all(triton_dst_k[valid_tokens:] == sentinel)
670 + assert torch.all(triton_dst_scale[valid_tokens:] == sentinel)
671 k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[
672 :valid_tokens
673 ].view(torch.float32)
vllm/models/deepseek_v4/amd/rocm.py
+48 -3
@@ -19,6 +19,7 @@ from vllm.models.deepseek_v4.sparse_mla import (
19 DeepseekV4SparseMLAMetadataBuilder,
20 )
21 from vllm.platforms import current_platform
22 +from vllm.platforms.rocm import _ON_GFX950
23 from vllm.triton_utils import tl, triton
24 from vllm.v1.attention.backend import (
25 CommonAttentionMetadata,
@@ -36,6 +37,19 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
37 from vllm.v1.worker.workspace import current_workspace_manager
38
39
40 +def _trust_dsv4_extra_cache_nan_free(
41 + kv_cache_dtype: str,
42 + has_kv_transfer: bool,
43 + has_extra_cache: bool,
44 +) -> bool:
45 + return (
46 + _ON_GFX950
47 + and kv_cache_dtype == "fp8_ds_mla"
48 + and not has_kv_transfer
49 + and has_extra_cache
50 + )
51 +
52 +
53 def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
54 lengths = lengths.to(dtype=torch.int32).contiguous()
55 indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
@@ -294,9 +308,13 @@ def _copy_ragged_to_graph_buffers(
308
309 max_entries = max(num_rows * max_entries_per_row, 1)
310 ragged_out = ragged_indices_buffer[:max_entries]
297 - nnz = ragged_indices.numel()
298 - if nnz > 0:
299 - ragged_out[:nnz].copy_(ragged_indices, non_blocking=True)
311 + source_entries = ragged_indices.numel()
312 + if source_entries > 0:
313 + ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True)
314 + if _ON_GFX950:
315 + # Preserve the graph-stable base pointer while exposing source capacity
316 + # to the sync-free split selector; indptr still carries the true NNZ.
317 + ragged_out = ragged_out[: max(source_entries, 1)]
318 return ragged_out, indptr_out
319
320
@@ -306,6 +324,7 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata):
324
325 c128a_decode_topk_ragged_indices: torch.Tensor | None = None
326 c128a_decode_topk_ragged_indptr: torch.Tensor | None = None
327 + for_cudagraph_capture: bool = False
328
329
330 @dataclass
@@ -370,6 +389,16 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBui
389 c128a_decode_topk_ragged_indptr=ragged_indptr,
390 )
391
392 + def build_for_cudagraph_capture(
393 + self, common_attn_metadata: CommonAttentionMetadata
394 + ) -> DeepseekV4ROCMAiterMLASparseMetadata:
395 + metadata = cast(
396 + DeepseekV4ROCMAiterMLASparseMetadata,
397 + super().build_for_cudagraph_capture(common_attn_metadata),
398 + )
399 + metadata.for_cudagraph_capture = _ON_GFX950
400 + return metadata
401 +
402
403 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
404 # Keep fused multi-step decode disabled until update_draft_decode_metadata()
@@ -451,7 +480,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
480 backend_cls = DeepseekV4ROCMAiterMLASparseBackend
481
482 def __init__(self, *args, **kwargs):
483 + vllm_config = args[0] if args else kwargs["vllm_config"]
484 super().__init__(*args, **kwargs)
485 + self._has_kv_transfer = vllm_config.kv_transfer_config is not None
486 # Block scale for the preshuffled weight; None = not preshuffled.
487 self._wqa_wkv_scale: torch.Tensor | None = None
488 self._wo_b_scale: torch.Tensor | None = None
@@ -605,6 +636,13 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
636 attn_metadata=rocm_metadata,
637 swa_only=swa_only,
638 output=output[:num_decode_tokens],
639 + adaptive_splits=(
640 + _ON_GFX950
641 + and not swa_only
642 + and self.compress_ratio == 128
643 + and rocm_metadata is not None
644 + and rocm_metadata.for_cudagraph_capture
645 + ),
646 )
647
648 def _forward_decode(
@@ -615,6 +653,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
653 attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
654 swa_only: bool,
655 output: torch.Tensor,
656 + adaptive_splits: bool,
657 ) -> None:
658 num_decodes = swa_metadata.num_decodes
659 num_decode_tokens = swa_metadata.num_decode_tokens
@@ -666,6 +705,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
705 nope_head_dim=self.nope_head_dim,
706 rope_head_dim=self.rope_head_dim,
707 output=output,
708 + adaptive_splits=adaptive_splits,
709 + extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
710 + self.kv_cache_dtype,
711 + self._has_kv_transfer,
712 + not swa_only and kv_cache is not None,
713 + ),
714 )
715
716 def _forward_prefill(
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+21 -2
@@ -24,8 +24,14 @@ from typing import Any
24
25 import torch
26
27 +from vllm.platforms import current_platform
28 from vllm.triton_utils import tl, triton
29
30 +if current_platform.is_rocm():
31 + from vllm.platforms.rocm import _ON_GFX950
32 +else:
33 + _ON_GFX950 = False
34 +
35 from .fused_indexer_q import _fp32x2_to_fp4x2
36
37
@@ -61,12 +67,15 @@ def compress_norm_rope_store_triton(
67 if head_dim == 512:
68 kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
69 num_warps = 4
70 + kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950}
71 elif use_fp4_cache:
72 kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
73 num_warps = 1
74 + kernel_kwargs = {}
75 else:
76 kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
77 num_warps = 1
78 + kernel_kwargs = {}
79
80 kernel[(num_actual,)](
81 # state cache
@@ -103,6 +112,7 @@ def compress_norm_rope_store_triton(
112 SCALE_DIM=scale_dim,
113 KV_BLOCK_STRIDE=kv_cache.stride(0),
114 num_warps=num_warps,
115 + **kernel_kwargs,
116 **pdl_kwargs,
117 )
118
@@ -145,6 +155,7 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn(
155 TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4
156 SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad)
157 KV_BLOCK_STRIDE: tl.constexpr,
158 + SANITIZE_CACHE_NANS: tl.constexpr,
159 ):
160 """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope).
161
@@ -261,7 +272,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn(
272
273 scale_idx = tl.arange(0, N_QUANT_BLOCKS)
274 encoded = exponents + 127.0
264 - encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0)
275 + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
276 + encoded = tl.maximum(tl.minimum(encoded, max_encoded), 0.0)
277 tl.store(
278 scale_ptr + scale_idx,
279 encoded.to(tl.uint8),
@@ -289,6 +301,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn(
301 new_even = even * cos_v - odd * sin_v
302 new_odd = odd * cos_v + even * sin_v
303 result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32
304 + if SANITIZE_CACHE_NANS:
305 + result = tl.where(result == result, result, 0.0)
306
307 # Store rotated rope portion as bf16 into the cache's bf16 area.
308 bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
@@ -417,6 +431,7 @@ def _finalize_norm_rope_quant_store_sparse_attn(
431 TOKEN_STRIDE: tl.constexpr,
432 SCALE_DIM: tl.constexpr,
433 KV_BLOCK_STRIDE: tl.constexpr,
434 + SANITIZE_CACHE_NANS: tl.constexpr,
435 ):
436 """Stage 2: read compressed_kv[512] from scratch buffer, then
437 RMSNorm + FP8 quant (nope) + RoPE + bf16 store
@@ -474,7 +489,8 @@ def _finalize_norm_rope_quant_store_sparse_attn(
489 tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM)
490
491 scale_idx = tl.arange(0, N_QUANT_BLOCKS)
477 - encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0)
492 + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
493 + encoded = tl.maximum(tl.minimum(exponents + 127.0, max_encoded), 0.0)
494 tl.store(
495 scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS
496 )
@@ -494,6 +510,8 @@ def _finalize_norm_rope_quant_store_sparse_attn(
510 new_even = even * cos_v - odd * sin_v
511 new_odd = odd * cos_v + even * sin_v
512 result = tl.interleave(new_even, new_odd)
513 + if SANITIZE_CACHE_NANS:
514 + result = tl.where(result == result, result, 0.0)
515 bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
516 rope_local = block - NOPE_HEAD_DIM
517 is_rope = (block >= NOPE_HEAD_DIM) & mask
@@ -564,6 +582,7 @@ def _launch_two_stage_sparse_attn_compressor(
582 TOKEN_STRIDE=token_stride,
583 SCALE_DIM=scale_dim,
584 KV_BLOCK_STRIDE=kv_cache.stride(0),
585 + SANITIZE_CACHE_NANS=_ON_GFX950,
586 )
587
588
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+784 -62
@@ -214,6 +214,87 @@ def _cp_gather_indexer_quant_cache_kernel(
214 tl.store(dst_k_ptr + offset, val, mask=valid_block)
215
216
217 +@triton.jit(do_not_specialize=["num_batches"])
218 +def _cp_gather_indexer_quant_cache_gfx950_kernel(
219 + kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B]
220 + # [n_blks, blk_size, head_dim]
221 + kv_cache_scale_ptr, # [n_blks, blk_size]
222 + k_fp8_ptr, # [num_tokens, head_dim]
223 + k_scale_ptr, # [num_tokens]
224 + block_table_ptr, # [batch_size, block_table_stride]
225 + cu_seqlen_ptr, # [batch_size + 1]
226 + token_to_seq_ptr, # [num_tokens]
227 + block_size,
228 + block_table_stride,
229 + kv_cache_stride,
230 + kv_cache_scale_stride,
231 + LAYOUT: tl.constexpr,
232 + HEAD_DIM: tl.constexpr,
233 + BLOCK_TILE_SIZE: tl.constexpr,
234 + HEAD_TILE_SIZE: tl.constexpr,
235 + num_batches,
236 + BLOCK_TABLE_WIDTH: tl.constexpr,
237 + NUM_BLOCKS: tl.constexpr,
238 +):
239 + tid = tl.program_id(0)
240 + offset = tl.arange(0, HEAD_DIM)
241 + batch_id = tl.load(token_to_seq_ptr + tid)
242 + valid_batch = (batch_id >= 0) & (batch_id < num_batches)
243 + safe_batch_id = tl.where(valid_batch, batch_id, 0)
244 + batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
245 + batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
246 + batch_offset = tid - batch_start
247 + valid_token = valid_batch & (tid >= batch_start) & (tid < batch_end)
248 + if not valid_token:
249 + return
250 + block_table_id = batch_offset // block_size
251 + block_offset = batch_offset % block_size
252 + valid_block_table = (
253 + valid_token
254 + & (block_table_id >= 0)
255 + & (block_table_id < BLOCK_TABLE_WIDTH)
256 + & (block_offset >= 0)
257 + & (block_offset < block_size)
258 + )
259 + safe_block_table_id = tl.where(valid_block_table, block_table_id, 0)
260 + block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id
261 + block_id = tl.load(
262 + block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
263 + )
264 + valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
265 + # The packed KV layout makes per-block strides large
266 + # enough that block_id * stride can exceed 32-bit range.
267 + safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
268 + safe_block_offset = tl.where(valid_block, block_offset, 0)
269 + tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
270 + if LAYOUT == "SHUFFLE":
271 + src_cache_offset = (
272 + safe_block_id * kv_cache_stride
273 + + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE
274 + + tiled_block_offset * HEAD_TILE_SIZE
275 + )
276 + else:
277 + src_cache_offset = (
278 + safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM
279 + )
280 + src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset
281 + dst_offset = tid * HEAD_DIM
282 + src_scale_ptr = kv_cache_scale_ptr + src_scale_offset
283 + src_cache_ptr = kv_cache_ptr + src_cache_offset
284 + dst_k_ptr = k_fp8_ptr + dst_offset
285 + scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0)
286 + tl.store(k_scale_ptr + tid, scale_val)
287 + if LAYOUT == "SHUFFLE":
288 + tiled_src_offset = (
289 + offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE
290 + + offset % HEAD_TILE_SIZE
291 + )
292 + else:
293 + tiled_src_offset = offset
294 + val = tl.load(src_cache_ptr + tiled_src_offset)
295 + tl.store(dst_k_ptr + offset, val, mask=valid_block)
296 +
297 +
298 def cp_gather_indexer_k_quant_cache_triton(
299 k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4]
300 k_fp8: torch.Tensor,
@@ -237,7 +318,7 @@ def cp_gather_indexer_k_quant_cache_triton(
318 grid = (num_tokens,)
319 k_fp8_scale = k_fp8_scale.view(torch.float32)
320 layout = "NORMAL" if block_size == 1 else "SHUFFLE"
240 - _cp_gather_indexer_quant_cache_kernel[grid](
321 + kernel_args = (
322 k_cache_value,
323 k_cache_scale,
324 k_fp8,
@@ -253,11 +334,22 @@ def cp_gather_indexer_k_quant_cache_triton(
334 head_dim,
335 block_tile_size,
336 head_tile_size,
256 - num_tokens,
257 - cu_seqlen.shape[0] - 1,
258 - block_table.shape[1],
259 - num_blocks,
337 )
338 + if _ON_GFX950:
339 + _cp_gather_indexer_quant_cache_gfx950_kernel[grid](
340 + *kernel_args,
341 + cu_seqlen.shape[0] - 1,
342 + block_table.shape[1],
343 + num_blocks,
344 + )
345 + else:
346 + _cp_gather_indexer_quant_cache_kernel[grid](
347 + *kernel_args,
348 + num_tokens,
349 + cu_seqlen.shape[0] - 1,
350 + block_table.shape[1],
351 + num_blocks,
352 + )
353
354
355 # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156
@@ -1211,6 +1303,90 @@ def _sparse_attn_prefill_ragged_kernel(
1303 )
1304
1305
1306 +@triton.jit
1307 +def _decode_e8m0_scales_triton(encoded_scales):
1308 + scale_bits = encoded_scales.to(tl.int32) << 23
1309 + scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits)
1310 + return scale_bits.to(tl.float32, bitcast=True)
1311 +
1312 +
1313 +@triton.jit
1314 +def _load_fp8_ds_mla_gfx950_nope_exact_chunk(
1315 + token_data_ptr,
1316 + token_scale_ptr,
1317 + valid,
1318 + CHUNK_START: tl.constexpr,
1319 + CHUNK_SIZE: tl.constexpr,
1320 + BLOCK_K: tl.constexpr,
1321 + IS_FNUZ: tl.constexpr,
1322 +):
1323 + offsets = CHUNK_START + tl.arange(0, CHUNK_SIZE)
1324 + x_uint8 = tl.load(
1325 + token_data_ptr[:, None] + offsets[None, :],
1326 + mask=valid[:, None],
1327 + other=0,
1328 + )
1329 + scale_offsets = CHUNK_START // 64 + tl.arange(0, CHUNK_SIZE // 64)
1330 + encoded_scales = tl.load(
1331 + token_scale_ptr[:, None] + scale_offsets[None, :],
1332 + mask=valid[:, None],
1333 + other=127,
1334 + )
1335 + scales = _decode_e8m0_scales_triton(encoded_scales)
1336 + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, CHUNK_SIZE // 64, 64))
1337 + scales = tl.reshape(scales, (BLOCK_K, CHUNK_SIZE))
1338 + if IS_FNUZ:
1339 + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32)
1340 + else:
1341 + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32)
1342 + value = (x_f32 * scales).to(tl.bfloat16)
1343 + zero = tl.zeros((BLOCK_K, CHUNK_SIZE), dtype=tl.bfloat16)
1344 + return tl.where(valid[:, None], value, zero)
1345 +
1346 +
1347 +@triton.jit
1348 +def _load_fp8_ds_mla_gfx950_tail128(
1349 + token_data_ptr,
1350 + token_scale_ptr,
1351 + valid,
1352 + NOPE_DIM: tl.constexpr,
1353 + BLOCK_K: tl.constexpr,
1354 + IS_FNUZ: tl.constexpr,
1355 +):
1356 + tail_offsets = 384 + tl.arange(0, 128)
1357 + nope_mask = tail_offsets < NOPE_DIM
1358 + x_uint8 = tl.load(
1359 + token_data_ptr[:, None] + tail_offsets[None, :],
1360 + mask=valid[:, None] & nope_mask[None, :],
1361 + other=0,
1362 + )
1363 + scale_offsets = 6 + tl.arange(0, 2)
1364 + scale_mask = scale_offsets < NOPE_DIM // 64
1365 + encoded_scales = tl.load(
1366 + token_scale_ptr[:, None] + scale_offsets[None, :],
1367 + mask=valid[:, None] & scale_mask[None, :],
1368 + other=127,
1369 + )
1370 + scales = _decode_e8m0_scales_triton(encoded_scales)
1371 + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 2, 64))
1372 + scales = tl.reshape(scales, (BLOCK_K, 128))
1373 + if IS_FNUZ:
1374 + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32)
1375 + else:
1376 + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32)
1377 + nope = (x_f32 * scales).to(tl.bfloat16)
1378 +
1379 + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
1380 + rope = tl.load(
1381 + rope_ptr[:, None] + (tail_offsets[None, :] - NOPE_DIM),
1382 + mask=valid[:, None] & ~nope_mask[None, :],
1383 + other=0.0,
1384 + )
1385 + value = tl.where(nope_mask[None, :], nope, rope)
1386 + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16)
1387 + return tl.where(valid[:, None], value, zero)
1388 +
1389 +
1390 @triton.jit
1391 def _sparse_attn_decode_ragged_kernel(
1392 q_ptr,
@@ -1685,6 +1861,407 @@ def _sparse_attn_decode_partial_kernel(
1861 )
1862
1863
1864 +@triton.jit
1865 +def _sparse_attn_decode_gfx950_partial_loaded_tile(
1866 + q_combined,
1867 + cache_ptr,
1868 + slot,
1869 + valid,
1870 + cache_stride0,
1871 + scale: tl.constexpr,
1872 + head_mask,
1873 + m_i,
1874 + l_i,
1875 + acc_nope_0a,
1876 + acc_nope_0b,
1877 + acc_nope_1,
1878 + acc_tail,
1879 + BLOCK_SIZE: tl.constexpr,
1880 + NOPE_DIM: tl.constexpr,
1881 + BLOCK_K: tl.constexpr,
1882 + IS_FNUZ: tl.constexpr,
1883 + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr,
1884 +):
1885 + safe_slot = tl.where(valid, slot, 0)
1886 + block_idx = safe_slot // BLOCK_SIZE
1887 + pos_in_block = safe_slot % BLOCK_SIZE
1888 + cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0
1889 + token_data_ptr = cache_block_ptr + pos_in_block * 576
1890 + token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8
1891 + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
1892 + token_data_ptr,
1893 + token_scale_ptr,
1894 + valid,
1895 + 0,
1896 + 128,
1897 + BLOCK_K,
1898 + IS_FNUZ,
1899 + )
1900 + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
1901 + token_data_ptr,
1902 + token_scale_ptr,
1903 + valid,
1904 + 128,
1905 + 128,
1906 + BLOCK_K,
1907 + IS_FNUZ,
1908 + )
1909 + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
1910 + token_data_ptr,
1911 + token_scale_ptr,
1912 + valid,
1913 + 256,
1914 + 128,
1915 + BLOCK_K,
1916 + IS_FNUZ,
1917 + )
1918 + k_tail = _load_fp8_ds_mla_gfx950_tail128(
1919 + token_data_ptr,
1920 + token_scale_ptr,
1921 + valid,
1922 + NOPE_DIM,
1923 + BLOCK_K,
1924 + IS_FNUZ,
1925 + )
1926 + if not TRUST_EXTRA_CACHE_NAN_FREE:
1927 + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16)
1928 + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero)
1929 + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero)
1930 + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero)
1931 + k_tail = tl.where(k_tail == k_tail, k_tail, zero)
1932 + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1)
1933 + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1)
1934 + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1)
1935 +
1936 + scores = tl.dot(q_combined, tl.trans(k_combined))
1937 + scores *= scale * 1.4426950408889634
1938 + scores = tl.where(
1939 + head_mask[:, None] & valid[None, :],
1940 + scores,
1941 + -3.4028234663852886e38,
1942 + )
1943 + m_block = tl.max(scores, axis=1)
1944 + m_new = tl.maximum(m_i, m_block)
1945 + alpha = tl.exp2(m_i - m_new)
1946 + p = tl.exp2(scores - m_new[:, None])
1947 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1948 + l_new = l_i * alpha + tl.sum(p, axis=1)
1949 + p_bf16 = p.to(k_nope_0a.dtype)
1950 + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a)
1951 + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b)
1952 + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1)
1953 + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail)
1954 + return (
1955 + m_new,
1956 + l_new,
1957 + acc_nope_0a,
1958 + acc_nope_0b,
1959 + acc_nope_1,
1960 + acc_tail,
1961 + )
1962 +
1963 +
1964 +@triton.jit
1965 +def _sparse_attn_decode_gfx950_partial_kernel(
1966 + q_ptr,
1967 + main_cache_ptr,
1968 + main_indices_ptr,
1969 + main_indptr_ptr,
1970 + extra_cache_ptr,
1971 + extra_indices_ptr,
1972 + extra_indptr_ptr,
1973 + part_m_ptr,
1974 + part_l_ptr,
1975 + part_acc_ptr,
1976 + q_stride0: tl.constexpr,
1977 + q_stride1: tl.constexpr,
1978 + main_cache_stride0: tl.constexpr,
1979 + extra_cache_stride0: tl.constexpr,
1980 + main_num_rows,
1981 + extra_num_rows,
1982 + MAIN_BLOCK_SIZE: tl.constexpr,
1983 + EXTRA_BLOCK_SIZE: tl.constexpr,
1984 + scale: tl.constexpr,
1985 + num_heads: tl.constexpr,
1986 + HAS_EXTRA: tl.constexpr,
1987 + NOPE_DIM: tl.constexpr,
1988 + ROPE_DIM: tl.constexpr,
1989 + IS_FNUZ_MAIN: tl.constexpr,
1990 + IS_FNUZ_EXTRA: tl.constexpr,
1991 + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr,
1992 + ADAPTIVE_SPLITS: tl.constexpr,
1993 + ONE_WAVE_SPLITS: tl.constexpr,
1994 + BLOCK_H: tl.constexpr,
1995 + BLOCK_K: tl.constexpr,
1996 + NUM_SPLITS: tl.constexpr,
1997 + NUM_STAGES: tl.constexpr,
1998 +):
1999 + query_idx = tl.program_id(0)
2000 + split_id = tl.program_id(1)
2001 + pid_h = tl.program_id(2)
2002 +
2003 + tl.static_assert(NOPE_DIM == 448)
2004 + tl.static_assert(ROPE_DIM == 64)
2005 +
2006 + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
2007 + if num_heads % BLOCK_H == 0:
2008 + head_mask = tl.full((BLOCK_H,), True, tl.int1)
2009 + else:
2010 + head_mask = head_offsets < num_heads
2011 + neg_large = -3.4028234663852886e38
2012 +
2013 + if ADAPTIVE_SPLITS:
2014 + main_start = tl.load(main_indptr_ptr + query_idx)
2015 + main_end = tl.load(main_indptr_ptr + query_idx + 1)
2016 + main_len = main_end - main_start
2017 + if HAS_EXTRA:
2018 + extra_start = tl.load(extra_indptr_ptr + query_idx)
2019 + extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
2020 + extra_len = extra_end - extra_start
2021 + else:
2022 + extra_start = 0
2023 + extra_len = 0
2024 + split4_span: tl.constexpr = 4 * BLOCK_K
2025 + split4_iters = (main_len + split4_span - 1) // split4_span
2026 + split4_iters += (extra_len + split4_span - 1) // split4_span
2027 + use_four_splits = split4_iters <= 3
2028 + work_splits = NUM_SPLITS
2029 + if ONE_WAVE_SPLITS > 4 and ONE_WAVE_SPLITS < NUM_SPLITS:
2030 + one_wave_span: tl.constexpr = ONE_WAVE_SPLITS * BLOCK_K
2031 + one_wave_iters = (main_len + one_wave_span - 1) // one_wave_span
2032 + one_wave_iters += (extra_len + one_wave_span - 1) // one_wave_span
2033 + work_splits = tl.where(one_wave_iters <= 3, ONE_WAVE_SPLITS, work_splits)
2034 + work_splits = tl.where(use_four_splits, 4, work_splits)
2035 + if split_id >= work_splits:
2036 + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets
2037 + tl.store(part_m_ptr + pm_base, neg_large, mask=head_mask)
2038 + tl.store(part_l_ptr + pm_base, 0.0, mask=head_mask)
2039 + return
2040 + else:
2041 + work_splits = NUM_SPLITS
2042 +
2043 + nope_offsets_0a = tl.arange(0, 128)
2044 + nope_offsets_0b = 128 + tl.arange(0, 128)
2045 + nope_offsets_0 = tl.arange(0, 256)
2046 + tail_offsets = 256 + tl.arange(0, 256)
2047 + nope_offsets_1 = 256 + tl.arange(0, 128)
2048 + tail_offsets_128 = 384 + tl.arange(0, 128)
2049 +
2050 + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
2051 + q_nope_0 = tl.load(
2052 + q_row_ptr + nope_offsets_0[None, :],
2053 + mask=head_mask[:, None],
2054 + other=0.0,
2055 + )
2056 + q_tail = tl.load(
2057 + q_row_ptr + tail_offsets[None, :],
2058 + mask=head_mask[:, None],
2059 + other=0.0,
2060 + )
2061 + q_combined = tl.cat(q_nope_0, q_tail, dim=1)
2062 +
2063 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
2064 + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
2065 + acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2066 + acc_nope_0b = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2067 + acc_nope_1 = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2068 + acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2069 + k_offsets = tl.arange(0, BLOCK_K)
2070 +
2071 + if not ADAPTIVE_SPLITS:
2072 + main_start = tl.load(main_indptr_ptr + query_idx)
2073 + main_end = tl.load(main_indptr_ptr + query_idx + 1)
2074 + main_len = main_end - main_start
2075 + main_chunk = (main_len + work_splits - 1) // work_splits
2076 + main_lo = split_id * main_chunk
2077 + main_hi = tl.minimum(main_lo + main_chunk, main_len)
2078 +
2079 + for k_start in tl.range(
2080 + main_lo,
2081 + main_hi,
2082 + BLOCK_K,
2083 + num_stages=NUM_STAGES,
2084 + ):
2085 + k_pos = k_start + k_offsets
2086 + in_range = k_pos < main_hi
2087 + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
2088 + valid = in_range & (slot >= 0) & (slot < main_num_rows)
2089 + (
2090 + m_i,
2091 + l_i,
2092 + acc_nope_0a,
2093 + acc_nope_0b,
2094 + acc_nope_1,
2095 + acc_tail,
2096 + ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
2097 + q_combined,
2098 + main_cache_ptr,
2099 + slot,
2100 + valid,
2101 + main_cache_stride0,
2102 + scale,
2103 + head_mask,
2104 + m_i,
2105 + l_i,
2106 + acc_nope_0a,
2107 + acc_nope_0b,
2108 + acc_nope_1,
2109 + acc_tail,
2110 + MAIN_BLOCK_SIZE,
2111 + NOPE_DIM,
2112 + BLOCK_K,
2113 + IS_FNUZ_MAIN,
2114 + False,
2115 + )
2116 +
2117 + if HAS_EXTRA:
2118 + if not ADAPTIVE_SPLITS:
2119 + extra_start = tl.load(extra_indptr_ptr + query_idx)
2120 + extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
2121 + extra_len = extra_end - extra_start
2122 + extra_chunk = (extra_len + work_splits - 1) // work_splits
2123 + extra_lo = split_id * extra_chunk
2124 + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len)
2125 +
2126 + outer_block_k: tl.constexpr = 2 * BLOCK_K
2127 + outer_k_offsets = tl.arange(0, outer_block_k)
2128 + extra_hi_full = (
2129 + extra_lo + ((extra_hi - extra_lo) // outer_block_k) * outer_block_k
2130 + )
2131 + for k_start in tl.range(
2132 + extra_lo,
2133 + extra_hi_full,
2134 + outer_block_k,
2135 + num_stages=NUM_STAGES,
2136 + ):
2137 + slot = tl.load(extra_indices_ptr + extra_start + k_start + outer_k_offsets)
2138 + valid = (slot >= 0) & (slot < extra_num_rows)
2139 + slot_pairs = tl.trans(tl.reshape(slot, (2, BLOCK_K)))
2140 + valid_pairs = tl.trans(tl.reshape(valid, (2, BLOCK_K)))
2141 + slot_lo, slot_hi = tl.split(slot_pairs)
2142 + valid_lo, valid_hi = tl.split(valid_pairs)
2143 + (
2144 + m_i,
2145 + l_i,
2146 + acc_nope_0a,
2147 + acc_nope_0b,
2148 + acc_nope_1,
2149 + acc_tail,
2150 + ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
2151 + q_combined,
2152 + extra_cache_ptr,
2153 + slot_lo,
2154 + valid_lo,
2155 + extra_cache_stride0,
2156 + scale,
2157 + head_mask,
2158 + m_i,
2159 + l_i,
2160 + acc_nope_0a,
2161 + acc_nope_0b,
2162 + acc_nope_1,
2163 + acc_tail,
2164 + EXTRA_BLOCK_SIZE,
2165 + NOPE_DIM,
2166 + BLOCK_K,
2167 + IS_FNUZ_EXTRA,
2168 + TRUST_EXTRA_CACHE_NAN_FREE,
2169 + )
2170 + (
2171 + m_i,
2172 + l_i,
2173 + acc_nope_0a,
2174 + acc_nope_0b,
2175 + acc_nope_1,
2176 + acc_tail,
2177 + ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
2178 + q_combined,
2179 + extra_cache_ptr,
2180 + slot_hi,
2181 + valid_hi,
2182 + extra_cache_stride0,
2183 + scale,
2184 + head_mask,
2185 + m_i,
2186 + l_i,
2187 + acc_nope_0a,
2188 + acc_nope_0b,
2189 + acc_nope_1,
2190 + acc_tail,
2191 + EXTRA_BLOCK_SIZE,
2192 + NOPE_DIM,
2193 + BLOCK_K,
2194 + IS_FNUZ_EXTRA,
2195 + TRUST_EXTRA_CACHE_NAN_FREE,
2196 + )
2197 + for tail_idx in tl.static_range(2):
2198 + tail_start = extra_hi_full + tail_idx * BLOCK_K
2199 + if tail_start < extra_hi:
2200 + k_pos = tail_start + k_offsets
2201 + in_range = k_pos < extra_hi
2202 + slot = tl.load(
2203 + extra_indices_ptr + extra_start + k_pos,
2204 + mask=in_range,
2205 + other=-1,
2206 + )
2207 + valid = in_range & (slot >= 0) & (slot < extra_num_rows)
2208 + (
2209 + m_i,
2210 + l_i,
2211 + acc_nope_0a,
2212 + acc_nope_0b,
2213 + acc_nope_1,
2214 + acc_tail,
2215 + ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
2216 + q_combined,
2217 + extra_cache_ptr,
2218 + slot,
2219 + valid,
2220 + extra_cache_stride0,
2221 + scale,
2222 + head_mask,
2223 + m_i,
2224 + l_i,
2225 + acc_nope_0a,
2226 + acc_nope_0b,
2227 + acc_nope_1,
2228 + acc_tail,
2229 + EXTRA_BLOCK_SIZE,
2230 + NOPE_DIM,
2231 + BLOCK_K,
2232 + IS_FNUZ_EXTRA,
2233 + TRUST_EXTRA_CACHE_NAN_FREE,
2234 + )
2235 +
2236 + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets
2237 + m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large)
2238 + tl.store(part_m_ptr + pm_base, m_store, mask=head_mask)
2239 + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask)
2240 + acc_base = part_acc_ptr + (
2241 + (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets[:, None]
2242 + ) * (NOPE_DIM + ROPE_DIM)
2243 + tl.store(
2244 + acc_base + nope_offsets_0a[None, :],
2245 + acc_nope_0a,
2246 + mask=head_mask[:, None],
2247 + )
2248 + tl.store(
2249 + acc_base + nope_offsets_0b[None, :],
2250 + acc_nope_0b,
2251 + mask=head_mask[:, None],
2252 + )
2253 + tl.store(
2254 + acc_base + nope_offsets_1[None, :],
2255 + acc_nope_1,
2256 + mask=head_mask[:, None],
2257 + )
2258 + tl.store(
2259 + acc_base + tail_offsets_128[None, :],
2260 + acc_tail,
2261 + mask=head_mask[:, None],
2262 + )
2263 +
2264 +
2265 @triton.jit
2266 def _sparse_attn_decode_reduce_kernel(
2267 part_m_ptr,
@@ -1701,6 +2278,7 @@ def _sparse_attn_decode_reduce_kernel(
2278 pa_stride_h,
2279 num_heads,
2280 HAS_ATTN_SINK: tl.constexpr,
2281 + ADAPTIVE_SPLITS: tl.constexpr,
2282 COMB_DIM: tl.constexpr,
2283 BLOCK_H: tl.constexpr,
2284 NUM_SPLITS: tl.constexpr,
@@ -1769,17 +2347,27 @@ def _sparse_attn_decode_reduce_kernel(
2347 other=neg_large,
2348 )
2349 w_s = tl.exp(m_s - m_final)
2350 + if ADAPTIVE_SPLITS:
2351 + active_split = m_s > neg_large
2352 + w_s = tl.where(head_mask & active_split, w_s, 0.0)
2353 acc_base = (
2354 part_acc_ptr
2355 + query_idx * pa_stride0
2356 + s * pa_stride_s
2357 + head_offsets[:, None] * pa_stride_h
2358 )
1778 - acc_s = tl.load(
1779 - acc_base + comb_offsets[None, :],
1780 - mask=head_mask[:, None],
1781 - other=0.0,
1782 - )
2359 + if ADAPTIVE_SPLITS:
2360 + acc_s = tl.load(
2361 + acc_base + comb_offsets[None, :],
2362 + mask=head_mask[:, None] & active_split[:, None],
2363 + other=0.0,
2364 + )
2365 + else:
2366 + acc_s = tl.load(
2367 + acc_base + comb_offsets[None, :],
2368 + mask=head_mask[:, None],
2369 + other=0.0,
2370 + )
2371 acc += w_s[:, None] * acc_s
2372
2373 out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0)
@@ -1986,6 +2574,49 @@ def _decode_num_splits(
2574 return best_splits
2575
2576
2577 +def _decode_gfx950_num_splits(
2578 + num_queries: int,
2579 + heads_blocks: int,
2580 + avg_main_len: float = 0.0,
2581 + avg_extra_len: float = 0.0,
2582 + block_k: int = 32,
2583 +) -> int:
2584 + base = max(1, num_queries * heads_blocks)
2585 + cu = max(1, _decode_cu_count())
2586 + target_workgroups = 2 * cu
2587 + num_splits = min(
2588 + 32,
2589 + max(
2590 + 1,
2591 + math.ceil(target_workgroups / base),
2592 + ),
2593 + )
2594 + if (
2595 + base >= 16
2596 + and num_splits > 4
2597 + and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3
2598 + ):
2599 + return 4
2600 + if 16 <= base < 64:
2601 + one_wave_splits = max(1, cu // base)
2602 + one_wave_iters = _decode_partial_iters(
2603 + avg_main_len, avg_extra_len, one_wave_splits, block_k
2604 + )
2605 + target_waves = 1 if one_wave_iters <= 9 else 2
2606 + num_splits = min(num_splits, max(1, target_waves * cu // base))
2607 + if base >= 16 and num_splits > 1:
2608 + target_waves = (base * num_splits + cu - 1) // cu
2609 + target_iters = _decode_partial_iters(
2610 + avg_main_len, avg_extra_len, num_splits, block_k
2611 + )
2612 + for splits in range(1, num_splits):
2613 + waves = (base * splits + cu - 1) // cu
2614 + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k)
2615 + if waves == target_waves and iters == target_iters:
2616 + return splits
2617 + return num_splits
2618 +
2619 +
2620 def _rocm_sparse_attn_decode_ragged_triton(
2621 q: torch.Tensor,
2622 main_cache: torch.Tensor,
@@ -1998,6 +2629,9 @@ def _rocm_sparse_attn_decode_ragged_triton(
2629 extra_cache: torch.Tensor | None = None,
2630 extra_indices: torch.Tensor | None = None,
2631 extra_indptr: torch.Tensor | None = None,
2632 + out: torch.Tensor | None = None,
2633 + extra_cache_nan_free: bool = False,
2634 + adaptive_splits: bool = False,
2635 ) -> torch.Tensor:
2636 assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
2637 assert main_cache.ndim == 3, (
@@ -2038,6 +2672,10 @@ def _rocm_sparse_attn_decode_ragged_triton(
2672 and extra_indices is not None
2673 and extra_indptr is not None
2674 )
2675 + assert not extra_cache_nan_free or (_ON_GFX950 and has_extra), (
2676 + "extra_cache_nan_free requires a gfx950 compressed cache with trusted "
2677 + "canonical-writer provenance"
2678 + )
2679 if has_extra:
2680 assert extra_cache is not None
2681 assert extra_indices is not None
@@ -2059,7 +2697,16 @@ def _rocm_sparse_attn_decode_ragged_triton(
2697 extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
2698
2699 block_h = 16
2062 - out = torch.empty_like(q, dtype=torch.bfloat16)
2700 + if out is None:
2701 + out = torch.empty_like(q, dtype=torch.bfloat16)
2702 + else:
2703 + assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}"
2704 + assert out.device == q.device, (
2705 + f"expected out on device {q.device}, got {out.device}"
2706 + )
2707 + assert out.dtype == torch.bfloat16, (
2708 + f"expected out dtype {torch.bfloat16}, got {out.dtype}"
2709 + )
2710 heads_blocks = triton.cdiv(num_heads, block_h)
2711 nope_block = triton.next_power_of_2(nope_head_dim)
2712 comb_dim = nope_head_dim + rope_head_dim
@@ -2103,14 +2750,35 @@ def _rocm_sparse_attn_decode_ragged_triton(
2750 return out
2751
2752 block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950.
2106 - # Average per-query segment lengths, read sync-free from the ragged index
2107 - # sizes, let the split heuristic avoid over-splitting
2108 - # main_indices/extra_indices are flat [nnz] int32.
2109 - inv_q = 1.0 / max(1, num_queries)
2110 - avg_main_len = main_indices.numel() * inv_q
2111 - avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
2112 - num_splits = _decode_num_splits(
2113 - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k
2753 + if _ON_GFX950:
2754 + inv_q = 1.0 / max(1, num_queries)
2755 + avg_main_len = main_indices.numel() * inv_q
2756 + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
2757 + num_splits = _decode_gfx950_num_splits(
2758 + num_queries,
2759 + heads_blocks,
2760 + avg_main_len,
2761 + avg_extra_len,
2762 + block_k,
2763 + )
2764 + else:
2765 + # Average per-query segment lengths, read sync-free from the ragged
2766 + # index sizes, let the split heuristic avoid over-splitting.
2767 + inv_q = 1.0 / max(1, num_queries)
2768 + avg_main_len = main_indices.numel() * inv_q
2769 + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
2770 + num_splits = _decode_num_splits(
2771 + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k
2772 + )
2773 +
2774 + base_workgroups = num_queries * heads_blocks
2775 + adaptive_splits = (
2776 + _ON_GFX950 and adaptive_splits and base_workgroups >= 16 and num_splits > 4
2777 + )
2778 + one_wave_splits = (
2779 + max(1, _decode_cu_count() // base_workgroups)
2780 + if adaptive_splits and 16 <= base_workgroups < 64
2781 + else num_splits
2782 )
2783
2784 part_m = torch.empty(
@@ -2123,48 +2791,88 @@ def _rocm_sparse_attn_decode_ragged_triton(
2791 device=q.device,
2792 )
2793
2126 - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)](
2127 - q,
2128 - main_cache,
2129 - main_indices,
2130 - main_indptr,
2131 - extra_cache,
2132 - extra_indices,
2133 - extra_indptr,
2134 - part_m,
2135 - part_l,
2136 - part_acc,
2137 - q.stride(0),
2138 - q.stride(1),
2139 - main_cache.stride(0),
2140 - extra_cache.stride(0),
2141 - part_m.stride(0),
2142 - part_m.stride(1),
2143 - part_acc.stride(0),
2144 - part_acc.stride(1),
2145 - part_acc.stride(2),
2146 - main_cache.shape[0] * main_cache.shape[1],
2147 - extra_cache.shape[0] * extra_cache.shape[1],
2148 - main_cache.shape[1],
2149 - extra_cache.shape[1],
2150 - scale,
2151 - num_heads,
2152 - HAS_EXTRA=has_extra,
2153 - NOPE_DIM=nope_head_dim,
2154 - NOPE_BLOCK=nope_block,
2155 - ROPE_DIM=rope_head_dim,
2156 - # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950).
2157 - # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere).
2158 - # Reading both with a single IS_FNUZ would decode one of them with the
2159 - # wrong FNUZ/OCP scale ratio (~1.87×).
2160 - IS_FNUZ_MAIN=is_fnuz,
2161 - IS_FNUZ_EXTRA=False,
2162 - BLOCK_H=block_h,
2163 - BLOCK_K=block_k,
2164 - NUM_SPLITS=num_splits,
2165 - NUM_STAGES=1,
2166 - num_warps=4,
2167 - )
2794 + if _ON_GFX950:
2795 + _sparse_attn_decode_gfx950_partial_kernel[
2796 + (num_queries, num_splits, heads_blocks)
2797 + ](
2798 + q,
2799 + main_cache,
2800 + main_indices,
2801 + main_indptr,
2802 + extra_cache,
2803 + extra_indices,
2804 + extra_indptr,
2805 + part_m,
2806 + part_l,
2807 + part_acc,
2808 + q.stride(0),
2809 + q.stride(1),
2810 + main_cache.stride(0),
2811 + extra_cache.stride(0),
2812 + main_cache.shape[0] * main_cache.shape[1],
2813 + extra_cache.shape[0] * extra_cache.shape[1],
2814 + main_cache.shape[1],
2815 + extra_cache.shape[1],
2816 + scale,
2817 + num_heads,
2818 + HAS_EXTRA=has_extra,
2819 + NOPE_DIM=nope_head_dim,
2820 + ROPE_DIM=rope_head_dim,
2821 + IS_FNUZ_MAIN=is_fnuz,
2822 + IS_FNUZ_EXTRA=False,
2823 + TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free,
2824 + ADAPTIVE_SPLITS=adaptive_splits,
2825 + ONE_WAVE_SPLITS=one_wave_splits,
2826 + BLOCK_H=block_h,
2827 + BLOCK_K=block_k,
2828 + NUM_SPLITS=num_splits,
2829 + NUM_STAGES=1,
2830 + num_warps=4,
2831 + waves_per_eu=0,
2832 + )
2833 + else:
2834 + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)](
2835 + q,
2836 + main_cache,
2837 + main_indices,
2838 + main_indptr,
2839 + extra_cache,
2840 + extra_indices,
2841 + extra_indptr,
2842 + part_m,
2843 + part_l,
2844 + part_acc,
2845 + q.stride(0),
2846 + q.stride(1),
2847 + main_cache.stride(0),
2848 + extra_cache.stride(0),
2849 + part_m.stride(0),
2850 + part_m.stride(1),
2851 + part_acc.stride(0),
2852 + part_acc.stride(1),
2853 + part_acc.stride(2),
2854 + main_cache.shape[0] * main_cache.shape[1],
2855 + extra_cache.shape[0] * extra_cache.shape[1],
2856 + main_cache.shape[1],
2857 + extra_cache.shape[1],
2858 + scale,
2859 + num_heads,
2860 + HAS_EXTRA=has_extra,
2861 + NOPE_DIM=nope_head_dim,
2862 + NOPE_BLOCK=nope_block,
2863 + ROPE_DIM=rope_head_dim,
2864 + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950).
2865 + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere).
2866 + # Reading both with a single IS_FNUZ would decode one of them with the
2867 + # wrong FNUZ/OCP scale ratio (~1.87×).
2868 + IS_FNUZ_MAIN=is_fnuz,
2869 + IS_FNUZ_EXTRA=False,
2870 + BLOCK_H=block_h,
2871 + BLOCK_K=block_k,
2872 + NUM_SPLITS=num_splits,
2873 + NUM_STAGES=1,
2874 + num_warps=4,
2875 + )
2876
2877 _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](
2878 part_m,
@@ -2181,6 +2889,7 @@ def _rocm_sparse_attn_decode_ragged_triton(
2889 part_acc.stride(2),
2890 num_heads,
2891 HAS_ATTN_SINK=has_attn_sink,
2892 + ADAPTIVE_SPLITS=adaptive_splits,
2893 COMB_DIM=comb_dim,
2894 BLOCK_H=1,
2895 NUM_SPLITS=num_splits,
@@ -2206,6 +2915,9 @@ def _rocm_sparse_attn_decode_triton(
2915 main_ragged_indptr: torch.Tensor | None = None,
2916 extra_ragged_indices: torch.Tensor | None = None,
2917 extra_ragged_indptr: torch.Tensor | None = None,
2918 + out: torch.Tensor | None = None,
2919 + extra_cache_nan_free: bool = False,
2920 + adaptive_splits: bool = False,
2921 ) -> torch.Tensor:
2922 if main_ragged_indices is None or main_ragged_indptr is None:
2923 main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
@@ -2241,6 +2953,9 @@ def _rocm_sparse_attn_decode_triton(
2953 extra_cache=extra_cache,
2954 extra_indices=extra_ragged_indices,
2955 extra_indptr=extra_ragged_indptr,
2956 + out=out,
2957 + extra_cache_nan_free=extra_cache_nan_free,
2958 + adaptive_splits=adaptive_splits,
2959 )
2960
2961
@@ -2312,6 +3027,8 @@ def rocm_sparse_attn_decode(
3027 nope_head_dim: int,
3028 rope_head_dim: int,
3029 output: torch.Tensor,
3030 + extra_cache_nan_free: bool = False,
3031 + adaptive_splits: bool = False,
3032 ) -> None:
3033 assert swa_k_cache.dtype == torch.uint8, (
3034 "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
@@ -2341,6 +3058,7 @@ def rocm_sparse_attn_decode(
3058 if topk_indices is not None:
3059 extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
3060
3061 + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
3062 attn_out = _rocm_sparse_attn_decode_triton(
3063 q=q,
3064 main_cache=swa_k_cache,
@@ -2357,5 +3075,9 @@ def rocm_sparse_attn_decode(
3075 main_ragged_indptr=swa_ragged_indptr,
3076 extra_ragged_indices=topk_ragged_indices,
3077 extra_ragged_indptr=topk_ragged_indptr,
3078 + out=direct_out,
3079 + extra_cache_nan_free=extra_cache_nan_free,
3080 + adaptive_splits=adaptive_splits,
3081 )
2361 - output.copy_(attn_out.to(output.dtype))
3082 + if direct_out is None:
3083 + output.copy_(attn_out.to(output.dtype))