From edfccc946925a8439e27abd76da8d1a640d1a1b7 Mon Sep 17 00:00:00 2001 From: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:53:47 +0800 Subject: [PATCH] [ROCM][Perf] Add ROCM_SEGMENTED_ATTN attention-backend works across RDNA3/3.5/4 (#59132) Signed-off-by: root Signed-off-by: big-yellow-duck Co-authored-by: Tan Pin Siang <1716735+tanpinsiang@users.noreply.github.com> Co-authored-by: vllmellm <190700713+vllmellm@users.noreply.github.com> Co-authored-by: Codex --- .../attention/test_segmented_attention.py | 1408 +++++++++++++++++ .../test_rocm_attention_backends_selection.py | 210 +++ vllm/config/kernel.py | 4 + vllm/model_executor/models/qwen3_dflash.py | 19 + vllm/model_executor/warmup/kernel_warmup.py | 7 + .../rocm_segmented_attn_autotune_warmup.py | 1335 ++++++++++++++++ vllm/v1/attention/backends/registry.py | 3 + .../attention/backends/rocm_segmented_attn.py | 371 +++++ vllm/v1/attention/ops/segmented_attention.py | 1074 +++++++++++++ 9 files changed, 4431 insertions(+) create mode 100644 tests/kernels/attention/test_segmented_attention.py create mode 100644 vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py create mode 100644 vllm/v1/attention/backends/rocm_segmented_attn.py create mode 100644 vllm/v1/attention/ops/segmented_attention.py diff --git a/tests/kernels/attention/test_segmented_attention.py b/tests/kernels/attention/test_segmented_attention.py new file mode 100644 index 0000000000..0fbf5fe894 --- /dev/null +++ b/tests/kernels/attention/test_segmented_attention.py @@ -0,0 +1,1408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Segmented attention accuracy, cache semantics, and startup tuning contracts.""" + +import math + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import triton + +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + +pytestmark = [ + pytest.mark.skip_global_cleanup, + pytest.mark.skipif( + not (current_platform.is_rocm() and on_gfx1x()), + reason="Segmented attention requires RDNA GPUs", + ), +] + + +def _make_paged_attention_case( + query_lens: list[int], + context_lens: list[int], + *, + num_heads: int, + num_kv_heads: int, + head_size: int, + block_size: int, + fp8: bool, + causal: bool = True, + sliding_window: int = 0, + softcap: float = 0.0, + sinks: torch.Tensor | None = None, +): + device = torch.device("cuda:0") + generator = torch.Generator(device=device).manual_seed(419) + total_queries = sum(query_lens) + qkv = ( + torch.randn( + total_queries, + num_heads + 2 * num_kv_heads, + head_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + * 0.25 + ) + query, key, value = qkv.split((num_heads, num_kv_heads, num_kv_heads), dim=1) + cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16 + k_scale = torch.tensor(0.125 if fp8 else 1.0, device=device) + v_scale = torch.tensor(0.25 if fp8 else 1.0, device=device) + max_seq_len = max(q + c for q, c in zip(query_lens, context_lens)) + blocks_per_seq = triton.cdiv(max_seq_len, block_size) + num_blocks = len(query_lens) * blocks_per_seq + key_cache = ( + torch.randn( + num_blocks, + block_size, + num_kv_heads, + head_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 0.25 + / k_scale + ).to(cache_dtype) + value_cache = ( + torch.randn( + num_blocks, + block_size, + num_kv_heads, + head_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 0.25 + / v_scale + ).to(cache_dtype) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32).view( + len(query_lens), blocks_per_seq + ) + starts = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + seq_lens = torch.tensor( + [q + c for q, c in zip(query_lens, context_lens)], + device=device, + dtype=torch.int32, + ) + + first = 0 + for seq, (query_len, context_len) in enumerate(zip(query_lens, context_lens)): + for query_pos in range(query_len): + cache_pos = context_len + query_pos + block = block_table[seq, cache_pos // block_size] + offset = cache_pos % block_size + key_cache[block, offset] = (key[first + query_pos].float() / k_scale).to( + cache_dtype + ) + value_cache[block, offset] = ( + value[first + query_pos].float() / v_scale + ).to(cache_dtype) + first += query_len + + references = [] + first = 0 + for seq, (query_len, context_len) in enumerate(zip(query_lens, context_lens)): + physical = block_table[seq].long() + seq_len = context_len + query_len + full_key = ( + key_cache[physical].flatten(0, 1)[:seq_len].float() * k_scale + ).repeat_interleave(num_heads // num_kv_heads, dim=1) + full_value = ( + value_cache[physical].flatten(0, 1)[:seq_len].float() * v_scale + ).repeat_interleave(num_heads // num_kv_heads, dim=1) + scores = torch.einsum( + "qhd,khd->hqk", query[first : first + query_len].float(), full_key + ) / math.sqrt(head_size) + if softcap > 0: + scores = softcap * torch.tanh(scores / softcap) + query_positions = context_len + torch.arange(query_len, device=device) + key_positions = torch.arange(seq_len, device=device) + causal_mask = key_positions[None, :] > query_positions[:, None] + if causal: + scores.masked_fill_(causal_mask[None], -float("inf")) + if sliding_window > 0: + window_mask = key_positions[None, :] < ( + query_positions[:, None] - sliding_window + 1 + ) + scores.masked_fill_(window_mask[None], -float("inf")) + if sinks is not None: + sink_scores = sinks.float()[:, None, None].expand(-1, query_len, 1) + probabilities = torch.cat((scores, sink_scores), dim=-1).softmax(-1) + probabilities = probabilities[..., :-1] + else: + probabilities = scores.softmax(-1) + references.append(torch.einsum("hqk,khd->qhd", probabilities, full_value)) + first += query_len + + return { + "query": query, + "key": key, + "value": value, + "key_cache": key_cache, + "value_cache": value_cache, + "block_table": block_table, + "starts": starts, + "seq_lens": seq_lens, + "k_scale": k_scale, + "v_scale": v_scale, + "reference": torch.cat(references), + "max_seq_len": max_seq_len, + } + + +def _run_segmented_case(case, output, max_query_len, *, force_splits=None, **kwargs): + from vllm.v1.attention.ops.segmented_attention import ( + run_segmented_attention, + select_segmented_config, + ) + + config = None + heads, dim = case["query"].shape[1:] + if force_splits is not None: + config = dict( + select_segmented_config( + case["seq_lens"].numel(), + max_query_len, + case["max_seq_len"], + heads, + case["key_cache"].shape[2], + dim, + case["key_cache"].element_size() == 1, + ), + splits=force_splits, + ) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + max_query_len, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + dim**-0.5, + config=config, + **kwargs, + ) + + +def test_segmented_tuning_candidates_preserve_workspace_bound(): + """D128/D256 candidates retain their existing split scratch bound.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + from vllm.v1.attention.ops import segmented_attention as segmented + + for batch, query_len, seq_len, heads, kv_heads, dim, fp8 in ( + (1, 1, 128, 8, 8, 128, True), + (32, 1, 262144, 16, 1, 128, True), + (4, 1024, 8192, 16, 4, 256, False), + (1, 1024, 131072, 6, 1, 256, True), + ): + default = segmented.select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + candidates = tuning._candidate_configs(default, batch, query_len, heads, dim) + assert candidates + assert ( + tuning._normalized_config(default, batch, query_len, heads, dim) + in candidates + ) + assert all(config["splits"] <= default["splits"] for config in candidates) + assert all(dim % config["bk"] == 0 for config in candidates) + + +@pytest.mark.parametrize("dim", (128, 256)) +def test_segmented_fp8_long_extend_split_workspace_bound(dim): + from vllm.v1.attention.ops.segmented_attention import ( + MAX_LONG_EXTEND_WORKSPACE_BYTES, + segmented_query_capacity, + select_segmented_config, + ) + + for batch, query_len, seq_len, heads in ( + (1, 256, 8192, 6), + (1, 256, 32768, 6), + (1, 256, 131072, 6), + (1, 1024, 32768, 1), + (1, 1024, 131072, 1), + (1, 1024, 131072, 6), + (1, 4096, 131072, 1), + (4, 1024, 131072, 6), + (5, 257, 131072, 6), + (21, 256, 131072, 6), + ): + config = select_segmented_config(batch, query_len, seq_len, heads, 1, dim, True) + if seq_len >= 32768: + scratch = ( + batch + * segmented_query_capacity(query_len) + * heads + * config["splits"] + * (dim + 1) + * 4 + ) + assert scratch <= MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + + for batch in (1, 4, 16): + for query_len in (257, 1025, 4095, 8192): + for ratio in (1, 6, 16): + for kv_heads in (1, 4): + for seq_len in (32768, 131072): + heads = ratio * kv_heads + config = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, True + ) + if config["splits"] == 1: + continue + scratch = ( + batch + * segmented_query_capacity(query_len) + * heads + * config["splits"] + * (dim + 1) + * 4 + ) + assert scratch <= MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + + +@torch.inference_mode() +def test_segmented_bf16_d256_long_extend_matches_dense_reference(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x segmented prefill") + + case = _make_paged_attention_case( + [1024], + [7168], + num_heads=6, + num_kv_heads=1, + head_size=256, + block_size=1568, + fp8=False, + ) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + 1024, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + 256**-0.5, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize( + "query_len,heads,qk_amplitude,dim", + [ + (256, 1, 1, 256), + (256, 6, 1, 256), + (512, 6, 1, 256), + (256, 6, 4, 256), + (256, 6, 16, 256), + (256, 6, 1, 128), + (512, 6, 4, 128), + (256, 6, 16, 128), + (8192, 1, 1, 128), + (8192, 1, 1, 256), + ], +) +@torch.inference_mode() +def test_segmented_fp8_long_extend_matches_dense_reference( + query_len, heads, qk_amplitude, dim +): + from vllm.platforms.rocm import on_gfx12x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx12x(): + pytest.skip("FP8 KV requires gfx12") + case = _make_paged_attention_case( + [query_len], + [8192 - query_len], + num_heads=heads, + num_kv_heads=1, + head_size=dim, + block_size=1568, + fp8=True, + ) + if qk_amplitude != 1: + # Check peaked logits too: one-term FP8 Q quantization passed only the + # nearly uniform fixture and deviated by 4-70% at these amplitudes. + case["query"].mul_(qk_amplitude) + case["k_scale"].mul_(qk_amplitude) + full_key = ( + case["key_cache"][case["block_table"][0].long()] + .flatten(0, 1)[:8192] + .float() + * case["k_scale"] + ).repeat_interleave(heads, dim=1) + full_value = ( + case["value_cache"][case["block_table"][0].long()] + .flatten(0, 1)[:8192] + .float() + * case["v_scale"] + ).repeat_interleave(heads, dim=1) + scores = torch.einsum( + "qhd,khd->hqk", case["query"].float(), full_key + ) / math.sqrt(dim) + future = torch.arange(8192, device=case["query"].device)[None, :] > ( + 8192 + - query_len + + torch.arange(query_len, device=case["query"].device)[:, None] + ) + scores.masked_fill_(future[None], -float("inf")) + case["reference"] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), full_value) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + query_len, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + dim**-0.5, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +def test_segmented_tuning_protects_static_incumbent(): + """Noise-sized gains must not replace the static configuration.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + incumbent = {"bm": 16} + challenger = {"bm": 32} + samples = { + tuning._config_key(incumbent): [100.0, 101.0, 99.0, 100.0, 100.0], + tuning._config_key(challenger): [99.0, 100.0, 98.0, 99.0, 99.0], + } + winner, comparisons = tuning._select_tuned_winner( + incumbent, [incumbent, challenger], samples + ) + assert winner["config"] == incumbent + assert comparisons[1]["paired_speedup_vs_default"] < 1.02 + + +def test_segmented_tuning_promotes_verified_challenger(): + """A finalist with a stable material gain should replace the incumbent.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + incumbent = {"bm": 16} + challenger = {"bm": 32} + samples = { + tuning._config_key(incumbent): [100.0, 102.0, 98.0, 101.0, 99.0], + tuning._config_key(challenger): [94.0, 96.0, 92.0, 95.0, 93.0], + } + winner, _ = tuning._select_tuned_winner(incumbent, [incumbent, challenger], samples) + assert winner["config"] == challenger + assert winner["paired_speedup_vs_default"] > 1.02 + + +def test_segmented_tuning_balances_tp_workloads_without_overlap(): + """Every missing bucket belongs to exactly one reasonably balanced rank.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + workloads = list(tuning._workloads(8192, 262144, 32)) + shards = tuning._shard_workloads( + workloads, + 4, + 8192, + 12, + 2, + 256, + True, + ) + flattened = [workload for shard in shards for workload in shard] + assert sorted(flattened) == sorted(workloads) + assert len(flattened) == len(set(flattened)) + owners = { + workload[:2]: rank for rank, shard in enumerate(shards) for workload in shard + } + assert all( + owners[workload[:2]] == rank + for rank, shard in enumerate(shards) + for workload in shard + ) + + weights = [ + sum( + tuning._workload_weight(workload, 8192, 12, 2, 256, True) + for workload in shard + ) + for shard in shards + ] + largest_group = max( + sum( + tuning._workload_weight(workload, 8192, 12, 2, 256, True) + for workload in workloads + if workload[:2] == group + ) + for group in {workload[:2] for workload in workloads} + ) + assert max(weights) - min(weights) <= largest_group + + +def test_segmented_tuning_prunes_scheduler_and_kv_limits(): + """Generated buckets must be reachable under scheduler and cache limits.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + limits = (8192, 262144, 32) + raw = set( + tuning._workloads( + *limits, + dtype=torch.bfloat16, + kv_dtype=torch.float8_e4m3fn, + heads=16, + kv_heads=1, + dim=128, + page=16, + ) + ) + assert {batch for batch, _, _ in raw} == {1, 2, 4, 8, 16, 32} + assert {query for _, query, _ in raw} == set(tuning._QUERY_BUCKETS) + assert (32, 1, 262144) in raw + assert (1, 8192, 262144) in raw + assert (1, 8192, 131072) in raw + assert (1, 8, 131072) in raw + assert all( + sum(tuning._query_lengths(batch, query, limits[0])) <= limits[0] + and batch <= limits[2] + and query <= min(limits[0], tuning.MAX_QUERY_LEN) + and query <= seq_len <= limits[1] + for batch, query, seq_len in raw + ) + + layouts = ((16, 16 * 128 * 2),) + bounded = set( + tuning._workloads( + *limits, + memory_budget_bytes=2**50, + dtype=torch.bfloat16, + kv_dtype=torch.float8_e4m3fn, + heads=16, + kv_heads=1, + dim=128, + page=16, + cache_layouts=layouts, + cache_budget_bytes=2**20, + ) + ) + assert bounded < raw + assert all( + batch * sum(((seq_len + block - 1) // block) * size for block, size in layouts) + <= 2**20 + for batch, _, seq_len in bounded + ) + + +@pytest.fixture +def segmented_tuner(tmp_path, monkeypatch): + """Exercise real cache files and lookup while replacing GPU timing only.""" + from types import SimpleNamespace + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + from vllm.v1.attention.ops import segmented_attention as segmented + + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _: SimpleNamespace( + name="test", gcnArchName="gfx1201", multi_processor_count=32 + ), + ) + monkeypatch.setattr(tuning, "_TABLES", {}) + calls = [] + + def tune(*args, **kwargs): + heads, kv_heads, dim = args[3:6] + batch, query_len, seq_len = args[-1] + default = segmented.select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, False + ) + best = tuning._candidate_configs(default, batch, query_len, heads, dim)[0] + calls.append(args[-1]) + return { + "workload": list(args[-1]), + "query_lengths": tuning._query_lengths(batch, query_len, args[-2]), + "default": default, + "best": best, + "results": [], + } + + monkeypatch.setattr(tuning, "_tune_workload", tune) + geometry = (torch.device("cuda:0"), torch.bfloat16, 4, 2, 128, 16, 128**-0.5) + return SimpleNamespace( + module=tuning, + warmup_args=(*geometry, 4, 8, 2), + lookup_args=(*geometry[:2], torch.bfloat16, *geometry[2:], 1, 1, 8), + calls=calls, + cache_dir=tmp_path, + ) + + +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("device", ["cuda", "cuda:0"]) +def test_segmented_tuning_releases_only_new_kernel_modules(monkeypatch, fail, device): + """Tuning releases candidate modules even on failure without invalidating graphs.""" + import weakref + from types import SimpleNamespace + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + class Module: + pass + + existing = Module() + other_device = Module() + caches = [{"existing": existing} for _ in range(2)] + for name, cache in zip( + ("_segmented_attention_stage", "_segmented_attention_reduce"), caches + ): + monkeypatch.setattr( + tuning, + name, + SimpleNamespace(device_caches={0: (cache,), 1: ({"other": other_device},)}), + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda _: None) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + references = [] + try: + with tuning._temporary_tuning_kernels(torch.device(device)): + for cache in caches: + cache["candidate"] = Module() + references.append(weakref.ref(cache["candidate"])) + if fail: + raise RuntimeError("tuning failed") + except RuntimeError: + assert fail + assert all(reference() is None for reference in references) + assert all(cache == {"existing": existing} for cache in caches) + for kernel in ( + tuning._segmented_attention_stage, + tuning._segmented_attention_reduce, + ): + assert kernel.device_caches[1][0] == {"other": other_device} + + +def test_segmented_tuning_cleanup_preserves_graphs_and_reloads_kernels(): + """Existing graphs survive cleanup, and released variants still run correctly.""" + import weakref + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + case = _make_paged_attention_case( + [3], + [11], + num_heads=10, + num_kv_heads=2, + head_size=128, + fp8=False, + block_size=16, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, 3, force_splits=1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_segmented_case(case, output, 3, force_splits=1) + caches = [ + kernel.device_caches[0][0] + for kernel in ( + tuning._segmented_attention_stage, + tuning._segmented_attention_reduce, + ) + ] + keys = [set(cache) for cache in caches] + with tuning._temporary_tuning_kernels(torch.device("cuda:0")): + _run_segmented_case(case, output, 3, force_splits=2) + references = [ + weakref.ref(cache[key]) + for cache, existing in zip(caches, keys) + for key in cache.keys() - existing + ] + assert references + assert all(reference() is None for reference in references) + output.fill_(float("nan")) + graph.replay() + torch.testing.assert_close(output.float(), case["reference"], atol=0.01, rtol=0.01) + output.fill_(float("nan")) + _run_segmented_case(case, output, 3, force_splits=2) + torch.testing.assert_close(output.float(), case["reference"], atol=0.01, rtol=0.01) + + +@pytest.mark.parametrize( + "mode", + [ + pytest.param({}, id="causal"), + pytest.param({"causal": False}, id="noncausal"), + pytest.param({"sliding_window": 7}, id="sliding-window"), + pytest.param({"sliding_window": 7, "causal": False}, id="draft-window"), + pytest.param({"has_sinks": True}, id="sinks"), + ], +) +def test_segmented_tuning_persists_without_retuning(segmented_tuner, monkeypatch, mode): + """A fresh in-memory table reloads the same winner without timing or writing.""" + state = segmented_tuner + tuning = state.module + tuning.warmup_segmented_attention( + *state.warmup_args, memory_budget_bytes=2**50, **mode + ) + assert state.calls + expected = tuning.get_segmented_config(*state.lookup_args, **mode) + assert expected is not None + (cache,) = state.cache_dir.rglob("*.json") + saved = cache.read_bytes(), cache.stat().st_mtime_ns + tuning._TABLES.clear() + + def forbidden(*args, **kwargs): + pytest.fail("Persistent segmented startup cache attempted to retune") + + monkeypatch.setattr(tuning, "_tune_workload", forbidden) + tuning.warmup_segmented_attention( + *state.warmup_args, memory_budget_bytes=2**50, **mode + ) + assert tuning.get_segmented_config(*state.lookup_args, **mode) == expected + assert (cache.read_bytes(), cache.stat().st_mtime_ns) == saved + + +@pytest.mark.parametrize( + "max_tokens,mode", + [ + pytest.param(3, {}, id="scheduler-token-budget"), + pytest.param(4, {"causal": False}, id="noncausal"), + pytest.param(4, {"sliding_window": 7}, id="sliding-window"), + pytest.param(4, {"sliding_window": 7, "causal": False}, id="draft-window"), + pytest.param(4, {"has_sinks": True}, id="sinks"), + ], +) +def test_segmented_tuning_does_not_reuse_incompatible_winners( + segmented_tuner, max_tokens, mode +): + """Changed attention semantics or query mix require their own measurements.""" + state = segmented_tuner + tuning = state.module + tuning.warmup_segmented_attention(*state.warmup_args, memory_budget_bytes=2**50) + previous_calls = len(state.calls) + if mode: + assert tuning.get_segmented_config(*state.lookup_args, **mode) is None + args = (*state.warmup_args[:7], max_tokens, *state.warmup_args[8:]) + tuning.warmup_segmented_attention(*args, memory_budget_bytes=2**50, **mode) + assert len(state.calls) > previous_calls + assert len(list(state.cache_dir.rglob("*.json"))) == 2 + assert tuning.get_segmented_config(*state.lookup_args, **mode) is not None + + +@pytest.mark.parametrize( + "dim,page,hq,hk,dtype,kv_dtype,byte_cache,force_splits", + [ + (128, 16, 4, 4, torch.bfloat16, torch.bfloat16, False, None), + (128, 32, 8, 2, torch.float16, torch.float16, False, None), + (256, 784, 12, 2, torch.bfloat16, torch.bfloat16, False, None), + (256, 1568, 12, 2, torch.bfloat16, torch.float8_e4m3fn, False, None), + (128, 32, 16, 1, torch.bfloat16, torch.float8_e4m3fn, True, None), + (256, 32, 4, 2, torch.float16, torch.float8_e4m3fn, False, None), + (128, 32, 8, 2, torch.bfloat16, torch.bfloat16, False, 32), + ], +) +@torch.inference_mode() +def test_segmented_attention_ragged_unified_graph_replay( + monkeypatch, dim, page, hq, hk, dtype, kv_dtype, byte_cache, force_splits +): + """Unified-cache segmented prefill follows metadata on graph replay.""" + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not ( + on_gfx12x() if kv_dtype.itemsize == 1 else on_gfx1x() + ): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + if force_splits is not None: + original = dispatcher.select_segmented_config + monkeypatch.setattr( + dispatcher, + "select_segmented_config", + lambda *args: dict(original(*args), splits=force_splits), + ) + query_lengths = [0, 1, 2, 7, 33, 129] + contexts = [0, 67, 0, page - 1, page + 1, 33] + device = torch.device("cuda:0") + total = sum(query_lengths) + q = torch.randn(total, hq + 1, dim, device=device, dtype=dtype)[:, :hq] + k = torch.randn(total, hk + 1, dim, device=device, dtype=dtype)[:, :hk] + v = torch.randn_like(k) + blocks_per_seq = triton.cdiv(max(map(sum, zip(query_lengths, contexts))), page) + num_blocks = len(contexts) * blocks_per_seq + page_elements = page * hk * dim + backing = torch.empty(num_blocks * 2 * page_elements, device=device, dtype=kv_dtype) + cache_shape = (num_blocks, page, hk, dim) + strides = (2 * page_elements, hk * dim, dim, 1) + kc = torch.as_strided(backing, cache_shape, strides) + vc = torch.as_strided(backing, cache_shape, strides, page_elements) + ks = torch.tensor(0.13 if kv_dtype.itemsize == 1 else 1.0, device=device) + vs = torch.tensor(0.27 if kv_dtype.itemsize == 1 else 1.0, device=device) + kc.copy_((torch.randn(cache_shape, device=device).float() * 0.25 / ks).to(kv_dtype)) + vc.copy_((torch.randn(cache_shape, device=device).float() * 0.25 / vs).to(kv_dtype)) + table = torch.randperm(num_blocks, device=device, dtype=torch.int32).view( + len(contexts), blocks_per_seq + ) + host_table = table.cpu().tolist() + starts = torch.empty(len(contexts) + 1, device=device, dtype=torch.int32) + lengths = torch.empty(len(contexts), device=device, dtype=torch.int32) + output = torch.full((total, hq + 1, dim), 17.0, device=device, dtype=dtype)[:, :hq] + + def metadata(qlens): + starts.copy_( + torch.tensor( + [0, *torch.tensor(qlens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + ) + lengths.copy_( + torch.tensor( + [c + n for c, n in zip(contexts, qlens)], + device=device, + dtype=torch.int32, + ) + ) + first = 0 + for seq, (nq, context) in enumerate(zip(qlens, contexts)): + for local in range(nq): + position = context + local + physical = host_table[seq][position // page] + offset = position % page + kc[physical, offset].copy_((k[first + local].float() / ks).to(kv_dtype)) + vc[physical, offset].copy_((v[first + local].float() / vs).to(kv_dtype)) + first += nq + + def run(): + dispatcher.segmented_attention( + query=q, + key=k, + value=v, + output=output, + kv_cache_dtype="fp8" if kv_dtype.itemsize == 1 else "auto", + key_cache=kc.view(torch.uint8) if byte_cache else kc, + value_cache=vc.view(torch.uint8) if byte_cache else vc, + block_table=table, + query_start_loc=starts, + seq_lens=lengths, + max_seq_len=max(contexts) + 129, + max_query_len=129, + k_scale=ks, + v_scale=vs, + sm_scale=dim**-0.5, + ) + + def check(qlens): + first = 0 + for seq, (nq, context) in enumerate(zip(qlens, contexts)): + if nq == 0: + first += nq + continue + physical = table[seq].long() + full_k = ( + kc[physical].reshape(-1, hk, dim)[: context + nq].float() * ks + ).repeat_interleave(hq // hk, 1) + full_v = ( + vc[physical].reshape(-1, hk, dim)[: context + nq].float() * vs + ).repeat_interleave(hq // hk, 1) + logits = torch.einsum( + "qhd,khd->hqk", q[first : first + nq].float(), full_k + ) / math.sqrt(dim) + mask = ( + torch.arange(context + nq, device=device)[None, :] + > context + torch.arange(nq, device=device)[:, None] + ) + logits.masked_fill_(mask[None], -float("inf")) + ref = torch.einsum("hqk,khd->qhd", logits.softmax(-1), full_v) + actual = output[first : first + nq].float() + assert torch.isfinite(actual).all() + relative = (actual - ref).norm(dim=-1) / ref.norm(dim=-1).clamp_min(1e-6) + assert relative.max().item() < 0.01 + first += nq + + metadata(query_lengths) + run() + torch.accelerator.synchronize() + check(query_lengths) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + replay_lengths = [1, 0, 7, 2, 34, 128] + q.mul_(0.75) + k.add_(0.25) + v.add_(0.5) + metadata(replay_lengths) + output.fill_(17.0) + graph.replay() + torch.accelerator.synchronize() + check(replay_lengths) + + +@pytest.mark.parametrize( + "fp8,force_splits", + [(False, None), (True, None), (False, 4), (True, 4)], +) +@torch.inference_mode() +def test_segmented_ragged_prefill_matches_dense_reference(fp8, force_splits): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + case = _make_paged_attention_case( + [1, 2, 7, 129], + [31, 64, 97, 128], + num_heads=12, + num_kv_heads=2, + head_size=256, + block_size=32, + fp8=fp8, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, max_query_len=129, force_splits=force_splits) + + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@torch.inference_mode() +def test_rocm_segmented_attn_layout_cache_update(fp8): + from types import SimpleNamespace + + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + cache_dtype = torch.uint8 if fp8 else torch.bfloat16 + kv_cache_dtype = "fp8" if fp8 else "auto" + impl = RocmSegmentedAttentionImpl( + 12, + 256, + 256**-0.5, + 2, + None, + None, + kv_cache_dtype, + attn_type=AttentionType.DECODER, + ) + cache = torch.zeros(3, 2, 32, 512, device="cuda", dtype=cache_dtype) + key = torch.randn(5, 2, 256, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + slots = torch.tensor([0, 7, 32, 65, 95], device="cuda", dtype=torch.int64) + k_scale = torch.tensor(0.125 if fp8 else 1.0, device="cuda") + v_scale = torch.tensor(0.25 if fp8 else 1.0, device="cuda") + layer = SimpleNamespace(_k_scale=k_scale, _v_scale=v_scale) + + impl.do_kv_cache_update(layer, key, value, cache, slots) + key_cache, value_cache = impl._split_kv_cache(cache) + if fp8: + key_cache = key_cache.view(impl.fp8_dtype) + value_cache = value_cache.view(impl.fp8_dtype) + cached_key = torch.stack( + [key_cache[int(slot) // 32, int(slot) % 32] for slot in slots] + ) + cached_value = torch.stack( + [value_cache[int(slot) // 32, int(slot) % 32] for slot in slots] + ) + expected_key = (key.float() / k_scale).to(key_cache.dtype) + expected_value = (value.float() / v_scale).to(value_cache.dtype) + torch.testing.assert_close(cached_key, expected_key, rtol=0, atol=0) + torch.testing.assert_close(cached_value, expected_value, rtol=0, atol=0) + + +@torch.inference_mode() +def test_segmented_mixed_causality_matches_dense_reference(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x unified attention per-request causal fallback") + + query_lens = [2, 3] + context_lens = [5, 6] + args = dict(num_heads=8, num_kv_heads=2, head_size=128, block_size=32) + case = _make_paged_attention_case( + query_lens, context_lens, fp8=False, causal=True, **args + ) + noncausal_case = _make_paged_attention_case( + query_lens, context_lens, fp8=False, causal=False, **args + ) + reference = torch.cat( + ( + case["reference"][: query_lens[0]], + noncausal_case["reference"][query_lens[0] :], + ) + ) + causal = torch.tensor([True, False], device=case["query"].device) + output = torch.empty_like(case["query"]) + dispatcher.segmented_attention( + query=case["query"], + key=case["key"], + value=case["value"], + output=output, + kv_cache_dtype="auto", + key_cache=case["key_cache"], + value_cache=case["value_cache"], + block_table=case["block_table"], + query_start_loc=case["starts"], + seq_lens=case["seq_lens"], + max_seq_len=case["max_seq_len"], + max_query_len=max(query_lens), + k_scale=case["k_scale"], + v_scale=case["v_scale"], + sm_scale=128**-0.5, + causal=causal, + ) + + relative = (output.float() - reference).norm(dim=-1) / reference.norm( + dim=-1 + ).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("window_size", [1, 2048]) +@torch.inference_mode() +def test_segmented_sliding_window_matches_dense_reference(fp8, causal, window_size): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + case = _make_paged_attention_case( + [8], + [4096], + num_heads=32, + num_kv_heads=8, + head_size=128, + block_size=32, + fp8=fp8, + causal=causal, + sliding_window=window_size, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case( + case, output, max_query_len=8, sliding_window=window_size - 1, causal=causal + ) + + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("window_size", [0, 128]) +@pytest.mark.parametrize("head_size", [64, 128]) +@torch.inference_mode() +def test_segmented_sinks_match_dense_reference(fp8, window_size, head_size): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + sinks = torch.linspace( + -1.0, + 8.0, + 16, + device="cuda:0", + dtype=torch.bfloat16 if head_size == 64 else torch.float32, + ) + case = _make_paged_attention_case( + [1, 7], + [127, 130], + num_heads=16, + num_kv_heads=2, + head_size=head_size, + block_size=32, + fp8=fp8, + sliding_window=window_size, + sinks=sinks, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case( + case, + output, + max_query_len=7, + sliding_window=window_size - 1 if window_size else -1, + sinks=sinks, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("window_size", [0, 128]) +@pytest.mark.parametrize("query_len", [1, 8]) +@torch.inference_mode() +def test_segmented_sink_split_reduce_matches_reference(fp8, window_size, query_len): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.ops.segmented_attention import ( + run_segmented_attention, + select_segmented_config, + ) + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + sinks = torch.linspace(-1.0, 8.0, 16, device="cuda:0") + case = _make_paged_attention_case( + [query_len], + [8192 - query_len], + num_heads=16, + num_kv_heads=2, + head_size=64, + block_size=64, + fp8=fp8, + sliding_window=window_size, + sinks=sinks, + ) + attention_span = min(8192, window_size + query_len) if window_size else 8192 + config = dict( + select_segmented_config(1, query_len, attention_span, 16, 2, 64, fp8), + splits=4, + ) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + query_len, + 8192, + case["k_scale"], + case["v_scale"], + 64**-0.5, + sliding_window=window_size - 1 if window_size else -1, + sinks=sinks, + config=config, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("query_len,context_len", [(1, 8191), (8, 32760), (1, 131071)]) +@torch.inference_mode() +def test_segmented_long_context_fp8_sinks_match_dense_reference(query_len, context_len): + from vllm.platforms.rocm import on_gfx12x + + if not current_platform.is_rocm() or not on_gfx12x(): + pytest.skip("FP8 KV requires gfx12") + sinks = torch.linspace(-1.0, 1.0, 16, device="cuda:0", dtype=torch.bfloat16) + case = _make_paged_attention_case( + [query_len], + [context_len], + num_heads=16, + num_kv_heads=2, + head_size=64, + block_size=64, + fp8=True, + sinks=sinks, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, max_query_len=query_len, sinks=sinks) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize( + "feature", + [ + "softcap", + "sliding_softcap", + "output_scale", + ], +) +@torch.inference_mode() +def test_segmented_softcap_and_output_scaling_match_dense_reference(feature): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x unified attention feature fallback") + + num_heads = 8 + configured_sliding_window = 16 if feature.startswith("sliding") else 0 + sliding_window = configured_sliding_window - 1 if configured_sliding_window else -1 + softcap = 5.0 if "softcap" in feature else 0.0 + case = _make_paged_attention_case( + [7, 2], + [40, 35], + num_heads=num_heads, + num_kv_heads=2, + head_size=128, + block_size=32, + fp8=False, + sliding_window=configured_sliding_window, + softcap=softcap, + sinks=None, + ) + output_scale = None + if feature == "output_scale": + output_scale = torch.tensor(0.5, device="cuda:0") + output = torch.empty_like(case["query"], dtype=current_platform.fp8_dtype()) + else: + output = torch.empty_like(case["query"]) + + dispatcher.segmented_attention( + query=case["query"], + key=case["key"], + value=case["value"], + output=output, + kv_cache_dtype="auto", + key_cache=case["key_cache"], + value_cache=case["value_cache"], + block_table=case["block_table"], + query_start_loc=case["starts"], + seq_lens=case["seq_lens"], + max_seq_len=case["max_seq_len"], + max_query_len=7, + k_scale=case["k_scale"], + v_scale=case["v_scale"], + sm_scale=128**-0.5, + sliding_window=sliding_window, + softcap=softcap, + output_scale=output_scale, + sinks=None, + ) + + actual = output.float() + if output_scale is not None: + actual *= output_scale + torch.testing.assert_close(actual, case["reference"], atol=0.2, rtol=0.2) + else: + relative = (actual - case["reference"]).norm(dim=-1) / case["reference"].norm( + dim=-1 + ).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@torch.inference_mode() +def test_segmented_attention_cache_offsets_cross_int32_boundary(): + """A small logical prefix can live beyond 2**31 elements in a strided cache.""" + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x segmented prefill") + torch.accelerator.empty_cache() + free, _ = torch.accelerator.get_memory_info() + if free < 10 * 2**30: + pytest.skip("The address-boundary fixture needs 10 GiB free VRAM") + device = torch.device("cuda:0") + stride = 2**30 + 4096 + kc = torch.empty_strided( + (3, 32, 1, 128), + (stride, 128, 128, 1), + device=device, + dtype=torch.bfloat16, + ) + vc = torch.empty_strided( + (3, 32, 1, 128), + (stride, 128, 128, 1), + device=device, + dtype=torch.bfloat16, + ) + dense_k = torch.randn(2, 32, 1, 128, device=device, dtype=torch.bfloat16) + dense_v = torch.randn_like(dense_k) + for logical, physical in enumerate((2, 1)): + kc[physical].copy_(dense_k[logical]) + vc[physical].copy_(dense_v[logical]) + q = torch.randn(2, 4, 128, device=device, dtype=torch.bfloat16) + out = torch.empty_like(q) + table = torch.tensor([[2, 1]], device=device, dtype=torch.int32) + starts = torch.tensor([0, 2], device=device, dtype=torch.int32) + lengths = torch.tensor([35], device=device, dtype=torch.int32) + one = torch.ones((), device=device) + run_segmented_attention( + q, out, kc, vc, table, starts, lengths, 2, 35, one, one, 128**-0.5 + ) + full_k = dense_k.flatten(0, 1)[:35].float().repeat_interleave(4, 1) + full_v = dense_v.flatten(0, 1)[:35].float().repeat_interleave(4, 1) + scores = torch.einsum("qhd,khd->hqk", q.float(), full_k) / math.sqrt(128) + scores[:, 0, -1] = -float("inf") + expected = torch.einsum("hqk,khd->qhd", scores.softmax(-1), full_v) + error = ((out.float() - expected).norm(dim=-1) / expected.norm(dim=-1)).max() + assert torch.isfinite(out).all() and error < 0.01 + + +@pytest.mark.parametrize("dim,hq,hk", [(128, 16, 1), (256, 12, 2)]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_segmented_attention_reservation_covers_ragged_query_caps( + monkeypatch, dim, hq, hk, fp8 +): + """Every supported query bucket fits the startup buffer after locking.""" + from vllm.v1.attention.ops import segmented_attention as segmented + + sizes = segmented.segmented_attention_workspace_size( + 32, + hq, + hk, + dim, + 65536, + fp8=fp8, + ) + workspace = tuple(torch.empty(size) for size in sizes) + query_lengths = { + query_len + for capacity in segmented._query_capacity_buckets() + for query_len in (max(1, capacity // 2 + 1), capacity) + } + for batch in range(1, 33): + for query_len in query_lengths: + qcap = segmented.segmented_query_capacity(query_len) + cfg = segmented.select_segmented_config( + batch, query_len, 65536, hq, hk, dim, fp8 + ) + shapes = segmented.segmented_workspace_shapes( + batch, qcap, hq, hk, dim, cfg["splits"] + ) + if shapes is not None: + assert math.prod(shapes[0]) <= workspace[0].numel() + assert math.prod(shapes[1]) <= workspace[1].numel() + + +def test_segmented_attention_reservation_respects_scheduler_token_limit( + monkeypatch, +) -> None: + """Workspace planning excludes batch/query pairs the scheduler cannot form.""" + from vllm.v1.attention.ops import segmented_attention as segmented + + calls = [] + + def record_config(batch, query_len, *_args): + calls.append((batch, query_len)) + return {"splits": 1} + + monkeypatch.setattr(segmented, "select_segmented_config", record_config) + segmented.segmented_attention_workspace_size.cache_clear() + segmented.segmented_attention_workspace_size( + 32, + 16, + 2, + 128, + 65536, + max_tokens=8, + ) + + assert calls + assert all(batch + query_len - 1 <= 8 for batch, query_len in calls) + assert (8, 1) in calls + assert (1, 8) in calls + + +def test_segmented_attention_query_capacity_buckets() -> None: + from vllm.v1.attention.ops.segmented_attention import ( + MAX_QUERY_LEN, + segmented_query_capacity, + ) + + assert [ + segmented_query_capacity(query_len) + for query_len in (1, 2, 3, 33, 129, 1025, 2049, MAX_QUERY_LEN) + ] == [1, 2, 4, 64, 256, 2048, 4096, 8192] + with pytest.raises(ValueError, match="Query length"): + segmented_query_capacity(MAX_QUERY_LEN + 1) + + +def test_segmented_scratch_is_stable_and_isolated_across_runtime_lanes(monkeypatch): + """Draft/DBO scratch never aliases target scratch or generic workspace.""" + import vllm.v1.worker.workspace as workspace + from vllm.v1.attention.ops import segmented_attention as segmented + + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + monkeypatch.setattr(segmented, "is_workspace_manager_initialized", lambda: True) + monkeypatch.setattr(segmented, "current_workspace_manager", lambda: manager) + buffers = {} + for ubatch in range(2): + active_ubatch[0] = ubatch + for lane in range(2): + with workspace.use_workspace_lane(lane): + buffers[ubatch, lane] = segmented.get_segmented_attention_workspace( + torch.device("cpu"), (1024, 16) + ) + manager.get_simultaneous(((2048,), torch.float32)) + manager.lock() + pointers: set[int] = set() + for (ubatch, lane), expected in buffers.items(): + active_ubatch[0] = ubatch + with workspace.use_workspace_lane(lane): + actual = segmented.get_segmented_attention_workspace( + torch.device("cpu"), (1024, 16) + ) + assert actual is expected + pointers.update(t.data_ptr() for t in actual) + assert ( + manager.get_simultaneous(((1,), torch.float32))[0].data_ptr() + not in pointers + ) + assert len(pointers) == 8 diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index b99a33f9f7..067d1dbc8b 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -61,6 +61,216 @@ def test_aiter_unified_attention_uses_dedicated_metadata_builder(): ) +def test_segmented_attention_uses_dedicated_backend_components(): + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionBackend, + RocmSegmentedAttentionImpl, + RocmSegmentedAttentionMetadataBuilder, + ) + + assert ( + RocmSegmentedAttentionBackend.get_builder_cls() + is RocmSegmentedAttentionMetadataBuilder + ) + assert RocmSegmentedAttentionBackend.get_impl_cls() is RocmSegmentedAttentionImpl + assert RocmSegmentedAttentionBackend.get_name() == "ROCM_SEGMENTED_ATTN" + assert RocmSegmentedAttentionBackend.supports_non_causal() + assert RocmSegmentedAttentionBackend.supports_sliding_window() + assert RocmSegmentedAttentionBackend.supports_sink() + assert not RocmSegmentedAttentionBackend.supports_mm_prefix() + + +@pytest.mark.parametrize( + "sliding_window,logits_soft_cap,has_sinks", + [(4096, None, False), (None, 50.0, False), (None, None, True)], +) +def test_segmented_attention_admits_unified_fallback_features( + monkeypatch, sliding_window, logits_soft_cap, has_sinks +): + from vllm.platforms import rocm + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setattr(rocm, "on_gfx12x", lambda: True) + sinks = torch.zeros(8) if has_sinks else None + + impl = RocmSegmentedAttentionImpl( + num_heads=8, + head_size=128, + scale=128**-0.5, + num_kv_heads=2, + alibi_slopes=None, + sliding_window=sliding_window, + kv_cache_dtype="auto", + logits_soft_cap=logits_soft_cap, + attn_type=AttentionType.DECODER, + sinks=sinks, + ) + + assert isinstance(impl, RocmSegmentedAttentionImpl) + + +@pytest.mark.parametrize( + "platform,rdna", + [("cuda", False), ("cpu", False), ("xpu", False), ("rocm", False), ("rocm", True)], +) +def test_segmented_attention_requires_rocm_rdna(monkeypatch, platform, rdna): + """Reject unsupported platforms before probing ROCm or constructing kernels.""" + from vllm.platforms import rocm + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backends import rocm_segmented_attn as segmented + + monkeypatch.setattr( + segmented, + "current_platform", + SimpleNamespace(is_rocm=lambda: platform == "rocm"), + ) + + def on_rdna(): + assert platform == "rocm", "Must not probe ROCm on another platform" + return rdna + + monkeypatch.setattr(rocm, "on_gfx1x", on_rdna) + reason = segmented.RocmSegmentedAttentionBackend.supports_combination( + 128, + torch.bfloat16, + "auto", + 32, + False, + False, + False, + False, + DeviceCapability(12, 0), + ) + if rdna: + assert reason is None + segmented.RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + else: + assert reason == "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm" + with pytest.raises(ValueError, match="requires AMD RDNA GPUs on ROCm"): + segmented.RocmSegmentedAttentionImpl( + 8, 128, 128**-0.5, 2, None, None, "auto" + ) + + +@pytest.mark.parametrize("explicit_cache", [False, True]) +def test_segmented_triton_cache_survives_standalone_compilation( + monkeypatch, tmp_path, explicit_cache +): + """Backend initialization preserves one cache across model compilations.""" + import os + + from vllm.compilation.compiler_interface import InductorStandaloneAdaptor + from vllm.platforms import rocm + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) + monkeypatch.delenv("TORCHINDUCTOR_CACHE_DIR", raising=False) + expected = tmp_path / "rocm_segmented_attention" / "triton" + if explicit_cache: + expected = tmp_path / "custom" + monkeypatch.setenv("TRITON_CACHE_DIR", str(expected)) + for model in ("target", "draft", "restart_target"): + RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + assert os.environ["TRITON_CACHE_DIR"] == str(expected) + assert expected.is_dir() + InductorStandaloneAdaptor("binary").initialize_cache(str(tmp_path / model)) + assert os.environ["TRITON_CACHE_DIR"] == str(expected) + + +def test_segmented_attention_is_opt_in(monkeypatch): + from vllm.platforms import rocm + from vllm.platforms.rocm import RocmPlatform, _get_backend_priorities + + assert AttentionBackendEnum.ROCM_SEGMENTED_ATTN not in _get_backend_priorities( + use_mla=False, use_sparse=False + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setattr(rocm, "on_gfx12x", lambda: True) + config = AttentionSelectorConfig( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=32, + ) + + path = RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN, + attn_selector_config=config, + ) + + assert path == AttentionBackendEnum.ROCM_SEGMENTED_ATTN.get_path() + + +def test_segmented_attention_rejects_mm_prefix_at_selection(monkeypatch): + from vllm.platforms import rocm + from vllm.platforms.rocm import RocmPlatform + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + config = AttentionSelectorConfig( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=32, + use_mm_prefix=True, + ) + + with pytest.raises(ValueError, match="multimodal token full attention"): + RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN, + attn_selector_config=config, + ) + + +def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch): + from vllm.config import KernelConfig + from vllm.platforms import rocm + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + impl = RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + impl._segmented_attention_config = SimpleNamespace( + kernel_config=KernelConfig(), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=3), + model_config=SimpleNamespace(max_model_len=16), + ) + query = torch.empty(3, 8, 128, dtype=torch.bfloat16) + key = torch.empty(3, 2, 128, dtype=torch.bfloat16) + value = torch.empty_like(key) + output = torch.empty_like(query) + kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16) + metadata = SimpleNamespace( + use_cascade=False, + num_actual_tokens=3, + block_table=torch.zeros(1, 1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 3], dtype=torch.int32), + seq_lens=torch.tensor([3], dtype=torch.int32), + max_seq_len=3, + max_query_len=3, + causal=True, + ) + layer = SimpleNamespace(_k_scale=torch.ones(()), _v_scale=torch.ones(())) + + with patch( + "vllm.v1.attention.backends.rocm_segmented_attn.segmented_attention" + ) as dispatch: + result = impl.forward(layer, query, key, value, kv_cache, metadata, output) + + assert result is output + dispatch.assert_called_once() + assert dispatch.call_args.kwargs["key_cache"].shape == (2, 16, 2, 128) + + def test_aiter_unified_attention_capture_preserves_query_start_locations(): from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( RocmAiterUnifiedAttentionMetadataBuilder, diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index df69f17149..b8bfca082a 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -278,6 +278,9 @@ class KernelConfig: enable_flashinfer_autotune: bool = None # type: ignore[assignment] """If True, run FlashInfer autotuning during kernel warmup.""" + enable_rocm_segmented_attn_autotune: bool = False + """If True, autotune ROCm segmented attention during kernel warmup on RDNA GPUs.""" + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated # to the shared JIT warmup infrastructure. # https://github.com/vllm-project/vllm/pull/47451 @@ -406,6 +409,7 @@ class KernelConfig: "enable_cutedsl_warmup", "enable_jit_warmup", "enable_flashinfer_autotune", + "enable_rocm_segmented_attn_autotune", "ir_op_priority", # handled separately below } if self.linear_backend_per_quant is None: diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3161cf2261..44b7865daa 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -34,6 +34,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType @@ -269,6 +270,24 @@ class DFlashQwen3Attention(nn.Module): sinks=self.attention_sink_bias, ) self.causal = causal + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + if on_gfx1x() and isinstance(self.attn.impl, RocmSegmentedAttentionImpl): + # Startup segmented tuning needs the draft's actual attention mode. + # Causality is normally supplied later through attention metadata. + self.attn.segmented_causal = causal + speculative_config = get_current_vllm_config().speculative_config + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + self.attn.segmented_query_limit = ( + speculative_config.num_speculative_tokens + 1 + ) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 228ce73803..55d5ea7a7d 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -220,6 +220,13 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): if process_local_only: return + if current_platform.is_rocm(): + from vllm.model_executor.warmup.rocm_segmented_attn_autotune_warmup import ( + rocm_segmented_attn_autotune_warmup, + ) + + rocm_segmented_attn_autotune_warmup(worker) + flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) diff --git a/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py b/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py new file mode 100644 index 0000000000..0ebcc58c3c --- /dev/null +++ b/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py @@ -0,0 +1,1335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup-only, persistent launch tuning for segmented ROCm attention.""" + +from __future__ import annotations + +import gc +import hashlib +import itertools +import json +import math +import os +import statistics +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +from filelock import FileLock + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.v1.attention.ops import segmented_attention +from vllm.v1.attention.ops.segmented_attention import ( + MAX_QUERY_LEN, + _segmented_attention_reduce, + _segmented_attention_stage, + compile_segmented_attention, + run_segmented_attention, + segmented_query_capacity, + segmented_workspace_shapes, + select_segmented_config, +) + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_TABLES: dict[tuple, dict] = {} +_QUERY_BUCKETS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) +_BATCH_BUCKETS = (1, 2, 4, 8, 16, 32) +_SEQUENCE_BUCKETS = (128, 512, 2048, 8192, 32768, 131072, 262144) +_COMPILE_WORKERS = 2 +_CONDITION_ROUNDS = 2 +_SCREEN_ROUNDS = 5 +_FINALIST_COUNT = 3 +_FINAL_ROUNDS = 7 +_MIN_PROMOTION_SPEEDUP = 1.02 + + +def _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window=-1, + causal=True, + has_sinks=False, +): + return ( + torch.device(device).index, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + + +def _normalized_config(config, batch, query_len, heads, dim): + qcap = segmented_query_capacity(query_len) + result = dict(config) + result.setdefault("waves_per_eu", 2) + result.setdefault("reduce_d", 64 if batch * qcap * heads < 64 else dim) + result.setdefault("reduce_warps", 4) + return result + + +def _candidate_configs(default, batch, query_len, heads, dim, seq_len=None): + """Search launch tiles and split counts within a bounded workspace.""" + base = _normalized_config(default, batch, query_len, heads, dim) + max_splits = base["splits"] + long_d64_decode = ( + dim == 64 + and batch == 1 + and query_len == 1 + and seq_len is not None + and seq_len >= 131072 + ) + if long_d64_decode: + max_splits = 512 + elif dim == 64 and query_len <= 512 and seq_len is not None and seq_len >= 4096: + max_splits = min(64, 2 * max_splits) + candidates: list[dict] = [] + seen = set() + + def add(**updates): + config = {**base, **updates} + if ( + config["splits"] > max_splits + or dim % config["bk"] + or config["reduce_d"] > dim + ): + return + identity = tuple(sorted(config.items())) + if identity not in seen: + seen.add(identity) + candidates.append(config) + + if long_d64_decode: + add() + for splits in (128, 256, 512, 64): + for bn, warps, stages in ( + (32, 1, 1), + (32, 2, 1), + (32, 2, 2), + (64, 1, 1), + (64, 2, 1), + (64, 2, 2), + ): + add( + bm=16, + bn=bn, + bk=64, + splits=splits, + warps=warps, + stages=stages, + ) + return candidates + + split_choices = [base["splits"]] + if max_splits > base["splits"]: + split_choices.append(max_splits) + for divisor in (2, 4): + split_choices.append(max(1, base["splits"] // divisor)) + split_choices.append(1) + for splits in split_choices: + add(splits=splits) + + tile_variants: tuple[tuple[int, int, int, int, int, int], ...] = ( + (16, 32, dim, 4, 1, 2), + (16, 64, min(128, dim), 4, 1, 2), + (32, 32, 64, 4, 2, 2), + (32, 64, 64, 4, 1, 2), + (64, 32, min(128, dim), 4, 1, 6), + (128, 32, min(128, dim), 8, 1, 6), + ) + if dim == 64: + tile_variants += ( + (16, 16, 64, 4, 1, 2), + (32, 16, 64, 4, 2, 2), + (64, 32, 64, 4, 1, 2), + (64, 64, 64, 4, 1, 2), + (128, 32, 64, 4, 1, 2), + (128, 64, 64, 4, 1, 2), + ) + for bm, bn, bk, warps, stages, waves in tile_variants: + add( + bm=bm, + bn=bn, + bk=bk, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + half_splits = max(1, base["splits"] // 2) + for bm, bn, bk, warps, stages, waves in tile_variants[:3]: + add( + bm=bm, + bn=bn, + bk=bk, + splits=half_splits, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + if dim == 64: + for splits in (max_splits, half_splits): + for bm, bn, bk, warps, stages, waves in tile_variants[2:]: + add( + bm=bm, + bn=bn, + bk=bk, + splits=splits, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + add(waves_per_eu=6) + add(stages=2 if base["stages"] == 1 else 1) + add(reduce_d=64 if base["reduce_d"] != 64 else dim) + return candidates + + +def _identity( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + sliding_window, + causal, + max_query_len, + has_sinks=False, + physical_max_len=None, +): + properties = torch.cuda.get_device_properties(device) + from triton._C.libtriton import get_cache_invalidating_env_vars + + source = Path(__file__) + kernel = Path(segmented_attention.__file__) + return { + "schema": 5, + "gpu": properties.name, + "arch": properties.gcnArchName, + "compute_units": properties.multi_processor_count, + "torch": torch.__version__, + "hip": torch.version.hip, + "compiler_env": get_cache_invalidating_env_vars(), + "kernel_sha256": hashlib.sha256(kernel.read_bytes()).hexdigest(), + "tuner_sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + "dtype": str(dtype), + "kv_dtype": str(kv_dtype), + "heads": heads, + "kv_heads": kv_heads, + "dim": dim, + "page": page, + "scale": scale, + "max_tokens": max_tokens, + "max_len": max_len, + "max_seqs": max_seqs, + "sliding_window": sliding_window, + "causal": causal, + "max_query_len": max_query_len, + "has_sinks": has_sinks, + "physical_max_len": physical_max_len, + "query_buckets": list(_QUERY_BUCKETS), + "batch_buckets": list(_BATCH_BUCKETS), + "sequence_buckets": list(_SEQUENCE_BUCKETS), + "compile_workers": _COMPILE_WORKERS, + "condition_rounds": _CONDITION_ROUNDS, + "screen_rounds": _SCREEN_ROUNDS, + "finalist_count": _FINALIST_COUNT, + "final_rounds": _FINAL_ROUNDS, + "min_promotion_speedup": _MIN_PROMOTION_SPEEDUP, + } + + +def _memory_budget(device): + free, total = torch.accelerator.get_memory_info(device) + available = int(min(free // 2, total // 4)) + return 1 << (available.bit_length() - 1) if available else 0 + + +def _query_lengths(batch, query_len, max_tokens): + if batch * query_len <= max_tokens: + return [query_len] * batch + return [min(query_len, max_tokens - batch + 1)] + [1] * (batch - 1) + + +def _scratch_bytes( + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_max_len=None, +): + batch, query_len, seq_len = workload + queries = _query_lengths(batch, query_len, max_tokens) + tokens = sum(queries) + scratch_seq_len = ( + max(seq_len, physical_max_len) + if physical_max_len is not None and batch == 1 and query_len <= 8 + else seq_len + ) + blocks = batch * math.ceil(scratch_seq_len / page) + cache = 2 * blocks * page * kv_heads * dim * kv_dtype.itemsize + tensors = tokens * dim * (4 * heads + 2 * kv_heads) * dtype.itemsize + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + candidates = _candidate_configs( + default, batch, query_len, heads, dim, seq_len=seq_len + ) + shapes = segmented_workspace_shapes( + batch, + segmented_query_capacity(query_len), + heads, + kv_heads, + dim, + max(config["splits"] for config in candidates), + ) + workspace = 0 if shapes is None else sum(math.prod(shape) * 4 for shape in shapes) + return cache + tensors + workspace + 384 * 1024**2 + + +def _workloads( + max_tokens, + max_len, + max_seqs, + *, + memory_budget_bytes=None, + dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + heads=8, + kv_heads=1, + dim=128, + page=16, + cache_layouts=(), + cache_budget_bytes=None, + max_query_len=MAX_QUERY_LEN, + physical_max_len=None, +): + query_limit = min(max_tokens, max_len, max_query_len, MAX_QUERY_LEN) + queries = {q for q in _QUERY_BUCKETS if q <= query_limit} + if query_limit: + queries.add(query_limit) + for query_len in sorted(queries): + batch_limit = min(max_seqs, max_tokens - query_len + 1) + if batch_limit < 1: + continue + if query_len <= 2: + batches = {min(batch, batch_limit) for batch in _BATCH_BUCKETS} + else: + batches = {1, min(4, batch_limit), batch_limit} + if query_len <= 2: + sequences = { + seq for seq in _SEQUENCE_BUCKETS if query_len <= seq <= max_len + } + sequences.update((query_len, max_len)) + else: + sequences = { + query_len, + max(query_len, min(max_len, 2048)), + max(query_len, min(max_len, 8192)), + max_len, + } + # Keep long-context verification and prefill buckets even when + # the model's maximum length is pruned by the KV memory budget. + sequences.update( + seq for seq in (32768, 131072) if query_len <= seq <= max_len + ) + for batch, seq_len in itertools.product(sorted(batches), sorted(sequences)): + workload = (batch, query_len, seq_len) + if ( + memory_budget_bytes is not None + and _scratch_bytes( + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_max_len, + ) + > memory_budget_bytes + ): + continue + if cache_budget_bytes is not None: + cache_bytes = batch * sum( + math.ceil(seq_len / block) * size if block else size + for block, size in cache_layouts + ) + if cache_bytes > cache_budget_bytes: + continue + yield workload + + +def _make_inputs( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_seq_len=None, +): + batch, query_len, seq_len = workload + physical_seq_len = max(seq_len, physical_seq_len or seq_len) + generator = torch.Generator(device=device).manual_seed(1234) + + def randn(*shape): + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + queries = _query_lengths(batch, query_len, max_tokens) + tokens = sum(queries) + q = randn(tokens, heads, dim) + k = randn(tokens, kv_heads, dim) + v = randn(tokens, kv_heads, dim) + blocks_per_seq = math.ceil(seq_len / page) + physical_blocks_per_seq = math.ceil(physical_seq_len / page) + cache_blocks_per_seq = ( + physical_blocks_per_seq if batch == 1 and query_len <= 8 else blocks_per_seq + ) + cache_shape = (batch * cache_blocks_per_seq, 2, page, kv_heads, dim) + scale_value = 0.125 if kv_dtype.itemsize == 1 else 1.0 + cache = (randn(*cache_shape).float() / scale_value).to(kv_dtype) + kc, vc = cache[:, 0], cache[:, 1] + table = ( + torch.arange(physical_blocks_per_seq, device=device, dtype=torch.int32)[None, :] + .remainder(cache_blocks_per_seq) + .expand(batch, -1) + .contiguous() + ) + table += ( + torch.arange(batch, device=device, dtype=torch.int32)[:, None] + * cache_blocks_per_seq + ) + starts = torch.tensor( + [0, *itertools.accumulate(queries)], device=device, dtype=torch.int32 + ) + lengths = torch.full((batch,), physical_seq_len, device=device, dtype=torch.int32) + scale = torch.full((), scale_value, device=device, dtype=torch.float32) + return q, k, v, kc, vc, table, starts, lengths, scale + + +def _config_key(config): + return tuple(sorted(config.items())) + + +def _rotated(configs, offset): + offset %= len(configs) + return configs[offset:] + configs[:offset] + + +def _measure_configs(configs, run, device, eviction, rounds, graph_calls=0): + """Interleave candidates so clock and thermal drift affect them evenly.""" + stream = torch.cuda.current_stream(device) + samples: dict[tuple, list[float]] = {_config_key(config): [] for config in configs} + graphs = {} + if graph_calls: + for config in configs: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(graph_calls): + run(config) + graphs[_config_key(config)] = graph + for graph in graphs.values(): + graph.replay() + torch.accelerator.synchronize(device) + for round_index in range(rounds): + events = [] + order = _rotated(configs, round_index) + if round_index % 2: + order = list(reversed(order)) + for config in order: + if not graph_calls: + eviction.zero_() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + if graph_calls: + graphs[_config_key(config)].replay() + else: + run(config) + end.record(stream) + events.append((_config_key(config), start, end)) + events[-1][2].synchronize() + for key, start, end in events: + samples[key].append(start.elapsed_time(end) * 1000 / max(1, graph_calls)) + return samples + + +def _condition_configs(configs, run, device): + for round_index in range(_CONDITION_ROUNDS): + for config in _rotated(configs, round_index): + run(config) + torch.accelerator.synchronize(device) + + +def _select_tuned_winner( + default, finalists, samples, min_promotion_speedup=_MIN_PROMOTION_SPEEDUP +): + """Return a finalist only when it reliably beats the static incumbent.""" + default_key = _config_key(default) + default_samples = samples[default_key] + comparisons = [] + for config in finalists: + key = _config_key(config) + config_samples = samples[key] + paired_speedup = statistics.median( + baseline / candidate + for baseline, candidate in zip(default_samples, config_samples, strict=True) + ) + comparisons.append( + { + "config": config, + "us": statistics.median(config_samples), + "paired_speedup_vs_default": paired_speedup, + } + ) + challenger = max( + comparisons, key=lambda result: result["paired_speedup_vs_default"] + ) + if challenger["paired_speedup_vs_default"] >= min_promotion_speedup: + return challenger, comparisons + return ( + next(result for result in comparisons if result["config"] == default), + comparisons, + ) + + +def _precompile_configs( + configs, + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + scale, + workspace, + sliding_window, + causal, + sinks=None, +): + workers = min(_COMPILE_WORKERS, len(configs)) + with ( + ThreadPoolExecutor(max_workers=workers) as executor, + triton.AsyncCompileMode(executor, ignore_errors=True), + ): + for config in configs: + compile_segmented_attention( + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + kv_scale, + scale, + config, + workspace, + sliding_window=sliding_window, + causal=causal, + sinks=sinks, + ) + + +def _tune_workload( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + workload, + *, + sliding_window=-1, + causal=True, + has_sinks=False, + physical_seq_len=None, +): + batch, query_len, seq_len = workload + q, _, _, kc, vc, table, starts, lengths, kv_scale = _make_inputs( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_seq_len, + ) + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + configs = _candidate_configs(default, batch, query_len, heads, dim, seq_len=seq_len) + sinks = ( + torch.linspace(-1, 1, heads, dtype=torch.float32, device=device) + if has_sinks + else None + ) + qcap = segmented_query_capacity(query_len) + shapes = segmented_workspace_shapes( + batch, qcap, heads, kv_heads, dim, max(c["splits"] for c in configs) + ) + workspace = ( + None + if shapes is None + else tuple( + torch.empty(shape, dtype=torch.float32, device=device) for shape in shapes + ) + ) + reference = torch.empty_like(q) + output = torch.empty_like(q) + + def run(config, out=output): + run_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + query_len, + seq_len, + kv_scale, + kv_scale, + scale, + sliding_window=sliding_window, + causal=causal, + config=config, + workspace=workspace, + sinks=sinks, + ) + + _precompile_configs( + configs, + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + scale, + workspace, + sliding_window, + causal, + sinks, + ) + incumbent = configs[0] + reference.fill_(float("nan")) + run(incumbent, reference) + torch.accelerator.synchronize(device) + if not torch.isfinite(reference).all(): + raise AssertionError("Default segmented attention output is not finite") + reference_norm = reference.float().norm().clamp_min(1e-6) + eviction = torch.empty(256 * 1024**2, dtype=torch.int8, device=device) + results = [] + valid_configs = [] + for config in configs: + try: + output.fill_(float("nan")) + run(config) + torch.accelerator.synchronize(device) + torch.testing.assert_close(output, reference, atol=0.01, rtol=0.01) + relative_l2 = ( + (output.float() - reference.float()).norm() / reference_norm + ).item() + if relative_l2 > 0.005: + raise AssertionError(f"Relative L2 error {relative_l2} exceeds 0.5%") + results.append({"config": config, "relative_l2": relative_l2}) + valid_configs.append(config) + except ( + torch.OutOfMemoryError, + triton.OutOfResources, + triton.CompilationError, + AssertionError, + RuntimeError, + ) as error: + logger.warning("Rejected segmented attention config %s: %s", config, error) + if not results: + raise RuntimeError("No valid segmented attention launch configurations") + if incumbent not in valid_configs: + raise RuntimeError("Static segmented attention configuration was rejected") + + _condition_configs(valid_configs, run, device) + graph_calls = 5 if dim == 64 and has_sinks else 0 + screen_samples = _measure_configs( + valid_configs, run, device, eviction, _SCREEN_ROUNDS, graph_calls + ) + by_key = {_config_key(result["config"]): result for result in results} + for key, samples in screen_samples.items(): + by_key[key]["screen_us"] = statistics.median(samples) + by_key[key]["screen_samples_us"] = samples + + challengers = sorted( + (config for config in valid_configs if config != incumbent), + key=lambda config: statistics.median(screen_samples[_config_key(config)]), + ) + finalists = [incumbent, *challengers[: _FINALIST_COUNT - 1]] + _condition_configs(finalists, run, device) + final_samples = _measure_configs( + finalists, run, device, eviction, _FINAL_ROUNDS, graph_calls + ) + min_promotion_speedup = ( + 1.10 + if dim == 64 and has_sinks and sliding_window >= 0 + else _MIN_PROMOTION_SPEEDUP + ) + winner, comparisons = _select_tuned_winner( + incumbent, finalists, final_samples, min_promotion_speedup + ) + for comparison in comparisons: + result = by_key[_config_key(comparison["config"])] + result["final_us"] = comparison["us"] + result["final_samples_us"] = final_samples[_config_key(comparison["config"])] + result["paired_speedup_vs_default"] = comparison["paired_speedup_vs_default"] + winner_result = by_key[_config_key(winner["config"])] + logger.info( + "ROCm segmented attention autotune: B=%d Q=%d seq=%d best=%s " + "%.2f us speedup=%.3fx%s", + batch, + query_len, + seq_len, + winner["config"], + winner["us"], + winner["paired_speedup_vs_default"], + " (static retained)" if winner["config"] == incumbent else "", + ) + return { + "workload": list(workload), + "query_lengths": _query_lengths(batch, query_len, max_tokens), + "default": default, + "best": winner["config"], + "selection": { + "incumbent": incumbent, + "finalists": finalists, + "min_promotion_speedup": min_promotion_speedup, + "winner_us": winner_result["final_us"], + "paired_speedup_vs_default": winner["paired_speedup_vs_default"], + }, + "results": results, + } + + +def _save(path, data): + with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as file: + temporary = Path(file.name) + try: + json.dump(data, file, indent=2) + file.flush() + os.fsync(file.fileno()) + except BaseException: + temporary.unlink(missing_ok=True) + raise + try: + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _load_records(path, identity, heads, kv_heads, dim, fp8): + """Load a canonical table and any crash-recovery TP shards.""" + records: dict[tuple, dict] = {} + merged_shards = [] + sources = [path, *sorted(path.parent.glob(f"{path.name}.tp*.part"))] + for source in sources: + if not source.exists(): + continue + try: + saved = json.loads(source.read_text()) + if saved["identity"] != identity or not all( + _valid_record(record, heads, kv_heads, dim, fp8) + for record in saved["records"] + ): + raise ValueError("identity or record validation failed") + records.update( + (tuple(record["workload"]), record) for record in saved["records"] + ) + if source != path: + merged_shards.append(source) + except (ValueError, KeyError, TypeError, AssertionError): + logger.warning( + "Ignoring invalid segmented attention tuning cache %s", source + ) + return {"identity": identity, "records": list(records.values())}, merged_shards + + +def _tp_context(): + """Return the initialized TP coordinator, or local-only tuning state.""" + from vllm.distributed.parallel_state import ( + get_tp_group, + model_parallel_is_initialized, + ) + + if not model_parallel_is_initialized(): + return None, 0, 1 + group = get_tp_group() + return group, group.rank_in_group, group.world_size + + +def _gather_tp(group, value): + if group is None or group.world_size == 1: + return [value] + gathered = [None] * group.world_size + torch.distributed.all_gather_object(gathered, value, group=group.cpu_group) + return gathered + + +def _workload_weight(workload, max_tokens, heads, kv_heads, dim, fp8): + """Estimate compile plus execution cost for TP workload balancing.""" + batch, query_len, seq_len = workload + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + candidates = len( + _candidate_configs(default, batch, query_len, heads, dim, seq_len=seq_len) + ) + query_tokens = sum(_query_lengths(batch, query_len, max_tokens)) + attention_work = query_tokens * seq_len * heads * dim + return candidates * (1 << 40) + attention_work + + +def _shard_workloads(workloads, world_size, max_tokens, heads, kv_heads, dim, fp8): + """Balance compile-affine ``(batch, query)`` groups across TP ranks.""" + assignments: list[list[tuple[int, tuple]]] = [[] for _ in range(world_size)] + loads = [0] * world_size + groups: dict[tuple, dict] = {} + for index, workload in enumerate(workloads): + compile_group = groups.setdefault(workload[:2], {"weight": 0, "items": []}) + compile_group["weight"] += _workload_weight( + workload, max_tokens, heads, kv_heads, dim, fp8 + ) + compile_group["items"].append((index, workload)) + for compile_group in sorted( + groups.values(), key=lambda item: item["weight"], reverse=True + ): + rank = min( + range(world_size), key=lambda candidate: (loads[candidate], candidate) + ) + assignments[rank].extend(compile_group["items"]) + loads[rank] += compile_group["weight"] + return [ + [workload for _, workload in sorted(assignment)] for assignment in assignments + ] + + +def _valid_record(record, heads, kv_heads, dim, fp8): + workload = record["workload"] + if len(workload) != 3 or any( + type(value) is not int or value < 1 for value in workload + ): + return False + batch, query_len, seq_len = workload + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + return record["best"] in _candidate_configs( + default, batch, query_len, heads, dim, seq_len=seq_len + ) + + +@contextmanager +def _temporary_tuning_kernels(device): + """Release tuning-only HIP modules without invalidating existing graphs.""" + device_index = torch.device(device).index + if device_index is None: + device_index = torch.accelerator.current_device_index() + caches = [] + for kernel in (_segmented_attention_stage, _segmented_attention_reduce): + cache = kernel.device_caches[device_index][0] + caches.append((cache, set(cache))) + try: + yield + finally: + torch.accelerator.synchronize(device) + for cache, existing_keys in caches: + for key in cache.keys() - existing_keys: + del cache[key] + # CompiledKernel destruction unloads the HIP module. Keep pre-existing + # kernels alive, since an earlier graph capture may reference them. + gc.collect() + torch.accelerator.empty_cache() + + +@torch.inference_mode() +def warmup_segmented_attention( + device, + dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + *, + memory_budget_bytes=None, + cache_layouts=(), + cache_budget_bytes=None, + kv_dtype=torch.bfloat16, + sliding_window=-1, + causal=True, + max_query_len=MAX_QUERY_LEN, + has_sinks=False, + physical_max_len=None, +): + """Tune reachable segmented buckets before graph capture.""" + if ( + dtype not in (torch.bfloat16, torch.float16) + or kv_dtype + not in ( + torch.bfloat16, + torch.float16, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + or dim not in (64, 128, 256) + or kv_heads < 1 + or heads % kv_heads + or not 1 <= heads // kv_heads <= 16 + ): + return + key = _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + identity = _identity( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + sliding_window, + causal, + max_query_len, + has_sinks, + physical_max_len, + ) + if memory_budget_bytes is None: + memory_budget_bytes = _memory_budget(device) + workload_args = dict( + memory_budget_bytes=memory_budget_bytes, + dtype=dtype, + kv_dtype=kv_dtype, + heads=heads, + kv_heads=kv_heads, + dim=dim, + page=page, + cache_layouts=cache_layouts, + cache_budget_bytes=cache_budget_bytes, + max_query_len=max_query_len, + physical_max_len=( + physical_max_len if has_sinks and sliding_window >= 0 else None + ), + ) + workloads = list(_workloads(max_tokens, max_len, max_seqs, **workload_args)) + active = _TABLES.get(key, {}) + warmed = ( + {tuple(record["workload"]) for record in active.get("records", [])} + if active.get("identity") == identity + else set() + ) + if all(workload in warmed for workload in workloads): + return + digest = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + path = Path(envs.VLLM_CACHE_ROOT) / "rocm_segmented_attention" / f"{digest}.json" + path.parent.mkdir(parents=True, exist_ok=True) + group, tp_rank, tp_size = _tp_context() + if group is not None: + digests = _gather_tp(group, digest) + if any(peer != digest for peer in digests): + logger.warning( + "ROCm segmented attention TP ranks have different tuning " + "identities; tuning this rank independently" + ) + group, tp_rank, tp_size = None, 0, 1 + + unpruned = len( + list(_workloads(max_tokens, max_len, max_seqs, max_query_len=max_query_len)) + ) + start = time.monotonic() + lock = FileLock(str(path) + ".lock") if tp_rank == 0 else None + if lock is not None: + lock.acquire() + try: + if tp_rank == 0: + data, merged_shards = _load_records( + path, identity, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + else: + data, merged_shards = None, [] + if group is not None: + data = group.broadcast_object(data, src=0) + assert data is not None + records = {tuple(record["workload"]): record for record in data["records"]} + missing = [workload for workload in workloads if workload not in records] + loaded = len(workloads) - len(missing) + assignments = _shard_workloads( + missing, + tp_size, + max_tokens, + heads, + kv_heads, + dim, + kv_dtype.itemsize == 1, + ) + local_workloads = assignments[tp_rank] + logger.info( + "ROCm segmented attention tuning plan: workloads=%d missing=%d " + "pruned=%d tp_rank=%d/%d assigned=%d scratch_budget=%d MiB " + "cache_budget=%s", + len(workloads), + len(missing), + unpruned - len(workloads), + tp_rank, + tp_size, + len(local_workloads), + memory_budget_bytes // 2**20, + cache_budget_bytes, + ) + + if not missing: + if tp_rank == 0 and merged_shards: + _save(path, data) + for shard in merged_shards: + shard.unlink(missing_ok=True) + _TABLES[key] = data + logger.info( + "ROCm segmented attention autotune ready: tp_rank=%d/%d " + "local_tuned=0 tuned=0 loaded=%d failed=0 elapsed=%.2fs cache=%s", + tp_rank, + tp_size, + loaded, + time.monotonic() - start, + path, + ) + return + + shard_path = path.with_suffix(path.suffix + f".tp{tp_rank}.part") + local_records = {} + local_failed = 0 + fatal = None + try: + with _temporary_tuning_kernels(device): + for workload in local_workloads: + try: + record = _tune_workload( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + workload, + sliding_window=sliding_window, + causal=causal, + has_sinks=has_sinks, + physical_seq_len=( + physical_max_len + if has_sinks and sliding_window >= 0 + else None + ), + ) + except ( + torch.OutOfMemoryError, + RuntimeError, + AssertionError, + ) as error: + local_failed += 1 + logger.warning( + "Skipping segmented attention tuning workload %s: %s", + workload, + error, + ) + else: + local_records[workload] = record + _save( + shard_path, + { + "identity": identity, + "records": list(local_records.values()), + }, + ) + finally: + gc.collect() + torch.accelerator.empty_cache() + except BaseException as error: + fatal = f"{type(error).__name__}: {error}" + + gathered = _gather_tp( + group, + { + "records": list(local_records.values()), + "failed": local_failed, + "fatal": fatal, + }, + ) + if tp_rank == 0: + fatal_errors = [] + for rank, contribution in enumerate(gathered): + records.update( + (tuple(record["workload"]), record) + for record in contribution["records"] + ) + if contribution["fatal"] is not None: + fatal_errors.append(f"TP rank {rank}: {contribution['fatal']}") + data = {"identity": identity, "records": list(records.values())} + _save(path, data) + for shard in path.parent.glob(f"{path.name}.tp*.part"): + shard.unlink(missing_ok=True) + result = { + "data": data, + "tuned": sum(len(item["records"]) for item in gathered), + "failed": sum(item["failed"] for item in gathered), + "fatal": fatal_errors, + } + else: + result = None + if group is not None: + result = group.broadcast_object(result, src=0) + assert result is not None + data = result["data"] + finally: + if lock is not None: + lock.release() + + _TABLES[key] = data + logger.info( + "ROCm segmented attention autotune ready: tp_rank=%d/%d " + "local_tuned=%d tuned=%d loaded=%d failed=%d elapsed=%.2fs cache=%s", + tp_rank, + tp_size, + len(local_records), + result["tuned"], + loaded, + result["failed"], + time.monotonic() - start, + path, + ) + if result["fatal"]: + raise RuntimeError( + "ROCm segmented attention distributed tuning failed: " + + "; ".join(result["fatal"]) + ) + + +def get_segmented_config( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + batch, + query_len, + seq_len, + sliding_window=-1, + causal=True, + has_sinks=False, +): + """Return a warmed ceiling bucket, or None for the static fallback.""" + data = _TABLES.get( + _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + ) + if data is None: + return None + records = data["records"] + eligible = [ + record + for record in records + if record["workload"][0] >= batch + and record["workload"][1] >= query_len + and record["workload"][2] >= seq_len + ] + if not eligible: + return None + winner = min( + eligible, + key=lambda record: ( + record["workload"][1], + record["workload"][0], + record["workload"][2], + ), + ) + return dict(winner["best"]) + + +def _warmup_segmented_attention(layer, device, **limits) -> None: + impl = layer.impl + config = impl._segmented_attention_config + assert config is not None + spec = layer.get_kv_cache_spec(config) + assert spec is not None + sliding_window = impl.sliding_window[0] + query_limit = getattr(layer, "segmented_query_limit", None) + if not isinstance(query_limit, int): + query_limit = MAX_QUERY_LEN + causal = getattr(layer, "segmented_causal", True) + if not isinstance(causal, bool): + causal = True + tuning_max_len = config.model_config.max_model_len + if sliding_window >= 0: + tuning_max_len = min(tuning_max_len, sliding_window + query_limit) + warmup_segmented_attention( + device, + config.model_config.dtype, + impl.num_heads, + impl.num_kv_heads, + impl.head_size, + spec.block_size, + impl.scale, + config.scheduler_config.max_num_batched_tokens, + tuning_max_len, + config.scheduler_config.max_num_seqs, + kv_dtype=spec.dtype if spec.dtype != torch.uint8 else impl.fp8_dtype, + sliding_window=sliding_window, + causal=causal, + has_sinks=impl.sinks is not None, + physical_max_len=config.model_config.max_model_len, + max_query_len=query_limit, + **limits, + ) + impl._segmented_attention_warmed_up = True + + +def rocm_segmented_attn_autotune_warmup(worker: Worker) -> None: + """Autotune RDNA segmented attention after KV allocation and before capture.""" + if ( + not current_platform.is_rocm() + or not worker.vllm_config.kernel_config.enable_rocm_segmented_attn_autotune + ): + return + + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec + + if not on_gfx1x(): + return + layers = [] + for layer in worker.vllm_config.compilation_config.static_forward_context.values(): + impl = getattr(layer, "impl", None) + if not isinstance(impl, RocmSegmentedAttentionImpl): + continue + config = impl._segmented_attention_config + if ( + config is not None + and config.kernel_config.enable_rocm_segmented_attn_autotune + and not impl._segmented_attention_warmed_up + and impl.alibi_slopes is None + and not impl.logits_soft_cap + ): + layers.append(layer) + if not layers: + return + + specs = worker.get_kv_cache_spec() + layouts = tuple( + ( + (spec.block_size, spec.page_size_bytes) + if isinstance(spec, FullAttentionSpec) + else (0, spec.max_memory_usage_bytes(worker.vllm_config)) + ) + for spec in specs.values() + ) + cache_budget = sum( + tensor.size + for tensor in worker.model_runner.kv_cache_config.kv_cache_tensors + if not tensor.host_resident + ) + budget = _memory_budget(worker.device) + for layer in layers: + _warmup_segmented_attention( + layer, + worker.device, + memory_budget_bytes=budget, + cache_layouts=layouts, + cache_budget_bytes=cache_budget, + ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 32e2fabeb5..f28a331854 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -53,6 +53,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend" ) ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" + ROCM_SEGMENTED_ATTN = ( + "vllm.v1.attention.backends.rocm_segmented_attn.RocmSegmentedAttentionBackend" + ) ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" ROCM_AITER_TRITON_MLA = ( "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend" diff --git a/vllm/v1/attention/backends/rocm_segmented_attn.py b/vllm/v1/attention/backends/rocm_segmented_attn.py new file mode 100644 index 0000000000..c541eed451 --- /dev/null +++ b/vllm/v1/attention/backends/rocm_segmented_attn.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm attention backend for token-major segmented Triton attention.""" + +import os +from pathlib import Path +from typing import TYPE_CHECKING, ClassVar + +import torch + +import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.backend import AttentionLayer, AttentionType +from vllm.v1.attention.backends.rocm_attn import ( + RocmAttentionBackend, + RocmAttentionImpl, + RocmAttentionMetadata, + RocmAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import get_num_attention_heads_from_layers +from vllm.v1.attention.ops.segmented_attention import ( + get_segmented_attention_workspace, + segmented_attention, + segmented_attention_workspace_size, +) +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +if TYPE_CHECKING: + from vllm.platforms.interface import DeviceCapability + +logger = init_logger(__name__) + + +def is_rdna() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx1x + + return on_gfx1x() + + +class RocmSegmentedAttentionMetadataBuilder(RocmAttentionMetadataBuilder): + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + model_config = vllm_config.model_config + self.num_heads_q = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or model_config.get_num_attention_heads(vllm_config.parallel_config) + self.num_heads_kv = kv_cache_spec.num_kv_heads + self.headdim = kv_cache_spec.head_size + + +class RocmSegmentedAttentionBackend(RocmAttentionBackend): + """Explicit token-major ROCm backend optimized for segmented prefill.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + @staticmethod + def get_name() -> str: + return "ROCM_SEGMENTED_ATTN" + + @staticmethod + def get_impl_cls() -> type["RocmSegmentedAttentionImpl"]: + return RocmSegmentedAttentionImpl + + @staticmethod + def get_builder_cls() -> type["RocmSegmentedAttentionMetadataBuilder"]: + return RocmSegmentedAttentionMetadataBuilder + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [64, 128, 256] + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + return attn_type == AttentionType.DECODER + + @classmethod + def supports_non_causal(cls) -> bool: + return True + + @classmethod + def supports_sliding_window(cls) -> bool: + return True + + @classmethod + def supports_sink(cls) -> bool: + return True + + @classmethod + def supports_mm_prefix(cls) -> bool: + return False + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: "DeviceCapability", + ) -> str | None: + del ( + head_size, + dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ) + if not is_rdna(): + return "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm" + + from vllm.platforms.rocm import on_gfx12x + + if kv_cache_dtype in ("fp8", "fp8_e4m3") and not on_gfx12x(): + return "FP8 segmented attention requires gfx12" + return None + + +class RocmSegmentedAttentionImpl(RocmAttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: int | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + if not is_rdna(): + raise ValueError("ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm") + + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + sinks, + ) + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + fp8 = kv_cache_dtype in ("fp8", "fp8_e4m3") + arch_supported = on_gfx12x() if fp8 else on_gfx1x() + gqa_ratio_supported = ( + num_kv_heads > 0 + and num_heads % num_kv_heads == 0 + and 1 <= self.num_queries_per_kv <= 16 + ) + if not ( + arch_supported + and attn_type == AttentionType.DECODER + and head_size in (64, 128, 256) + and gqa_ratio_supported + and kv_cache_dtype in ("auto", "float16", "bfloat16", "fp8", "fp8_e4m3") + and alibi_slopes is None + ): + raise ValueError( + "ROCM_SEGMENTED_ATTN requires decoder attention on gfx1x " + "(gfx12 for FP8), head size 64, 128 or 256, GQA ratio 1-16, and " + "does not support ALiBi." + ) + # Startup tuning and target/draft compilation must share a Triton cache. + cache_dir = os.environ.setdefault( + "TRITON_CACHE_DIR", + str(Path(envs.VLLM_CACHE_ROOT) / "rocm_segmented_attention" / "triton"), + ) + Path(cache_dir).mkdir(parents=True, exist_ok=True) + logger.info_once("Using token-major ROCm segmented Triton attention") + self._segmented_attention_warmed_up = False + self._segmented_attention_config: VllmConfig | None = None + + def process_weights_after_loading(self, act_dtype: torch.dtype): + super().process_weights_after_loading(act_dtype) + config = get_current_vllm_config() + self._segmented_attention_config = config + + def _autotune_enabled(self) -> bool: + config = self._segmented_attention_config + return ( + config is not None + and config.kernel_config.enable_rocm_segmented_attn_autotune + ) + + def _get_workspace(self, device): + config = self._segmented_attention_config + assert config is not None + sizes = segmented_attention_workspace_size( + config.scheduler_config.max_num_seqs, + self.num_heads, + self.num_kv_heads, + self.head_size, + config.model_config.max_model_len, + max_tokens=config.scheduler_config.max_num_batched_tokens, + fp8=self.kv_cache_dtype in ("fp8", "fp8_e4m3"), + autotune=self._autotune_enabled(), + ) + return get_segmented_attention_workspace(device, sizes) + + def _split_kv_cache( + self, + kv_cache: torch.Tensor, + num_kv_heads: int | None = None, + head_size: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_kv_heads = num_kv_heads or self.num_kv_heads + head_size = head_size or self.head_size + return ( + kv_cache[:, 0].unflatten(-1, (num_kv_heads, head_size)), + kv_cache[:, 1].unflatten(-1, (num_kv_heads, head_size)), + ) + + def forward( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: RocmAttentionMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + if output_block_scale is not None: + raise NotImplementedError( + "fused block_scale output quantization is not supported" + ) + # Profile runs allocate the same fixed scratch used during graph replay. + workspace = self._get_workspace(query.device) + if attn_metadata is None: + return output.fill_(0) + assert not attn_metadata.use_cascade + + num_actual_tokens = attn_metadata.num_actual_tokens + key_cache, value_cache = self._split_kv_cache(kv_cache) + if is_quantized_kv_cache(self.kv_cache_dtype): + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0: + raise NotImplementedError( + "A non-1.0 q_scale with an FP8 query is not supported" + ) + + segmented_attention( + query=query[:num_actual_tokens], + key=key[:num_actual_tokens], + value=value[:num_actual_tokens], + output=output[:num_actual_tokens], + kv_cache_dtype=self.kv_cache_dtype, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + query_start_loc=attn_metadata.query_start_loc, + seq_lens=attn_metadata.seq_lens, + max_seq_len=attn_metadata.max_seq_len, + max_query_len=attn_metadata.max_query_len, + k_scale=layer._k_scale, + v_scale=layer._v_scale, + sm_scale=self.scale, + sliding_window=self.sliding_window[0], + output_scale=output_scale, + sinks=self.sinks, + causal=attn_metadata.causal, + softcap=self.logits_soft_cap, + workspace=workspace, + use_tuned_config=self._autotune_enabled(), + ) + return output + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ): + key_cache, value_cache = self._split_kv_cache(kv_cache) + if is_quantized_kv_cache(self.kv_cache_dtype): + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + triton_reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + def do_rope_and_kv_cache_update( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + key_cache, value_cache = self._split_kv_cache( + kv_cache, + layer.num_kv_heads, # type: ignore[attr-defined] + layer.head_size, # type: ignore[attr-defined] + ) + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) + if is_fp8_kv_cache: + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + + rocm_aiter_ops.triton_rope_and_cache( + query, + key, + value, + positions, + cos_sin_cache, + is_neox, + key_cache, + value_cache, + layer_slot_mapping, + layer._k_scale, + layer._v_scale, + True, + is_fp8_kv_cache, + ) diff --git a/vllm/v1/attention/ops/segmented_attention.py b/vllm/v1/attention/ops/segmented_attention.py new file mode 100644 index 0000000000..f98b07df2a --- /dev/null +++ b/vllm/v1/attention/ops/segmented_attention.py @@ -0,0 +1,1074 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ROCm token-major segmented attention kernels and dispatch.""" + +from __future__ import annotations + +from functools import lru_cache + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.kv_cache_interface import get_kv_quant_mode +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +from .triton_unified_attention import unified_attention + +MAX_QUERY_LEN = 8192 +MAX_SPLITS = 64 +# D128 scratch cap; wider outputs scale this budget with head dimension. +MAX_LONG_EXTEND_WORKSPACE_BYTES = 32 * 1024 * 1024 + + +def segmented_query_capacity(max_query_len: int) -> int: + """Round a supported query length to its compiled power-of-two capacity.""" + if not 1 <= max_query_len <= MAX_QUERY_LEN: + raise ValueError(f"Query length must be in [1, {MAX_QUERY_LEN}]") + return 1 << (max_query_len - 1).bit_length() + + +def _query_capacity_buckets() -> tuple[int, ...]: + return tuple(1 << exponent for exponent in range(MAX_QUERY_LEN.bit_length())) + + +def _long_extend_splits( + base_splits: int, + batch: int, + query_len: int, + seq_len: int, + hq: int, + hk: int, + dim: int, + bm: int, +) -> int: + """Add long-prefix parallelism while bounding the extra split workspace.""" + if seq_len < 32768: + return base_splits + groups = batch * hk * triton.cdiv(query_len * (hq // hk), bm) + split_cap = 16 if seq_len < 131072 else 32 + scratch_per_split = batch * segmented_query_capacity(query_len) * hq * (dim + 1) * 4 + scratch_limit = MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + scratch_splits = scratch_limit // scratch_per_split + scratch_cap = 1 << (scratch_splits.bit_length() - 1) if scratch_splits else 1 + occupancy_splits = triton.next_power_of_2(triton.cdiv(512, groups)) + return max(base_splits, min(split_cap, scratch_cap, occupancy_splits)) + + +@triton.jit +def _segmented_attention_stage( + Q, + K, + V, + TABLE, + STARTS, + LENS, + KS, + VS, + SINKS, + PART, + LSE, + OUT, + Q0: tl.constexpr, + Q1: tl.constexpr, + O0: tl.constexpr, + O1: tl.constexpr, + K0: tl.constexpr, + K1: tl.constexpr, + K2: tl.constexpr, + K3: tl.constexpr, + V0: tl.constexpr, + V1: tl.constexpr, + V2: tl.constexpr, + V3: tl.constexpr, + T0: tl.constexpr, + T1: tl.constexpr, + PAGE: tl.constexpr, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + QCAP: tl.constexpr, + LIMIT_Q: tl.constexpr, + SCALE: tl.constexpr, + FP8: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + SPLITS: tl.constexpr, + WINDOW: tl.constexpr, + CAUSAL: tl.constexpr, + HAS_SINKS: tl.constexpr, + PREFIX_FAST: tl.constexpr, + PV_SPLIT: tl.constexpr, + QK_PIPELINE: tl.constexpr, +): + item, mb, split = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq, kh = item // HK, item % HK + first = tl.load(STARTS + seq) + nq = tl.load(STARTS + seq + 1) - first + if nq < 1 or nq > LIMIT_Q: + return + G: tl.constexpr = HQ // HK + ROWS: tl.constexpr = QCAP * G + if mb * BM >= nq * G: + return + sequence = tl.load(LENS + seq) + prefix = sequence - nq + m = mb * BM + tl.arange(0, BM) + qpos, qh = m // G, kh * G + m % G + n = tl.arange(0, BN) + d = tl.arange(0, D) + if BK == D: + q_full = tl.load( + Q + (first + qpos[:, None]) * Q0 + qh[:, None] * Q1 + d[None, :], + mask=(m < nq * G)[:, None], + other=0.0, + ) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denom = tl.zeros((BM,), tl.float32) + if HAS_SINKS: + sink = tl.load(SINKS + qh, mask=m < nq * G, other=0).to(tl.float32) + maximum = tl.where(split == 0, sink * 1.4426950408889634, maximum) + denom = tl.where(split == 0, 1.0, denom) + if PV_SPLIT: + half = tl.arange(0, D // 2) + acc_low = tl.zeros((BM, D // 2), tl.float32) + acc_high = tl.zeros((BM, D // 2), tl.float32) + else: + acc = tl.zeros((BM, D), tl.float32) + tile_rows = tl.minimum((mb + 1) * BM, nq * G) + range_begin = 0 + if WINDOW >= 0: + first_qpos = (mb * BM) // G + range_begin = tl.maximum(0, prefix + first_qpos - WINDOW) + range_begin = (range_begin // BN) * BN + range_end = sequence + if CAUSAL: + range_end = tl.minimum(range_end, prefix + tl.cdiv(tile_rows, G)) + span = tl.maximum(0, range_end - range_begin) + segment = tl.cdiv(span, SPLITS * BN) * BN + begin = range_begin + split * segment + end = tl.minimum(begin + segment, range_end) + k_scale = 1.0 + v_scale = 1.0 + if FP8: + k_scale = tl.load(KS) + v_scale = tl.load(VS) + for tail in tl.static_range(2): + if tail: + loop_begin = tl.maximum(begin, (end // BN) * BN) + loop_end = end + else: + loop_begin = begin + loop_end = (end // BN) * BN + for start in tl.range(loop_begin, loop_end, BN): + ns = start + n + valid = tl.full((BN,), True, tl.int1) + if tail: + valid = ns < end + if FP8 and PAGE >= BN: + page = start // PAGE + inside = start - page * PAGE + n + crossed = inside >= PAGE + page += crossed + inside = tl.where(crossed, inside - PAGE, inside) + else: + page = ns // PAGE + inside = ns % PAGE + block = tl.load(TABLE + seq * T0 + page * T1, mask=valid, other=0).to( + tl.int64 + ) + scores = tl.zeros((BM, BN), tl.float32) + for ki in tl.range(D // BK, num_stages=QK_PIPELINE): + kd = ki * BK + tl.arange(0, BK) + if BK == D: + q = q_full + else: + q = tl.load( + Q + + (first + qpos[:, None]) * Q0 + + qh[:, None] * Q1 + + kd[None, :], + mask=(m < nq * G)[:, None], + other=0.0, + ) + k = tl.load( + K + + block[None, :] * K0 + + inside[None, :] * K1 + + kh * K2 + + kd[:, None] * K3, + mask=valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k.to(q.dtype), scores) + scores *= SCALE * 1.4426950408889634 * k_scale + # A complete prefix tile is visible to every valid causal row when + # sliding-window masking is disabled. + if not ( + PREFIX_FAST + and CAUSAL + and WINDOW < 0 + and start + BN <= prefix + and mb * BM + BM <= nq * G + ): + score_mask = valid[None, :] & (m[:, None] < nq * G) + if CAUSAL: + score_mask = score_mask & (ns[None, :] <= prefix + qpos[:, None]) + if WINDOW >= 0: + score_mask = score_mask & ( + ns[None, :] >= prefix + qpos[:, None] - WINDOW + ) + scores = tl.where(score_mask, scores, -float("inf")) + new_max = tl.maximum(maximum, tl.max(scores, 1)) + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + alpha = tl.exp2(maximum - safe_max) + p = tl.exp2(scores - safe_max[:, None]) + if PV_SPLIT: + acc_low *= alpha[:, None] + acc_high *= alpha[:, None] + v_low = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + half[None, :] * V3, + mask=valid[:, None], + other=0.0, + ) + v_high = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + (half[None, :] + D // 2) * V3, + mask=valid[:, None], + other=0.0, + ) + prob = p.to(Q.dtype.element_ty) + acc_low = tl.dot(prob, v_low.to(prob.dtype), acc_low) + acc_high = tl.dot(prob, v_high.to(prob.dtype), acc_high) + else: + acc *= alpha[:, None] + v = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + d[None, :] * V3, + mask=valid[:, None], + other=0.0, + ) + if FP8 and D == 256 and BM >= 32: + # Defer the scalar V dequantization to shorten the D-wide + # accumulator live range for the register-heavy 256-D tiles. + acc = tl.dot( + p.to(Q.dtype.element_ty), v.to(Q.dtype.element_ty), acc + ) + elif FP8: + acc += ( + tl.dot(p.to(Q.dtype.element_ty), v.to(Q.dtype.element_ty)) + * v_scale + ) + else: + acc = tl.dot(p.to(Q.dtype.element_ty), v, acc) + denom = denom * alpha + tl.sum(p, 1) + maximum = safe_max + norm = tl.where(denom > 0.0, denom, 1.0) + if PV_SPLIT: + result = tl.cat(acc_low, acc_high, dim=1) / norm[:, None] + else: + result = acc / norm[:, None] + if FP8 and D == 256 and BM >= 32: + result *= v_scale + if SPLITS == 1: + tl.store( + OUT + (first + qpos[:, None]) * O0 + qh[:, None] * O1 + d[None, :], + result, + mask=(m < nq * G)[:, None], + ) + else: + row = ((seq * HK + kh) * ROWS + m) * SPLITS + split + tl.store( + PART + row[:, None] * D + d[None, :], + result, + mask=(m < nq * G)[:, None], + ) + tl.store( + LSE + row, + maximum + tl.log2(denom), + mask=m < nq * G, + ) + + +@triton.jit +def _segmented_attention_reduce( + PART, + LSE, + OUT, + STARTS, + O0: tl.constexpr, + O1: tl.constexpr, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + QCAP: tl.constexpr, + LIMIT_Q: tl.constexpr, + SPLITS: tl.constexpr, + REDUCE_D: tl.constexpr, +): + m, item, db = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq, kh = item // HK, item % HK + first = tl.load(STARTS + seq) + nq = tl.load(STARTS + seq + 1) - first + G: tl.constexpr = HQ // HK + ROWS: tl.constexpr = QCAP * G + if nq < 1 or nq > LIMIT_Q or m >= nq * G: + return + s = tl.arange(0, SPLITS) + d = db * REDUCE_D + tl.arange(0, REDUCE_D) + row = ((seq * HK + kh) * ROWS + m) * SPLITS + lse = tl.load(LSE + row + s) + weights = tl.exp2(lse - tl.max(lse, 0)) + part = tl.load( + PART + (row + s)[:, None] * D + d[None, :], + mask=d[None, :] < D, + other=0.0, + ) + result = tl.sum(part * weights[:, None], 0) / tl.sum(weights, 0) + tl.store( + OUT + (first + m // G) * O0 + (kh * G + m % G) * O1 + d, + result, + mask=d < D, + ) + + +@lru_cache(maxsize=512) +def select_segmented_config(batch, max_query_len, max_seq_len, hq, hk, dim, fp8): + """Select the token-major cache configuration for segmented attention.""" + qcap = min(max_query_len, MAX_QUERY_LEN) + if batch < 1 or qcap < 1: + raise ValueError("Batch and query capacity must be positive") + + if qcap <= 2: + if dim == 64 and qcap == 1 and batch == 1 and max_seq_len >= 131072: + return dict( + bm=16, + bn=32, + bk=64, + splits=256, + warps=1, + stages=1, + ) + if dim == 64 and max_seq_len <= 256: + return dict( + bm=16, + bn=32 if fp8 else 64, + bk=64, + splits=1, + warps=4, + stages=1, + ) + if dim == 64 and max_seq_len >= 32768: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 16) + return dict( + bm=16, + bn=32, + bk=64, + splits=min(64, triton.next_power_of_2(triton.cdiv(128, groups))), + warps=4, + stages=1, + ) + if batch * hk == 1: + cfg = dict( + bm=16, + bn=64, + bk=min(128, dim), + splits=32, + warps=4, + stages=1, + ) + else: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 16) + cfg = dict( + bm=16, + bn=32, + bk=dim, + splits=min( + MAX_SPLITS, + triton.next_power_of_2(triton.cdiv(64, groups)), + ), + warps=4, + stages=1, + ) + while cfg["splits"] > 1 and max_seq_len < cfg["splits"] * cfg["bn"] * 2: + cfg["splits"] //= 2 + return cfg + + if qcap <= 8: + bm, bn = 16, 64 + bk, stages = 64, 1 + elif qcap <= 32: + bm, bk = 32, 64 + bn, stages = (32, 2) if fp8 else (64, 1) + else: + bm, bn, bk = 32, 64, 64 + stages = 1 if fp8 else 2 + groups = batch * hk * triton.cdiv(qcap * (hq // hk), bm) + target = 96 + if fp8 and qcap > 32: + target = 192 + splits = min(32, triton.next_power_of_2(triton.cdiv(target, groups))) + while splits > 1 and max_seq_len < splits * 128: + splits //= 2 + cfg = dict( + bm=bm, + bn=bn, + bk=bk, + splits=splits, + warps=4, + stages=stages, + ) + if dim == 64 and qcap >= 64 and qcap < max_seq_len <= qcap + 256: + cfg.update( + bm=64 if fp8 or qcap <= 256 else 128, + bn=32, + bk=64, + splits=1, + warps=4, + stages=1, + ) + elif dim == 64 and fp8 and max_seq_len <= 256: + cfg.update(bn=32, splits=1) + elif dim == 64 and qcap == 8 and max_seq_len >= 32768: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 32) + cfg.update( + bm=32, + bn=32 if fp8 else 64, + splits=min(64, triton.next_power_of_2(triton.cdiv(256, groups))), + ) + elif dim == 64 and fp8 and qcap == 8 and max_seq_len >= 4096: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 32) + cfg.update( + bm=32, + bn=64, + splits=min(32, triton.next_power_of_2(triton.cdiv(128, groups))), + ) + elif dim == 64 and fp8 and qcap >= 64 and max_seq_len >= 4096: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 64) + cfg.update( + bm=64, + bn=32, + splits=min(8, triton.next_power_of_2(triton.cdiv(1024, groups))), + stages=1, + ) + elif dim == 128 and fp8 and qcap <= 256 and max_seq_len <= qcap: + cfg.update(splits=1) + elif ( + dim == 128 + and fp8 + and qcap >= 256 + and max_seq_len >= 8192 + and ( + hq // hk >= 4 + or (max_seq_len >= 32768 and (hq // hk >= 2 or qcap >= 512)) + or max_seq_len >= 131072 + ) + ): + # Shorter prefixes need sufficient grouped-query reuse; at GQA1/Q256 + # the wider tile pays off only once the prefix reaches 131K tokens. + cfg.update( + bm=64, + bn=128, + bk=64, + warps=8, + stages=3, + waves_per_eu=6, + prefix_fast=True, + ) + if qcap >= 512 or max_seq_len >= 32768: + cfg["qk_pipeline"] = 3 + cfg["splits"] = _long_extend_splits( + cfg["splits"], batch, qcap, max_seq_len, hq, hk, dim, cfg["bm"] + ) + elif dim == 256 and fp8 and qcap >= 256 and max_seq_len >= 8192: + # Wider tiles amortize each FP8 KV load over twice as many queries and + # keys. Six waves is the best gfx1201 pressure/occupancy tradeoff. + cfg.update( + bm=64, + bn=128, + bk=64, + warps=8, + stages=1, + waves_per_eu=6, + ) + cfg["splits"] = _long_extend_splits( + cfg["splits"], batch, qcap, max_seq_len, hq, hk, dim, cfg["bm"] + ) + elif dim == 256 and not fp8 and qcap >= 1024 and max_seq_len >= 8192: + # Two half-width PV accumulators avoid spilling this BF16 D256 tile. + # Keep Q256 on the incumbent: the wider tile loses at that capacity. + cfg.update( + bm=64, + bn=64, + bk=64, + warps=8, + stages=3, + waves_per_eu=6, + pv_split=True, + qk_pipeline=3, + ) + elif dim == 128 and not fp8 and qcap > 128: + if max_seq_len >= 8192: + cfg.update( + bm=128, + bn=32, + bk=128, + warps=8, + stages=1, + waves_per_eu=6, + ) + groups = batch * hk * triton.cdiv(qcap * (hq // hk), cfg["bm"]) + cfg["splits"] = min(16, triton.next_power_of_2(triton.cdiv(192, groups))) + elif qcap >= 4096: + cfg.update( + bm=128, + bn=32, + bk=128, + splits=1, + warps=8, + stages=1, + waves_per_eu=6, + ) + elif qcap > 256 or hq > 10: + cfg.update( + bm=64, + bn=32, + bk=128, + splits=1, + warps=4, + stages=1, + waves_per_eu=6, + ) + else: + cfg.update(bm=32, bn=64, bk=64, splits=1, warps=4, stages=2) + return cfg + + +def segmented_workspace_shapes(batch, query_cap, hq, hk, dim, splits): + if splits == 1: + return None + partial = (batch, hk, query_cap * (hq // hk), splits, dim) + return partial, partial[:-1] + + +@lru_cache(maxsize=128) +def segmented_attention_workspace_size( + max_batch, + hq, + hk, + dim, + max_seq_len, + *, + max_tokens=None, + fp8=False, + autotune=False, +): + """Size the largest selected attention workspace before graph capture. + + ``max_tokens`` bounds reachable ``(batch, query_len)`` pairs using one + longest query and one token for each remaining sequence. Omitting it keeps + the legacy reservation behavior for callers without scheduler limits. + ``autotune`` reserves additional splits for tuned D64 candidates. + """ + largest = 0 + previous_capacity = 0 + for query_capacity in _query_capacity_buckets(): + query_lengths = {previous_capacity + 1, query_capacity} + for query_len in query_lengths: + if max_tokens is not None and query_len > max_tokens: + continue + batch_limit = max_batch + if max_tokens is not None: + batch_limit = min(batch_limit, max_tokens - query_len + 1) + for batch in range(1, batch_limit + 1): + cfg = select_segmented_config( + batch, query_len, max_seq_len, hq, hk, dim, fp8 + ) + splits = cfg["splits"] + if dim == 64 and query_len <= 512 and max_seq_len >= 4096 and autotune: + split_limit = ( + 512 + if batch == 1 and query_len == 1 and max_seq_len >= 131072 + else MAX_SPLITS + ) + splits = min(split_limit, 2 * splits) + if splits > 1: + largest = max( + largest, + splits * batch * query_capacity * hq, + ) + previous_capacity = query_capacity + return largest * dim, largest + + +def get_segmented_attention_workspace(device, sizes): + """Keep scratch storage stable and separate for each runtime lane/ubatch.""" + + def allocate(): + return tuple( + torch.empty(size, device=device, dtype=torch.float32) for size in sizes + ) + + if is_workspace_manager_initialized(): + return current_workspace_manager().get_persistent_resource( + ("rocm_segmented_attention", device, sizes), allocate + ) + return allocate() + + +def can_use_segmented_attention( + q, k, v, out, kc, vc, table, starts, lengths, k_scale, v_scale +): + """Check tensor metadata; unsupported feature checks live in the caller.""" + if ( + k is None + or v is None + or q.ndim != 3 + or kc.ndim != 4 + or vc.ndim != 4 + or q.dtype not in (torch.bfloat16, torch.float16) + or out.dtype != q.dtype + or q.shape != out.shape + or q.shape[-1] not in (64, 128, 256) + or kc.dtype != vc.dtype + or k.dtype != q.dtype + or v.dtype != q.dtype + or kc.dtype not in (q.dtype, torch.float8_e4m3fn, torch.float8_e4m3fnuz) + or table.ndim != 2 + or starts.ndim != 1 + or lengths.ndim != 1 + or table.shape[0] != lengths.numel() + or starts.numel() != lengths.numel() + 1 + or table.dtype != torch.int32 + or starts.dtype != torch.int32 + or lengths.dtype != torch.int32 + or starts.stride(0) != 1 + or lengths.stride(0) != 1 + or not q.is_cuda + ): + return False + hkv = kc.shape[2] + dim = q.shape[-1] + page = kc.shape[1] + if hkv < 1 or q.shape[1] % hkv or not 1 <= q.shape[1] // hkv <= 16: + return False + if ( + k.shape != (q.shape[0], hkv, dim) + or v.shape != k.shape + or kc.shape != (kc.shape[0], page, hkv, dim) + or vc.shape != kc.shape + or any( + t.device != q.device for t in (k, v, out, kc, vc, table, starts, lengths) + ) + or any(t.stride(-1) != 1 for t in (q, k, v, out, kc, vc)) + ): + return False + if kc.element_size() == 1: + return all( + isinstance(s, torch.Tensor) + and s.numel() == 1 + and s.dtype == torch.float32 + and s.device == q.device + for s in (k_scale, v_scale) + ) + return True + + +def run_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + max_query_len, + max_seq_len, + k_scale, + v_scale, + scale, + *, + sliding_window=-1, + causal=True, + sinks=None, + config=None, + workspace=None, +): + """Write eligible unified-cache attention rows; leave other requests untouched.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + if batch == 0 or max_query_len == 0 or q.numel() == 0: + return out + query_len = min(max_query_len, MAX_QUERY_LEN) + qcap = segmented_query_capacity(query_len) + fp8 = kc.element_size() == 1 + hk = kc.shape[2] + attention_span = max_seq_len + if sliding_window >= 0: + attention_span = min(max_seq_len, sliding_window + max_query_len) + cfg = config or select_segmented_config( + batch, query_len, attention_span, hq, hk, dim, fp8 + ) + splits = cfg["splits"] + shapes = segmented_workspace_shapes(batch, qcap, hq, hk, dim, splits) + if shapes is None: + partial = lse = out + elif workspace is not None: + partial, lse = ( + buffer.view(-1)[: shape.numel()].view(shape) + for buffer, shape in zip(workspace, map(torch.Size, shapes)) + ) + elif is_workspace_manager_initialized(): + partial, lse = current_workspace_manager().get_simultaneous( + (shapes[0], torch.float32), (shapes[1], torch.float32) + ) + else: + partial = torch.empty(shapes[0], device=q.device, dtype=torch.float32) + lse = torch.empty(shapes[1], device=q.device, dtype=torch.float32) + _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + cfg, + partial, + lse, + qcap, + sliding_window, + causal, + compile_only=False, + ) + return out + + +def compile_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + max_query_len, + k_scale, + v_scale, + scale, + config, + workspace, + *, + sliding_window=-1, + causal=True, + sinks=None, +): + """Compile one segmented attention configuration without launching it.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + qcap = segmented_query_capacity(min(max_query_len, MAX_QUERY_LEN)) + hk = kc.shape[2] + shapes = segmented_workspace_shapes(batch, qcap, hq, hk, dim, config["splits"]) + if shapes is None: + partial = lse = out + else: + partial, lse = workspace + _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + config, + partial, + lse, + qcap, + sliding_window, + causal, + compile_only=True, + ) + + +def _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + cfg, + partial, + lse, + qcap, + sliding_window, + causal, + *, + compile_only, +): + """Launch or compile the exact stage and reduction specializations.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + hk = kc.shape[2] + fp8 = kc.element_size() == 1 + splits = cfg["splits"] + stage_grid = ( + batch * hk, + triton.cdiv(qcap * (hq // hk), cfg["bm"]), + splits, + ) + stage_args = ( + q, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + partial, + lse, + out, + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + *kc.stride(), + *vc.stride(), + *table.stride(), + kc.shape[1], + hq, + hk, + dim, + qcap, + MAX_QUERY_LEN, + scale, + fp8, + cfg["bm"], + cfg["bn"], + cfg["bk"], + splits, + sliding_window, + causal, + sinks is not None, + cfg.get("prefix_fast", False), + cfg.get("pv_split", False), + cfg.get("qk_pipeline"), + ) + stage_options = dict( + num_warps=cfg["warps"], + num_stages=cfg["stages"], + waves_per_eu=cfg.get("waves_per_eu", 2), + ) + if compile_only: + _segmented_attention_stage.warmup(*stage_args, grid=stage_grid, **stage_options) + else: + _segmented_attention_stage[stage_grid](*stage_args, **stage_options) + if splits > 1: + reduce_d = cfg.get("reduce_d", 64 if batch * qcap * hq < 64 else dim) + reduce_grid = ( + qcap * (hq // hk), + batch * hk, + triton.cdiv(dim, reduce_d), + ) + reduce_args = ( + partial, + lse, + out, + starts, + out.stride(0), + out.stride(1), + hq, + hk, + dim, + qcap, + MAX_QUERY_LEN, + splits, + reduce_d, + ) + reduce_options = {"num_warps": cfg.get("reduce_warps", 4)} + if compile_only: + _segmented_attention_reduce.warmup( + *reduce_args, grid=reduce_grid, **reduce_options + ) + else: + _segmented_attention_reduce[reduce_grid](*reduce_args, **reduce_options) + + +def get_segmented_config(*args, **kwargs): + """Load the tuner on demand to avoid a kernel/tuner import cycle.""" + from vllm.model_executor.warmup.rocm_segmented_attn_autotune_warmup import ( + get_segmented_config as tuned_config, + ) + + return tuned_config(*args, **kwargs) + + +def segmented_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + kv_cache_dtype: str, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + max_query_len: int, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + *, + sliding_window: int = -1, + output_scale: torch.Tensor | None = None, + sinks: torch.Tensor | None = None, + causal: bool | torch.Tensor = True, + softcap: float = 0.0, + workspace=None, + use_tuned_config: bool = True, +) -> None: + """Run segmented attention when eligible, otherwise use unified Triton.""" + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + fp8_dtype = current_platform.fp8_dtype() + if key_cache.dtype == torch.uint8: + key_cache = key_cache.view(fp8_dtype) + value_cache = value_cache.view(fp8_dtype) + + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + segmented_pattern = ( + isinstance(causal, bool) + and (on_gfx12x() if key_cache.element_size() == 1 else on_gfx1x()) + and not softcap + and output_scale is None + and ( + sinks is None + or ( + sinks.ndim == 1 + and sinks.shape[0] == query.shape[1] + and sinks.device == query.device + and sinks.dtype in (torch.float16, torch.bfloat16, torch.float32) + ) + ) + ) + if ( + segmented_pattern + and 0 < max_query_len <= MAX_QUERY_LEN + and ( + can_use_segmented_attention( + query, + key, + value, + output, + key_cache, + value_cache, + block_table, + query_start_loc, + seq_lens, + k_scale, + v_scale, + ) + ) + ): + attention_span = max_seq_len + if sliding_window >= 0: + attention_span = min(max_seq_len, sliding_window + max_query_len) + config = None + if use_tuned_config: + config = get_segmented_config( + query.device, + query.dtype, + key_cache.dtype, + query.shape[1], + key_cache.shape[2], + query.shape[2], + key_cache.shape[1], + sm_scale, + len(seq_lens), + max_query_len, + attention_span, + sliding_window, + causal, + has_sinks=sinks is not None, + ) + if config is None: + config = select_segmented_config( + len(seq_lens), + max_query_len, + attention_span, + query.shape[1], + key_cache.shape[2], + query.shape[2], + key_cache.element_size() == 1, + ) + run_segmented_attention( + query, + output, + key_cache, + value_cache, + block_table, + query_start_loc, + seq_lens, + max_query_len, + max_seq_len, + k_scale, + v_scale, + sm_scale, + sliding_window=sliding_window, + causal=causal, + sinks=sinks, + config=config, + workspace=workspace, + ) + return + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=query_start_loc, + max_seqlen_q=max_query_len, + seqused_k=seq_lens, + max_seqlen_k=max_seq_len, + softmax_scale=sm_scale, + causal=causal, + window_size=(-1, -1) if sliding_window < 0 else (sliding_window, 0), + block_table=block_table, + softcap=softcap, + q_descale=None, + k_descale=k_scale, + v_descale=v_scale, + output_scale=output_scale, + sinks=sinks, + kv_quant_mode=get_kv_quant_mode(kv_cache_dtype), + )