Skip to content

[Triton/Gluon] [Config] Add a tuned gfx1101 MHA config, split small_head/default - #4493

Open
Ragua1 wants to merge 1 commit into
ROCm:mainfrom
Ragua1:mha-config-gfx1101
Open

Ragua1 wants to merge 1 commit into
ROCm:mainfrom
Ragua1:mha-config-gfx1101

Conversation

@Ragua1

@Ragua1 Ragua1 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Rewritten on 2026-09-10, after the branch was force-pushed from c27cabd7 to ec3af272. The earlier description covered the flat-layout file and values that no longer hold; that revision and the review anchored to it are unchanged in the history. What moved and why is in the comment below.

Motivation

aiter.ops.triton.attention.mha.flash_attn_func cannot run at all on gfx1101 (RDNA3). _get_config resolves configs/{arch}/triton/attention/mha/DEFAULT.json, and with no per-arch file and no fallback the call raises

FileNotFoundError: Required config file doesn't exist:
  aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json

before any kernel is compiled. gfx1101 is in RDNA_ARCHS, so the architecture is declared supported while the default MHA path is unreachable on it. This PR is the narrow remedy for the one architecture I can measure on.

What changed since the approval on 2026-08-06

This PR was approved at c27cabd7, and that revision no longer holds, so the branch has been rewritten rather than merged as approved. Two things moved underneath it: the config layout (#5018 / #5019 deleted the flat path without a fallback) and the kernel (#3936, #4414).

Re-measured against today's main, the approved fwd/default is 1.274x slower on head_dim 128, and on head_dim 64 it no longer establishes a gain — the intervals overlap the donor's. What moved is the donor, not the candidate: the inherited gfx1151 entry became faster on today's kernel while the tuned one stood still. Numerics did not move at all, so a correctness re-verification passes with digit-identical residuals — re-verifying rather than re-measuring would have merged the regression.

Rather than withdraw the tuning, this revision splits it in two, which removes the regression instead of trading one shape against another.

Technical Details

One new file, aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json, in the nested layout from #5019. 101 lines added, nothing removed, no code touched.

Nine of the eleven entries are byte-identical to gfx1151/triton/attention/mha/DEFAULT.json, the nearest tuned architecture. Two forward entries are tuned on gfx1101 instead of inherited:

fwd/default      BLOCK_M 64 -> 128, num_warps 4 -> 8, num_stages 2 -> 3   (vs. the gfx1151 donor)
fwd/small_head   BLOCK_M 128, num_warps 4, num_stages 1                   (new key, no donor equivalent)

fwd/small_head is the bucket added in #4414 for 16 < head_dim_v <= 64. Because it is opt-in per architecture by the presence of the key, this remains a data-only change that cannot affect any other architecture.

The split is needed because the two head_dim ranges want a different num_warps, and _get_config already discriminates on exactly that axis. Two independent coarse sweeps, one per range, selected these two entries separately — the numbers are below.

📌 The comment on that bucket does not describe this architecture, and these values should not be read as agreeing with it. It states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which num_stages=3 cures. On gfx1101 the ordering is the opposite: on the tuned M128 / N32 / w4 tile, num_stages 1 / 2 / 3 measured 2.346 / 2.429 / 2.450 ms. If that comment is meant as an invariant rather than as an observation about the architectures tuned so far, it changes the answer here, and I would rather be told than assume.

⚠️ num_stages is Triton-version-dependent on this architecture, so these values should not be assumed to carry backwards — click for why

Under the Triton used here, use_async_copy defaults to arch in ["gfx950", "gfx1250"] and is_pingpong_schedule_enabled admits only gfx942 and gfx950 (triton/backends/amd/compiler.py:43 and :33), so gfx1101 gets neither. num_stages therefore means classic software pipelining into registers only, which is why it can trade against BLOCK_M and why the answer differs between the two head-dim ranges. Under Triton 3.7.1 the same knob had a different gate — use_async_copy read the env knob without testing the arch — so these numbers should not be assumed to carry backwards.

fwd/dropout_or_fp32, both pe entries and both backward sections are inherited unchanged; only the forward path was exercised, so no claim is made about backward.

Test Plan

Windows native ROCm, RX 7800 XT (gfx1101, RDNA3, 60 CU), AITER checkout on PYTHONPATH with no install, AITER_TRITON_ONLY implicit on win32.

Every number below was measured on one software stack: torch 2.15.0a0+rocm10.1.0a20260822, HIP 7.16.26332, triton-windows 3.8.0.post28, a Windows build of Triton 3.8.0 and the newest release available. Older stacks are deliberately not reported, so that no row here describes a configuration that is no longer current.

fp16, seed 424242. Each point is the median of 20 timed iterations after 5 warmups with torch.cuda.synchronize() per iteration. Sweeps are single-shot rankings over a grid; every claimed ratio comes from a series of 10 independent repeats, and is claimed only where the [min, max] intervals across those repeats are disjoint. Where they overlap this is stated and nothing is claimed. Ratios are against the gfx1151 donor entry, since inheriting is the alternative to this file.

Two shape sets. Diffusion attention: SDXL self-attention (2, 4096, 4096, 10, 64), head_dim 64, and a Flux joint-attention shape (1, 4608, 4608, 24, 128), head_dim 128. LLM prefill, causal=True, GQA, batch 1, head dims from op_tests/op_benchmarks/triton/utils/model_configs.json (hidden_size / num_attention_heads); the sequence lengths are not in that file and are my choice, 4096 and 16384.

⚠️ The upstream MHA test and benchmark harness cannot run on this host, so its shape set is used with my own measuring path — click for the chain

op_tests/triton_tests/attention/test_mha.py and bench_mha.py both reach aiter.dtypes, which aiter/__init__.py binds only in the non-AITER_TRITON_ONLY branch, and the JIT path below it ends in RuntimeError("ROCm version file not found"). The upstream shape set is therefore used with my own measuring path rather than the upstream harness. This is independent of the present change, and I preferred not to introduce a workaround into the verification being reported.

Test Result

The defect is still present on the current base

With the new file removed from the checkout, on main @ 12620102 and the stack above:

FileNotFoundError: Required config file doesn't exist:
  .../aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json
  _triton_kernels/attention/mha.py:955  ->  utils/config_utils.py:50

It is raised on the first call and on every later call in the same process. On that base the architectures that do have an MHA config are gfx942, gfx950, gfx1151 and gfx1250.

Two escape hatches work on the same host and the same shape: passing config= explicitly, and mha_set_impl("dao_ai"). Both return max abs 6.470e-05 against the fp32 reference, against 6.013e-05 for torch's own fp16 SDPA measured in the same process. The kernel therefore runs correctly on this architecture, and it is only the config lookup that makes the default path unreachable.

Routing and correctness

Verified on main @ 12620102 with this file applied, each entry in a separate process because _get_config memoises. head_dim 64 routing to small_head and head_dim 128 routing to default is the behaviour the split depends on.

All six forward entries resolve as intended, max abs 6.469e-05 to 7.008e-05 — click for the per-entry output
v_head_dim 128, fp16            -> default              max abs vs fp32 reference 7.008e-05
v_head_dim  64, fp16            -> small_head           max abs vs fp32 reference 6.469e-05
v_head_dim  64, fp16 + dropout  -> dropout_or_fp32      finite, dropout active
v_head_dim  64, fp32            -> dropout_or_fp32      max abs 8.643e-07
v_head_dim  64, pe              -> pe                   max abs 6.685e-05
v_head_dim  64, pe + dropout    -> pe_dropout_or_fp32   finite, dropout active

Every configuration measured in this campaign produced an identical maximum absolute difference per shape, the inherited one included, so tile parameters have no numerical effect on this kernel and the choice is purely a performance decision.

How the two entries were selected

One coarse sweep per head-dim range, 72 candidates each over BLOCK_M x BLOCK_N x num_warps x num_stages, ranked on that range's shape. The two sweeps are independent and each selected the entry this PR ships for its own range.

sweep shape best of 72 ms vs. donor approved M128 w4 s1 on the same shape
head_dim 64 M128 N32 w4 s1 = this PR's small_head 2.346 0.974 it is that entry
head_dim 128 M128 N32 w8 s3 = this PR's default 6.444 0.933 8.619, i.e. 1.25

The separating knob is num_warps: at head_dim 128 the same tile measures 8.619 ms at w4 s1 against 6.657 ms at w8 s1, while at head_dim 64 the ordering reverses. A refine stage over PRELOAD_V and waves_per_eu on the three best candidates produced no cell that beat M128 N32 w8 s3 pv0 we1 on head_dim 128.

Three of the 72 candidates on the head_dim 128 shape are unusable on this architecture, and the failure is the expected one: out of resource: shared memory, Required: 69632 (resp. 73728), Hardware limit: 65536.

head_dim 128 — the regression and its removal

10 repeats per configuration, and the whole series run twice. [min, max] over the repeats in ms, ratio of medians against the donor in brackets.

series donor M64 w4 s2 approved M128 w4 s1 this PR's default, M128 w8 s3
run 1 [6.805, 7.028] [8.575, 8.776] (1.273) ❌ [6.437, 6.503] (0.944)
run 2 [6.787, 7.034] [8.698, 8.789] (1.274) ❌ [6.411, 6.513] (0.943)

✅ / ❌ = intervals disjoint from the donor's. The two runs agree to the third significant figure, and the same M128 w8 s3 entry costs 1.070–1.077 on head_dim 64 — which is the reason one fwd/default cannot serve both ranges.

LLM prefill shapes, causal=True, GQA, batch 1

Ratios against the donor, whose median in ms is given for scale.

All ten shapes: the entry this file routes to is the faster candidate in every row, significantly so in all ten — click for the table
shape head_dim donor ms M128 w4 s1 M128 w8 s3 routed by this file to
llama3-8B, sq 4096 128 3.885 1.331 ❌ 0.951 (overlap) default
llama3-8B, sq 16384 128 56.949 1.316 ❌ 0.904 default
mixtral-7B, sq 4096 128 3.950 1.331 ❌ 0.935 default
mixtral-7B, sq 16384 128 57.483 1.285 ❌ 0.895 default
kimi-k2.5 tp4, sq 4096 112 8.344 1.298 ❌ 0.892 default
kimi-k2.5 tp4, sq 16384 112 143.159 1.136 ❌ 0.775 default
deepseek-V3, sq 4096 56 8.746 0.828 0.924 ✅ small_head
deepseek-V3, sq 16384 56 167.621 0.644 0.714 ✅ small_head
glm-4.7-fp8 tp4, sq 4096 53 5.964 0.915 1.019 (overlap) small_head
glm-4.7-fp8 tp4, sq 16384 53 93.844 0.852 0.922 small_head

In every row the entry this file routes the shape to is the faster of the two candidates, and in nine of the ten it is significantly faster than the donor. Comparing the two candidates directly, their [min, max] intervals are disjoint in all ten rows — the routing decision is significant on every shape in the set, even where the winner's margin over the donor is not.

This set matters because #3560 tuned the donor on n=112 LLM shapes, so choosing fwd/default on diffusion shapes replaces an LLM-derived value with a diffusion-derived one. The measurement says the two families do not disagree on this card: both want num_warps 8 at head_dim 112–128, and the axis that separates them is head_dim, not workload.

head_dim 64, fwd/small_head

The LLM rows above are the evidence for this entry, and they are the only significant evidence for it.

On the diffusion shape it ranked first of 72 candidates in the sweep, but across 10 repeats it measures 0.933 and 0.936 with intervals that overlap the donor's in both runs, so no gain is claimed there.

The entry is nonetheless the right one for its range, and the comparison that establishes that is against the alternative rather than against the donor. Giving fwd/default to both ranges measures 1.070–1.077 on the diffusion head_dim 64 shape, and on the four head_dim 53–56 LLM shapes it is slower than M128 w4 s1 by 0.924 vs 0.828, 0.714 vs 0.644, 1.019 vs 0.915 and 0.922 vs 0.852 — disjoint intervals in all four.

What this does not establish

Batch > 1, the varlen / thd layout, sliding window, and the decode phase (seqlen_q = 1) were not measured. Forward path only; both backward sections are inherited. One card, one host and one software stack — no other RDNA or CDNA part is available to me, and nothing here is a claim about gfx1151, whose values this PR does not change.

A cross-attention shape (2, 4096, 77, 10, 64) was measured and is deliberately not reported: the donor's own timing on it varies by 31 % across repeats on this host, which is far above every other shape in the set, so it discriminates nothing in either direction.

A cross-architecture fallback for the same reachability defect is a separate change and is deliberately not bundled here. I am happy to file it separately if that is wanted.

@Ragua1
Ragua1 requested a review from a team July 31, 2026 21:42
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every PR:

  • ✅ Pre-checks (submodule verification, code formatting)
  • ✅ Aiter op tests (gfx942 + gfx950)
  • ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label Tests
ci:triton-300x Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4493 --add-label <label>

@zufayu
zufayu requested a review from vgokhale August 5, 2026 02:48
@vgokhale
vgokhale requested a review from carlushuang August 5, 2026 15:54
carlushuang
carlushuang previously approved these changes Aug 6, 2026

@carlushuang carlushuang left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

gfx1101 (RDNA3, e.g. RX 7800 XT) ships no MHA config, so `_get_config` in
`_triton_kernels/attention/mha.py` finds no
`configs/gfx1101/triton/attention/mha/DEFAULT.json` and every call to
`aiter.ops.triton.attention.mha.flash_attn_func` fails before a kernel runs.
Of the architectures in RDNA_ARCHS, only gfx1151 ships one today.

Nine of the eleven entries are taken verbatim from the gfx1151 donor
(RDNA3.5, added in ROCm#3423, tuned in ROCm#3560), which is the nearest tuned
architecture. Two forward entries are tuned on gfx1101 instead of inherited,
and they differ from each other only in `num_warps` and `num_stages`:

  fwd/default     BLOCK_M 128, num_warps 8, num_stages 3   (donor: 64 / 4 / 2)
  fwd/small_head  BLOCK_M 128, num_warps 4, num_stages 1

The split uses the `small_head` bucket added in ROCm#4414, which is opt-in per
architecture by the mere presence of the key, so this stays a data-only
change. It is needed because one `fwd/default` cannot serve both halves on
this card: measured against the donor, `M128 w4 s1` is 0.935x on head_dim 64
but 1.269x on head_dim 128, while `M128 w8 s3` is 0.914x on head_dim 128 but
1.062x on head_dim 64.

Measured on Windows native ROCm, triton 3.8.0, fp16, 10 independent repeats
of 20 iterations after 5 warmups, `torch.cuda.synchronize()` per iteration;
a result counts only when the [min, max] intervals across repeats are
disjoint. Ratios are against the gfx1151 donor entry, i.e. against what a
donor-inherited config would do.

  head_dim 128, `default`     Flux joint 0.914x / 0.891x / 0.892x on three
                              independent torch+ROCm stacks (2.11/7.15,
                              2.11/10.1, 2.15/10.1), all pinned to the same
                              triton; llama3-8B 0.882x, mixtral-7B 0.880x,
                              kimik25-tp4 0.849x at seqlen 16384
  head_dim <= 64, `small_head` SDXL self-attn 0.935x; deepseek-V3 0.658x and
                              glm47fp8-tp4 0.857x at seqlen 16384

The LLM shapes come from `op_tests/op_benchmarks/triton/utils/model_configs.json`
(prefill, causal, GQA, batch 1); the sequence lengths are not in that file and
are chosen here. Not covered: batch > 1, varlen/thd, sliding window, decode.

Note that the `small_head` comment in `_get_config` does not describe gfx1101.
It states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which
num_stages=3 cures. On this card the ordering is the opposite -- on the tuned
M128/N32/w4 tile, num_stages 1/2/3 measured 2.300 / 2.399 / 2.492 ms. The
bucket is still the right mechanism here, for a different reason: the two
head_dim ranges want a different `num_warps`, not a different `num_stages`.

Signed-off-by: Martin Domanský <ragua@email.cz>
@github-actions github-actions Bot changed the title [triton-mha] add gfx1101 tuning config [Triton/Gluon] [Config] add gfx1101 tuning config Sep 10, 2026
@Ragua1

Ragua1 commented Sep 10, 2026

Copy link
Copy Markdown
Contributor Author

The branch has been force-updated from c27cabd7 to ec3af272 and the description rewritten. GitHub dismissed the earlier approval on that push, and correctly so: it does not cover what is on the branch now, so this needs a fresh look rather than a merge on the strength of it. Apologies for the churn.

Two independent reasons, neither of which was visible when it was approved.

First, the path. The approved revision added configs/gfx1101-MHA-DEFAULT.json in the flat layout, which #5018 / #5019 removed without a fallback. _get_config now resolves configs/gfx1101/triton/attention/mha/DEFAULT.json, so merging the approved revision as-is would have added a file that nothing reads.

Second, and the reason this took longer: on re-measurement against today's main the approved values no longer hold either. The tuned fwd/default is 1.274x slower on head_dim 128, and on head_dim 64 it no longer establishes a gain at all — its intervals overlap the donor's. The approved values were tuned against an older kernel and an older Triton, and what moved since is the donor rather than the candidate: the inherited gfx1151 entry got faster on the current kernel while the tuned one stood still. Numerics did not move at all, so a correctness re-verification passes with digit-identical residuals; re-verifying rather than re-measuring would have merged the regression.

Rather than withdraw the tuning, this revision splits it across the small_head bucket added in #4414 (16 < head_dim_v <= 64) and fwd/default, which removes the regression instead of trading one shape against another. Two independent coarse sweeps, one per head-dim range, each selected the entry the file now ships for that range.

Both halves are then confirmed on ten LLM prefill shapes taken from op_tests/op_benchmarks/triton/utils/model_configs.json, which were not used for tuning. At head_dim 112–128 the default values give 0.775–0.935x against the donor on five of the six shapes, and 0.951x with overlapping intervals on the sixth. At head_dim 53–56 the small_head values give 0.644–0.915x on all four. Comparing the two candidates against each other rather than against the donor, their intervals are disjoint on all ten of those shapes, so the routing decision is supported on every shape in the set.

⚠️ Every number in the description comes from one current software stack — torch 2.15 on ROCm 10.1 with Triton 3.8.0 — and the whole series was run twice on it. Earlier stacks are deliberately not reported, so that no row describes a configuration that is no longer current. Protocol, the full tables and the shapes where nothing is claimed are in the description.

Still one new file, data only, no code touched, and no other architecture is affected — small_head is opt-in by the presence of the key.

One thing in the description I would rather flag here as well: the comment on that bucket in _get_config states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which num_stages=3 cures. On gfx1101 the measured ordering is the opposite, so this file uses the bucket for a different reason than the comment gives. If the comment is meant as an invariant rather than as an observation about the architectures tuned so far, please say so — that would change the answer.

@mgehre-amd — you added the gfx1151 MHA config in #3423 and tuned it in #3560, and nine of the eleven entries here are byte-identical to that file, so you are likely the right person to judge the two that are not. I have deliberately not asked @carlushuang to re-review: the values now on the branch are not the ones that were approved, and I would rather have that judged fresh than carried over.

@Ragua1 Ragua1 changed the title [Triton/Gluon] [Config] add gfx1101 tuning config [Triton/Gluon] [Config] Add a tuned gfx1101 MHA config, split small_head/default Sep 10, 2026
@Boss2002n
Boss2002n requested a lite review from Copilot September 11, 2026 01:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Unresolved correctness, backward-coverage, and bf16-validation findings must be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a gfx1101 Triton MHA configuration, making default dispatch reachable and splitting forward tuning by head dimension.

Changes:

  • Adds tuned fwd/default and fwd/small_head entries.
  • Adds inherited PE and backward configurations.
  • Adds architecture-specific MHA configuration data.

Review findings:

  • Critical (2 votes): Add d=8/16 correctness coverage or a safe route.
  • Moderate (2 votes): Validate or defer/gate gfx1101 backward support.
  • Moderate (1 vote): Measure bf16 or add dtype-specific selection.
File summaries
File Description
aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json Adds gfx1101 MHA tuning and dispatch entries.
Review details

Suppressed comments (1)

aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json:19

  • [verified] The new tiles are selected for both fp16 and bf16 because _get_config has no dtype-specific config branch, while every reported tuning run is fp16 and the benchmark defaults to bf16. A config that wins fp16 is not established to be a good bf16 choice, so this can regress the default model dtype on gfx1101. Author must measure the routed d64/d128 cases in bf16 as well, or add dtype-specific selection before shipping these values.
      "BLOCK_M": 128,
      "BLOCK_N": 32,
      "PRELOAD_V": false,
      "waves_per_eu": 1,
      "num_warps": 8,
      "num_ctas": 1,
      "num_stages": 3
  • Files reviewed: 1/1 changed files
  • Comments generated: 2
  • Review effort level: Lite

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json
Comment thread aiter/ops/triton/configs/gfx1101/triton/attention/mha/DEFAULT.json
@Ragua1

Ragua1 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Both review findings are answered in their own threads, with measurements on gfx1101 rather than from the source alone. Neither finding leads to a change in the file. The values stand as submitted.

What remains unmeasured is listed here rather than left to be discovered:

  • fp8, on any section.
  • The three inherited forward sections fwd/dropout_or_fp32, fwd/pe and fwd/pe_dropout_or_fp32. They are byte-identical to the gfx1151 donor and were not measured on gfx1101.
  • Backward performance. The backward arm is a correctness smoke test only.

If any of these is required before approval, please say which and it will be measured.

One question is left open deliberately, because it is a decision about _get_config rather than about this file. It is raised in a separate comment, so it can be answered or declined independently of this PR.

@Ragua1

Ragua1 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

A question about _get_config's bucketing rather than about this PR's file, raised separately so it can be answered independently.

The small_head branch selects on a band, 16 < head_dim_v <= 64. The hazard its comment documents is attached to a value instead:

Using num_stages=3 recovers performance and is numerically verified for these dims, but regresses d128 and miscompiles d<=16, so only 16 < d <= 64 uses this path.

The band and the value are equivalent only as long as fwd/default does not itself carry num_stages 3.

Reading the four architectures that ship an MHA table on main at 12620102:

arch ships fwd/small_head fwd/default fwd/small_head
gfx942 no BLOCK_M 128, num_warps 4, num_stages 1
gfx1151 no BLOCK_M 64, num_warps 4, num_stages 2
gfx1250 no BLOCK_M 128, num_warps 4, num_stages 1
gfx950 yes BLOCK_M 128, num_warps 4, num_stages 1 BLOCK_M 128, num_warps 4, num_stages 3

On gfx950, the only architecture that has the bucket, fwd/default is at s1 and small_head at s3. The bucket is what raises num_stages for the middle band, and the d <= 16 clause is what stops that rise from reaching the small dimensions. The band and the hazard coincide there.

Tuning gfx1101 produced the opposite polarity. d64 selected num_warps 4 with num_stages 1, and d128 selected num_warps 8 with num_stages 3, so in this PR small_head holds s1 and default holds s3.

That the two architectures tune differently is not itself surprising. gfx950 is a CDNA part with Wave64 and matrix cores, gfx1101 is RDNA3 with Wave32, and different optima are expected. What does not carry across is the protection: with the polarity reversed, the same lower bound routes d <= 16 toward s3 rather than away from it.

So the observation is not that the values differ. It is that the bucket expresses a band, while the constraint it was introduced to enforce is a property of a value. Those two agree for only one of the two polarities.

This is not a defect report. On gfx1101, num_stages showed no numerical effect at any head dimension, in fp16 or bf16, including under causal masking and at dimensions that are not powers of two. The question is whether the dispatch is intended to express the constraint at all.

Three readings, and the code does not tell us which is intended:

  1. The band is the whole intent, and it is each architecture's own responsibility to keep num_stages 3 out of fwd/default. If so, that requirement is today discoverable only from a comment on the branch it does not apply to. A line in configs/CLAUDE.md would make it checkable by whoever adds the next table.
  2. The d <= 16 clause is meant to bind generally. A lower bucket would then express it in the dispatch, rather than relying on every architecture's default staying below s3.
  3. The clause is specific to gfx950 and the compiler it was measured on, and is not meant to constrain other architectures.

Which of the three is it? If the answer is the second, we are willing to prepare that change as a separate PR. It touches the elif chain and every architecture's table, so it should not ride along with a config-only change, and it needs your view on the intended shape first.

@Ragua1
Ragua1 requested a review from carlushuang September 11, 2026 20:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants