From c9e45c6a376fccf7f051fde4b0db480809d69913 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Mon, 5 Oct 2026 19:08:22 +0000 Subject: [PATCH 01/10] [ROCm][DSv4.1][Perf] Route gfx942 sparse MLA prefill to AITER pa_prefill_sparse gfx942 sparse MLA prefill stays on the in-tree Triton kernel. When the installed aiter has the gfx942 pa_prefill_sparse launch, call that kernel and keep the in-tree path otherwise. Signed-off-by: Frida Andersson --- .../attention/test_rocm_triton_attn_dsv4.py | 62 ++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 80 +++++++++++++++++++ 2 files changed, 142 insertions(+) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 3b200be514d5..70c7035baf89 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -834,6 +834,7 @@ def fake_dense_fallback(*args, out, **kwargs): monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True) monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None) + monkeypatch.setattr(mod, "_get_aiter_pa_prefill_sparse", lambda: None) monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_triton", fake_dense_fallback) monkeypatch.setattr( mod, @@ -858,6 +859,67 @@ def fake_dense_fallback(*args, out, **kwargs): assert torch.count_nonzero(output) == 0 +def test_sparse_attn_prefill_aiter_gfx942_routing(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16) + kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16) + indices = torch.tensor([[0], [1]], dtype=torch.int32) + topk_length = torch.ones(2, dtype=torch.int32) + attn_sink = torch.empty(1, dtype=torch.float32) + output = torch.empty_like(q) + pa_calls = 0 + + def fake_pa( + q, kv, indices, indptr, kv_ext, idx_ext, indptr_ext, sink, scale, out=None + ): + nonlocal pa_calls + pa_calls += 1 + assert kv_ext is None and idx_ext is None and indptr_ext is None + assert out is output + assert indices.dtype == torch.int32 + assert indptr.dtype == torch.int32 + out.zero_() + return out + + monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: False) + monkeypatch.setattr(mod, "_get_aiter_pa_prefill_sparse", lambda: fake_pa) + monkeypatch.setattr( + mod, + "build_ragged_indices_from_dense", + lambda *args, **kwargs: ( + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([0, 1, 2], dtype=torch.int32), + ), + ) + monkeypatch.setattr( + mod, + "_rocm_sparse_attn_prefill_triton", + lambda *args, **kwargs: pytest.fail("unexpected dense Triton fallback"), + ) + monkeypatch.setattr( + mod, + "_rocm_sparse_attn_prefill_ragged_triton", + lambda *args, **kwargs: pytest.fail("unexpected ragged Triton fallback"), + ) + + mod.rocm_sparse_attn_prefill( + q=q, + kv=kv, + indices=indices, + topk_length=topk_length, + scale=HEAD_DIM**-0.5, + head_dim=HEAD_DIM, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + attn_sink=attn_sink, + output=output, + ) + + assert pa_calls == 1 + assert torch.count_nonzero(output) == 0 + + @requires_gfx950 @torch.inference_mode() def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None: diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index a3412aa83508..81330ba53e12 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools import importlib +import inspect import math from collections.abc import Callable from importlib.util import find_spec @@ -46,6 +47,31 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: return pa_sparse_prefill_opus +@functools.cache +def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: + """gfx942 prefill kernel from aiter#6002. Missing that launch keeps the in-tree kernel.""" + if not _ON_GFX942: + return None + try: + from aiter.ops.triton.attention import pa_prefill_sparse as pa_mod + except ImportError: + return None + try: + source = inspect.getsource(pa_mod) + except OSError: + return None + # aiter#6002 is the gfx942 launch. An older pa_prefill_sparse still imports + # and would run a different kernel, so stay on the in-tree path. + if 'DEVICE_ARCH == "gfx942"' not in source: + logger.info_once( + "AITER pa_prefill_sparse has no gfx942 kernel; " + "using the in-tree sparse prefill kernel" + ) + return None + logger.info_once("Using AITER pa_prefill_sparse for sparse MLA prefill on gfx942") + return pa_mod.pa_prefill_sparse + + # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -3526,6 +3552,38 @@ def _rocm_sparse_attn_prefill_ragged_aiter_opus( return True +def _rocm_sparse_attn_prefill_aiter( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + indptr: torch.Tensor, + scale: float, + attn_sink: torch.Tensor | None, + output: torch.Tensor, +) -> bool: + """gfx942 path for aiter#6002. One KV pool; the extend sources stay unset.""" + pa_prefill_sparse = _get_aiter_pa_prefill_sparse() + if pa_prefill_sparse is None: + return False + indices = _as_int32_contiguous_1d(indices) + indptr = _as_int32_contiguous_1d(indptr) + written = pa_prefill_sparse( + q, + kv, + indices, + indptr, + None, + None, + None, + attn_sink, + float(scale), + out=output if output.shape == q.shape and output.dtype == q.dtype else None, + ) + if written.data_ptr() != output.data_ptr(): + output.copy_(written[..., : output.shape[-1]].to(output.dtype)) + return True + + @functools.lru_cache def _decode_cu_count() -> int: try: @@ -4115,6 +4173,28 @@ def rocm_sparse_attn_prefill( ): return + if _get_aiter_pa_prefill_sparse() is not None: + if ragged_indices is None or ragged_indptr is None: + assert indices is not None + indices_2d = indices.reshape(indices.shape[0], -1) + ragged_indices, ragged_indptr = build_ragged_indices_from_dense( + indices_2d, + topk_length + if topk_length is not None + else (indices_2d >= 0).sum(dim=-1, dtype=torch.int32), + num_rows=kv.shape[0], + ) + if _rocm_sparse_attn_prefill_aiter( + q=q, + kv=kv.squeeze(1), + indices=ragged_indices, + indptr=ragged_indptr, + scale=scale, + attn_sink=opus_attn_sink, + output=output, + ): + return + if ragged_indices is not None and ragged_indptr is not None: _rocm_sparse_attn_prefill_ragged_triton( q=q, From bc916c3e820dde670a58b3b8035669ffd5eab3ee Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Mon, 5 Oct 2026 19:39:30 +0000 Subject: [PATCH 02/10] [ROCm][DSv4.1] Wrap the gfx942 prefill getter docstring ruff E501 rejects the one-line docstring at 94 characters. Signed-off-by: Frida Andersson --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 81330ba53e12..c96f83dfd6b2 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -49,7 +49,10 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: @functools.cache def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: - """gfx942 prefill kernel from aiter#6002. Missing that launch keeps the in-tree kernel.""" + """gfx942 prefill kernel from aiter#6002. + + Missing that launch keeps the in-tree kernel. + """ if not _ON_GFX942: return None try: From 8d52d751fa7070b4afc43152c70f1740a66f025c Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Tue, 6 Oct 2026 15:40:00 +0000 Subject: [PATCH 03/10] [ROCm][DSv4.1] Check the gfx942 sparse-prefill route The routing test accepted any int32 tensor, so the dense table still passed. Assert the packed indices and cover the source gate that selects the kernel. Signed-off-by: Frida Andersson --- .../attention/test_rocm_triton_attn_dsv4.py | 66 ++++++++++++++++++- 1 file changed, 63 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 70c7035baf89..11e43eb37378 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest import torch @@ -875,10 +875,12 @@ def fake_pa( ): nonlocal pa_calls pa_calls += 1 + assert kv.ndim == 2 assert kv_ext is None and idx_ext is None and indptr_ext is None assert out is output - assert indices.dtype == torch.int32 - assert indptr.dtype == torch.int32 + assert indices.ndim == 1 and indptr.ndim == 1 + assert torch.equal(indices, torch.tensor([0, 1], dtype=torch.int32)) + assert torch.equal(indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) out.zero_() return out @@ -920,6 +922,64 @@ def fake_pa( assert torch.count_nonzero(output) == 0 +def test_pa_prefill_sparse_getter_requires_gfx942_source(monkeypatch) -> None: + """The gfx942 launch is a string in the installed module, not the import.""" + import builtins + + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + sentinel = object() + source_text = "" + raise_import = False + raise_source = False + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if raise_import and name == "aiter.ops.triton.attention": + raise ImportError("aiter is not installed") + if name == "aiter.ops.triton.attention": + attention = ModuleType("aiter.ops.triton.attention") + pa = ModuleType("aiter.ops.triton.attention.pa_prefill_sparse") + pa.pa_prefill_sparse = sentinel + attention.pa_prefill_sparse = pa + return attention + return real_import(name, globals, locals, fromlist, level) + + def fake_getsource(module): + if raise_source: + raise OSError("source is missing") + return source_text + + monkeypatch.setattr(builtins, "__import__", fake_import) + monkeypatch.setattr(mod.inspect, "getsource", fake_getsource) + + def call(): + mod._get_aiter_pa_prefill_sparse.cache_clear() + try: + return mod._get_aiter_pa_prefill_sparse() + finally: + mod._get_aiter_pa_prefill_sparse.cache_clear() + + monkeypatch.setattr(mod, "_ON_GFX942", False) + source_text = 'DEVICE_ARCH == "gfx942"' + assert call() is None + + monkeypatch.setattr(mod, "_ON_GFX942", True) + raise_import = True + assert call() is None + raise_import = False + + raise_source = True + assert call() is None + raise_source = False + + source_text = "older pa_prefill_sparse" + assert call() is None + + source_text = 'DEVICE_ARCH == "gfx942"' + assert call() is sentinel + + @requires_gfx950 @torch.inference_mode() def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None: From be985832bd068ccfe4f764a33c15f94e3fd30c98 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Tue, 6 Oct 2026 15:46:11 +0000 Subject: [PATCH 04/10] [ROCm][DSv4.1] Drop the gfx942 source-gate unit test The routing test already checks the dispatch. The source check fails closed, so it does not need its own import mock. Signed-off-by: Frida Andersson --- .../attention/test_rocm_triton_attn_dsv4.py | 60 +------------------ 1 file changed, 1 insertion(+), 59 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 11e43eb37378..11d214d3c755 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools -from types import ModuleType, SimpleNamespace +from types import SimpleNamespace import pytest import torch @@ -922,64 +922,6 @@ def fake_pa( assert torch.count_nonzero(output) == 0 -def test_pa_prefill_sparse_getter_requires_gfx942_source(monkeypatch) -> None: - """The gfx942 launch is a string in the installed module, not the import.""" - import builtins - - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - sentinel = object() - source_text = "" - raise_import = False - raise_source = False - real_import = builtins.__import__ - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - if raise_import and name == "aiter.ops.triton.attention": - raise ImportError("aiter is not installed") - if name == "aiter.ops.triton.attention": - attention = ModuleType("aiter.ops.triton.attention") - pa = ModuleType("aiter.ops.triton.attention.pa_prefill_sparse") - pa.pa_prefill_sparse = sentinel - attention.pa_prefill_sparse = pa - return attention - return real_import(name, globals, locals, fromlist, level) - - def fake_getsource(module): - if raise_source: - raise OSError("source is missing") - return source_text - - monkeypatch.setattr(builtins, "__import__", fake_import) - monkeypatch.setattr(mod.inspect, "getsource", fake_getsource) - - def call(): - mod._get_aiter_pa_prefill_sparse.cache_clear() - try: - return mod._get_aiter_pa_prefill_sparse() - finally: - mod._get_aiter_pa_prefill_sparse.cache_clear() - - monkeypatch.setattr(mod, "_ON_GFX942", False) - source_text = 'DEVICE_ARCH == "gfx942"' - assert call() is None - - monkeypatch.setattr(mod, "_ON_GFX942", True) - raise_import = True - assert call() is None - raise_import = False - - raise_source = True - assert call() is None - raise_source = False - - source_text = "older pa_prefill_sparse" - assert call() is None - - source_text = 'DEVICE_ARCH == "gfx942"' - assert call() is sentinel - - @requires_gfx950 @torch.inference_mode() def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None: From d7f230fb0edf73e23142ffe88cab67d02f8c2ce0 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Tue, 6 Oct 2026 21:25:34 +0000 Subject: [PATCH 05/10] [ROCm][DSv4.1] Call the gfx942 prefill getter once The dispatch looked the function up twice. Pass it in, and drop the pull-request number from the comments. Signed-off-by: Frida Andersson --- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index c96f83dfd6b2..8e03c843e068 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -49,7 +49,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: @functools.cache def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: - """gfx942 prefill kernel from aiter#6002. + """gfx942 prefill kernel, when that launch is installed. Missing that launch keeps the in-tree kernel. """ @@ -63,8 +63,8 @@ def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: source = inspect.getsource(pa_mod) except OSError: return None - # aiter#6002 is the gfx942 launch. An older pa_prefill_sparse still imports - # and would run a different kernel, so stay on the in-tree path. + # An older pa_prefill_sparse still imports and would run a different + # kernel, so stay on the in-tree path. if 'DEVICE_ARCH == "gfx942"' not in source: logger.info_once( "AITER pa_prefill_sparse has no gfx942 kernel; " @@ -3556,6 +3556,7 @@ def _rocm_sparse_attn_prefill_ragged_aiter_opus( def _rocm_sparse_attn_prefill_aiter( + pa_prefill_sparse: Callable[..., torch.Tensor], q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, @@ -3563,11 +3564,8 @@ def _rocm_sparse_attn_prefill_aiter( scale: float, attn_sink: torch.Tensor | None, output: torch.Tensor, -) -> bool: - """gfx942 path for aiter#6002. One KV pool; the extend sources stay unset.""" - pa_prefill_sparse = _get_aiter_pa_prefill_sparse() - if pa_prefill_sparse is None: - return False +) -> None: + """gfx942 path. One KV pool; the extend sources stay unset.""" indices = _as_int32_contiguous_1d(indices) indptr = _as_int32_contiguous_1d(indptr) written = pa_prefill_sparse( @@ -3584,7 +3582,6 @@ def _rocm_sparse_attn_prefill_aiter( ) if written.data_ptr() != output.data_ptr(): output.copy_(written[..., : output.shape[-1]].to(output.dtype)) - return True @functools.lru_cache @@ -4176,7 +4173,8 @@ def rocm_sparse_attn_prefill( ): return - if _get_aiter_pa_prefill_sparse() is not None: + pa_prefill_sparse = _get_aiter_pa_prefill_sparse() + if pa_prefill_sparse is not None: if ragged_indices is None or ragged_indptr is None: assert indices is not None indices_2d = indices.reshape(indices.shape[0], -1) @@ -4187,7 +4185,8 @@ def rocm_sparse_attn_prefill( else (indices_2d >= 0).sum(dim=-1, dtype=torch.int32), num_rows=kv.shape[0], ) - if _rocm_sparse_attn_prefill_aiter( + _rocm_sparse_attn_prefill_aiter( + pa_prefill_sparse, q=q, kv=kv.squeeze(1), indices=ragged_indices, @@ -4195,8 +4194,8 @@ def rocm_sparse_attn_prefill( scale=scale, attn_sink=opus_attn_sink, output=output, - ): - return + ) + return if ragged_indices is not None and ragged_indptr is not None: _rocm_sparse_attn_prefill_ragged_triton( From 1f774e571174f3d9a71e189c0113cc6a0ac1c8f0 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Wed, 7 Oct 2026 17:39:10 +0000 Subject: [PATCH 06/10] [ROCm][DSv4.1] Drop the gfx942 routing monkeypatch test The stand-in only checked that the dispatch called it. Kernel accuracy stays in aiter. Signed-off-by: Frida Andersson --- .../attention/test_rocm_triton_attn_dsv4.py | 63 ------------------- 1 file changed, 63 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 11d214d3c755..7bd0837d14e8 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -859,69 +859,6 @@ def fake_dense_fallback(*args, out, **kwargs): assert torch.count_nonzero(output) == 0 -def test_sparse_attn_prefill_aiter_gfx942_routing(monkeypatch) -> None: - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - q = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16) - kv = torch.empty(2, 1, HEAD_DIM, dtype=torch.bfloat16) - indices = torch.tensor([[0], [1]], dtype=torch.int32) - topk_length = torch.ones(2, dtype=torch.int32) - attn_sink = torch.empty(1, dtype=torch.float32) - output = torch.empty_like(q) - pa_calls = 0 - - def fake_pa( - q, kv, indices, indptr, kv_ext, idx_ext, indptr_ext, sink, scale, out=None - ): - nonlocal pa_calls - pa_calls += 1 - assert kv.ndim == 2 - assert kv_ext is None and idx_ext is None and indptr_ext is None - assert out is output - assert indices.ndim == 1 and indptr.ndim == 1 - assert torch.equal(indices, torch.tensor([0, 1], dtype=torch.int32)) - assert torch.equal(indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) - out.zero_() - return out - - monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: False) - monkeypatch.setattr(mod, "_get_aiter_pa_prefill_sparse", lambda: fake_pa) - monkeypatch.setattr( - mod, - "build_ragged_indices_from_dense", - lambda *args, **kwargs: ( - torch.tensor([0, 1], dtype=torch.int32), - torch.tensor([0, 1, 2], dtype=torch.int32), - ), - ) - monkeypatch.setattr( - mod, - "_rocm_sparse_attn_prefill_triton", - lambda *args, **kwargs: pytest.fail("unexpected dense Triton fallback"), - ) - monkeypatch.setattr( - mod, - "_rocm_sparse_attn_prefill_ragged_triton", - lambda *args, **kwargs: pytest.fail("unexpected ragged Triton fallback"), - ) - - mod.rocm_sparse_attn_prefill( - q=q, - kv=kv, - indices=indices, - topk_length=topk_length, - scale=HEAD_DIM**-0.5, - head_dim=HEAD_DIM, - nope_head_dim=NOPE_HEAD_DIM, - rope_head_dim=ROPE_HEAD_DIM, - attn_sink=attn_sink, - output=output, - ) - - assert pa_calls == 1 - assert torch.count_nonzero(output) == 0 - - @requires_gfx950 @torch.inference_mode() def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None: From 2645a398d8bcf7667cf025fef5219a7e17aa75ac Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Wed, 7 Oct 2026 21:06:30 +0000 Subject: [PATCH 07/10] [ROCm][DSv4.1] Import pa_prefill_sparse when AITER is enabled The gfx942 launch ships with the aiter this depends on, so the getter no longer inspects the kernel source. Signed-off-by: Frida Andersson --- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 28 ++++--------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 8e03c843e068..4ebca3a01002 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools import importlib -import inspect import math from collections.abc import Callable from importlib.util import find_spec @@ -49,30 +48,15 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: @functools.cache def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: - """gfx942 prefill kernel, when that launch is installed. + """gfx942 prefill kernel when AITER is enabled.""" + from vllm._aiter_ops import rocm_aiter_ops - Missing that launch keeps the in-tree kernel. - """ - if not _ON_GFX942: - return None - try: - from aiter.ops.triton.attention import pa_prefill_sparse as pa_mod - except ImportError: - return None - try: - source = inspect.getsource(pa_mod) - except OSError: - return None - # An older pa_prefill_sparse still imports and would run a different - # kernel, so stay on the in-tree path. - if 'DEVICE_ARCH == "gfx942"' not in source: - logger.info_once( - "AITER pa_prefill_sparse has no gfx942 kernel; " - "using the in-tree sparse prefill kernel" - ) + if not _ON_GFX942 or not rocm_aiter_ops.is_enabled(): return None + from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse + logger.info_once("Using AITER pa_prefill_sparse for sparse MLA prefill on gfx942") - return pa_mod.pa_prefill_sparse + return pa_prefill_sparse # Conservative perf gate, not a correctness bound: OPUS is correct for any query From 9e72cb898146e5c41032413ac2a4e481516f856e Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Thu, 8 Oct 2026 14:07:36 +0000 Subject: [PATCH 08/10] [ROCm][DSv4.1] Skip pa_prefill_sparse for short or mismatched calls GLM sends decode rows through this function. Reuse the OPUS query cutoff and dtype checks so those calls stay on the in-tree kernel. Signed-off-by: Frida Andersson --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 4ebca3a01002..07c5cbd1d281 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -4158,7 +4158,12 @@ def rocm_sparse_attn_prefill( return pa_prefill_sparse = _get_aiter_pa_prefill_sparse() - if pa_prefill_sparse is not None: + if ( + pa_prefill_sparse is not None + and q.shape[0] >= _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES + and q.dtype in (torch.bfloat16, torch.float16) + and kv.dtype == q.dtype + ): if ragged_indices is None or ragged_indptr is None: assert indices is not None indices_2d = indices.reshape(indices.shape[0], -1) From e2a8220318e8c547b0ed70f97c94e30efd5e8fc6 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Thu, 8 Oct 2026 20:52:58 +0000 Subject: [PATCH 09/10] [ROCm][DSv4.1] Gate gfx942 pa_prefill_sparse before the import The getter imported and logged before the size check, so a short call could fail on a missing symbol. Use a gfx942 cutoff of 512, measured against the in-tree ragged prefill, and log only when the kernel launches. Signed-off-by: Frida Andersson --- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 07c5cbd1d281..de6c8f51b089 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -53,15 +53,30 @@ def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: if not _ON_GFX942 or not rocm_aiter_ops.is_enabled(): return None - from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse - - logger.info_once("Using AITER pa_prefill_sparse for sparse MLA prefill on gfx942") + try: + from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse + except ImportError: + return None return pa_prefill_sparse # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 +# gfx942 gate against the in-tree ragged prefill, not the gfx950 OPUS +# crossover. MI325, H=16, D=512, bf16: AITER is behind through 256 rows at +# fanout 128 and ahead from 512 rows at fanouts 128, 512, and 2048. +_GFX942_AITER_PA_PREFILL_SPARSE_MIN_QUERIES = 512 + + +def _can_use_aiter_pa_prefill_sparse(q: torch.Tensor, kv: torch.Tensor) -> bool: + """Cheap gate. The AITER import runs only after this passes.""" + return ( + _ON_GFX942 + and q.shape[0] >= _GFX942_AITER_PA_PREFILL_SPARSE_MIN_QUERIES + and q.dtype in (torch.bfloat16, torch.float16) + and kv.dtype == q.dtype + ) def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool: @@ -3550,6 +3565,7 @@ def _rocm_sparse_attn_prefill_aiter( output: torch.Tensor, ) -> None: """gfx942 path. One KV pool; the extend sources stay unset.""" + logger.info_once("Using AITER pa_prefill_sparse for sparse MLA prefill on gfx942") indices = _as_int32_contiguous_1d(indices) indptr = _as_int32_contiguous_1d(indptr) written = pa_prefill_sparse( @@ -4157,13 +4173,11 @@ def rocm_sparse_attn_prefill( ): return - pa_prefill_sparse = _get_aiter_pa_prefill_sparse() - if ( - pa_prefill_sparse is not None - and q.shape[0] >= _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES - and q.dtype in (torch.bfloat16, torch.float16) - and kv.dtype == q.dtype - ): + if _can_use_aiter_pa_prefill_sparse(q, kv): + pa_prefill_sparse = _get_aiter_pa_prefill_sparse() + else: + pa_prefill_sparse = None + if pa_prefill_sparse is not None: if ragged_indices is None or ragged_indptr is None: assert indices is not None indices_2d = indices.reshape(indices.shape[0], -1) From 167f303b65f4991c717e2c9f7301e6bfe9e379c7 Mon Sep 17 00:00:00 2001 From: Frida Andersson Date: Thu, 8 Oct 2026 21:29:19 +0000 Subject: [PATCH 10/10] [ROCm][DSv4.1] Raise when pa_prefill_sparse is missing A short call already skips the import. An eligible call should raise if the symbol is missing, instead of staying on the in-tree kernel with no log. Match the cutoff comment to the measured table. Signed-off-by: Frida Andersson --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index de6c8f51b089..1a0bdfa1c49c 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -53,10 +53,8 @@ def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: if not _ON_GFX942 or not rocm_aiter_ops.is_enabled(): return None - try: - from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse - except ImportError: - return None + from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse + return pa_prefill_sparse @@ -64,8 +62,9 @@ def _get_aiter_pa_prefill_sparse() -> Callable[..., torch.Tensor] | None: # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 # gfx942 gate against the in-tree ragged prefill, not the gfx950 OPUS -# crossover. MI325, H=16, D=512, bf16: AITER is behind through 256 rows at -# fanout 128 and ahead from 512 rows at fanouts 128, 512, and 2048. +# crossover. MI325, H=16, D=512, bf16: from 512 rows AITER is ahead at +# fanouts 128, 512, and 2048. At 256 rows and fanout 128 the ratio is 1.04, +# on a flat ~0.025 ms. _GFX942_AITER_PA_PREFILL_SPARSE_MIN_QUERIES = 512