[Triton/Gluon] [gfx950] pa_decode_sparse: pick BLOCK_K by occupancy, not a constant (up to +43 %) - #5294
Conversation
BLOCK_K was fixed at 64. That is the right choice only while every CTA
still gets a CU of its own -- the wider tile (4 warps, 64 KB LDS) does
more work per CU. Once the grid needs more than one wave, the narrower
tile (2 warps, 32 KB LDS) packs several CTAs per CU and hides the gather
latency, which dominates this kernel: it reaches only ~10 % of peak HBM
bandwidth on MI355X even with fully contiguous indices, and ~7 % of peak
flops, so neither memory nor math is the limit.
Deciding on the grid size rather than the query count matters because the
two are not proportional: num_splits varies, so 34 queries can launch more
CTAs than 64 do. Sorted by CTA count the picture is monotonic and the
crossover sits exactly at the CU count.
Measured on MI355X (gfx950), 64 heads, head_dim 512 (448 NoPE + 64 RoPE),
packed fp8 cache, top-k 2048, min of 5 rounds x 15 iters:
queries fixed K=64 this patch delta vs per-point optimum
16 63.2 us 62.6 us +0.9% -0.9%
24 78.3 72.9 +7.4% +0.4%
32 73.4 73.4 +0.1% -0.1%
48 121.3 109.7 +10.6% +0.0%
64 131.0 131.1 -0.1% +0.1%
96 218.0 193.8 +12.5% +0.1%
128 264.6 185.4 +42.7% +0.2%
192 402.8 355.2 +13.4% +0.0%
256 550.2 392.4 +40.2% -2.6%
Never worse than -0.6 % anywhere on the curve, and within +-0.9 % of the
per-point optimum throughout. Numerics are unchanged: max abs deviation
against a torch reference is 0.000061 for both tile widths.
The choice is made after num_splits is known, so it costs one comparison
and no extra work. Also applies unchanged to a 4-bit packed cache variant
we tested (128 queries: 289.9 -> 220.2 us).
Refs: ROCm#5293
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
|
Some context I found after opening this, which may help place the change: this is the measure #2416 recommended in March, applied to the sparse decode kernel. #2416 documents decode attention on gfx950 reaching 0.3–6 % of HBM bandwidth across Llama3-405B, Gemma2-9B, DeepSeek-R1 and Kimi-K2.5, and lists as root cause "No gfx950-specific decode attention tuning" with the recommendation "Consider decode-specific tile sizes". That recommendation was never carried out for Two notes on how our numbers relate to that issue:
The LDS numbers back the mechanism: gfx950 has 160 KB addressable LDS per workgroup (aiter's own This also lines up with AMD's own guidance in the ATOM decode blog post (2026-09-03), which names "a small-M GEMM or single-query attention operation may not produce enough workgroups to occupy all compute units" as a root cause and prescribes "smaller tile heights" for small-M paths — there applied to GEMM, here to sparse attention decode. |
There was a problem hiding this comment.
🟡 Changes recommended
The new BLOCK_K switching logic can leave num_splits computed for the old tile width, making split selection inconsistent with the chosen tile and potentially impacting performance across shapes.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Updates the gfx950 (CDNA4) Gluon driver for pa_decode_sparse to dynamically select BLOCK_K based on estimated occupancy, aiming to improve latency-bound performance at higher CTA counts.
Changes:
- Adds an occupancy-based heuristic that narrows
BLOCK_Kfrom 64 → 32 when the launch grid exceeds the device CU count. - Documents the rationale and measured performance impact for the new heuristic.
File summaries
| File | Description |
|---|---|
aiter/ops/triton/attention/pa_decode_sparse.py |
Adds occupancy-based BLOCK_K selection for the gfx950 sparse-decode Gluon driver. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # Measured on MI355X: predicts the faster tile on 21/21 load points, worth | ||
| # up to +43 % at 128 queries (264.5 -> 186.9 us) with identical numerics. | ||
| if num_queries * num_splits * heads_blocks > get_num_sms(): | ||
| BLOCK_K = 32 | ||
| num_warps = BLOCK_K // 16 |
Fixes the fixed
BLOCK_K = 64in the gfx950 sparse-decode driver, as described in #5293.What changes:
BLOCK_Kis chosen afternum_splitsis known — 32 when the launch produces more CTAs than there are CUs, 64 otherwise. One comparison, no extra work.Why the grid size and not the query count: the two are not proportional, because
num_splitsvaries. 34 queries can launch more CTAs than 64 do, which is why the fastest tile width is not monotonic innum_queries. Sorted by CTA count it is monotonic, and the crossover sits exactly atget_num_sms().Why it helps: this kernel is neither bandwidth- nor compute-bound — on MI355X it reaches ~10 % of peak HBM bandwidth (even with fully contiguous indices, so it is not the gather pattern) and ~7 % of peak flops. It is latency-bound, so occupancy is what matters. The narrow tile uses 32 KB LDS instead of 64 KB and lets several CTAs share a CU.
Measurements (MI355X/gfx950, 64 heads, head_dim 512 = 448 NoPE + 64 RoPE, packed fp8 cache, top-k 2048, minimum of 5 rounds × 15 iterations):
Never worse than −0.6 % anywhere on the curve. Numerics are unchanged — max abs deviation against a torch reference is 0.000061 for both tile widths.
The rule also transfers unchanged to a 4-bit packed cache variant we tested locally (128 queries: 289.9 → 220.2 µs), so it is not specific to the fp8 layout.
Things I measured that did not help, in case it saves someone the search:
GSPT8 or 32 instead of 16 (−14 % / −45 %),waves_per_eu1/2/4 (−4 % to −69 %), cache modifier.csinstead of.cg(does not compile), and forcingkv_splits(the automatic choice is already optimal at 64+ queries; more splits cost up to 45 %).Happy to re-measure on other shapes — different head counts, top-k values, or a bf16 cache — if that would help review.