[ROCm][Attention] Wide MHA decode attention kernel for gfx1151 - #1087
Draft
roberteg16 wants to merge 3 commits into
Draft
[ROCm][Attention] Wide MHA decode attention kernel for gfx1151#1087roberteg16 wants to merge 3 commits into
roberteg16 wants to merge 3 commits into
Conversation
roberteg16
force-pushed
the
rogarcia.wide-decode-attn
branch
4 times, most recently
from
August 12, 2026 14:43
8fafabb to
c9c3258
Compare
Adds a split/reduce decode-attention kernel for MHA (num_heads == num_kv_heads)
on AMD RDNA3.5, a faster alternative to Triton's unified attention for that
shape. Enabled by default on gfx1151 (VLLM_ROCM_WIDE_DECODE_ATTN=0 forces
Triton, =1 forces on).
Measured against Triton over 120 cells (head counts 8/16/32/64 x head_dim
128/256/512 x M 1/4 x ctx 128..65536): 1.0-4.7x on the 2D path, 0.98-6.0x on
the 3D one, reaching 90-95% of the 238 GiB/s streaming peak at long context.
One wave owns each (sequence, KV head, KV segment) and writes fp32 partials; a
reduce pass merges them with the usual log-sum-exp rescale.
- csrc/rocm/wide_decode_attn.cu: the kernels, ported from the standalone
attn_decode_hip/ harness where they were developed and tuned. fp16 and bf16
(v_dot2_f32_f16 / v_dot2_f32_bf16); the body is gfx11-only with stubs on
other arches so the TU links in multi-arch builds. Registered as
torch.ops._rocm_C.wide_decode_attn; every constraint is a TORCH_CHECK, since
the alternative failure mode is an unwritten output tensor.
- vllm/envs.py: VLLM_ROCM_WIDE_DECODE_ATTN tri-state (unset = default-on on
gfx1151).
- vllm/v1/attention/ops/rocm_wide_decode_attn.py: the host-side predicate, a
graph-safe custom op wrapper with a no-op fake, and the varlen -> fixed-M
gather/scatter.
- vllm/v1/attention/backends/triton_attn.py: workspace allocation in the
builder, and a fast path in forward() that falls through to
unified_attention for anything the predicate declines.
- tests/kernels/attention/test_rocm_wide_decode_attn.py: kernel vs an explicit
fp32 paged reference across shapes, dtypes, batches and context edges, plus
the predicate and the gather/scatter roundtrip.
The predicate enforces every restriction the kernel has, because violating one
produces wrong output rather than an error: MHA only (one query head per KV
head is baked into the wave mapping), head_size in {64,128,256,512}, fp16/bf16
with an unquantized cache (there is no scale in the ABI), causal with no
sliding window / ALiBi / sinks, no fused output quant, and the NHD cache layout
(HND permutes block_size and num_kv_heads inside a block).
Only M=1 and M=4 are instantiated. Query lengths 2 and 3 are padded up to the
M=4 kernel rather than built: padding costs a measured geomean 1.024x (worst
1.081x, and in 8 of 120 cells M=4 is outright faster), and it keeps every
launch on one of the two M values the tuning sweep actually covered -- the
tuned rule steps UNROLL between M=2 and M=3 on 7 of 12 shapes, so instantiating
them would mean shipping configs fitted to no measurement. Padding goes at the
front because the mask is `pos > ctx_len + m`: real token j lands at row
M-len+j and sees exactly [0, num_computed+j], so seq_lens passes through
unmodified.
Padding also removes the need for a uniform-length decode block, so the batch
does not have to be reordered. That matters because reorder_batch_threshold is
taken as a min across every attention group -- setting it here would change
batch ordering for every TRITON_ATTN user on every platform, and would make
AttentionCGSupport.ALWAYS untrue. The fast path is confined to decode-only
batches and leaves all of that alone.
MEMORY AND CAPTURE CORRECTNESS. Four things the integration had to get right,
each of which fails only on hardware:
The M=1 path must not slice the M dimension off the shared M=4 workspace:
the shape is right but the strides are not (131072 where a dense block wants
32768), and the kernels index the partials with raw pointer arithmetic, so
it would read the wrong elements silently rather than fault. A dense view is
carved over the front of the storage instead. Contiguity of the partials,
block_table and seq_lens is checked at the op boundary; the KV caches
deliberately are NOT required to be contiguous (they are k/v views into one
[num_blocks, 2, ...] allocation), only dense within a block -- which also
rejects HND.
No device sync in the guard. Comparing num_actual_tokens against
query_start_loc[-1] would call __bool__ on a device tensor, syncing per layer
per forward and making the branch uncapturable. max_query_len <= 4 already
bounds every request in the batch.
No data-dependent shapes in the scatter. Selecting real rows with a boolean
mask materializes via nonzero(), which syncs and cannot be captured, and
TRITON_ATTN declares AttentionCGSupport.ALWAYS. build_gather_index returns a
fixed num_seqs*M index that aims padding rows at a scratch slot instead.
Padding counts are independent. gpu_model_runner pads seq_lens and
num_actual_tokens separately under capture, so num_reqs is passed in
explicitly and drives every slice; the workspace is sized for the larger of
max_num_seqs and max_cudagraph_capture_size.
Also: ENCODER_DECODER reaches forward() and its cross-attention is not
causal, so the path is gated on attn_type == DECODER; and a sliding_window
of 1 is stored as (0, 0) by this Impl, so the predicate receives W itself
rather than testing the stored extent against 0.
TUNING. Three rules ported from the standalone harness, each fitted on
measurement:
Unrolled reduce loops. Making the segment count a runtime argument dropped
the unroll the templated version got for free -- geomean 1.011 over 240
cells, up to 1.16x at short context where the reduce is 13% of the traffic.
Unrolling both loops by 4 with independent accumulators brings it to 1.0026.
KV segments capped at 32 when head_dim is 64. The cap of 16 was fitted over
head_dim 128/256/512 where the quotient rarely reaches it; at head_dim 64
vec2 is 1, so it binds on every shape. Over 136 cells the conditional cap
takes the worst case from 1.424x to 1.123x, while raising it everywhere
would regress Hq=8 at head_dim 128.
Head grouping (HEADS_PER_WAVE). At head_dim 64 a lane holds 4 bytes and the
wave reduces one dot2 per token, so the 5-step DPP chain becomes most of the
inner loop. Adjacent KV heads are contiguous in the paged layout, so a wave
takes 2 or 4 as one run: b128 loads and a 3-step reduction. The rule is 2
heads per wave below 16 heads and 4 at or above, with a segment numerator of
256 at head_dim 64 against 512 above -- both conditioned on head_dim 64,
since applying either everywhere costs geomean 1.070 at head_dim >= 128
against 1.012.
P.V accumulation goes through elem_traits::accum_pv, emitting VOPD-pairable
dot2 rather than v_fma_mix_f32. It removes all 64 fma_mix and changes the
runtime by nothing -- the kernel sits at 97% of this machine's streaming
ceiling and waits on memory, not issue -- but costs one VGPR less.
The C++ and Python copies of the tuned rule agree on all 4096
(M, head_size, num_kv_heads) combinations; the partials are indexed
[num_seqs, num_heads, nseg, M, head_size], so a disagreement would be a memory
fault, and the op re-derives the count and TORCH_CHECKs it.
Verified so far without a GPU: the source compiles for gfx1151 with both dtypes
instantiated symmetrically (112 kernels, zero scratch, occupancy 7-16), bf16
lowers to v_dot2_f32_bf16 rather than an fp32 fallback, the predicate rejects
all twelve unsupported cases, the M=4-padded mask reproduces Triton's
context_len = seq_len - cur_batch_query_len exactly for every real query
length, and the fixed-shape scatter roundtrips exactly on irregular lengths.
The test suite last ran green at 82 passed; it and an end-to-end A/B need to
run again on hardware.
AI assistance (Claude) was used.
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Robert Esclapez-Garcia <rogarcia@amd.com>
…ters
Review feedback. No behaviour change except the last item.
NAMING. The kernel, the registered op, the source file, the Python module and
the env var are named for what they are: causal MHA decode attention on
RDNA3.5. VLLM_ROCM_WIDE_DECODE_ATTN becomes VLLM_ROCM_RDNA35_CAUSAL_MHA.
COMMENTS. Those that narrated how the code got here -- what was tried, what
regressed, what was removed and re-added -- are gone. What remains explains the
code as it stands, keeping the measurements that justify a constant or a
restriction. The file header now walks the seven phases of the split kernel in
order, from task decomposition to the partial write.
DEAD PARAMETERS. Three mode switches had a default that nothing ever overrode,
so each carried an unused arm through every instantiation:
PV_DOT2_MODE chose dot2 against a one-hot weight over the fma_mix form.
On by default; the fma_mix arm is deleted.
BURST_MAX_MODE chose rescaling once per burst over once per token. Swept over
40 cells with both values compiled and timed, per-burst wins
38 of 40: forcing it on costs geomean 1.0005 against the best
config per cell, forcing it off costs 1.028-1.038 with a worst
case of 1.11.
FAST_ADDR_MODE chose whether to emit a second address chain for bursts inside
one page. Its -1 mode resolved to VEC2_PER_LANE <= 8, and every
instantiated combination satisfies that -- head_dim 64 reaches
4 at most, and head_dim >= 128 only runs ungrouped, so
VEC2_PER_LANE is 2, 4 or 8. The condition was constant-true in
every build; the sweep agrees it should be (on wins 35 of 40).
The template goes from eight parameters to five. Resource usage is unchanged
across all 160 instantiations: no scratch, VGPR 18-196, occupancy 7-16.
ROUTING. head_size 64 with 32 heads at M=1 is no longer declined. It runs
0.90-0.96x of Triton's 2D path there (1.05-1.07x of the 3D one) and 1.19-1.42x
at M=4; one predicate covering every supported shape is easier to reason about
than a table of exceptions, and the loss is bounded and small.
Verified on gfx1151: pytest 82/82; the fast path matches unified_attention to
1.2e-04 over decode-only, irregular and uniform M=4 batches; all three capture
and replay under cudagraph with zero deviation; Llama-2-7B-AWQ end-to-end decode
44.7 tok/s against Triton's 42.2.
Signed-off-by: Robert Esclapez-Garcia <rogarcia@amd.com>
6672106e1 dropped the losing-shape table, so can_run() no longer looks at max_query_len beyond the range check and its verdict is uniform over M=1..4 (checked over the 80 (num_heads, head_size) combinations). The builder's any() probe existed only to catch a shape declined at M=1 and accepted at M=4 -- head_size 64 with 32 heads was the one entry -- so it collapses to a single call, and the comment saying otherwise goes with it. No behaviour change: the allocation is made for exactly the same layers. Not run on device -- this session has no GPU access (not in the render group, so /dev/kfd is unreadable). Verified by calling can_run() directly over the shape matrix. Signed-off-by: Robert Esclapez-Garcia <rogarcia@amd.com> Co-Authored-By: Claude <noreply@anthropic.com>
roberteg16
force-pushed
the
rogarcia.wide-decode-attn
branch
from
August 12, 2026 15:30
c9c3258 to
59e75e0
Compare
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
1 participant
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
Adds
rdna35_causal_mha_attn, a split/reduce decode-attention kernel for MHA(
num_heads == num_kv_heads) on gfx1151, as a fast path insideTRITON_ATTN.Enabled by default on gfx1151;
VLLM_ROCM_RDNA35_CAUSAL_MHA=0forces Triton and=1forces it on. Anything the host predicate declines falls through tounified_attentionunchanged, so the path is A/B-able at runtime.One wave owns each (sequence, KV-head-group, KV segment) and writes fp32
partials; a reduce pass merges them with the usual log-sum-exp rescale. fp16 and
bf16, using
v_dot2_f32_f16/v_dot2_f32_bf16. The body is gfx11-only withstubs elsewhere, so the TU links in multi-arch and CDNA builds.
End-to-end
TheBloke/Llama-2-7B-AWQ(32×32×128), input 1920 / output 128,--max-num-seqs 1,three repetitions per mode — all three identical, so this is not run-to-run noise:
Prefill is untouched by construction — the fast path is decode-only. The fast
path was confirmed to actually fire (64 invocations with it on, 0 with it off,
identical generated text).
Kernel benchmarks
312 cells: head counts 8/16/32/64 × head_dim 64/128/256/512 × M 1/4 × ctx
128..65536, fp16, block_size 16, against
unified_attention's 2D and 3D pathson the same inputs. Geomean 1.87× vs 2D and 1.79× vs 3D; 283/312 cells beat
both paths; median 88% of the 238 GiB/s streaming peak.
Per-cell detail (160 rows, ctx 128 / 1024 / 4096 / 16384 / 65536)
The 29 cells that do not beat both paths are all head_dim 64, where a lane holds
the least work. The worst is 32 heads at M=1: 0.90–0.96× of the 2D path, though
still 1.05–1.07× of the 3D one, and 1.19–1.42× at M=4. Those shapes are routed
here anyway — one predicate covering every supported shape is easier to reason
about than a table of exceptions, and the loss is bounded and small.
Restrictions
Every one is a property of the kernel, and routing a case that violates one
produces wrong output rather than an error, so the host predicate
(
rocm_wide_decode_attn.can_run) enforces all of them:head_sizein {64, 128, 256, 512}AttentionType.DECODERonly —ENCODER_DECODERcross-attention is not causalmax_query_len <= 4)The kernel's own knobs (KV segment count, unroll, KV heads per wave) are chosen
by a rule in the op, keyed on head size, head count and query length; callers
never pick a configuration.
Query lengths 2 and 3 are padded up to the M=4 kernel rather than instantiated.
Padding costs a measured geomean 1.024× and keeps every launch on one of the two
M values the tuning sweep covered; it also removes any need to reorder the
batch, which would have changed batch ordering for every
TRITON_ATTNuser onevery platform.
Test plan
pytest tests/kernels/attention/test_rdna35_causal_mha_attn.py— 82 passedon gfx1151 (shapes, dtypes, batches with unequal lengths, context edges,
softcap, the predicate, and the gather/scatter roundtrip)
unified_attentionon identical state: 1.2e-04 max abs diffover decode-only, irregular
[1,3,2,4]and uniform M=4 batchesbit-identical to Triton
(M, head_size, num_kv_heads) combinations
AI assistance (Claude) was used; every number above comes from a run on this
machine, and a human has reviewed the change.