Repository navigation
Conversation
alexnails
requested review from
BBuf,
DarkSharpness,
HydraQYH,
celve and
yuan-luo
as code owners
August 11, 2026 23:18
2de2324 bounded the extend stage of _fwd_kernel for sliding-window layers but left stage 1 iterating every BLOCK_N tile of the prefix KV, relying on SKIP_TILE = tl.max(tl.max(final_mask, 1), 0) == 0 to suppress the MMA -- exactly the waste that commit removed from stage 2. Apply the same floor here. Stage 1's window mask keeps (q, kv) iff q <= kv + SLIDING_WINDOW_SIZE with q = cur_seq_len_prefix + cur_block_m*BLOCK_M + i and kv = start_n + j; row i = 0 always exists (the early return above guarantees cur_block_m*BLOCK_M < cur_seq_len_extend) and max j = BLOCK_N - 1, so a tile can hold an unmasked element only if cur_seq_len_prefix + cur_block_m*BLOCK_M <= start_n + BLOCK_N - 1 + SLIDING_WINDOW_SIZE, whose smallest BLOCK_N-aligned solution is (A // BLOCK_N) * BLOCK_N for A = cur_seq_len_prefix + cur_block_m*BLOCK_M - SLIDING_WINDOW_SIZE (A is exactly the lowest kv index row 0 of the block can attend to). The ceil-to-floor identity ceil((A - BLOCK_N + 1)/BLOCK_N) == A // BLOCK_N holds for every residue, so the bound is tight and conservative for any BLOCK_M/BLOCK_N; brute-forcing the mask as written over 8 tilings x 20 windows x 17 prefixes x 10 extend lengths visits 794,088 tiles below the floor and finds none holding an unmasked element. Once cur_block_m*BLOCK_M >= SLIDING_WINDOW_SIZE the floor reaches cur_seq_len_prefix and the loop empties out: no query in that m-block can see any prefix token. Only cur_seq_len_prefix > 0 is affected -- chunked prefill or a radix-cache hit. At prefix 0 the loop already ran zero iterations. TritonAttnBackend hands SWA layers window_kv_indptr/indices, so cur_seq_len_prefix is min(prefix, window) and the work removed is ceil(window/BLOCK_N) tiles per m-block: on gpt-oss-120b at TP4 (window 128, BLOCK_M 128 / BLOCK_N 64) an 8192-token chunk over a cached prefix drops from 382 to 256 iterated tiles per (seq, head). A caller that hands the kernel the FULL prefix -- the layout the repo's own test_extend_attention_sliding_window uses -- instead saves O(prefix/BLOCK_N) per m-block, which is what makes this more than tidiness. Bit-identical: max_abs exactly 0.0 over 139 shape/config combinations -- prefix straddling every BLOCK_M/BLOCK_N boundary, windows 1/15/16/63/64/65/127/128/129/ 191/192/255/256/257/1023/4096, prefix < window, prefix == window, window > seq_len, ragged batches mixing prefixes and extends, window-bounded and unbounded prefix kv_indices, head_dim 64/128/256, fp16 and bf16, with and without attention sinks, causal and non-causal, the compact and legacy tile grids, the (64,64) (64,128) (128,128) (32,64) (16,64) (32,32) tilings other backends select, the DCP skip_prefix / skip_extend / lse_extend paths (where the loop bound meets a constant-zero upper bound), and USE_CUSTOM_MASK + sliding window (EAGLE target-verify). An independent fp32 torch reference agrees with both sides on nine prefix-carrying shapes at max_abs <= 5.2e-3 against tol 2e-2, so the bit-equality is equality of correct results. For a full-attention layer (SLIDING_WINDOW_SIZE <= 0) the constexpr branch folds away and the emitted AMDGCN is identical -- 1313 instruction lines on both sides, zero instruction differences, only assembler/DWARF metadata moves. Microbench (MI350X, gpt-oss shapes), sliding-window / full-attention layers, both A/B orders agreeing: prefix8192 extend8192 x8 -17.0% / +0.3%, prefix4096 extend4096 x8 -17.1% / -0.4%, prefix7168 extend1024 x8 -14.7% / -0.5%, prefix8192 extend8192 x1 -15.4% / -0.2%. Unbounded-prefix layout: prefix1024 -57.9%, prefix4096 -83.9%, prefix8192 -91.2%, prefix16384 -94.9%. Self-correction worth recording: at prefix 0, where the loop cannot execute a single iteration, sliding-window layers are 1.15-1.25% SLOWER. That is not noise -- a same-variant control (the identical file loaded twice, compiled separately) spreads only 0.03-0.19% on the same shapes. The mechanism is in the AMDGCN metadata: the dynamic loop lower bound raises sgpr_spill_count from 92 to 102 for the sliding-window kernel (vgpr_count 114 and vgpr_spill_count 0 unchanged, and both counts identical for full attention). Clamping the floor to prefix_end was tried and measures the same, so the cost is inherent to a runtime lower bound. In situ on a served gpt-oss-120b at TP4, ISL 8192 with --chunked-prefill-size 4096 so the second chunk of every request carries a prefix, at 8-way concurrency, over 216 matched _fwd_kernel calls per rank (6 forwards x 36 layers, 6210 GPU events, corrupt-row guard 1.000x, sum/union 1.000): sliding-window layers 36.5 -> 34.3 us/call, -5.9% (per rank -6.1 / -5.5 / -6.9 / -5.2%, same sign on all four), _fwd_kernel total 33.9 -> 33.5 ms (-1.0%). The full-attention layers, which the assembly proves cannot change, land at 277.3 -> 276.2 us/call (-0.4%, per-rank -2.6% to +0.6%) -- that is this instrument's noise floor, and the sliding-window signal sits at twice its worst excursion. The tile arithmetic predicts -7.9% for this mix (chunk 1 at prefix 0 costing +1.2%, chunk 2 at prefix 128 saving 17%). Serving latency does not resolve it, and should not be expected to: SWA extend attention is ~2% of prefill GPU time, so -6% of it is ~0.1%. An 8k1k deck sweep at OSL 1024, run as an A/B/A with a replicated baseline in one window (all 15 points completed in full), moves TTFT p50 by -1.0 / -2.7 / -2.1 / -1.5 / +3.1% at concurrency 1 / 2 / 4 / 8 / 16 against a base-vs-base spread of +4.9 / +1.6 / +3.7 / +2.3 / +8.8% at the same points; output throughput moves -0.2 / -0.3 / +0.6 / +1.5 / -0.6% against a base spread of -0.2 / +0.9 / -0.6 / -1.3 / +0.1%. A radix-cache-hit prefill probe (4096-token shared prefix, server-reported cached_tokens=4096 on every request, /flush_cache never called) moves p50 by -0.1 / -0.2 / +0.2 / -0.6% at concurrency 1 / 2 / 4 / 8 against a base spread of -0.2% to +0.5%. Both are no-regression checks, not wins. The case for this change is that it removes provably no-op work and makes the two stages consistent. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_fwd_kernel_unified (deterministic extend attention) walks its whole prefix+extend
kv range for sliding-window layers and leans on the same SKIP_TILE reduction the
two-stage kernel used to. Its window mask compares absolute positions,
q_abs = cur_window_start + cur_seq_prefix_len + cur_block_m*BLOCK_M + i
k_abs = cur_window_start + start_n + j
keep iff q_abs <= k_abs + SLIDING_WINDOW_SIZE
so cur_window_start cancels and the condition is identical to the two-stage
kernel's with cur_seq_prefix_len added. Row i = 0 always exists (the early return
above guarantees cur_block_m*BLOCK_M < cur_seq_q_len) and max j = BLOCK_N - 1, so
the first tile that can hold an unmasked element is (A // BLOCK_N) * BLOCK_N for
A = cur_seq_prefix_len + cur_block_m*BLOCK_M - SLIDING_WINDOW_SIZE. Because this
loop spans prefix and extend KV together, the single floor removes the fully
masked prefix tiles and the fully masked extend tiles below the window at once.
Brute-forcing the mask as written (including the causal-over-extend term, and with
cur_window_start at 0 and at 7777 so the cancellation is checked rather than
assumed) visits 4,100,280 tiles below the floor over 8 tilings x 20 windows x 17
prefixes x 10 extend lengths and finds none holding an unmasked element.
Bit-identical: max_abs exactly 0.0 over 68 shape/config combinations -- windows
1/15/16/63/64/65/127/128/129/191/192/255/256/257/1023/4096, prefix and window
straddling every BLOCK_M/BLOCK_N boundary, window > seq_len, ragged batches,
head_dim 64/128/256, fp16 and bf16, with and without sinks, causal and non-causal,
window_start_pos 0 / 7777 / 99999, and the (64,64) (64,128) (128,128) (32,64)
(16,64) (32,32) tilings. An independent fp32 torch reference agrees with both
sides on nine shapes at max_abs <= 5.2e-3 against tol 2e-2. For a full-attention
layer the emitted AMDGCN is identical -- 939 instruction lines on both sides,
zero instruction differences.
Microbench (MI350X, gpt-oss shapes), sliding-window / full-attention layers, with
a same-variant control (the identical file loaded twice) spreading 0.06-0.54%:
prefix0 extend8192 x8 -46.5% / +0.3%, prefix128 extend8192 x8 -47.4%,
prefix4096 extend4096 x8 -70.0% / -0.0%, prefix1024 extend1024 x8 -55.6%,
prefix128 extend1024 x8 -31.3%. The win is large even at prefix 0 because this
loop spans the extend KV as well, so the one floor also delivers the saving
2de2324 gave the two-stage kernel and the unified kernel never received --
deterministic mode has had no windowing optimisation at all until now. Unlike the
two-stage prefix trim, this costs no extra scalar spills (sgpr_spill_count 95 on
both sides at W=128, 5 on both at W=-1), which is why it has no prefix-0 penalty.
This path has no repo coverage today: test_extend_attention_unified_vs_regular
never passes sliding_window_size, so the numerics above are the only correctness
check the unified kernel's sliding-window path has anywhere. Worth adding a
windowed case to that test separately.
Deliberately no serving evidence: deterministic mode is off by default, so no deck
arm was run and no deck number is claimed for it. The justification is the same
bit-identity plus strictly-fewer-iterations argument as the two-stage bound.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
alexnails
force-pushed
the
perf/gfx950-swa-prefix-trim
branch
from
August 24, 2026 08:09
a7aed00 to
4f2cd17
Compare
This branch has not been deployed
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.
Read this first
Neither commit here produces a measurable end-to-end improvement, and the first carries a small cost in one regime. Both are exhaustively proven bit-identical. I am submitting them as correctness-and-efficiency changes with full disclosure, not as a performance claim — a maintainer may reasonably decide the first is not worth its cost.
Commit 1 — bound the stage-1 prefix loop
#34462 bounded the stage-2 extend loop. Stage 1 has the identical structure: its mask keeps
(q, kv)iffWp + m0 + i <= start_n + j + W, so the same floor withcur_seq_len_prefixadded empties the loop oncem0 >= W.Bound verified exhaustively, not sampled.
A = Wp + m0 - W, floor(A // BLOCK_N) * BLOCK_N;ceil((A-BN+1)/BN) == A//BNwas checked for every residue class, then the mask as written was brute-forced over 8 tilings × 20 windows × 17 prefixes × 10 extend lengths: 794,088 tiles below the floor, zero holding an unmasked element, with 50,560 tightness witnesses. Extra conjuncts (causal, custom mask) can only empty a tile further, and were checked anyway.Isolated kernel
In a served profile
gpt-oss-120b TP4, ISL 8192 with
--chunked-prefill-size 4096, cc8, 216 matched calls per rank:_fwd_kernel, SWA layers_fwd_kernel, full-attention layersPredicted −7.9% from tile arithmetic, so the mechanism behaves as modelled.
Served: no resolvable change
run_sweep.sh, 8k/1k, OSL 1024, arms A → B → A. All 15 points reported 100% completion; every delta sits inside the base-vs-base spread. This is exactly what the ~0.1%-of-prefill-GPU prediction requires — the affected regime (prefix > 0, i.e. chunked prefill and radix-cache hits) is not what a flush-between-points sweep measures.The cost, stated plainly
At prefix 0 the loop cannot execute a single iteration, yet SWA layers measure 1.15–1.25% slower. Confirmed genuine by a same-variant control (the identical file loaded twice reproduced to ±0.2%). The mechanism is visible in the AMDGCN metadata:
sgpr_spill_countrises 92 → 102. A clamped reformulation measures the same, so it appears inherent to introducing a runtime lower bound at all.So commit 1 is a trade: it helps chunked-prefill and cache-hit traffic and costs ~1.2% on prefix-0 traffic.
Commit 2 — the same bug in
_fwd_kernel_unifiedThe unified kernel carries the identical one-sided window mask.
cur_window_startcancels on both sides, giving the same floor withcur_seq_prefix_lenadded; 4,100,280 tiles checked below it, zero violations.Microbench: −31% to −70% — larger than commit 1 because the unified loop spans extend KV as well, so this floor also delivers the stage-2 saving that kernel never received.
No serving evidence exists and none is claimed. Deterministic mode is off by default, so no deck arm was run. This commit carries no prefix-0 cost (
sgpr_spill_count95 → 95).Numerics
max_absexactly 0.0 — 139 configurations for commit 1, 68 for commit 2 — each with an independent fp32 torch reference confirming both arms compute real windowed attention (base and candidate both ≤5.2e-3 against tol 2e-2), so this is not equality of degenerate output.Coverage includes both grid modes, unbounded prefix, non-causal, 6 tilings, head_dim 64/128/256, fp16/bf16, and sinks on/off. Two coverage gaps found and closed during validation: the DCP
skip_prefix/skip_extend/lse_extendpaths (a runtime lower bound meeting a constant-zero upper bound) andUSE_CUSTOM_MASKcombined with sliding window (the EAGLE target-verify route).Full-attention layers are identical at the assembly level: 1313 and 939 AMDGCN instruction lines respectively, zero instruction diffs.
Coverage gap worth a follow-up
test_extend_attention_unified_vs_regularnever passessliding_window_size, so the numerics in this PR are currently the only correctness check that path has anywhere in the repo.Hardware: MI350X (gfx950), ROCm 7.2 userspace.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #32709182831
Latest PR Test (Extra): ❌ Run #32709182601
Latest PR Test (AMD ROCm 7.2): ❌ Run #32709182701