@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 98b29aa99a

[ROCm][Perf][GLM-5.3-Flash] Add AITER topk backend for decodes (#58167)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Simon Danielsson committed Oct 2, 2026 at 12:47 UTC 98b29aa99aee8fe58a21b168b45a5dcb910f63c0
8 files changed +705 -240
benchmarks/kernels/benchmark_top_k_per_row_decode.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark the ROCm sparse-indexer decode top-k in isolation.
4 +
5 +Example usage:
6 +python3 benchmarks/kernels/benchmark_top_k_per_row_decode.py \
7 + --rows 4 16 32 --lens 2048 8192 16384 65536 --capacity 262144
8 +"""
9 +
10 +import argparse
11 +import itertools
12 +import json
13 +
14 +import torch
15 +
16 +import vllm._custom_ops # noqa: F401
17 +from vllm._aiter_ops import rocm_aiter_ops
18 +from vllm.triton_utils import triton
19 +
20 +DISTRIBUTIONS = ("indexer", "fp8", "randn")
21 +
22 +
23 +def make_logits(rows: int, capacity: int, dist: str, spread: float):
24 + normal = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
25 + if dist == "randn":
26 + return normal
27 + if dist == "fp8":
28 + return normal.to(torch.float8_e4m3fnuz).to(torch.float32)
29 + return normal * spread + 3.0
30 +
31 +
32 +def make_inputs(rows: int, live: int, capacity: int, k: int, args):
33 + torch.manual_seed(args.seed)
34 + logits = make_logits(rows, capacity, args.dist, args.spread)
35 + seq_lens = torch.full((rows,), live, device="cuda", dtype=torch.int32)
36 + invalid = torch.arange(capacity, device="cuda")[None] >= seq_lens[:, None]
37 + logits.masked_fill_(invalid, -float("inf"))
38 + indices = torch.empty((rows, k), device="cuda", dtype=torch.int32)
39 + return logits, seq_lens, indices
40 +
41 +
42 +def run_native(logits, seq_lens, indices, k):
43 + torch.ops._C.top_k_per_row_decode(
44 + logits,
45 + 1,
46 + seq_lens,
47 + indices,
48 + logits.shape[0],
49 + logits.stride(0),
50 + logits.stride(1),
51 + k,
52 + )
53 +
54 +
55 +def run_aiter(logits, seq_lens, indices, k):
56 + rocm_aiter_ops.indexer_top_k_decode(logits, 1, seq_lens, indices, k)
57 +
58 +
59 +BACKENDS = {"native": run_native, "aiter": run_aiter}
60 +
61 +
62 +def check(logits, indices, k):
63 + selected = logits.gather(1, indices.long()).sort(dim=1, descending=True).values
64 + expected = logits.topk(k, dim=1).values
65 + torch.testing.assert_close(selected, expected, atol=0, rtol=0)
66 +
67 +
68 +def capture(run):
69 + side = torch.cuda.Stream()
70 + side.wait_stream(torch.cuda.current_stream())
71 + with torch.cuda.stream(side):
72 + for _ in range(3):
73 + run()
74 + torch.cuda.current_stream().wait_stream(side)
75 + graph = torch.cuda.CUDAGraph()
76 + with torch.cuda.graph(graph):
77 + run()
78 + return graph
79 +
80 +
81 +def benchmark(backend: str, rows: int, live: int, args) -> dict:
82 + logits, seq_lens, indices = make_inputs(rows, live, args.capacity, args.top_k, args)
83 + fn = BACKENDS[backend]
84 +
85 + def run():
86 + fn(logits, seq_lens, indices, args.top_k)
87 +
88 + run()
89 + torch.accelerator.synchronize()
90 + check(logits, indices, args.top_k)
91 + timed = run
92 + if args.graph:
93 + graph = capture(run)
94 + indices.zero_()
95 + graph.replay()
96 + torch.accelerator.synchronize()
97 + check(logits, indices, args.top_k)
98 + timed = graph.replay
99 + ms = triton.testing.do_bench(timed, warmup=args.warmup, rep=args.rep)
100 + return dict(
101 + backend=backend,
102 + mode="graph" if args.graph else "eager",
103 + dist=args.dist,
104 + rows=rows,
105 + live=live,
106 + capacity=args.capacity,
107 + k=args.top_k,
108 + us=ms * 1000,
109 + live_gbps=rows * live * 4 / (ms * 1e6),
110 + )
111 +
112 +
113 +def main():
114 + parser = argparse.ArgumentParser(description=__doc__)
115 + parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8])
116 + parser.add_argument(
117 + "--lens",
118 + type=int,
119 + nargs="+",
120 + default=[32768, 65536, 98304, 131072],
121 + help="Valid columns per row; must be at least --top-k",
122 + )
123 + parser.add_argument("--capacity", type=int, default=262144)
124 + parser.add_argument("--top-k", type=int, default=512)
125 + parser.add_argument(
126 + "--backends", nargs="+", choices=list(BACKENDS), default=["native", "aiter"]
127 + )
128 + parser.add_argument(
129 + "--graph",
130 + action="store_true",
131 + help="Time a captured graph replay, matching how decode actually runs",
132 + )
133 + parser.add_argument(
134 + "--dist",
135 + choices=DISTRIBUTIONS,
136 + default="indexer",
137 + help="Logit value distribution; see the module docstring",
138 + )
139 + parser.add_argument(
140 + "--spread",
141 + type=float,
142 + default=0.01,
143 + help="Relative width of the --dist indexer cluster",
144 + )
145 + parser.add_argument("--warmup", type=int, default=25)
146 + parser.add_argument("--rep", type=int, default=200)
147 + parser.add_argument("--seed", type=int, default=123)
148 + args = parser.parse_args()
149 +
150 + print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args))))
151 + for backend, rows, live in itertools.product(args.backends, args.rows, args.lens):
152 + if live < args.top_k:
153 + continue
154 + print(json.dumps(benchmark(backend, rows, live, args)), flush=True)
155 +
156 +
157 +if __name__ == "__main__":
158 + main()
tests/kernels/test_top_k_per_row.py
+225 -118
@@ -636,8 +636,8 @@ def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None:
636 @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
637 @torch.inference_mode()
638 def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
639 + from vllm._aiter_ops import rocm_aiter_ops
640 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
640 - _get_aiter_top_k_kernel,
641 _launch_aiter_top_k_per_row_prefill,
642 )
643
@@ -647,17 +647,14 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
647 top_k = 4
648 indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda")
649
650 - aiter_topk_kernel = _get_aiter_top_k_kernel(
650 + if not rocm_aiter_ops.is_indexer_top_k_supported(
651 is_prefill=True,
652 compress_ratio=4,
653 num_rows=logits.shape[0],
654 - on_gfx950=True,
655 - )
656 - if aiter_topk_kernel is None:
654 + ):
655 pytest.skip("AITER top-k is unavailable")
656 assert (
657 _launch_aiter_top_k_per_row_prefill(
660 - aiter_topk_kernel,
658 logits,
659 row_starts,
660 row_ends,
@@ -683,10 +680,7 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
680 def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
681 num_speculative_tokens: int,
682 ) -> None:
686 - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
687 - _get_aiter_top_k_kernel,
688 - _launch_aiter_top_k_per_row_decode,
689 - )
683 + from vllm._aiter_ops import rocm_aiter_ops
684
685 query_tokens = num_speculative_tokens + 1
686 offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda")
@@ -701,20 +695,18 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
695 top_k = 4
696 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
697
704 - aiter_topk_kernel = _get_aiter_top_k_kernel(
698 + if not rocm_aiter_ops.is_indexer_top_k_supported(
699 is_prefill=False,
700 compress_ratio=4,
701 num_rows=num_rows,
702 max_valid_seq_len=seq_lens.max().item(),
709 - on_gfx950=True,
710 - )
711 - if aiter_topk_kernel is None:
703 + ):
704 pytest.skip("AITER top-k is unavailable")
705 assert (
714 - _launch_aiter_top_k_per_row_decode(
715 - aiter_topk_kernel,
706 + rocm_aiter_ops.indexer_top_k_decode(
707 logits,
717 - seq_lens,
708 + 1,
709 + seq_lens.reshape(-1),
710 indices,
711 top_k,
712 )
@@ -729,130 +721,107 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
721 assert torch.all(indices[row_idx, num_valid:] == -1)
722
723
732 -def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
733 - from vllm.v1.attention.ops import rocm_aiter_mla_sparse
724 +@requires_gfx950
725 +@torch.inference_mode()
726 +def test_aiter_decode_topk_matches_per_row_under_mtp() -> None:
727 + """The aiter backend expands (B, 1) seq_lens into per-row ends before
728 + launching, so next_n > 1 must agree with the in-tree kernel, including
729 + the rows that clamp to a zero-length prefix.
730 + """
731 + from vllm.model_executor.layers.indexer_topk import get_indexer_topk
732
735 - def prefill_kernel(*args, **kwargs) -> None:
736 - pass
733 + next_n = 4
734 + top_k = 512
735 + num_columns = 1024
736 + seq_lens = torch.tensor([[2], [513], [700]], dtype=torch.int32, device="cuda")
737 + num_rows = seq_lens.shape[0] * next_n
738 +
739 + torch.manual_seed(0)
740 + logits = torch.randn(num_rows, num_columns, dtype=torch.float32, device="cuda")
741 +
742 + outputs = {}
743 + for backend in ("aiter", "per_row"):
744 + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
745 + get_indexer_topk(backend)(logits, seq_lens, next_n, indices, top_k, num_columns)
746 + outputs[backend] = indices
747 +
748 + row_ends = (
749 + (seq_lens - next_n + 1 + torch.arange(next_n, device="cuda"))
750 + .clamp_(min=0)
751 + .reshape(-1)
752 + .tolist()
753 + )
754 + assert row_ends[:next_n] == [0, 0, 1, 2]
755 +
756 + for row_idx, row_end in enumerate(row_ends):
757 + num_valid = min(top_k, row_end)
758 + aiter_row = outputs["aiter"][row_idx]
759 + assert set(aiter_row[:num_valid].tolist()) == set(
760 + outputs["per_row"][row_idx][:num_valid].tolist()
761 + )
762 + assert torch.all(aiter_row[num_valid:] == -1)
763
738 - def decode_kernel(*args, **kwargs) -> None:
739 - pass
764 +
765 +def _enable_aiter_topk(monkeypatch, enabled: bool = True) -> None:
766 + from vllm._aiter_ops import rocm_aiter_ops
767
768 monkeypatch.setattr(
742 - rocm_aiter_mla_sparse,
743 - "_get_aiter_topk_ops",
744 - lambda: (prefill_kernel, decode_kernel),
769 + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: enabled, raising=True
770 )
746 - get_kernel = rocm_aiter_mla_sparse._get_aiter_top_k_kernel
771
748 - eligible_cases = [
749 - (
750 - dict(is_prefill=True, compress_ratio=4, num_rows=1, on_gfx950=True),
751 - prefill_kernel,
752 - ),
753 - (
754 - dict(
755 - is_prefill=False,
756 - compress_ratio=4,
757 - num_rows=1,
758 - max_valid_seq_len=65_536,
759 - on_gfx950=True,
760 - ),
761 - decode_kernel,
762 - ),
763 - (
764 - dict(
765 - is_prefill=False,
766 - compress_ratio=4,
767 - num_rows=257,
768 - max_valid_seq_len=250_000,
769 - num_columns=524_288,
770 - on_gfx950=True,
771 - ),
772 - decode_kernel,
773 - ),
774 - (
775 - dict(
776 - is_prefill=False,
777 - compress_ratio=2,
778 - num_rows=385,
779 - max_valid_seq_len=250_000,
780 - num_columns=524_288,
781 - topk_tokens=512,
782 - on_gfx950=True,
783 - ),
784 - decode_kernel,
785 - ),
786 - (
787 - dict(
788 - is_prefill=False,
789 - compress_ratio=2,
790 - num_rows=384,
791 - max_valid_seq_len=50_000,
792 - num_columns=1_048_577,
793 - topk_tokens=512,
794 - on_gfx950=True,
795 - ),
796 - decode_kernel,
797 - ),
798 - ]
799 - for kwargs, expected_kernel in eligible_cases:
800 - assert get_kernel(**kwargs) is expected_kernel
772
802 - def fail_if_imported() -> None:
803 - pytest.fail("ineligible shapes must not import AITER top-k")
773 +def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
774 + from vllm._aiter_ops import rocm_aiter_ops
775
805 - monkeypatch.setattr(
806 - rocm_aiter_mla_sparse,
807 - "_get_aiter_topk_ops",
808 - fail_if_imported,
809 - )
810 - ineligible_cases = [
811 - dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True),
776 + _enable_aiter_topk(monkeypatch)
777 + is_supported = rocm_aiter_ops.is_indexer_top_k_supported
778 +
779 + eligible_cases = [
780 + dict(is_prefill=True, compress_ratio=4, num_rows=1),
781 + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_536),
782 dict(
783 is_prefill=False,
814 - compress_ratio=2,
815 - num_rows=384,
816 - max_valid_seq_len=1_048_576,
817 - num_columns=1_048_576,
818 - topk_tokens=512,
819 - on_gfx950=True,
784 + compress_ratio=4,
785 + num_rows=257,
786 + max_valid_seq_len=250_000,
787 ),
788 dict(
789 is_prefill=False,
823 - compress_ratio=4,
824 - num_rows=1,
825 - max_valid_seq_len=65_537,
826 - on_gfx950=True,
790 + compress_ratio=2,
791 + num_rows=384,
792 + max_valid_seq_len=1_048_576,
793 ),
794 + ]
795 + for kwargs in eligible_cases:
796 + assert is_supported(**kwargs) is True
797 +
798 + ineligible_cases = [
799 + dict(is_prefill=True, compress_ratio=1, num_rows=1),
800 + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_537),
801 dict(
802 is_prefill=False,
803 compress_ratio=4,
804 num_rows=256,
805 max_valid_seq_len=125_000,
833 - on_gfx950=True,
834 - ),
835 - dict(
836 - is_prefill=False,
837 - compress_ratio=4,
838 - num_rows=320,
839 - max_valid_seq_len=2_500,
840 - on_gfx950=False,
806 ),
807 ]
808 for kwargs in ineligible_cases:
844 - assert get_kernel(**kwargs) is None
809 + assert is_supported(**kwargs) is False
810 +
811 + _enable_aiter_topk(monkeypatch, enabled=False)
812 + assert is_supported(is_prefill=True, compress_ratio=4, num_rows=1) is False
813
846 - monkeypatch.setattr(rocm_aiter_mla_sparse, "_get_aiter_topk_ops", lambda: None)
847 - assert (
848 - get_kernel(
849 - is_prefill=True,
850 - compress_ratio=4,
851 - num_rows=1,
852 - on_gfx950=True,
853 - )
854 - is None
855 - )
814 +
815 +def test_dsv4_native_top_k_window() -> None:
816 + """The in-tree decode kernel's measured advantage band over AITER, which
817 + applies to the DSV4 indexer only."""
818 + from vllm._aiter_ops import rocm_aiter_ops
819 +
820 + prefers_native = rocm_aiter_ops.dsv4_indexer_prefers_native_top_k
821 + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=512) is True
822 + assert prefers_native(num_rows=385, num_columns=524_288, topk_tokens=512) is False
823 + assert prefers_native(num_rows=128, num_columns=1_048_577, topk_tokens=512) is False
824 + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=1024) is False
825
826
827 @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:
1896 if not is_sm100:
1897 with pytest.raises(RuntimeError, match="SM100a/SM103a"):
1898 resolve("deep_select")
1899 +
1900 +
1901 +def _patch_aiter_topk(monkeypatch, decode_kernel, enabled: bool = True) -> None:
1902 + from vllm._aiter_ops import rocm_aiter_ops
1903 +
1904 + _enable_aiter_topk(monkeypatch, enabled=enabled)
1905 + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: enabled, raising=True)
1906 + monkeypatch.setattr(rocm_aiter_ops, "indexer_top_k_decode", decode_kernel)
1907 +
1908 +
1909 +def _make_topk(backend: str):
1910 + from vllm.config import VllmConfig, set_current_vllm_config
1911 + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
1912 +
1913 + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
1914 + with set_current_vllm_config(cfg):
1915 + return SparseIndexerTopk()
1916 +
1917 +
1918 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
1919 +def test_sparse_indexer_topk_aiter_requires_rocm() -> None:
1920 + """Explicitly requesting the ROCm-only backend must fail fast elsewhere."""
1921 + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda")
1922 + with pytest.raises(RuntimeError, match="ROCm"):
1923 + _make_topk("aiter").resolve_backend(logits, 2048, 8)
1924 +
1925 +
1926 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
1927 +def test_sparse_indexer_topk_backend_resolution_aiter(monkeypatch) -> None:
1928 + """Aiter is opt-in only: "auto" never selects it (the shape-dependent perf
1929 + gate lives at the call site), and an explicit request validates hard
1930 + capability only."""
1931 +
1932 + def decode_kernel(*args, **kwargs) -> None:
1933 + pass
1934 +
1935 + _patch_aiter_topk(monkeypatch, decode_kernel)
1936 + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda")
1937 +
1938 + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row"
1939 + assert _make_topk("aiter").resolve_backend(logits, 2048, 8) == "aiter"
1940 + # Shapes the call-site perf gate would decline still resolve.
1941 + assert _make_topk("aiter").resolve_backend(logits, 512, 8) == "aiter"
1942 + with pytest.raises(RuntimeError, match="topk_tokens must be in"):
1943 + _make_topk("aiter").resolve_backend(logits, 3000, 8)
1944 +
1945 + _enable_aiter_topk(monkeypatch, enabled=False)
1946 + with pytest.raises(RuntimeError, match="no tuned indexer top-k kernels"):
1947 + _make_topk("aiter").resolve_backend(logits, 2048, 8)
1948 +
1949 + _patch_aiter_topk(monkeypatch, decode_kernel, enabled=False)
1950 + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row"
1951 + with pytest.raises(RuntimeError, match="VLLM_ROCM_USE_AITER"):
1952 + _make_topk("aiter").resolve_backend(logits, 2048, 8)
1953 +
1954 +
1955 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
1956 +@torch.inference_mode()
1957 +def test_sparse_indexer_topk_aiter_applies_no_shape_gate(monkeypatch) -> None:
1958 + """Shape-dependent perf gating belongs to the call site: once "aiter" is
1959 + requested the dispatcher must run it for every supported k, including the
1960 + k=512 band the in-tree kernel wins on DSV4's compressed-KV logits."""
1961 + from vllm.model_executor.layers import indexer_topk
1962 +
1963 + calls: list[tuple] = []
1964 + _patch_aiter_topk(monkeypatch, lambda *args: calls.append(args))
1965 + monkeypatch.setattr(
1966 + indexer_topk.ops,
1967 + "top_k_per_row_decode",
1968 + lambda *args: pytest.fail("an explicit aiter request must not fall back"),
1969 + )
1970 +
1971 + num_rows = 8
1972 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda")
1973 + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda")
1974 + indices = torch.empty((num_rows, 512), dtype=torch.int32, device="cuda")
1975 + op = _make_topk("aiter")
1976 +
1977 + op(logits, seq_lens, 1, indices, 512, 4096)
1978 + # Past the compressed-context boundary the call-site gate declines on.
1979 + op(logits, seq_lens, 1, indices, 512, 4 * (64 * 1024 + 1))
1980 + assert len(calls) == 2
1981 +
1982 +
1983 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
1984 +@torch.inference_mode()
1985 +def test_sparse_indexer_topk_auto_stays_on_per_row(monkeypatch) -> None:
1986 + """The "auto" chain must not reach AITER on its own: the decode top-k runs
1987 + through top_k_per_row_decode unless the caller asks for the aiter backend."""
1988 + from vllm.model_executor.layers import indexer_topk
1989 +
1990 + def decode_kernel(*args, **kwargs) -> None:
1991 + pytest.fail("AITER is opt-in only")
1992 +
1993 + _patch_aiter_topk(monkeypatch, decode_kernel)
1994 + calls: list[tuple] = []
1995 + monkeypatch.setattr(
1996 + indexer_topk.ops, "top_k_per_row_decode", lambda *args: calls.append(args)
1997 + )
1998 +
1999 + num_rows = 8
2000 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda")
2001 + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda")
2002 + indices = torch.empty((num_rows, 1024), dtype=torch.int32, device="cuda")
2003 +
2004 + _make_topk("auto")(logits, seq_lens, 1, indices, 1024, 4096)
2005 + assert len(calls) == 1
2006 +
2007 +
2008 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm")
2009 +@torch.inference_mode()
2010 +def test_sparse_indexer_topk_aiter_expands_mtp_row_ends(monkeypatch) -> None:
2011 + """AITER's decode kernel only implements the next_n == 1 form, so spec
2012 + decode batches must be handed the equivalent expanded per-row ends."""
2013 + from vllm.model_executor.layers import indexer_topk
2014 +
2015 + recorded: dict = {}
2016 +
2017 + def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None:
2018 + recorded["next_n"] = next_n
2019 + recorded["row_ends"] = seq_lens.tolist()
2020 +
2021 + _patch_aiter_topk(monkeypatch, decode_kernel)
2022 + monkeypatch.setattr(
2023 + indexer_topk.ops,
2024 + "top_k_per_row_decode",
2025 + lambda *args: pytest.fail("an explicit aiter request must not fall back"),
2026 + )
2027 +
2028 + next_n = 2
2029 + logits = torch.randn(4, 4096, dtype=torch.float32, device="cuda")
2030 + seq_lens = torch.tensor([[100], [200]], dtype=torch.int32, device="cuda")
2031 + indices = torch.empty((4, 1024), dtype=torch.int32, device="cuda")
2032 +
2033 + _make_topk("aiter")(logits, seq_lens, next_n, indices, 1024, 800)
2034 +
2035 + assert recorded["next_n"] == 1
2036 + assert recorded["row_ends"] == [99, 100, 199, 200]
tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
+72 -14
@@ -8,10 +8,6 @@ import torch
8 from vllm.config import VllmConfig, set_current_vllm_config
9 from vllm.platforms import current_platform
10
11 -pytestmark = pytest.mark.skipif(
12 - not current_platform.is_cuda(), reason="CUDA-only dispatch"
13 -)
14 -
11
12 def _require_deep_gemm() -> None:
13 from vllm.utils.deep_gemm import has_deep_gemm
@@ -20,17 +16,10 @@ def _require_deep_gemm() -> None:
16 pytest.skip("kpool indexer requires DeepGEMM")
17
18
23 -@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
24 -def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
25 - """The kpool indexer must read kernel_config.sparse_indexer_topk_backend
26 - and hand it to the shared SparseIndexerTopk dispatcher, rather than
27 - hard-coding its own cooperative/persistent/per_row choice."""
28 - _require_deep_gemm()
29 - from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
30 -
19 +def _build(indexer_cls, backend: str):
20 cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
21 with set_current_vllm_config(cfg):
33 - op = SparseAttnIndexerKpool(
22 + return indexer_cls(
23 k_cache=None,
24 quant_block_size=128,
25 scale_fmt="ue8m0",
@@ -40,4 +29,73 @@ def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> N
29 max_total_seq_len=8192,
30 topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"),
31 )
43 - assert op.topk_backend == backend
32 +
33 +
34 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only dispatch")
35 +@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
36 +def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
37 + """The kpool indexer must read kernel_config.sparse_indexer_topk_backend
38 + and hand it to the shared SparseIndexerTopk dispatcher, rather than
39 + hard-coding its own cooperative/persistent/per_row choice."""
40 + _require_deep_gemm()
41 + from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
42 +
43 + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend
44 +
45 +
46 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch")
47 +@pytest.mark.parametrize("backend", ["auto", "aiter", "per_row", "torch"])
48 +def test_kpool_indexer_dispatches_through_shared_topk_backend_rocm(
49 + backend: str,
50 +) -> None:
51 + """The AMD kpool indexer must go through the same dispatcher, so the
52 + AITER decode top-k is reachable from kernel_config."""
53 + from vllm.models.glm5next.amd.sparse_indexer import SparseAttnIndexerKpool
54 +
55 + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend
56 +
57 +
58 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch")
59 +@pytest.mark.parametrize(
60 + "configured,full_cudagraph,expected",
61 + [
62 + ("auto", False, "aiter"),
63 + ("auto", True, "auto"),
64 + ("aiter", True, "aiter"),
65 + ("per_row", False, "per_row"),
66 + ],
67 +)
68 +def test_kpool_decode_topk_backend_never_auto_selects_aiter_under_full_cudagraph(
69 + monkeypatch: pytest.MonkeyPatch,
70 + configured: str,
71 + full_cudagraph: bool,
72 + expected: str,
73 +) -> None:
74 + """A FULL cudagraph bakes in one kernel for every replay, so "auto" cannot
75 + use the context length to pick AITER and must stay on the in-tree kernel.
76 + Explicit backends keep working, and outside FULL the long-context heuristic
77 + still narrows to AITER."""
78 + from vllm._aiter_ops import rocm_aiter_ops
79 + from vllm.models.glm5next.amd import sparse_indexer
80 +
81 + monkeypatch.setattr(
82 + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: True, raising=False
83 + )
84 + monkeypatch.setattr(
85 + rocm_aiter_ops,
86 + "is_indexer_top_k_supported",
87 + lambda **kwargs: True,
88 + raising=False,
89 + )
90 +
91 + assert (
92 + sparse_indexer._kpool_decode_topk_backend(
93 + configured,
94 + num_rows=8,
95 + max_valid_seq_len=64 * 1024,
96 + select_k=512,
97 + index_kpool=4,
98 + full_cudagraph=full_cudagraph,
99 + )
100 + == expected
101 + )
vllm/_aiter_ops.py
+99
@@ -1913,6 +1913,11 @@ def _sync_aiter_situv2_moe_env() -> None:
1913 os.environ.pop(name, None)
1914
1915
1916 +_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
1917 +_GFX950_C4A_NATIVE_MAX_ROWS = 256
1918 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
1919 +
1920 +
1921 class rocm_aiter_ops:
1922 """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
1923
@@ -2546,6 +2551,100 @@ class rocm_aiter_ops:
2551 """Whether AITER has DSV4 native-I384 configs through the given M."""
2552 return cls._probe_dsv4_i384_fhmoe_capability(num_tokens)
2553
2554 + @classmethod
2555 + @if_aiter_supported
2556 + def is_indexer_top_k_enabled(cls) -> bool:
2557 + """gfx950 is the only arch with tuned AITER indexer top-k kernels."""
2558 + from vllm.platforms.rocm import on_gfx950
2559 +
2560 + return cls._AITER_ENABLED and on_gfx950()
2561 +
2562 + @classmethod
2563 + def is_indexer_top_k_supported(
2564 + cls,
2565 + *,
2566 + is_prefill: bool,
2567 + compress_ratio: int,
2568 + num_rows: int,
2569 + max_valid_seq_len: int | None = None,
2570 + ) -> bool:
2571 + """Whether AITER's sparse indexer top-k beats the in-tree kernel for
2572 + this shape."""
2573 + if compress_ratio <= 1 or not cls.is_indexer_top_k_enabled():
2574 + return False
2575 +
2576 + if not is_prefill:
2577 + assert max_valid_seq_len is not None
2578 + # AITER v0.1.19 decode is one-block only. This measured gfx950
2579 + # FP32/k=1024 compressed-row boundary is independent of the native
2580 + # split-count boundary in sampler.cu.
2581 + if (
2582 + num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS
2583 + and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN
2584 + ):
2585 + return False
2586 +
2587 + return True
2588 +
2589 + @staticmethod
2590 + def dsv4_indexer_prefers_native_top_k(
2591 + *,
2592 + num_rows: int,
2593 + num_columns: int,
2594 + topk_tokens: int,
2595 + ) -> bool:
2596 + """The in-tree decode kernel's measured advantage window over AITER.
2597 + Tuned on DSV4's compressed-KV logits, so it applies to that indexer
2598 + only."""
2599 + return (
2600 + topk_tokens == 512
2601 + and 0 < num_rows <= 384
2602 + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS
2603 + )
2604 +
2605 + @staticmethod
2606 + def indexer_top_k_decode(
2607 + logits: torch.Tensor,
2608 + next_n: int,
2609 + seq_lens: torch.Tensor,
2610 + topk_indices: torch.Tensor,
2611 + topk_tokens: int,
2612 + ) -> None:
2613 + from aiter.ops.topk import top_k_per_row_decode
2614 +
2615 + top_k_per_row_decode(
2616 + logits,
2617 + next_n,
2618 + seq_lens,
2619 + topk_indices,
2620 + logits.shape[0],
2621 + logits.stride(0),
2622 + logits.stride(1),
2623 + k=topk_tokens,
2624 + )
2625 +
2626 + @staticmethod
2627 + def indexer_top_k_prefill(
2628 + logits: torch.Tensor,
2629 + row_starts: torch.Tensor,
2630 + row_ends: torch.Tensor,
2631 + indices: torch.Tensor,
2632 + topk_tokens: int,
2633 + ) -> None:
2634 + from aiter.ops.topk import top_k_per_row_prefill
2635 +
2636 + top_k_per_row_prefill(
2637 + logits,
2638 + row_starts,
2639 + row_ends,
2640 + indices,
2641 + None,
2642 + logits.shape[0],
2643 + logits.stride(0),
2644 + logits.stride(1),
2645 + k=topk_tokens,
2646 + )
2647 +
2648 @staticmethod
2649 def register_ops_once() -> None:
2650 global _OPS_REGISTERED
vllm/config/kernel.py
+3
@@ -167,6 +167,7 @@ SparseIndexerTopkBackend = Literal[
167 "per_row",
168 "flashinfer",
169 "torch",
170 + "aiter",
171 ]
172
173 # Architectures whose model code wires up the flashinfer moe_ep experts. MTP
@@ -289,6 +290,8 @@ class KernelConfig:
290 - "per_row": Use vLLM's top_k_per_row_decode kernel
291 - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel
292 - "torch": Use a plain torch.topk implementation (debug reference)
293 + - "aiter": Use AITER's top_k_per_row_decode kernel (ROCm gfx950 only,
294 + requires VLLM_ROCM_USE_AITER)
295
296 Explicit values raise RuntimeError when their constraints are not met.
297 """
vllm/model_executor/layers/indexer_topk.py
+57 -10
@@ -163,6 +163,14 @@ class SparseIndexerTopk(torch.nn.Module):
163 current_platform.has_device_capability(90)
164 and not current_platform.is_device_capability_family(120)
165 )
166 + self._is_rocm = current_platform.is_rocm()
167 + self._aiter_enabled = False
168 + self._aiter_capable = False
169 + if self._is_rocm:
170 + from vllm._aiter_ops import rocm_aiter_ops
171 +
172 + self._aiter_enabled = bool(rocm_aiter_ops.is_enabled())
173 + self._aiter_capable = bool(rocm_aiter_ops.is_indexer_top_k_enabled())
174
175 def resolve_backend(
176 self, logits: torch.Tensor, topk_tokens: int, num_rows: int
@@ -193,6 +201,19 @@ class SparseIndexerTopk(torch.nn.Module):
201 f"inputs violate DeepSelect's constraints: dtype={logits.dtype},"
202 f" stride={logits.stride()}, topk={topk_tokens}"
203 )
204 + elif self._backend == "aiter":
205 + if not self._is_rocm:
206 + failures.append("requires a ROCm platform")
207 + elif not self._aiter_enabled:
208 + failures.append("requires AITER (VLLM_ROCM_USE_AITER=1)")
209 + elif not self._aiter_capable:
210 + failures.append(
211 + "AITER has no tuned indexer top-k kernels for this GPU arch"
212 + )
213 + if topk_tokens not in (512, 1024, 2048, 4096):
214 + failures.append(
215 + f"topk_tokens must be in (512, 1024, 2048, 4096), got {topk_tokens}"
216 + )
217 elif self._backend == "flashinfer":
218 if not self._is_cuda:
219 failures.append("requires a CUDA platform")
@@ -216,11 +237,14 @@ class SparseIndexerTopk(torch.nn.Module):
237 self, logits: torch.Tensor, topk_tokens: int, num_rows: int
238 ) -> str:
239 """The priority chain: cooperative -> persistent -> per_row.
219 - deep_select/flashinfer/torch are opt-in only.
240 + aiter/deep_select/flashinfer/torch are opt-in only.
241
242 cooperative_topk is preferred whenever it is applicable, i.e. within
243 its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to
244 persistent_topk.
245 +
246 + aiter not auto-resolved as it does not outperform the per-row on every shape,
247 + hence we delegate to the caller to decide when to use it.
248 """
249 if not self._cooperative_constraints(logits, topk_tokens, num_rows):
250 return "cooperative"
@@ -285,7 +309,8 @@ class SparseIndexerTopk(torch.nn.Module):
309 max_seq_len: int,
310 ) -> None:
311 """Run the resolved decode top-k implementation, writing into
288 - topk_indices (int32, -1 fill for rows shorter than topk_tokens)."""
312 + topk_indices (int32, -1 fill for rows shorter than topk_tokens).
313 + """
314 backend = self.resolve_backend(logits, topk_tokens, logits.shape[0])
315 if backend == "deep_select":
316 row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
@@ -344,15 +369,37 @@ class SparseIndexerTopk(torch.nn.Module):
369 ) < row_ends.unsqueeze(1)
370 indices = torch.where(in_range, indices, -1)
371 topk_indices.copy_(indices)
347 - else:
348 - assert backend == "per_row", f"unknown topk backend: {backend}"
349 - ops.top_k_per_row_decode(
372 + elif backend == "aiter":
373 + from vllm._aiter_ops import rocm_aiter_ops
374 +
375 + # AITER's decode kernel only implements the 1D/next_n == 1 form,
376 + # which is equivalent once the ragged ends are expanded.
377 + rocm_aiter_ops.indexer_top_k_decode(
378 logits,
351 - next_n,
352 - seq_lens,
379 + 1,
380 + self._row_ends(seq_lens, next_n, logits.shape[0]),
381 topk_indices,
354 - logits.shape[0],
355 - logits.stride(0),
356 - logits.stride(1),
382 topk_tokens,
383 )
384 + else:
385 + assert backend == "per_row", f"unknown topk backend: {backend}"
386 + self._per_row(logits, seq_lens, next_n, topk_indices, topk_tokens)
387 +
388 + @staticmethod
389 + def _per_row(
390 + logits: torch.Tensor,
391 + seq_lens: torch.Tensor,
392 + next_n: int,
393 + topk_indices: torch.Tensor,
394 + topk_tokens: int,
395 + ) -> None:
396 + ops.top_k_per_row_decode(
397 + logits,
398 + next_n,
399 + seq_lens,
400 + topk_indices,
401 + logits.shape[0],
402 + logits.stride(0),
403 + logits.stride(1),
404 + topk_tokens,
405 + )
vllm/models/glm5next/amd/sparse_indexer.py
+76 -6
@@ -8,10 +8,11 @@ import vllm.envs as envs
8 from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels
9 from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11 -from vllm.config import get_current_vllm_config_or_none
11 +from vllm.config import CUDAGraphMode, get_current_vllm_config_or_none
12 from vllm.forward_context import get_forward_context
13 from vllm.logger import init_logger
14 from vllm.model_executor.custom_op import CustomOp
15 +from vllm.model_executor.layers.indexer_topk import get_indexer_topk
16 from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
17 from vllm.models.glm5next.common.sparse_indexer import (
18 RADIX_TOPK_WORKSPACE_SIZE,
@@ -24,6 +25,7 @@ from vllm.models.glm5next.common.sparse_indexer import (
25 kv_cache_as_quant_view,
26 )
27 from vllm.platforms import current_platform
28 +from vllm.utils.math_utils import cdiv
29 from vllm.utils.torch_utils import (
30 LayerNameType,
31 _encode_layer_name,
@@ -37,6 +39,9 @@ from vllm.v1.worker.workspace import current_workspace_manager
39
40 logger = init_logger(__name__)
41
42 +_GFX950_KPOOL_AITER_MAX_POOLS = 256 * 1024 # (1M ctx)
43 +_GFX950_KPOOL_AITER_MIN_POOLS = 16 * 1024
44 +
45 # kpool write helper: form pools from the current token batch and compress them
46 # into the index K cache via the fused Triton kernel.
47
@@ -87,6 +92,51 @@ def _kpool_compress_insert(
92 )
93
94
95 +def _kpool_decode_topk_backend(
96 + configured: str,
97 + *,
98 + num_rows: int,
99 + max_valid_seq_len: int,
100 + select_k: int,
101 + index_kpool: int,
102 + full_cudagraph: bool,
103 +) -> str:
104 + """Select the topk backend for decodes.
105 +
106 + Heuristic based on ctx lengths:
107 + - <16k pools: in-tree hip kernel
108 + - 16-256k pools: use aiter
109 + - >256k pools: in-tree because not measured on >1m ctx
110 +
111 + Since under FULL cudagraphs we cannot access the context length,
112 + and the aiter kernel does not always outperform the in-tree kernel,
113 + we default to in-tree kernel under FULL cudagraphs. Users can opt in
114 + if they know their context length is long enough to benefit from aiter.
115 + """
116 + if (
117 + configured != "auto"
118 + or full_cudagraph
119 + or not rocm_aiter_ops.is_indexer_top_k_enabled()
120 + ):
121 + return configured
122 + if rocm_aiter_ops.is_indexer_top_k_supported(
123 + is_prefill=False,
124 + compress_ratio=index_kpool,
125 + num_rows=num_rows,
126 + max_valid_seq_len=max_valid_seq_len,
127 + ):
128 + return "aiter"
129 + if (
130 + index_kpool > 1
131 + and select_k == 512
132 + and _GFX950_KPOOL_AITER_MIN_POOLS
133 + <= max_valid_seq_len
134 + <= _GFX950_KPOOL_AITER_MAX_POOLS
135 + ):
136 + return "aiter"
137 + return configured
138 +
139 +
140 @eager_break_during_capture
141 def sparse_attn_indexer_kpool(
142 hidden_states: torch.Tensor,
@@ -117,6 +167,7 @@ def sparse_attn_indexer_kpool(
167 # path and when the tail cache is disabled.
168 tail_kv_cache: torch.Tensor | None = None,
169 tail_prefix: str | None = None,
170 + topk_backend: str = "auto",
171 ) -> torch.Tensor:
172 # careful! this will be None in dummy run
173 attn_metadata = get_forward_context().attn_metadata
@@ -580,15 +631,29 @@ def sparse_attn_indexer_kpool(
631 else:
632 topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
633
583 - torch.ops._C.top_k_per_row_decode(
634 + full_cudagraph = (
635 + get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL
636 + )
637 + if full_cudagraph:
638 + topk_max_seq_len = max_pool_len * index_kpool
639 + else:
640 + topk_max_seq_len = attn_metadata_narrowed.max_seq_len
641 +
642 + resolved_backend = _kpool_decode_topk_backend(
643 + topk_backend,
644 + num_rows=num_rows,
645 + max_valid_seq_len=cdiv(topk_max_seq_len, index_kpool),
646 + select_k=select_k,
647 + index_kpool=index_kpool,
648 + full_cudagraph=full_cudagraph,
649 + )
650 + get_indexer_topk(resolved_backend)(
651 logits,
585 - next_n,
652 seq_lens,
653 + next_n,
654 topk_dst,
588 - num_rows,
589 - logits.stride(0),
590 - logits.stride(1),
655 select_k,
656 + topk_max_seq_len,
657 )
658
659 # Resolve to token-level indices in the output buffer.
@@ -664,6 +729,10 @@ class SparseAttnIndexerKpool(CustomOp):
729 self.topk_indices_buffer = topk_indices_buffer
730 self.skip_k_cache_insert = skip_k_cache_insert
731 self.use_fp4_cache = use_fp4_cache
732 + cfg = get_current_vllm_config_or_none()
733 + self.topk_backend = (
734 + cfg.kernel_config.sparse_indexer_topk_backend if cfg is not None else "auto"
735 + )
736
737 def forward_hip(
738 self,
@@ -726,6 +795,7 @@ class SparseAttnIndexerKpool(CustomOp):
795 positions,
796 self.tail_cache.kv_cache if self.tail_cache is not None else None,
797 self.tail_cache.prefix if self.tail_cache is not None else None,
798 + self.topk_backend,
799 )
800
801 def forward_native(
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+15 -92
@@ -32,18 +32,6 @@ logger = init_logger(__name__)
32 FP8_DTYPE = current_platform.fp8_dtype()
33
34
35 -@functools.cache
36 -def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None:
37 - try:
38 - from aiter.ops.topk import (
39 - top_k_per_row_decode,
40 - top_k_per_row_prefill,
41 - )
42 - except ImportError:
43 - return None
44 - return top_k_per_row_prefill, top_k_per_row_decode
45 -
46 -
35 @functools.cache
36 def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
37 from vllm._aiter_ops import rocm_aiter_ops
@@ -58,9 +46,6 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
46 return pa_sparse_prefill_opus
47
48
61 -_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
62 -_GFX950_C4A_NATIVE_MAX_ROWS = 256
63 -_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
49 # Conservative perf gate, not a correctness bound: OPUS is correct for any query
50 # count, but Triton stays faster below this measured crossover.
51 _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
@@ -71,43 +56,6 @@ def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:
56 return compress_ratio == 4
57
58
74 -def _get_aiter_top_k_kernel(
75 - *,
76 - is_prefill: bool,
77 - compress_ratio: int,
78 - num_rows: int,
79 - max_valid_seq_len: int | None = None,
80 - num_columns: int | None = None,
81 - topk_tokens: int = 1024,
82 - on_gfx950: bool = _ON_GFX950,
83 -) -> Callable[..., None] | None:
84 - if compress_ratio <= 1 or not on_gfx950:
85 - return None
86 -
87 - if not is_prefill:
88 - assert max_valid_seq_len is not None
89 - if (
90 - topk_tokens == 512
91 - and 0 < num_rows <= 384
92 - and num_columns is not None
93 - and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS
94 - ):
95 - return None
96 - # AITER v0.1.19 decode is one-block only. This measured gfx950
97 - # FP32/k=1024 compressed-row boundary is independent of the native
98 - # split-count boundary in sampler.cu.
99 - if (
100 - num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS
101 - and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN
102 - ):
103 - return None
104 -
105 - topk_ops = _get_aiter_topk_ops()
106 - if topk_ops is None:
107 - return None
108 - return topk_ops[0] if is_prefill else topk_ops[1]
109 -
110 -
59 @triton.jit
60 def _localize_aiter_prefill_topk_kernel(
61 indices_ptr,
@@ -145,44 +93,18 @@ def _localize_aiter_prefill_topk(
93
94
95 def _launch_aiter_top_k_per_row_prefill(
148 - top_k_per_row_prefill: Callable[..., None],
96 logits: torch.Tensor,
97 row_starts: torch.Tensor,
98 row_ends: torch.Tensor,
99 indices: torch.Tensor,
100 topk_tokens: int,
101 ) -> None:
155 - top_k_per_row_prefill(
156 - logits,
157 - row_starts,
158 - row_ends,
159 - indices,
160 - None,
161 - logits.shape[0],
162 - logits.stride(0),
163 - logits.stride(1),
164 - k=topk_tokens,
165 - )
166 - _localize_aiter_prefill_topk(indices, row_starts)
167 -
102 + from vllm._aiter_ops import rocm_aiter_ops
103
169 -def _launch_aiter_top_k_per_row_decode(
170 - top_k_per_row_decode: Callable[..., None],
171 - logits: torch.Tensor,
172 - seq_lens: torch.Tensor,
173 - indices: torch.Tensor,
174 - topk_tokens: int,
175 -) -> None:
176 - top_k_per_row_decode(
177 - logits,
178 - 1,
179 - seq_lens.reshape(-1),
180 - indices,
181 - logits.shape[0],
182 - logits.stride(0),
183 - logits.stride(1),
184 - k=topk_tokens,
104 + rocm_aiter_ops.indexer_top_k_prefill(
105 + logits, row_starts, row_ends, indices, topk_tokens
106 )
107 + _localize_aiter_prefill_topk(indices, row_starts)
108
109
110 @triton.jit
@@ -1195,6 +1117,8 @@ def rocm_aiter_sparse_attn_indexer(
1117 candidate_block_size: int = 0,
1118 candidate_write: bool = False,
1119 ) -> torch.Tensor:
1120 + from vllm._aiter_ops import rocm_aiter_ops
1121 +
1122 # careful! this will be None in dummy run
1123 forward_context = get_forward_context()
1124 attn_metadata = forward_context.attn_metadata
@@ -1347,14 +1271,12 @@ def rocm_aiter_sparse_attn_indexer(
1271
1272 num_rows = logits.shape[0]
1273
1350 - aiter_topk_kernel = _get_aiter_top_k_kernel(
1274 + if rocm_aiter_ops.is_indexer_top_k_supported(
1275 is_prefill=True,
1276 compress_ratio=compress_ratio,
1277 num_rows=num_rows,
1354 - )
1355 - if aiter_topk_kernel is not None:
1278 + ):
1279 _launch_aiter_top_k_per_row_prefill(
1357 - aiter_topk_kernel,
1280 logits,
1281 chunk.cu_seqlen_ks,
1282 chunk.cu_seqlen_ke,
@@ -1448,19 +1370,20 @@ def rocm_aiter_sparse_attn_indexer(
1370 max_compressed_seq_len = max_model_len
1371 else:
1372 max_compressed_seq_len = layer_attn_metadata.max_seq_len // compress_ratio
1451 - aiter_topk_kernel = _get_aiter_top_k_kernel(
1373 + if rocm_aiter_ops.is_indexer_top_k_supported(
1374 is_prefill=False,
1375 compress_ratio=compress_ratio,
1376 num_rows=num_rows,
1377 max_valid_seq_len=max_compressed_seq_len,
1378 + ) and not rocm_aiter_ops.dsv4_indexer_prefers_native_top_k(
1379 + num_rows=num_rows,
1380 num_columns=logits.shape[1],
1381 topk_tokens=topk_tokens,
1458 - )
1459 - if aiter_topk_kernel is not None:
1460 - _launch_aiter_top_k_per_row_decode(
1461 - aiter_topk_kernel,
1382 + ):
1383 + rocm_aiter_ops.indexer_top_k_decode(
1384 logits,
1463 - decode_metadata.seq_lens,
1385 + 1,
1386 + decode_metadata.seq_lens.reshape(-1),
1387 topk_indices,
1388 topk_tokens,
1389 )