@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7e91760650

[ROCm][Perf] Route large DSV4 sparse prefill to AITER OPUS (#54855)

Signed-off-by: jiacao-amd <jiahui.cao@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

jiacao-amd committed Sep 9, 2026 at 23:39 UTC 7e917606504f48f22491eb9d04ce50c28618d45f
2 files changed +318
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+211
@@ -381,6 +381,217 @@ def test_sparse_attn_prefill_ragged_kernel() -> None:
381 torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
382
383
384 +@pytest.mark.parametrize(
385 + ("num_queries", "on_gfx950", "expected"),
386 + [(1023, True, False), (1024, True, True), (1024, False, False)],
387 +)
388 +def test_aiter_sparse_prefill_opus_selection(
389 + num_queries: int, on_gfx950: bool, expected: bool
390 +) -> None:
391 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
392 + _can_use_aiter_sparse_prefill_opus,
393 + )
394 +
395 + device = torch.device("cuda")
396 + q = torch.empty(num_queries, 16, HEAD_DIM, dtype=torch.bfloat16, device=device)
397 + kv = torch.empty(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device)
398 + attn_sink = torch.empty(16, dtype=torch.float32, device=device)
399 + output = torch.empty_like(q)
400 +
401 + assert (
402 + _can_use_aiter_sparse_prefill_opus(
403 + q, kv, attn_sink, output, on_gfx950=on_gfx950
404 + )
405 + is expected
406 + )
407 +
408 +
409 +def test_aiter_sparse_prefill_opus_selection_rejects_incompatible_inputs() -> None:
410 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
411 + _can_use_aiter_sparse_prefill_opus,
412 + )
413 +
414 + device = torch.device("cuda")
415 + q = torch.empty(1024, 16, HEAD_DIM, dtype=torch.bfloat16, device=device)
416 + kv = torch.empty(1024, HEAD_DIM, dtype=torch.bfloat16, device=device)
417 + attn_sink = torch.empty(16, dtype=torch.float32, device=device)
418 + output = torch.empty_like(q)
419 +
420 + assert _can_use_aiter_sparse_prefill_opus(q, kv, attn_sink, output, on_gfx950=True)
421 + assert not _can_use_aiter_sparse_prefill_opus(
422 + q,
423 + kv,
424 + attn_sink,
425 + torch.empty(16, 1024, HEAD_DIM, dtype=q.dtype, device=device).transpose(0, 1),
426 + on_gfx950=True,
427 + )
428 + assert not _can_use_aiter_sparse_prefill_opus(
429 + q, kv, attn_sink[:-1], output, on_gfx950=True
430 + )
431 + assert not _can_use_aiter_sparse_prefill_opus(
432 + q.cpu(), kv.cpu(), attn_sink.cpu(), output.cpu(), on_gfx950=True
433 + )
434 +
435 +
436 +def test_sparse_attn_prefill_aiter_opus_routing(monkeypatch) -> None:
437 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
438 +
439 + q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
440 + kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
441 + indices = torch.tensor([[0], [1]], dtype=torch.int32)
442 + topk_length = torch.ones(2, dtype=torch.int32)
443 + attn_sink = torch.empty(1, dtype=torch.float32)
444 + output = torch.empty_like(q)
445 + opus_calls = 0
446 +
447 + def fake_opus(*args, out):
448 + nonlocal opus_calls
449 + opus_calls += 1
450 + assert args[2].dtype == torch.int32
451 + assert args[3].dtype == torch.int32
452 + assert args[5].numel() == 0
453 + assert torch.count_nonzero(args[6]) == 0
454 + out.zero_()
455 + return out
456 +
457 + monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
458 + monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: fake_opus)
459 + monkeypatch.setattr(
460 + mod,
461 + "build_ragged_indices_from_dense",
462 + lambda *args, **kwargs: (
463 + torch.tensor([0, 1], dtype=torch.int32),
464 + torch.tensor([0, 1, 2], dtype=torch.int32),
465 + ),
466 + )
467 + monkeypatch.setattr(
468 + mod,
469 + "_rocm_sparse_attn_prefill_triton",
470 + lambda *args, **kwargs: pytest.fail("unexpected dense Triton fallback"),
471 + )
472 + monkeypatch.setattr(
473 + mod,
474 + "_rocm_sparse_attn_prefill_ragged_triton",
475 + lambda *args, **kwargs: pytest.fail("unexpected ragged Triton fallback"),
476 + )
477 +
478 + mod.rocm_sparse_attn_prefill(
479 + q=q,
480 + kv=kv,
481 + indices=indices,
482 + topk_length=topk_length,
483 + scale=HEAD_DIM**-0.5,
484 + head_dim=HEAD_DIM,
485 + nope_head_dim=NOPE_HEAD_DIM,
486 + rope_head_dim=ROPE_HEAD_DIM,
487 + attn_sink=attn_sink,
488 + output=output,
489 + )
490 +
491 + assert opus_calls == 1
492 + assert torch.count_nonzero(output) == 0
493 +
494 +
495 +def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> None:
496 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
497 +
498 + q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
499 + kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16)
500 + indices = torch.tensor([[0], [1]], dtype=torch.int32)
501 + topk_length = torch.ones(2, dtype=torch.int32)
502 + attn_sink = torch.empty(1, dtype=torch.float32)
503 + output = torch.empty_like(q)
504 + dense_fallback_calls = 0
505 +
506 + def fake_dense_fallback(*args, **kwargs):
507 + nonlocal dense_fallback_calls
508 + dense_fallback_calls += 1
509 + return torch.zeros_like(q)
510 +
511 + monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
512 + monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
513 + monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_triton", fake_dense_fallback)
514 + monkeypatch.setattr(
515 + mod,
516 + "_rocm_sparse_attn_prefill_ragged_triton",
517 + lambda *args, **kwargs: pytest.fail("unexpected ragged Triton fallback"),
518 + )
519 +
520 + mod.rocm_sparse_attn_prefill(
521 + q=q,
522 + kv=kv,
523 + indices=indices,
524 + topk_length=topk_length,
525 + scale=HEAD_DIM**-0.5,
526 + head_dim=HEAD_DIM,
527 + nope_head_dim=NOPE_HEAD_DIM,
528 + rope_head_dim=ROPE_HEAD_DIM,
529 + attn_sink=attn_sink,
530 + output=output,
531 + )
532 +
533 + assert dense_fallback_calls == 1
534 + assert torch.count_nonzero(output) == 0
535 +
536 +
537 +@requires_gfx950
538 +@torch.inference_mode()
539 +def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None:
540 + opus_mod = pytest.importorskip("aiter.ops.pa_sparse_prefill_opus")
541 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
542 +
543 + monkeypatch.setattr(
544 + mod, "_get_aiter_sparse_prefill_opus", lambda: opus_mod.pa_sparse_prefill_opus
545 + )
546 +
547 + device = torch.device("cuda")
548 + torch.manual_seed(4)
549 + num_queries = 8
550 + num_heads = 16
551 + q = (
552 + torch.randn(
553 + num_queries,
554 + num_heads,
555 + HEAD_DIM,
556 + dtype=torch.bfloat16,
557 + device=device,
558 + )
559 + * 0.125
560 + )
561 + kv = torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
562 + rows = [list(range(query_idx + 1)) for query_idx in range(num_queries)]
563 + indices = torch.tensor(
564 + [index for row in rows for index in row], dtype=torch.int32, device=device
565 + )
566 + indptr = torch.tensor(
567 + [0]
568 + + [
569 + sum(len(row) for row in rows[: query_idx + 1])
570 + for query_idx in range(num_queries)
571 + ],
572 + dtype=torch.int32,
573 + device=device,
574 + )
575 + attn_sink = torch.linspace(
576 + -0.25, 0.25, num_heads, dtype=torch.float32, device=device
577 + )
578 + output = torch.empty_like(q)
579 + scale = HEAD_DIM**-0.5
580 +
581 + assert mod._rocm_sparse_attn_prefill_ragged_aiter_opus(
582 + q=q,
583 + kv=kv,
584 + indices=indices,
585 + indptr=indptr,
586 + scale=scale,
587 + attn_sink=attn_sink,
588 + output=output,
589 + )
590 + expected = _ref_sparse_prefill_ragged(q, kv, rows, scale, attn_sink)
591 +
592 + torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
593 +
594 +
595 @torch.inference_mode()
596 def test_sparse_attn_decode_ragged_kernel() -> None:
597 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+107
@@ -13,6 +13,7 @@ import vllm.envs as envs
13 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14 from vllm.config import CUDAGraphMode
15 from vllm.forward_context import get_forward_context
16 +from vllm.logger import init_logger
17 from vllm.platforms import current_platform
18 from vllm.triton_utils import tl, triton
19 from vllm.utils.torch_utils import LayerNameType
@@ -26,6 +27,8 @@ else:
27 _ON_GFX942 = False
28 _ON_GFX950 = False
29
30 +logger = init_logger(__name__)
31 +
32
33 @functools.cache
34 def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None:
@@ -39,8 +42,25 @@ def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | N
42 return top_k_per_row_prefill, top_k_per_row_decode
43
44
45 +@functools.cache
46 +def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
47 + from vllm._aiter_ops import rocm_aiter_ops
48 +
49 + if not rocm_aiter_ops.is_mla_enabled():
50 + return None
51 + try:
52 + from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_opus
53 + except ImportError:
54 + return None
55 + logger.info_once("Using AITER OPUS for large sparse MLA prefill on gfx950")
56 + return pa_sparse_prefill_opus
57 +
58 +
59 _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
60 _GFX950_C4A_NATIVE_MAX_ROWS = 256
61 +# Conservative perf gate, not a correctness bound: OPUS is correct for any query
62 +# count, but Triton stays faster below this measured crossover.
63 +_GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
64
65
66 def _get_aiter_top_k_kernel(
@@ -2667,6 +2687,66 @@ def _rocm_sparse_attn_prefill_triton(
2687 )
2688
2689
2690 +def _can_use_aiter_sparse_prefill_opus(
2691 + q: torch.Tensor,
2692 + kv: torch.Tensor,
2693 + attn_sink: torch.Tensor | None,
2694 + output: torch.Tensor,
2695 + on_gfx950: bool = _ON_GFX950,
2696 +) -> bool:
2697 + return (
2698 + on_gfx950
2699 + and q.shape[0] >= _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES
2700 + and q.is_cuda
2701 + and output.shape == q.shape
2702 + and kv.shape[-1] == q.shape[-1]
2703 + and q.dtype in (torch.bfloat16, torch.float16)
2704 + and kv.dtype == q.dtype
2705 + and output.dtype == q.dtype
2706 + and kv.device == q.device
2707 + and output.device == q.device
2708 + and q.stride(-1) == 1
2709 + and kv.stride(-1) == 1
2710 + and output.stride() == q.stride()
2711 + and attn_sink is not None
2712 + and attn_sink.shape == (q.shape[1],)
2713 + and attn_sink.dtype == torch.float32
2714 + and attn_sink.device == q.device
2715 + )
2716 +
2717 +
2718 +def _rocm_sparse_attn_prefill_ragged_aiter_opus(
2719 + q: torch.Tensor,
2720 + kv: torch.Tensor,
2721 + indices: torch.Tensor,
2722 + indptr: torch.Tensor,
2723 + scale: float,
2724 + attn_sink: torch.Tensor,
2725 + output: torch.Tensor,
2726 +) -> bool:
2727 + pa_sparse_prefill_opus = _get_aiter_sparse_prefill_opus()
2728 + if pa_sparse_prefill_opus is None:
2729 + return False
2730 +
2731 + indices = _as_int32_contiguous_1d(indices)
2732 + indptr = _as_int32_contiguous_1d(indptr)
2733 + empty_indices = indices[:0]
2734 + empty_indptr = torch.zeros_like(indptr)
2735 + pa_sparse_prefill_opus(
2736 + q,
2737 + kv,
2738 + indices,
2739 + indptr,
2740 + kv[:1],
2741 + empty_indices,
2742 + empty_indptr,
2743 + attn_sink.contiguous(),
2744 + float(scale),
2745 + out=output,
2746 + )
2747 + return True
2748 +
2749 +
2750 @functools.lru_cache
2751 def _decode_cu_count() -> int:
2752 try:
@@ -3170,6 +3250,33 @@ def rocm_sparse_attn_prefill(
3250 rope_head_dim,
3251 "rocm_sparse_attn_prefill",
3252 )
3253 + opus_attn_sink = None if attn_sink is None else attn_sink[: q.shape[1]]
3254 + if (
3255 + _can_use_aiter_sparse_prefill_opus(q, kv.squeeze(1), opus_attn_sink, output)
3256 + and _get_aiter_sparse_prefill_opus() is not None
3257 + ):
3258 + if ragged_indices is None or ragged_indptr is None:
3259 + assert indices is not None
3260 + indices_2d = indices.reshape(indices.shape[0], -1)
3261 + ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
3262 + indices_2d,
3263 + topk_length
3264 + if topk_length is not None
3265 + else (indices_2d >= 0).sum(dim=-1, dtype=torch.int32),
3266 + num_rows=kv.shape[0],
3267 + )
3268 + assert opus_attn_sink is not None
3269 + if _rocm_sparse_attn_prefill_ragged_aiter_opus(
3270 + q=q,
3271 + kv=kv.squeeze(1),
3272 + indices=ragged_indices,
3273 + indptr=ragged_indptr,
3274 + scale=scale,
3275 + attn_sink=opus_attn_sink,
3276 + output=output,
3277 + ):
3278 + return
3279 +
3280 if ragged_indices is not None and ragged_indptr is not None:
3281 output_chunk = _rocm_sparse_attn_prefill_ragged_triton(
3282 q=q,