From 98b29aa99aee8fe58a21b168b45a5dcb910f63c0 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:47:41 +0200 Subject: [PATCH] [ROCm][Perf][GLM-5.3-Flash] Add AITER topk backend for decodes (#58167) Signed-off-by: simondanielsson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../kernels/benchmark_top_k_per_row_decode.py | 158 ++++++++ tests/kernels/test_top_k_per_row.py | 343 ++++++++++++------ .../test_sparse_indexer_topk_dispatch.py | 86 ++++- vllm/_aiter_ops.py | 99 +++++ vllm/config/kernel.py | 3 + vllm/model_executor/layers/indexer_topk.py | 67 +++- vllm/models/glm5next/amd/sparse_indexer.py | 82 ++++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 107 +----- 8 files changed, 705 insertions(+), 240 deletions(-) create mode 100644 benchmarks/kernels/benchmark_top_k_per_row_decode.py diff --git a/benchmarks/kernels/benchmark_top_k_per_row_decode.py b/benchmarks/kernels/benchmark_top_k_per_row_decode.py new file mode 100644 index 0000000000..11337ca110 --- /dev/null +++ b/benchmarks/kernels/benchmark_top_k_per_row_decode.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the ROCm sparse-indexer decode top-k in isolation. + +Example usage: +python3 benchmarks/kernels/benchmark_top_k_per_row_decode.py \ + --rows 4 16 32 --lens 2048 8192 16384 65536 --capacity 262144 +""" + +import argparse +import itertools +import json + +import torch + +import vllm._custom_ops # noqa: F401 +from vllm._aiter_ops import rocm_aiter_ops +from vllm.triton_utils import triton + +DISTRIBUTIONS = ("indexer", "fp8", "randn") + + +def make_logits(rows: int, capacity: int, dist: str, spread: float): + normal = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if dist == "randn": + return normal + if dist == "fp8": + return normal.to(torch.float8_e4m3fnuz).to(torch.float32) + return normal * spread + 3.0 + + +def make_inputs(rows: int, live: int, capacity: int, k: int, args): + torch.manual_seed(args.seed) + logits = make_logits(rows, capacity, args.dist, args.spread) + seq_lens = torch.full((rows,), live, device="cuda", dtype=torch.int32) + invalid = torch.arange(capacity, device="cuda")[None] >= seq_lens[:, None] + logits.masked_fill_(invalid, -float("inf")) + indices = torch.empty((rows, k), device="cuda", dtype=torch.int32) + return logits, seq_lens, indices + + +def run_native(logits, seq_lens, indices, k): + torch.ops._C.top_k_per_row_decode( + logits, + 1, + seq_lens, + indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k, + ) + + +def run_aiter(logits, seq_lens, indices, k): + rocm_aiter_ops.indexer_top_k_decode(logits, 1, seq_lens, indices, k) + + +BACKENDS = {"native": run_native, "aiter": run_aiter} + + +def check(logits, indices, k): + selected = logits.gather(1, indices.long()).sort(dim=1, descending=True).values + expected = logits.topk(k, dim=1).values + torch.testing.assert_close(selected, expected, atol=0, rtol=0) + + +def capture(run): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + return graph + + +def benchmark(backend: str, rows: int, live: int, args) -> dict: + logits, seq_lens, indices = make_inputs(rows, live, args.capacity, args.top_k, args) + fn = BACKENDS[backend] + + def run(): + fn(logits, seq_lens, indices, args.top_k) + + run() + torch.accelerator.synchronize() + check(logits, indices, args.top_k) + timed = run + if args.graph: + graph = capture(run) + indices.zero_() + graph.replay() + torch.accelerator.synchronize() + check(logits, indices, args.top_k) + timed = graph.replay + ms = triton.testing.do_bench(timed, warmup=args.warmup, rep=args.rep) + return dict( + backend=backend, + mode="graph" if args.graph else "eager", + dist=args.dist, + rows=rows, + live=live, + capacity=args.capacity, + k=args.top_k, + us=ms * 1000, + live_gbps=rows * live * 4 / (ms * 1e6), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument( + "--lens", + type=int, + nargs="+", + default=[32768, 65536, 98304, 131072], + help="Valid columns per row; must be at least --top-k", + ) + parser.add_argument("--capacity", type=int, default=262144) + parser.add_argument("--top-k", type=int, default=512) + parser.add_argument( + "--backends", nargs="+", choices=list(BACKENDS), default=["native", "aiter"] + ) + parser.add_argument( + "--graph", + action="store_true", + help="Time a captured graph replay, matching how decode actually runs", + ) + parser.add_argument( + "--dist", + choices=DISTRIBUTIONS, + default="indexer", + help="Logit value distribution; see the module docstring", + ) + parser.add_argument( + "--spread", + type=float, + default=0.01, + help="Relative width of the --dist indexer cluster", + ) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--seed", type=int, default=123) + args = parser.parse_args() + + print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args)))) + for backend, rows, live in itertools.product(args.backends, args.rows, args.lens): + if live < args.top_k: + continue + print(json.dumps(benchmark(backend, rows, live, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index b6c469181f..05fda51b7c 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -636,8 +636,8 @@ def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: + from vllm._aiter_ops import rocm_aiter_ops from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _get_aiter_top_k_kernel, _launch_aiter_top_k_per_row_prefill, ) @@ -647,17 +647,14 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: top_k = 4 indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda") - aiter_topk_kernel = _get_aiter_top_k_kernel( + if not rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=True, compress_ratio=4, num_rows=logits.shape[0], - on_gfx950=True, - ) - if aiter_topk_kernel is None: + ): pytest.skip("AITER top-k is unavailable") assert ( _launch_aiter_top_k_per_row_prefill( - aiter_topk_kernel, logits, row_starts, row_ends, @@ -683,10 +680,7 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( num_speculative_tokens: int, ) -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _get_aiter_top_k_kernel, - _launch_aiter_top_k_per_row_decode, - ) + from vllm._aiter_ops import rocm_aiter_ops query_tokens = num_speculative_tokens + 1 offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda") @@ -701,20 +695,18 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( top_k = 4 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") - aiter_topk_kernel = _get_aiter_top_k_kernel( + if not rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=False, compress_ratio=4, num_rows=num_rows, max_valid_seq_len=seq_lens.max().item(), - on_gfx950=True, - ) - if aiter_topk_kernel is None: + ): pytest.skip("AITER top-k is unavailable") assert ( - _launch_aiter_top_k_per_row_decode( - aiter_topk_kernel, + rocm_aiter_ops.indexer_top_k_decode( logits, - seq_lens, + 1, + seq_lens.reshape(-1), indices, top_k, ) @@ -729,130 +721,107 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( assert torch.all(indices[row_idx, num_valid:] == -1) -def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: - from vllm.v1.attention.ops import rocm_aiter_mla_sparse +@requires_gfx950 +@torch.inference_mode() +def test_aiter_decode_topk_matches_per_row_under_mtp() -> None: + """The aiter backend expands (B, 1) seq_lens into per-row ends before + launching, so next_n > 1 must agree with the in-tree kernel, including + the rows that clamp to a zero-length prefix. + """ + from vllm.model_executor.layers.indexer_topk import get_indexer_topk - def prefill_kernel(*args, **kwargs) -> None: - pass + next_n = 4 + top_k = 512 + num_columns = 1024 + seq_lens = torch.tensor([[2], [513], [700]], dtype=torch.int32, device="cuda") + num_rows = seq_lens.shape[0] * next_n - def decode_kernel(*args, **kwargs) -> None: - pass + torch.manual_seed(0) + logits = torch.randn(num_rows, num_columns, dtype=torch.float32, device="cuda") + + outputs = {} + for backend in ("aiter", "per_row"): + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + get_indexer_topk(backend)(logits, seq_lens, next_n, indices, top_k, num_columns) + outputs[backend] = indices + + row_ends = ( + (seq_lens - next_n + 1 + torch.arange(next_n, device="cuda")) + .clamp_(min=0) + .reshape(-1) + .tolist() + ) + assert row_ends[:next_n] == [0, 0, 1, 2] + + for row_idx, row_end in enumerate(row_ends): + num_valid = min(top_k, row_end) + aiter_row = outputs["aiter"][row_idx] + assert set(aiter_row[:num_valid].tolist()) == set( + outputs["per_row"][row_idx][:num_valid].tolist() + ) + assert torch.all(aiter_row[num_valid:] == -1) + + +def _enable_aiter_topk(monkeypatch, enabled: bool = True) -> None: + from vllm._aiter_ops import rocm_aiter_ops monkeypatch.setattr( - rocm_aiter_mla_sparse, - "_get_aiter_topk_ops", - lambda: (prefill_kernel, decode_kernel), + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: enabled, raising=True ) - get_kernel = rocm_aiter_mla_sparse._get_aiter_top_k_kernel + + +def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + _enable_aiter_topk(monkeypatch) + is_supported = rocm_aiter_ops.is_indexer_top_k_supported eligible_cases = [ - ( - dict(is_prefill=True, compress_ratio=4, num_rows=1, on_gfx950=True), - prefill_kernel, + dict(is_prefill=True, compress_ratio=4, num_rows=1), + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_536), + dict( + is_prefill=False, + compress_ratio=4, + num_rows=257, + max_valid_seq_len=250_000, ), - ( - dict( - is_prefill=False, - compress_ratio=4, - num_rows=1, - max_valid_seq_len=65_536, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=4, - num_rows=257, - max_valid_seq_len=250_000, - num_columns=524_288, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=2, - num_rows=385, - max_valid_seq_len=250_000, - num_columns=524_288, - topk_tokens=512, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=2, - num_rows=384, - max_valid_seq_len=50_000, - num_columns=1_048_577, - topk_tokens=512, - on_gfx950=True, - ), - decode_kernel, - ), - ] - for kwargs, expected_kernel in eligible_cases: - assert get_kernel(**kwargs) is expected_kernel - - def fail_if_imported() -> None: - pytest.fail("ineligible shapes must not import AITER top-k") - - monkeypatch.setattr( - rocm_aiter_mla_sparse, - "_get_aiter_topk_ops", - fail_if_imported, - ) - ineligible_cases = [ - dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), dict( is_prefill=False, compress_ratio=2, num_rows=384, max_valid_seq_len=1_048_576, - num_columns=1_048_576, - topk_tokens=512, - on_gfx950=True, - ), - dict( - is_prefill=False, - compress_ratio=4, - num_rows=1, - max_valid_seq_len=65_537, - on_gfx950=True, ), + ] + for kwargs in eligible_cases: + assert is_supported(**kwargs) is True + + ineligible_cases = [ + dict(is_prefill=True, compress_ratio=1, num_rows=1), + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_537), dict( is_prefill=False, compress_ratio=4, num_rows=256, max_valid_seq_len=125_000, - on_gfx950=True, - ), - dict( - is_prefill=False, - compress_ratio=4, - num_rows=320, - max_valid_seq_len=2_500, - on_gfx950=False, ), ] for kwargs in ineligible_cases: - assert get_kernel(**kwargs) is None + assert is_supported(**kwargs) is False - monkeypatch.setattr(rocm_aiter_mla_sparse, "_get_aiter_topk_ops", lambda: None) - assert ( - get_kernel( - is_prefill=True, - compress_ratio=4, - num_rows=1, - on_gfx950=True, - ) - is None - ) + _enable_aiter_topk(monkeypatch, enabled=False) + assert is_supported(is_prefill=True, compress_ratio=4, num_rows=1) is False + + +def test_dsv4_native_top_k_window() -> None: + """The in-tree decode kernel's measured advantage band over AITER, which + applies to the DSV4 indexer only.""" + from vllm._aiter_ops import rocm_aiter_ops + + prefers_native = rocm_aiter_ops.dsv4_indexer_prefers_native_top_k + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=512) is True + assert prefers_native(num_rows=385, num_columns=524_288, topk_tokens=512) is False + assert prefers_native(num_rows=128, num_columns=1_048_577, topk_tokens=512) is False + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=1024) is False @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @@ -1927,3 +1896,141 @@ def test_sparse_indexer_topk_backend_resolution() -> None: if not is_sm100: with pytest.raises(RuntimeError, match="SM100a/SM103a"): resolve("deep_select") + + +def _patch_aiter_topk(monkeypatch, decode_kernel, enabled: bool = True) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + _enable_aiter_topk(monkeypatch, enabled=enabled) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: enabled, raising=True) + monkeypatch.setattr(rocm_aiter_ops, "indexer_top_k_decode", decode_kernel) + + +def _make_topk(backend: str): + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk + + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + return SparseIndexerTopk() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +def test_sparse_indexer_topk_aiter_requires_rocm() -> None: + """Explicitly requesting the ROCm-only backend must fail fast elsewhere.""" + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="ROCm"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +def test_sparse_indexer_topk_backend_resolution_aiter(monkeypatch) -> None: + """Aiter is opt-in only: "auto" never selects it (the shape-dependent perf + gate lives at the call site), and an explicit request validates hard + capability only.""" + + def decode_kernel(*args, **kwargs) -> None: + pass + + _patch_aiter_topk(monkeypatch, decode_kernel) + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") + + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" + assert _make_topk("aiter").resolve_backend(logits, 2048, 8) == "aiter" + # Shapes the call-site perf gate would decline still resolve. + assert _make_topk("aiter").resolve_backend(logits, 512, 8) == "aiter" + with pytest.raises(RuntimeError, match="topk_tokens must be in"): + _make_topk("aiter").resolve_backend(logits, 3000, 8) + + _enable_aiter_topk(monkeypatch, enabled=False) + with pytest.raises(RuntimeError, match="no tuned indexer top-k kernels"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + _patch_aiter_topk(monkeypatch, decode_kernel, enabled=False) + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" + with pytest.raises(RuntimeError, match="VLLM_ROCM_USE_AITER"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_aiter_applies_no_shape_gate(monkeypatch) -> None: + """Shape-dependent perf gating belongs to the call site: once "aiter" is + requested the dispatcher must run it for every supported k, including the + k=512 band the in-tree kernel wins on DSV4's compressed-KV logits.""" + from vllm.model_executor.layers import indexer_topk + + calls: list[tuple] = [] + _patch_aiter_topk(monkeypatch, lambda *args: calls.append(args)) + monkeypatch.setattr( + indexer_topk.ops, + "top_k_per_row_decode", + lambda *args: pytest.fail("an explicit aiter request must not fall back"), + ) + + num_rows = 8 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") + indices = torch.empty((num_rows, 512), dtype=torch.int32, device="cuda") + op = _make_topk("aiter") + + op(logits, seq_lens, 1, indices, 512, 4096) + # Past the compressed-context boundary the call-site gate declines on. + op(logits, seq_lens, 1, indices, 512, 4 * (64 * 1024 + 1)) + assert len(calls) == 2 + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_auto_stays_on_per_row(monkeypatch) -> None: + """The "auto" chain must not reach AITER on its own: the decode top-k runs + through top_k_per_row_decode unless the caller asks for the aiter backend.""" + from vllm.model_executor.layers import indexer_topk + + def decode_kernel(*args, **kwargs) -> None: + pytest.fail("AITER is opt-in only") + + _patch_aiter_topk(monkeypatch, decode_kernel) + calls: list[tuple] = [] + monkeypatch.setattr( + indexer_topk.ops, "top_k_per_row_decode", lambda *args: calls.append(args) + ) + + num_rows = 8 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") + indices = torch.empty((num_rows, 1024), dtype=torch.int32, device="cuda") + + _make_topk("auto")(logits, seq_lens, 1, indices, 1024, 4096) + assert len(calls) == 1 + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_aiter_expands_mtp_row_ends(monkeypatch) -> None: + """AITER's decode kernel only implements the next_n == 1 form, so spec + decode batches must be handed the equivalent expanded per-row ends.""" + from vllm.model_executor.layers import indexer_topk + + recorded: dict = {} + + def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None: + recorded["next_n"] = next_n + recorded["row_ends"] = seq_lens.tolist() + + _patch_aiter_topk(monkeypatch, decode_kernel) + monkeypatch.setattr( + indexer_topk.ops, + "top_k_per_row_decode", + lambda *args: pytest.fail("an explicit aiter request must not fall back"), + ) + + next_n = 2 + logits = torch.randn(4, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.tensor([[100], [200]], dtype=torch.int32, device="cuda") + indices = torch.empty((4, 1024), dtype=torch.int32, device="cuda") + + _make_topk("aiter")(logits, seq_lens, next_n, indices, 1024, 800) + + assert recorded["next_n"] == 1 + assert recorded["row_ends"] == [99, 100, 199, 200] diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py index 9f94da459c..8b4a09f42f 100644 --- a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -8,10 +8,6 @@ import torch from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda(), reason="CUDA-only dispatch" -) - def _require_deep_gemm() -> None: from vllm.utils.deep_gemm import has_deep_gemm @@ -20,17 +16,10 @@ def _require_deep_gemm() -> None: pytest.skip("kpool indexer requires DeepGEMM") -@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) -def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: - """The kpool indexer must read kernel_config.sparse_indexer_topk_backend - and hand it to the shared SparseIndexerTopk dispatcher, rather than - hard-coding its own cooperative/persistent/per_row choice.""" - _require_deep_gemm() - from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool - +def _build(indexer_cls, backend: str): cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) with set_current_vllm_config(cfg): - op = SparseAttnIndexerKpool( + return indexer_cls( k_cache=None, quant_block_size=128, scale_fmt="ue8m0", @@ -40,4 +29,73 @@ def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> N max_total_seq_len=8192, topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"), ) - assert op.topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only dispatch") +@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: + """The kpool indexer must read kernel_config.sparse_indexer_topk_backend + and hand it to the shared SparseIndexerTopk dispatcher, rather than + hard-coding its own cooperative/persistent/per_row choice.""" + _require_deep_gemm() + from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool + + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +@pytest.mark.parametrize("backend", ["auto", "aiter", "per_row", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend_rocm( + backend: str, +) -> None: + """The AMD kpool indexer must go through the same dispatcher, so the + AITER decode top-k is reachable from kernel_config.""" + from vllm.models.glm5next.amd.sparse_indexer import SparseAttnIndexerKpool + + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +@pytest.mark.parametrize( + "configured,full_cudagraph,expected", + [ + ("auto", False, "aiter"), + ("auto", True, "auto"), + ("aiter", True, "aiter"), + ("per_row", False, "per_row"), + ], +) +def test_kpool_decode_topk_backend_never_auto_selects_aiter_under_full_cudagraph( + monkeypatch: pytest.MonkeyPatch, + configured: str, + full_cudagraph: bool, + expected: str, +) -> None: + """A FULL cudagraph bakes in one kernel for every replay, so "auto" cannot + use the context length to pick AITER and must stay on the in-tree kernel. + Explicit backends keep working, and outside FULL the long-context heuristic + still narrows to AITER.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.glm5next.amd import sparse_indexer + + monkeypatch.setattr( + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: True, raising=False + ) + monkeypatch.setattr( + rocm_aiter_ops, + "is_indexer_top_k_supported", + lambda **kwargs: True, + raising=False, + ) + + assert ( + sparse_indexer._kpool_decode_topk_backend( + configured, + num_rows=8, + max_valid_seq_len=64 * 1024, + select_k=512, + index_kpool=4, + full_cudagraph=full_cudagraph, + ) + == expected + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index f2d313bdd0..ca98e63b39 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1913,6 +1913,11 @@ def _sync_aiter_situv2_moe_env() -> None: os.environ.pop(name, None) +_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 +_GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 + + class rocm_aiter_ops: """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM. @@ -2546,6 +2551,100 @@ class rocm_aiter_ops: """Whether AITER has DSV4 native-I384 configs through the given M.""" return cls._probe_dsv4_i384_fhmoe_capability(num_tokens) + @classmethod + @if_aiter_supported + def is_indexer_top_k_enabled(cls) -> bool: + """gfx950 is the only arch with tuned AITER indexer top-k kernels.""" + from vllm.platforms.rocm import on_gfx950 + + return cls._AITER_ENABLED and on_gfx950() + + @classmethod + def is_indexer_top_k_supported( + cls, + *, + is_prefill: bool, + compress_ratio: int, + num_rows: int, + max_valid_seq_len: int | None = None, + ) -> bool: + """Whether AITER's sparse indexer top-k beats the in-tree kernel for + this shape.""" + if compress_ratio <= 1 or not cls.is_indexer_top_k_enabled(): + return False + + if not is_prefill: + assert max_valid_seq_len is not None + # AITER v0.1.19 decode is one-block only. This measured gfx950 + # FP32/k=1024 compressed-row boundary is independent of the native + # split-count boundary in sampler.cu. + if ( + num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS + and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN + ): + return False + + return True + + @staticmethod + def dsv4_indexer_prefers_native_top_k( + *, + num_rows: int, + num_columns: int, + topk_tokens: int, + ) -> bool: + """The in-tree decode kernel's measured advantage window over AITER. + Tuned on DSV4's compressed-KV logits, so it applies to that indexer + only.""" + return ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ) + + @staticmethod + def indexer_top_k_decode( + logits: torch.Tensor, + next_n: int, + seq_lens: torch.Tensor, + topk_indices: torch.Tensor, + topk_tokens: int, + ) -> None: + from aiter.ops.topk import top_k_per_row_decode + + top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k=topk_tokens, + ) + + @staticmethod + def indexer_top_k_prefill( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + topk_tokens: int, + ) -> None: + from aiter.ops.topk import top_k_per_row_prefill + + top_k_per_row_prefill( + logits, + row_starts, + row_ends, + indices, + None, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k=topk_tokens, + ) + @staticmethod def register_ops_once() -> None: global _OPS_REGISTERED diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index a42ad62047..7fe868ae5a 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -167,6 +167,7 @@ SparseIndexerTopkBackend = Literal[ "per_row", "flashinfer", "torch", + "aiter", ] # Architectures whose model code wires up the flashinfer moe_ep experts. MTP @@ -289,6 +290,8 @@ class KernelConfig: - "per_row": Use vLLM's top_k_per_row_decode kernel - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel - "torch": Use a plain torch.topk implementation (debug reference) + - "aiter": Use AITER's top_k_per_row_decode kernel (ROCm gfx950 only, + requires VLLM_ROCM_USE_AITER) Explicit values raise RuntimeError when their constraints are not met. """ diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index c3422445b2..f0da0312e1 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -163,6 +163,14 @@ class SparseIndexerTopk(torch.nn.Module): current_platform.has_device_capability(90) and not current_platform.is_device_capability_family(120) ) + self._is_rocm = current_platform.is_rocm() + self._aiter_enabled = False + self._aiter_capable = False + if self._is_rocm: + from vllm._aiter_ops import rocm_aiter_ops + + self._aiter_enabled = bool(rocm_aiter_ops.is_enabled()) + self._aiter_capable = bool(rocm_aiter_ops.is_indexer_top_k_enabled()) def resolve_backend( self, logits: torch.Tensor, topk_tokens: int, num_rows: int @@ -193,6 +201,19 @@ class SparseIndexerTopk(torch.nn.Module): f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," f" stride={logits.stride()}, topk={topk_tokens}" ) + elif self._backend == "aiter": + if not self._is_rocm: + failures.append("requires a ROCm platform") + elif not self._aiter_enabled: + failures.append("requires AITER (VLLM_ROCM_USE_AITER=1)") + elif not self._aiter_capable: + failures.append( + "AITER has no tuned indexer top-k kernels for this GPU arch" + ) + if topk_tokens not in (512, 1024, 2048, 4096): + failures.append( + f"topk_tokens must be in (512, 1024, 2048, 4096), got {topk_tokens}" + ) elif self._backend == "flashinfer": if not self._is_cuda: failures.append("requires a CUDA platform") @@ -216,11 +237,14 @@ class SparseIndexerTopk(torch.nn.Module): self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: """The priority chain: cooperative -> persistent -> per_row. - deep_select/flashinfer/torch are opt-in only. + aiter/deep_select/flashinfer/torch are opt-in only. cooperative_topk is preferred whenever it is applicable, i.e. within its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to persistent_topk. + + aiter not auto-resolved as it does not outperform the per-row on every shape, + hence we delegate to the caller to decide when to use it. """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" @@ -285,7 +309,8 @@ class SparseIndexerTopk(torch.nn.Module): max_seq_len: int, ) -> None: """Run the resolved decode top-k implementation, writing into - topk_indices (int32, -1 fill for rows shorter than topk_tokens).""" + topk_indices (int32, -1 fill for rows shorter than topk_tokens). + """ backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) if backend == "deep_select": row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) @@ -344,15 +369,37 @@ class SparseIndexerTopk(torch.nn.Module): ) < row_ends.unsqueeze(1) indices = torch.where(in_range, indices, -1) topk_indices.copy_(indices) - else: - assert backend == "per_row", f"unknown topk backend: {backend}" - ops.top_k_per_row_decode( + elif backend == "aiter": + from vllm._aiter_ops import rocm_aiter_ops + + # AITER's decode kernel only implements the 1D/next_n == 1 form, + # which is equivalent once the ragged ends are expanded. + rocm_aiter_ops.indexer_top_k_decode( logits, - next_n, - seq_lens, + 1, + self._row_ends(seq_lens, next_n, logits.shape[0]), topk_indices, - logits.shape[0], - logits.stride(0), - logits.stride(1), topk_tokens, ) + else: + assert backend == "per_row", f"unknown topk backend: {backend}" + self._per_row(logits, seq_lens, next_n, topk_indices, topk_tokens) + + @staticmethod + def _per_row( + logits: torch.Tensor, + seq_lens: torch.Tensor, + next_n: int, + topk_indices: torch.Tensor, + topk_tokens: int, + ) -> None: + ops.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_tokens, + ) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index 761f79d95d..c357608c47 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -8,10 +8,11 @@ import vllm.envs as envs from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import get_current_vllm_config_or_none +from vllm.config import CUDAGraphMode, get_current_vllm_config_or_none from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -24,6 +25,7 @@ from vllm.models.glm5next.common.sparse_indexer import ( kv_cache_as_quant_view, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -37,6 +39,9 @@ from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) +_GFX950_KPOOL_AITER_MAX_POOLS = 256 * 1024 # (1M ctx) +_GFX950_KPOOL_AITER_MIN_POOLS = 16 * 1024 + # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -87,6 +92,51 @@ def _kpool_compress_insert( ) +def _kpool_decode_topk_backend( + configured: str, + *, + num_rows: int, + max_valid_seq_len: int, + select_k: int, + index_kpool: int, + full_cudagraph: bool, +) -> str: + """Select the topk backend for decodes. + + Heuristic based on ctx lengths: + - <16k pools: in-tree hip kernel + - 16-256k pools: use aiter + - >256k pools: in-tree because not measured on >1m ctx + + Since under FULL cudagraphs we cannot access the context length, + and the aiter kernel does not always outperform the in-tree kernel, + we default to in-tree kernel under FULL cudagraphs. Users can opt in + if they know their context length is long enough to benefit from aiter. + """ + if ( + configured != "auto" + or full_cudagraph + or not rocm_aiter_ops.is_indexer_top_k_enabled() + ): + return configured + if rocm_aiter_ops.is_indexer_top_k_supported( + is_prefill=False, + compress_ratio=index_kpool, + num_rows=num_rows, + max_valid_seq_len=max_valid_seq_len, + ): + return "aiter" + if ( + index_kpool > 1 + and select_k == 512 + and _GFX950_KPOOL_AITER_MIN_POOLS + <= max_valid_seq_len + <= _GFX950_KPOOL_AITER_MAX_POOLS + ): + return "aiter" + return configured + + @eager_break_during_capture def sparse_attn_indexer_kpool( hidden_states: torch.Tensor, @@ -117,6 +167,7 @@ def sparse_attn_indexer_kpool( # path and when the tail cache is disabled. tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -580,15 +631,29 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - torch.ops._C.top_k_per_row_decode( + full_cudagraph = ( + get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL + ) + if full_cudagraph: + topk_max_seq_len = max_pool_len * index_kpool + else: + topk_max_seq_len = attn_metadata_narrowed.max_seq_len + + resolved_backend = _kpool_decode_topk_backend( + topk_backend, + num_rows=num_rows, + max_valid_seq_len=cdiv(topk_max_seq_len, index_kpool), + select_k=select_k, + index_kpool=index_kpool, + full_cudagraph=full_cudagraph, + ) + get_indexer_topk(resolved_backend)( logits, - next_n, seq_lens, + next_n, topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), select_k, + topk_max_seq_len, ) # Resolve to token-level indices in the output buffer. @@ -664,6 +729,10 @@ class SparseAttnIndexerKpool(CustomOp): self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache + cfg = get_current_vllm_config_or_none() + self.topk_backend = ( + cfg.kernel_config.sparse_indexer_topk_backend if cfg is not None else "auto" + ) def forward_hip( self, @@ -726,6 +795,7 @@ class SparseAttnIndexerKpool(CustomOp): positions, self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, + self.topk_backend, ) def forward_native( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 505c595b41..a3412aa835 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -32,18 +32,6 @@ logger = init_logger(__name__) FP8_DTYPE = current_platform.fp8_dtype() -@functools.cache -def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None: - try: - from aiter.ops.topk import ( - top_k_per_row_decode, - top_k_per_row_prefill, - ) - except ImportError: - return None - return top_k_per_row_prefill, top_k_per_row_decode - - @functools.cache def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: from vllm._aiter_ops import rocm_aiter_ops @@ -58,9 +46,6 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: return pa_sparse_prefill_opus -_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 -_GFX950_C4A_NATIVE_MAX_ROWS = 256 -_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # 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 @@ -71,43 +56,6 @@ def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool: return compress_ratio == 4 -def _get_aiter_top_k_kernel( - *, - is_prefill: bool, - compress_ratio: int, - num_rows: int, - max_valid_seq_len: int | None = None, - num_columns: int | None = None, - topk_tokens: int = 1024, - on_gfx950: bool = _ON_GFX950, -) -> Callable[..., None] | None: - if compress_ratio <= 1 or not on_gfx950: - return None - - if not is_prefill: - assert max_valid_seq_len is not None - if ( - topk_tokens == 512 - and 0 < num_rows <= 384 - and num_columns is not None - and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS - ): - return None - # AITER v0.1.19 decode is one-block only. This measured gfx950 - # FP32/k=1024 compressed-row boundary is independent of the native - # split-count boundary in sampler.cu. - if ( - num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS - and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN - ): - return None - - topk_ops = _get_aiter_topk_ops() - if topk_ops is None: - return None - return topk_ops[0] if is_prefill else topk_ops[1] - - @triton.jit def _localize_aiter_prefill_topk_kernel( indices_ptr, @@ -145,46 +93,20 @@ def _localize_aiter_prefill_topk( def _launch_aiter_top_k_per_row_prefill( - top_k_per_row_prefill: Callable[..., None], logits: torch.Tensor, row_starts: torch.Tensor, row_ends: torch.Tensor, indices: torch.Tensor, topk_tokens: int, ) -> None: - top_k_per_row_prefill( - logits, - row_starts, - row_ends, - indices, - None, - logits.shape[0], - logits.stride(0), - logits.stride(1), - k=topk_tokens, + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.indexer_top_k_prefill( + logits, row_starts, row_ends, indices, topk_tokens ) _localize_aiter_prefill_topk(indices, row_starts) -def _launch_aiter_top_k_per_row_decode( - top_k_per_row_decode: Callable[..., None], - logits: torch.Tensor, - seq_lens: torch.Tensor, - indices: torch.Tensor, - topk_tokens: int, -) -> None: - top_k_per_row_decode( - logits, - 1, - seq_lens.reshape(-1), - indices, - logits.shape[0], - logits.stride(0), - logits.stride(1), - k=topk_tokens, - ) - - @triton.jit def _indexer_k_quant_and_cache_kernel( k_ptr, # [num_tokens, head_dim] @@ -1195,6 +1117,8 @@ def rocm_aiter_sparse_attn_indexer( candidate_block_size: int = 0, candidate_write: bool = False, ) -> torch.Tensor: + from vllm._aiter_ops import rocm_aiter_ops + # careful! this will be None in dummy run forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata @@ -1347,14 +1271,12 @@ def rocm_aiter_sparse_attn_indexer( num_rows = logits.shape[0] - aiter_topk_kernel = _get_aiter_top_k_kernel( + if rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=True, compress_ratio=compress_ratio, num_rows=num_rows, - ) - if aiter_topk_kernel is not None: + ): _launch_aiter_top_k_per_row_prefill( - aiter_topk_kernel, logits, chunk.cu_seqlen_ks, chunk.cu_seqlen_ke, @@ -1448,19 +1370,20 @@ def rocm_aiter_sparse_attn_indexer( max_compressed_seq_len = max_model_len else: max_compressed_seq_len = layer_attn_metadata.max_seq_len // compress_ratio - aiter_topk_kernel = _get_aiter_top_k_kernel( + if rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=False, compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + ) and not rocm_aiter_ops.dsv4_indexer_prefers_native_top_k( + num_rows=num_rows, num_columns=logits.shape[1], topk_tokens=topk_tokens, - ) - if aiter_topk_kernel is not None: - _launch_aiter_top_k_per_row_decode( - aiter_topk_kernel, + ): + rocm_aiter_ops.indexer_top_k_decode( logits, - decode_metadata.seq_lens, + 1, + decode_metadata.seq_lens.reshape(-1), topk_indices, topk_tokens, )