Repository navigation
[ROCm][Bugfix] Fix ROCM_ATTN sliding-window boundary - #59550
Conversation
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>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
✅ @tangzzycc, CI is now available for this PR.
|
Please provide the vllm launch command, lmeval score of the model to show that acc is fixed. (show before and after this PR) |
Signed-off-by: tangzzycc <3081129260@qq.com>
|
Hi @tjtanaa ,I evaluated Gemma3 1B IT on Radeon 890M (
The before/after accuracy change is small on this subset: one additional correct answer (+0.78 percentage points). The main purpose of the fix is to pass the full window length Launch command (set PYTHONPATH=/path/to/vllm-checkout .venv/bin/vllm serve \
models/gemma-3-1b-it --served-model-name gemma3-1b \
--host 127.0.0.1 --port 8055 --attention-backend ROCM_ATTN \
--max-model-len 4096 --kv-cache-memory-bytes 134217728 --seed 1234Evaluation command (use a separate output directory for each run; .venv/bin/python -m lm_eval run --model local-completions \
--model_args 'model=gemma3-1b,base_url=http://127.0.0.1:8055/v1/completions,tokenizer=models/gemma-3-1b-it,tokenizer_backend=huggingface,tokenized_requests=True,add_bos_token=True,max_length=4096,max_gen_toks=512' \
--tasks gsm8k --num_fewshot 5 --limit 128 --seed 1234 \
--log_samples --output_path results-before |
|
/ci run |
|
✅ Triggered Buildkite CI #92784 for commit |
CI selector (shadow): 5 test steps (12 jobs) instead of 74 (98 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (5)
Would skip (today's rules run them) (70)
Would add (today's rules do not run them) (1)
AMD mirrors: would skip (18)
AMD mirrors: would add (42)
2 changed files · base |
|
/ci retry |
|
✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/92784 |
|
/ci cancel |
|
✅ No cancelable CI build is running for branch |
|
/ci retry |
|
✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/92784 |
|
/ci run |
|
✅ Triggered Buildkite CI #92796 for commit |
3 participants
Purpose
ROCM_ATTNstores the decoder window as(W - 1, 0), whereW - 1is the left span. The underlying decode and prefix-prefill kernels usequery_position - key_position < SLIDING_WINDOW, so their argument must be the total window lengthW, including the current token. PassingW - 1excludes the oldest valid token; forW=1, passing zero disables the window entirely.This PR passes
1 + self.sliding_window[0]tochunked_prefill_paged_decode. This matches the existing KV-cache docstring example: withsliding_window=4and seven computed tokens, the next token at position 7 must attend to tokens 4–7. The old argument is 3 and attends to 5–7; the corrected argument is 4 and attends to 4–7.TRITON_ATTNstores the same(W - 1, 0)tuple and already passesSLIDING_WINDOW=(1 + window_size[0]). Applying the same conversion makes the causal window boundaries consistent between the two backends. The no-window sentinel maps to zero, which the helper treats as disabled; encoder attention uses a separate path.Open-PR searches on 2026-10-01 found no duplicate. Related #49588, #53628, and #59132 do not modify this backend argument conversion.
Test Plan
Run the focused backend regression and the existing cached-KV boundary test:
.venv/bin/python -m pytest tests/kernels/attention/test_prefix_prefill.py \ -k 'rocm_backend_sliding_window or cached_kv_block_table_boundary' -qTest Result
Before the fix, all four windowed regression cases fail and both no-window controls pass. After the fix: 7 passed, including the existing boundary test. All applicable pre-commit checks passed.
Gemma3 1B IT (BF16, window 512) completed inference on Radeon 890M (
gfx1150) with inputs within and beyond the window, including chunked prefill and HIP graph decode. Compared before/after ROCM_ATTN with TRITON_ATTN and Transformers FP32. Before/after prompt logprobs were identical for input lengths up to 512; longer inputs showed changes, with errors against the reference implementations increasing or decreasing depending on input length.AI assistance was used during implementation, testing, and preparation of this PR. vLLM's
/pr-checklistskill was used to review the change.