iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Simon Danielsson
2026-10-02 11:47:41 +01:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 751fa96a49
commit 98b29aa99a
8 changed files with 705 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_top_k_per_row_decode.py
+158
View File
@@ -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()
tests/kernels/test_top_k_per_row.py
+225 -118
View File
@@ -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]
tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
+72 -14
View File
@@ -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
)
vllm/_aiter_ops.py
+99
View File
@@ -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
vllm/config/kernel.py
+3
View File
@@ -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.
"""
vllm/model_executor/layers/indexer_topk.py
+57 -10
View File
@@ -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,
)
vllm/models/glm5next/amd/sparse_indexer.py
+76 -6
View File
@@ -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(
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+15 -92
View File
@@ -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,
)