Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
23f9b05 to
b253193
Compare
ea2905f to
dc77055
Compare
|
The one red check here isn't the kernel, it's prepare_triton_wheel/prepare. It can't resolve the triton-kernels wheel off the AMD pip index, and since that job gates the Standard Tests, test_kvcache gets skipped and never runs, so a branch push won't clear it. @vgokhale, could someone on @ROCm/team_aiter who owns the CI re-run or patch that wheel step so the real tests fire? A review pass after would be great. This also unblocks #3962 (stacked on it) and the vLLM side in #46847. |
|
Still blocked on CI here. The only red is prepare_triton_wheel/prepare, which can't resolve the triton-kernels wheel from the AMD pip index, so the Standard Tests including test_kvcache stay skipped and a branch push won't clear it. Could someone on @ROCm/team_aiter re-run or fix that gate so the tests run, then take a review pass? @vgokhale, #3962 is stacked on this. |
|
@vgokhale this one is just waiting on the prepare_triton_wheel CI gate. Could you re-run it or point me at whats needed to get it green? I can rebase if you want. |
|
@vgokhale following up on this one, still blocked on the prepare_triton_wheel gate rather than the code. Could you re-run it or tell me what to fix? |
|
@cagrikymk @azaidy could one of you take a look at this one? It only adds aiter/ops/triton/attention/pa_decode_shuffle.py, a sliding-window decode kernel over the shuffled fp8 paged KV layout, so nothing existing changes. The only red is prepare_triton_wheel from the July run (it couldnt pull the triton-kernels wheel), which means the Standard Tests never ran. The branch is far behind main by now, if a fresh CI run is easier than a re-run I can rebase it, just tell me. |
|
@azaidy @vgokhale this one still hasnt had a real CI run, the July prepare_triton_wheel failure skipped the Standard Tests and nothing ran after that. It only adds pa_decode_shuffle.py and still has no conflicts with main. If one of you can pick it up I'll rebase it first so the tests actually run on current main. |
The shuffled fp8 KV layout is written by reshape_and_cache_shuffle and read at decode by the asm paged-attention path, which has no window argument. So a model with sliding-window attention can't use the shuffled cache: the only windowed fp8 decode is the NHD paged_attention_v1 path, on a different layout. There was no kernel that reads the shuffled layout and honors a window. Add paged_attention_decode_shuffle_swa, a triton paged decode that mirrors the shuffled K/V indexing, dequantizes fp8 KV to fp32 with per-tensor or per-token scales (A16W8, query stays 16-bit), runs an online-softmax decode, and masks keys outside [seq_len - window, seq_len). One program per (sequence, kv head), decode only. Validated against an independent fp8 reference on MI300X (gfx942) and MI355X (gfx950), both e4m3fn and e4m3fnuz, across head 128, GQA 32/8 and MQA, block 16/32, windowed and full, per-tensor and per-token scales. Max abs error stays below the fp8 quantization floor on both arches. Signed-off-by: adlashab <adlashab@amd.com>
dc77055 to
42387e2
Compare
github-actions
Bot
changed the title
|
@azaidy @vgokhale I rebased this onto current main (e8bdaf8), new head 42387e2, so CI finally runs on a fresh base. The July run never got past prepare_triton_wheel. The file itself is byte identical to before. I re-checked the kernel against an fp8 reference on MI355X today (torch 2.14.0, triton 3.8.0 CI wheel): all 12 cases pass for e4m3fn and e4m3fnuz, including both shuffled KV layouts vLLM main hands it. The caller is vllm-project/vllm#46847. Could one of you review once CI is through? Edit: the build_aiter_wheels (3.12) red on this run is the runner, it failed with "Cannot connect to the Docker daemon" before the build started (the 3.10 one passed). It needs a rerun of the failed jobs, I cant trigger that from my side. |
There was a problem hiding this comment.
🟡 Changes recommended
Unused V slots can propagate NaNs into output, and required tests and a benchmark are missing.
2 open findings
What changed in this PR
Adds Triton sliding-window decode for the shuffled fp8 KV cache, without changing existing attention paths.
Changes:
- Adds a decode kernel and public wrapper.
- Supports full or windowed attention with scalar or per-token K/V scales.
| File | Description |
|---|---|
aiter/ops/triton/attention/pa_decode_shuffle.py |
Implements windowed decode over shuffled fp8 KV caches. |
🧠 Review effort: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| p = tl.exp(scores - m_new[:, None]) | ||
| alpha = tl.exp(m_i - m_new) | ||
| l_i = l_i * alpha + tl.sum(p, axis=1) | ||
| acc = acc * alpha[:, None] + tl.dot(p.to(tl.float32), v) |
| tl.store(o_ptrs, out.to(out_ptr.dtype.element_ty), mask=q_mask) | ||
|
|
||
|
|
||
| def paged_attention_decode_shuffle_swa( |
|
I validated this shuffled FP8 contract on MI455/gfx1250 for MiniMax-M3 and opened stacked #6625. It addresses the two outstanding code-review gaps here: invalid/NaN V padding is zeroed before PV, and an in-tree producer-to-consumer test covers both FP8 page-16 and page-128 layouts, per-token scales, partial pages, GQA16, and zero-length graph rows. The stacked operator tests pass on gfx1250 with cosine >= 0.999. The companion ATOM integration is ROCm/ATOM #2614. I kept it stacked so this PR retains authorship and review history rather than copying the kernel into an unrelated standalone submission. |
…d test and bench Unused slots in the first and last block of the window get p = 0, but a NaN in their V or V scale still went through p @ v as 0 * NaN and made the whole row NaN (seq_len 1 with 16-slot pages is enough). Zero those V rows. Only the edge blocks have unused slots, so interior blocks skip the select. Rows with seq_len 0 (cudagraph padding) now return zeros, not 0 / 0. V scales use their own strides, so K and V scale tensors can differ. Log with % args and give the kernel a make_kernel_repr name. Add op_tests/triton_tests/attention/test_pa_decode_shuffle.py (both fp8 types, scalar and per-token scales, 16 and 32 slot pages, GQA, window edges, NaN in every unused slot, a seq_len 0 row) and op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py. Signed-off-by: adlashab <adlashab@amd.com>
7098885 to
3ea37a3
Compare
|
@XiaobingSuper thanks for validating the shuffled fp8 contract on MI455/gfx1250 with MiniMax-M3, thats the part I couldnt run myself. I fixed both review points in #3959 itself, new head 3ea37a3:
How the two PRs relate now: the V masking, the zero-length rows and the separate V scale strides are in both, so #6625 will conflict with the new #3959 head on pa_decode_shuffle.py, and on the test file because both use op_tests/triton_tests/attention/test_pa_decode_shuffle.py. What only #6625 has: the producer-to-consumer test through reshape_and_cache_with_pertoken_quant (my test writes the shuffled layout by hand, it doesnt go through the aiter cache writer), page size 128, the all-dead online-softmax guard, the contract asserts, and the gfx1250 validation. |
|
No activity for 15 days, so this is now labelled |
Summary
Adds
aiter/ops/triton/attention/pa_decode_shuffle.py, a paged decode kernelthat reads the shuffled fp8 KV layout used by the asm
pa_fwd/paged_attention_commonpath and applies a sliding-window mask.
Motivation
The shuffled fp8 layout is written by
reshape_and_cache_shuffleand consumed atdecode by the asm paged-attention entry point, which takes no window argument. So a
model with sliding-window attention can't use the shuffled cache: the only windowed
fp8 decode is the NHD
paged_attention_v1path, which uses a different (non-shuffled)layout. There was no kernel that both reads the shuffled layout and honors a window,
so SWA layers either fall back to NHD or silently compute full attention. This is the
correctness gap the kernel closes; the vLLM side that routes to it is a companion PR.
Changes
New
pa_decode_shuffle.pywithpaged_attention_decode_shuffle_swa(...). Itmirrors the shuffled indexing exactly (x = 16 // elem_size, = 16 for fp8):
dequantizes fp8 KV to fp32 with the per-tensor or per-token K/V scales (A16W8,
same as the asm path: query stays 16-bit, only KV is fp8), runs an online-softmax
decode, and masks keys outside
[seq_len - window, seq_len). One program per(sequence, kv head); decode only (one query token per sequence).
No change to existing kernels or call sites; this only adds the new file.
Performance
Not a perf change. This is a correctness/enablement kernel (windowed shuffled-fp8
decode that did not exist before). The split-K performance variant is a separate PR.
Testing
Standalone numerics harness against an independent fp8 reference, both fp8 dtypes
(
e4m3fn,e4m3fnuz), across the full case matrix: head 128, GQA 32/8 and MQA,block 16 and 32, windowed and full, per-tensor and per-token scales.
kernel-vs-fp8-reference 7.7e-3 to 1.81e-2, while the fp8 quantization noise itself
is 0.19 to 0.35, so the kernel error sits well under the quant floor.
<= 1.55e-2 against the same 0.18-0.35 quant floor.
The kernel is plain triton and arch-generic; the two arches cover both fp8 encodings.
Documentation
Layout and masking are documented in the module docstring. No user-guide change.
Dependencies
No new third-party dependencies (triton + torch only).
Breaking Changes
None. New file only.