Repository navigation
[ROCm][gfx1151] ROCM_RDNA_SINGLESEQ_ATTN: HIP attention on gfx1151 - #1346
Closed
roberteg16 wants to merge 3 commits into
Closed
roberteg16 wants to merge 3 commits into
roberteg16 wants to merge 3 commits into
Conversation
roberteg16
force-pushed
the
rogarcia.gfx1151-rdna35-upstream
branch
14 times, most recently
from
October 11, 2026 21:21
77871bd to
0099419
Compare
roberteg16
marked this pull request as ready for review
October 12, 2026 10:27
roberteg16
requested review from
AndreasKaratzas and
dllehr-amd
as code owners
October 12, 2026 10:27
7 of 8 tasks
mgehre-amd
reviewed
Oct 12, 2026
| if ( | ||
| current_platform.is_rocm() | ||
| and vllm_config.attention_config.backend is None | ||
| and vllm_config.scheduler_config.max_num_seqs == 1 |
There was a problem hiding this comment.
max_num_seq is guarantees that you only see seq size = 1, but even if that is not set, you might see seq size=1 (if there are no parallel requests). I wonder whether you should dynamically dispatch to this kernel even in this case? It would need to share the KV cache layout then with another attention algorihtm that can handle the seq size > 1 cases
Author
There was a problem hiding this comment.
It would need to share the KV cache layout then with another attention algorihtm that can handle the seq size > 1 cases
That's the major problem I am seeing. This kernel is not compatible either with ROCM_ATTN nor ROCM_SEGMENTED_ATTN. I could pair it with TRITON_ATTN which supports it, although it is horribly slow at upstream state. My goal now is to get single seq merged (which generally runs close to roofline at any config) and then upgrade it to support arbitrary seq_len. I am already working on it meanwhile.
Author
There was a problem hiding this comment.
Another temporal solution is to use a loop at the host side and always dispatch to our kernels. This should be "not that bad" for prefill, but for decode is not a good idea because of the kernel launch tax. Perhaps there is a way to make it cuda graph friendly.
roberteg16
changed the title
mgehre-amd
reviewed
Oct 12, 2026
| # TODO: Make this explicit in the selector in a future PR. | ||
| if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities: | ||
| backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN) | ||
| # The RDNA3.5 kernels are decoder self-attention only. |
There was a problem hiding this comment.
Can we find a better name for those kernels? An aspect that defines them without specifically tying them to a hardware revision? Those kernels might also work on RDNA 3 or RDNA 4
Author
There was a problem hiding this comment.
Do you have any suggestion? I would suggest to keep it like this until we open to other archs. Wdyt?
There was a problem hiding this comment.
I'd call it ROCM_RDNA_SINGLESEQ_ATTN or similar
mgehre-amd
reviewed
Oct 12, 2026
mgehre-amd
reviewed
Oct 12, 2026
| @@ -0,0 +1,673 @@ | |||
| HEAD_DIM,NUM_Q_HEADS,NUM_KV_HEADS,WINDOW,MAX_QUERY_LEN,BF16,PAGE_SIZE,MAX_SEGMENTS,ROW_GROUPS,MIN_SEGMENT_BLOCKS,WAVES,HEAD_DIM_SPLIT,PREFETCH,V_IN_LDS,DOT_PRODUCT | |||
There was a problem hiding this comment.
What happens if a model needs a config that is not in here?
Author
There was a problem hiding this comment.
A warning is emitted and we fallback to triton attn. The same is done on ROCM_ATTN and ROCM_SEGMENTED_ATTN
Author
There was a problem hiding this comment.
In case you want to add a new row, you have to follow benchmarks/kernels/gfx1151_decode_attn/README.md
Author
Full decode table, fp16, part 1/3: 14/2/64 … 32/8/128Ratio = baseline µs / ours µs.
|
Author
Full decode table, fp16, part 2/3: 32/32/128 … 16/8/256Ratio = baseline µs / ours µs.
|
Author
Full decode table, fp16, part 3/3: 16/8/256 w1024 … 32/4/512Ratio = baseline µs / ours µs.
|
Author
Full decode table, bf16, part 1/3: 14/2/64 … 32/8/128Ratio = baseline µs / ours µs.
|
Author
Full decode table, bf16, part 2/3: 32/32/128 … 16/8/256Ratio = baseline µs / ours µs.
|
Author
Full decode table, bf16, part 3/3: 16/8/256 w1024 … 32/4/512Ratio = baseline µs / ours µs.
|
Author
Full prefill table, fp16, part 1/2: 14/2/64 … 8/4/256Ratio = baseline µs / ours µs.
|
Author
Full prefill table, fp16, part 2/2: 8/4/256 w1024 … 32/4/512Ratio = baseline µs / ours µs.
|
Author
Full prefill table, bf16, part 1/2: 14/2/64 … 8/4/256Ratio = baseline µs / ours µs.
|
Author
Full prefill table, bf16, part 2/2: 8/4/256 w1024 … 32/4/512Ratio = baseline µs / ours µs.
|
roberteg16
force-pushed
the
rogarcia.gfx1151-rdna35-upstream
branch
2 times, most recently
from
October 13, 2026 09:27
30d97c3 to
a78dce0
Compare
…ache - --dtype (float16 or bfloat16) sets the Q/K/V and output dtype. - BenchmarkConfig.sliding_window: windowed impl, KV cache spec and block table. - --min-working-set-mb raises --num-layers until the KV working set no longer fits the last-level cache. - The KV cache holds random values instead of zeros. - process_weights_after_loading is called, max_model_len covers the longest sequence, and NAME@autotune runs the backend's startup autotuning before timing. - The stand-in model is the ungated HuggingFaceTB/SmolLM2-135M. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
roberteg16
force-pushed
the
rogarcia.gfx1151-rdna35-upstream
branch
from
October 13, 2026 11:46
a78dce0 to
ff70b52
Compare
…ntion
An opt-in attention backend for gfx1151 with max_num_seqs=1
(--attention-backend ROCM_RDNA_SINGLESEQ_ATTN). Calls of up to 8 query
tokens run a HIP decode kernel, longer ones a HIP prefill kernel, on
TRITON_ATTN's KV cache and metadata. Everything else runs TRITON_ATTN.
- Decode (csrc/rocm/rdna_singleseq_decode_attn.cu): GQA-packed WMMA,
fp16 and bf16, D in {64, 96, 128, 256, 512}, full and sliding-window
attention, split KV merged in-kernel. Full CUDA graphs for uniform
decode batches.
- Prefill (csrc/rocm/rdna_singleseq_prefill_attn.cu): FlashAttention-2
on WMMA, with split KV when the tiles alone leave the GPU idle. Runs
outside CUDA graphs.
- Each row of vllm/v1/attention/ops/
rocm_rdna_singleseq_{decode,prefill}_variants.csv is one build;
csrc/rocm/generate_rdna_singleseq_attn.py generates them at configure
time, for gfx1151 only.
- Strided queries, unbuilt shapes, alibi, sinks, softcap, quantized KV
and adaptive verification fall back to Triton.
gfx1151, fp16: prefill 8.3x geomean over TRITON_ATTN and 1.64x over
ROCM_SEGMENTED_ATTN@autotune (33 shapes x 21 cells); decode at D=96
(Phi-3.5-vision) 1.81x over TRITON_ATTN.
Tests: tests/kernels/attention/test_rocm_rdna_singleseq_hip.py.
AI assistance was used for this patch.
Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
Tools for the ROCM_RDNA_SINGLESEQ_ATTN kernels, in benchmarks/kernels/rocm_rdna_singleseq_hip_attn/tools: - jit.py: builds a variant standalone and makes the backend use it. - tune.py: searches the launch knobs and writes them into the CSVs. - matrix.py: times the kernels on shapes.csv, optionally against other backends. - utils.py: rooflines and shared options. The attention_benchmarks README notes that this backend's prefill is timed with --no-cuda-graphs. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
roberteg16
force-pushed
the
rogarcia.gfx1151-rdna35-upstream
branch
from
October 13, 2026 12:39
ff70b52 to
7f220bb
Compare
roberteg16
changed the title
6 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
New opt-in attention backend
ROCM_RDNA_SINGLESEQ_ATTN: two HIP kernels for gfx1151 (RDNA3.5) tuned for one sequence per step.--attention-backend ROCM_RDNA_SINGLESEQ_ATTNon gfx1151 with--max-num-seqs 1. It is never selected automatically, so default backend selection is unchanged; on other devices or withmax_num_seqs > 1it is rejected at selection.TRITON_ATTN's packed HND cache (LBHNC) and metadata; anything the kernels do not serve runsTRITON_ATTN's forward.UNIFORM_BATCH;cudagraph_mode=FULLfalls back toFULL_AND_PIECEWISE). Prefill runs outside graphs.csrc/rocm/rdna_singleseq_decode_attn.cucsrc/rocm/rdna_singleseq_prefill_attn.cuBoth: fp16/bf16, D ∈ {64, 96, 128, 256, 512},
Hq % Hkv == 0, causal full or sliding window, pages a multiple of 16 (hybrid pages included), unquantized KV. Not supported (Triton runs them): sinks, alibi, softcap, non-causal, quantized KV, multimodal bidirectional prefix, and a strided query (in-place RoPE on the QKV projection in eager mode). KV connectors and adaptive verification are rejected at selection.Every build is one CSV row (
vllm/v1/attention/ops/rocm_rdna_singleseq_{decode,prefill}_variants.csv, 672 + 430 rows over 33 configurations);csrc/rocm/generate_rdna_singleseq_attn.pyturns them into translation units at configure time, only when the build targets gfx1151. The_rocm_Cops arerdna_singleseq_{decode,prefill}_{attn,variant,table}.Commits: (1)
attention_benchmarksfix (dtype, sliding window, KV out of cache), (2) backend + kernels, (3) tools.Tools (
benchmarks/kernels/rocm_rdna_singleseq_hip_attn/tools)Every shape and launch knob is a compile-time define, so adding or retuning a model means new CSV rows. These tools produce and check them:
tune.py: searches the knob grid for one configuration (decode per M, prefill per step of M), checks each point against a float reference, times it, A/Bs the finalists against the current row and writes a decision;--applywrites decisions into the CSV (refuses ones measured on another kernel source).matrix.py: every configuration inshapes.csvagainst baseline backends (decode under CUDA graphs, prefill eagerly, as vLLM runs them), with the path taken (kernel/prefill/triton) and % of roofline.jit.py: builds single variants with clang outside_rocm_Cand hooks them into the backend, so the tuner tries knob sets without rebuilding vLLM.utils.py: the roofline and knob helpers both share.Usage in
benchmarks/kernels/rocm_rdna_singleseq_hip_attn/README.md. Withbenchmarks/attention_benchmarks, time this backend's prefill with--no-cuda-graphs(under capture it runs Triton).Build cost
Clean build of the 1102 variants (132 compile units), 32-thread host, no ccache: 25.2 s wall / 332 s CPU at
-j16, 19.6 s / 538 s at-j32._rocm_C.abi3.so42.1 → 59.3 MB (+17.1 MB); wheel 22.0 → 26.3 MB (+4.2 MB).Accuracy
GSM8K, 500 questions, 5-shot, greedy, seed 42 (strict / flexible). Ours at
max_num_seqs=1, reference at 16.Kernel performance (gfx1151)
Ratio = baseline time / ours (> 1: ours faster), geomean per shape (Hq/Hkv/D,
w= sliding window). Baselines:TRITON_ATTN,ROCM_SEGMENTED_ATTN(no autotune),ROCM_ATTN.%roof: median of the roofline reached.–: the baseline does not support the shape.Decode (M = 1..8, S = 128..32k)
3696 cells. Geomean: Triton 7.15x (min 1.01), segmented 1.61x (min 0.97), ROCM_ATTN 22.41x (min 1.08); median roofline 87 %.
Full decode table
One row per cell (µs, ratios, %roof), by shape: fp16: 1 2 3 · bf16: 1 2 3
Prefill (M = 16..16k after a cached prefix of 0, 4k, 16k)
2178 cells. Geomean: Triton 8.63x (min 1.60), segmented 2.75x (min 1.03), ROCM_ATTN 8.29x (min 1.22); median roofline 76 %.
Full prefill table
One row per cell (µs, ratios, %roof), by shape: fp16: 1 2 · bf16: 1 2
End to end (gfx1151,
vllm bench serve,--max-num-seqs 1, ours with--attention-backend ROCM_RDNA_SINGLESEQ_ATTN)–: ROCM_ATTN exceeded the 200 min timeout.Gemma 4 (D=512 layers:
ROCM_ATTNandROCM_SEGMENTED_ATTNdo not start), againstTRITON_ATTN:Tests
pytest tests/kernels/attention/test_rocm_rdna_singleseq_hip.py tests/v1/attention/test_rocm_attention_backends_selection.py: 228 passed.--max-num-seqs 1: same greedy output asTRITON_ATTNin eager mode, compiled with CUDA graphs,cudagraph_mode=FULL, and ngram speculative decoding.pre-commit(ruff, mypy, clang-format, typos, …) on every changed file: passed.Not a duplicate: no other PR adds RDNA3.5 HIP attention kernels or a single-sequence backend. AI assistance (Claude Code) was used; the submitter reviews every line.