perf(gqa): support sliding windows in the fused flash prefill kernel (-65.7% TTFT cumulative) - #622
Closed
BoarQing wants to merge 7 commits into
Closed
perf(gqa): support sliding windows in the fused flash prefill kernel (-65.7% TTFT cumulative)#622BoarQing wants to merge 7 commits into
BoarQing wants to merge 7 commits into
Conversation
gpt-oss carries a learned per-head sink logit on every layer. The decomposed prefill path handles it but the fused v5 kernel did not, so all 24 layers fell back to materialising the score matrix. Fold the sink into the v5 softmax denominator in the epilogue, where it costs one exp2f and one add per row. m_reg already holds the row max scaled by log2(e), so the natural-space term exp(sink - m) becomes exp2f(sink * kLog2e - m_reg). Rows that are entirely masked keep m_reg at -inf and are skipped, so their denominator stays at the clamped floor. Expose it through a new hip_gqa_flash_prefill_v3 entry point that also carries local_window_size for later use, and make v2 a thin wrapper so existing callers and test harnesses keep their signature. v3 declines (-1) a sink request at d != 64 and any window request, so those shapes keep falling back to the decomposed path instead of silently computing the wrong thing. Tuning launches deliberately omit the sink: one FMA per row cannot reorder the (M_TILES, BKV) candidates, and leaving it out keeps tuned configs comparable with the non-sink case. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
|
Thanks for opening a PR! This project follows LLVM's incremental-development and AI-tool-use Before requesting review, please check that:
Reviewers are assigned through |
Closed
The fused-path gate required head_sink == nullptr, so gpt-oss-20b, which sets a sink on all 24 layers, never reached the flash kernel. Every layer went through the decomposed path and materialised an fp32 score matrix. Thread head_sink and smooth_softmax through gqa_forward_fused to both the decode and prefill entry points, and replace the head_sink == nullptr term with a sink_ok gate that admits sinks on decode and on d == 64 prefill, matching what the kernels actually implement. The window rejection stays: v3 declines it. On gpt-oss-20b at 16k this moves the 12 full-attention layers onto the fused kernel; the 12 sliding-window layers still fall back because v3 declines the window. Interleaved A/B at seqlen 16384, alternating prebuilt DLLs across three rounds so machine drift hits both arms equally: 32,346 ms base against 23,108 ms, -9,238 ms or -28.6%, with a 102-115 ms run-to-run spread. RGP at chunk 16 agrees: a fused full-attention layer is one 5.6 ms dispatch in place of the 26.0 ms of score GEMM, causal mask, softmax and PV GEMM it replaces. Scaled over 12 layers and 32 chunks that predicts roughly 7.6 s, close to the 9.2 s measured. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
Extend the standalone prefill harness with the gpt-oss geometry (H=64, G=8, d=64) under three sink modes: none, per-head logits, and smooth_softmax. The CPU fp32 reference folds the sink into the denominator the same way the kernel does, and the sink vector round-trips through fp16 so the comparison is not polluted by a dtype difference the kernel does not have. Cases cover pure prefill and chunked prefill at past_len 512 and 8192, since the sink interacts with the running max across KV tiles. One d=128 case asserts that v3 declines the sink with -1 rather than computing it wrongly. All 13 cases pass at relL2 3.2e-04 to 5.7e-04; the sink cases sit within 2e-05 of their no-sink counterparts, which is fp16 rounding rather than a systematic shift. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
The v5 config cache was keyed on skv, which advances by one chunk on every call of a chunked prefill. A 16k prefill therefore saw 32 distinct keys and ran 32 tuning sweeps of 526 launches each, 16,832 tuning launches in total. That cost was invisible before because nothing on the gpt-oss path reached v5; with sinks now supported it dominates the first prefill. skv only scales the length of the KV loop, uniformly across candidates, so it cannot change which candidate wins. Confirmed rather than assumed: with HIPDNN_PREFILL_TUNE_DEBUG=1 over a full 16k prefill, M_TILES=2/BKV=32 wins at all 32 skv values from 512 to 16384, and the ranking of all four candidates is identical at every one. Reusing the first sweep's winner everywhere costs 0.000 ms of measured regret. Keying on (d, sq, Hq, G) leaves one sweep per process, and it lands at the smallest skv, where a sweep is cheapest. Measured back to back on the same build, cold TTFT on gpt-oss-20b at 16k drops from 139,582 ms to 21,905 ms. Warm TTFT is unchanged within run-to-run spread, as expected given the chosen config is identical. The remaining first-prefill cost is the single sweep: on the final build, cold 23,942 ms against warm 23,108 ms. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
gpt-oss alternates full-attention and sliding-window layers, 12 of each. The window was previously only ever applied as extra masking work on the decomposed path, so a windowed layer cost the same as a full one and in fact slightly more, because the mask kernel writes more -inf. Apply the window in v5 in two places. The per-row mask gains the lower bound that causal_mask_kernel_impl already uses, k < past_len + q - window + 1. More importantly the KV loop now starts at the first tile that any row of the Q tile can reach, rather than at tile 0: every row in a Q tile has a lower bound at least as high as the first row's, so that start covers the tile and the per-row mask trims the remainder. This is what makes a window cheaper instead of merely masked, since the number of tiles walked stops growing with skv once the window is filled. Skipping tiles introduces a case causal-only attention never produces: a row can have every key in a tile masked while it has still seen nothing, leaving both m_old and m_new at -inf, and exp2f(-inf - -inf) is NaN. Guard it by contributing nothing and leaving the running state untouched. This was not theoretical; it reproduced immediately as NaN output on window=128, and it does not reproduce at window=1 or at a window wider than the sequence, which is why the harness covers all three. The window is added to the autotune key. Unlike skv, it does not merely scale the KV loop, it caps how many tiles a block walks, which is a genuinely different work profile, so it can change which candidate wins. Measured in the standalone harness at gpt-oss geometry, per call: with window=128 the kernel takes 0.157 ms at past_len 512 and 0.134 ms at past_len 8192, i.e. flat in KV depth, against 6.27 ms for the same shape unwindowed. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 3527 - Commit: |
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
Run: 3527 - Commit: |
The fused gate rejected any local_window_size > 0, and the prefill call site hardcoded -1, so gpt-oss's 12 sliding-window layers stayed on the decomposed path even after sinks were supported. Thread the real window through gqa_forward_fused to hip_gqa_flash_prefill_v3 and replace the blanket rejection with a window_ok term admitting it on d == 64 prefill, which is what v5 implements. Decode is deliberately left out. Its kernel does clamp kv_lo for a window, so enabling it would be a one-word change, but nothing here verifies that path and prefill is where the window was costing us. Windowed decode continues to go decomposed. With this, every gpt-oss prefill layer reaches the fused kernel: a 2048-token prefill logs 96 of 96 GQA calls (4 chunks x 24 layers) as fused prefill and none as decomposed. Interleaved A/B at seqlen 16384, three rounds alternating prebuilt DLLs: 29,483 ms base, 21,049 ms with sinks alone, 10,113 ms with sinks and windows. That is -19,370 ms against base (-65.7%), of which this change contributes -10,936 ms. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
Ten window cases on top of the sink set, with the CPU reference gaining the same lower bound the kernel uses. The cases are chosen so that a window bug cannot hide. Window alone comes before window plus sink, so the two features fail separately. past_len 8192 with window 128 is the case where whole KV tiles must be skipped rather than merely masked. window=100 against BKV of 32 and 64 catches an off-by-one in the start-tile clamp. A window wider than the sequence must reproduce full attention exactly, and window=1 must reduce to each query seeing only itself; both of those pass even with the NaN bug present, which is precisely why the mid-range cases are needed. A d=128 case asserts v3 declines the window rather than silently ignoring it. All 23 cases pass, windowed ones at relL2 3.3e-04 to 3.5e-04, in line with the no-window cases. Co-Authored-By: Cursor <cursoragent@cursor.com> Made-with: Cursor
BoarQing
force-pushed
the
perf/gqa-flash-prefill-window
branch
from
August 1, 2026 03:25
a1ea5bf to
4fa4edd
Compare
Contributor
Author
9 tasks
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
Teaches the fused flash prefill kernel about sliding windows, which puts the
last of gpt-oss-20b's attention layers on the fused path.
Includes #598 (attention sinks).
This PR targets
main, so the diff is #598's four commits plus three windowcommits on top; the window commits are the last three and are self-contained.
The window is not observable end-to-end without #598, because the sink gate
otherwise rejects all 24 gpt-oss layers. Merging #598 first will shrink this
diff to the window commits alone.
TTFT at seqlen 16384, interleaved: 29,483 ms base -> 10,113 ms, -65.7%.
#598 gets the first 8,434 ms; this PR adds the other 10,936 ms.
After this, a 2048-token prefill logs 96 of 96 GQA calls (4 chunks x 24
layers) as fused prefill and none as decomposed.
Related issue or design
None. Third in a series on the gpt-oss-20b prefill bottleneck, after the
no-expand default (#597, independent
and performance-neutral) and sinks (#598,
which this one requires).
Why
gpt-oss alternates full-attention and sliding-window layers, 12 of each. The
window was only ever applied as extra masking on the decomposed path, so a
windowed layer cost the same as a full one ΓÇö in fact marginally more, because
the mask kernel writes more
-inf. A window is an opportunity to do less work,and none of it was being taken.
The saving is not in the mask, it is in the loop bound. Applying the window
in v5 has two parts:
causal_mask_kernel_implalreadydefines,
k < past_len + q - window + 1.of at tile 0. Every row in a Q tile has a lower bound at least as high as the
first row's, so that start covers the tile and the per-row mask trims the
remainder.
The second part is what makes a window cheaper rather than merely masked: the
number of tiles a block walks stops growing with
skvonce the window isfilled. In the standalone harness at gpt-oss geometry,
window=128costs0.157 ms at
past_len512 and 0.134 ms atpast_len8192 ΓÇö flat in KV depth ΓÇöagainst 6.27 ms for the same shape unwindowed.
A NaN the window introduces
Skipping tiles creates a case causal-only attention never produces: a row can
have every key in a tile masked while it has still seen nothing, leaving both
m_oldandm_newat-inf, andexp2f(-inf - -inf)is NaN, which thenpropagates through the correction factor into the accumulator.
This was not hypothetical ΓÇö it reproduced immediately as NaN output on the
window=128cases. It is worth noting that it does not reproduce atwindow=1or at a window wider than the sequence, both of which passed whilethe bug was live. Those two cases alone would have signed off a broken kernel,
which is why the harness covers the mid-range and chunked cases too.
The guard contributes nothing for such a row and leaves its running state
untouched.
What
gqa_kernel.hip: window lower bound in the per-row mask; KV loop start-tileclamp; the all-masked-tile guard above;
local_window_sizethreaded throughdispatchPrefillV5, the tuner andprefillV5Run;v3now accepts a windowat
d == 64and still declines elsewhere.skv(removed from the key intesting oss20b perf #598 because it only scales the KV loop uniformly), the window caps how many
tiles a block walks, which is a different work profile and can genuinely
change which candidate wins.
gqa.cpp: the prefill call site was passing a hardcoded-1; it now passesthe real window, and the blanket
local_window_size <= 0rejection becomes awindow_okterm admittingd == 64prefill.Performance
Hardware and method: gfx1151,
model_benchmark -l 16384 --use_random_tokens -g 1 -r 2 -w 1 -b 1 -ml 0,HIPDNN_EP_AUTOTUNE=1, noHIPDNN_EP_PERForHIPDNN_EP_DEBUG. Three rounds, interleaved: prebuilt DLL sets per arm areswapped in and alternated so machine drift hits every arm equally. All three
arms were rebuilt from the same
main(f5ac3547).Interleaving is not optional on this box: absolute TTFT drifts about 10% over a
session. The sink arm measured -28.6% both in this session (21,049 against
29,483) and in an earlier one (23,108 against 32,346), i.e. the ratio reproduces
even though the absolutes moved by 3 s.
Test plan
pass. Ten are new window cases, chosen so a window bug cannot hide:
past_len8192 withwindow128, where whole KV tiles must be skippedrather than merely masked;
window=100against BKV of 32 and 64, to catch an off-by-one in thestart-tile clamp;
window=1, which must reduce to each query seeing only itself;d=128case assertingv3declines the window rather than ignoring it.cases.
the sink branch (22 tests compared one by one).
HIPDNN_EP_DEBUG=1: 96 of 96 prefill GQA calls fused.Pre-existing, not caused by this change: the same four
test_gqa_decode_fixed_cache_llama_shapefailures present at base.Notes for reviewers
The clamp is only correct because rows within a Q tile have monotonically
non-decreasing lower bounds; if that ever stopped holding, the kernel would
silently drop keys.
kv_lo, so enabling it is nearly a one-word change, but nothing here verifiesit and prefill is where the window was costing us.
gpt-oss: there is no longer a score matrix in its prefill. That work would now
only help models still on the decomposed path.
cases, the measurement harnesses and this description. I reviewed the change
and the evidence, and every number above is a measurement I ran.