From 7cb4b59ddbdc4e0aaa92e42778e36a7dd8289ced Mon Sep 17 00:00:00 2001 From: tangzzycc <3081129260@qq.com> Date: Thu, 1 Oct 2026 12:59:38 +0800 Subject: [PATCH 1/2] [ROCm][Bugfix] Fix ROCM_ATTN sliding-window boundary Convert the stored left window span to the window length expected by the paged attention kernels. Add regression coverage for decode and chunked prefill. Signed-off-by: tangzzycc <3081129260@qq.com> --- .../kernels/attention/test_prefix_prefill.py | 58 +++++++++++++++++++ vllm/v1/attention/backends/rocm_attn.py | 2 +- 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index 290d645e0793..b2efa76da89a 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -5,6 +5,7 @@ import random import time from collections.abc import Callable +from types import SimpleNamespace import pytest import torch @@ -323,6 +324,63 @@ def test_contexted_kv_attention( torch.testing.assert_close(output, output_ref, atol=atol, rtol=0) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend") +@pytest.mark.parametrize("sliding_window", [None, 1, 4]) +@pytest.mark.parametrize("query_len", [1, 3]) +@torch.inference_mode() +def test_rocm_backend_sliding_window_includes_boundary_token( + sliding_window: int | None, query_len: int +) -> None: + """A window of W must include the current token and W - 1 preceding tokens.""" + from vllm.v1.attention.backends.rocm_attn import RocmAttentionImpl + from vllm.v1.attention.ops.paged_attn import PagedAttention + + device, dtype = "cuda:0", torch.float16 + num_heads, head_size, seq_len = 4, 128, 8 + query = torch.zeros(query_len, num_heads, head_size, dtype=dtype, device=device) + key = torch.zeros(seq_len, 1, head_size, dtype=dtype, device=device) + values = torch.arange(1, seq_len + 1, dtype=dtype, device=device) + value = values[:, None, None].expand_as(key) + kv_cache = torch.zeros(2, 1, 16, head_size, dtype=dtype, device=device) + _, value_cache = PagedAttention.split_kv_cache(kv_cache, 1, head_size) + value_cache[..., :seq_len] = value.permute(1, 2, 0) + scale = torch.tensor(1.0, device=device) + metadata = SimpleNamespace( + use_cascade=False, + num_actual_tokens=query_len, + query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device), + max_query_len=query_len, + max_seq_len=seq_len, + block_table=torch.zeros(1, 1, dtype=torch.int32, device=device), + causal=True, + ) + impl = RocmAttentionImpl( + num_heads=num_heads, + head_size=head_size, + scale=head_size**-0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=sliding_window, + kv_cache_dtype="auto", + ) + output = torch.empty_like(query) + impl.forward( + SimpleNamespace(_k_scale=scale, _v_scale=scale), + query, + key[-query_len:], + value[-query_len:], + kv_cache.transpose(0, 1), + metadata, + output, + ) + mask = create_causal_attention_mask_for_sdpa( + [query_len], [seq_len], sliding_window or 0, device=device, dtype=dtype + ) + expected = (mask.float().softmax(dim=-1) @ values.float())[:, None, None] + torch.testing.assert_close(output, expected.to(dtype).expand_as(output)) + + @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 453f49b46813..e79173bfb279 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -481,7 +481,7 @@ def forward( k_scale=layer._k_scale, v_scale=layer._v_scale, alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window[0], + sliding_window=1 + self.sliding_window[0], sm_scale=self.scale, output_scale=output_scale, sinks=self.sinks, From 4531f646667b2d67d54460ae5ac1801edecdd0b7 Mon Sep 17 00:00:00 2001 From: tangzzycc <3081129260@qq.com> Date: Sat, 3 Oct 2026 01:21:02 +0800 Subject: [PATCH 2/2] Clarify sliding-window parameter conversion Signed-off-by: tangzzycc <3081129260@qq.com> --- vllm/v1/attention/backends/rocm_attn.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index e79173bfb279..3ccc1e5d5783 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -481,6 +481,8 @@ def forward( k_scale=layer._k_scale, v_scale=layer._v_scale, alibi_slopes=self.alibi_slopes, + # self.sliding_window[0] is the FlashAttention-style left span (W - 1). + # chunked_prefill_paged_decode expects the full window length W. sliding_window=1 + self.sliding_window[0], sm_scale=self.scale, output_scale=output_scale,