feat(ck-tile): make dispatcher LDS capacity budget architecture-aware - #12173
Conversation
The codegen rejects any GEMM tile whose LDS staging footprint exceeds a cap keyed on the pipeline alone. With no architecture term, gfx950 (160 KB of LDS) and gfx1250 (320 KB) were both held to gfx942's 64 KB, so the largest and deepest tiles were never generated, benchmarked or selectable on those targets. Resolve the budget per architecture from each entry's lds_capacity_kb, a field that arch_specs.json already carried and ADDING_NEW_GPU.md already documented as mandatory, but that no code read. Pipelines now declare a basis (fraction of capacity, or an absolute limit) instead of a byte count, so a new GPU needs one number rather than ten. compv4 and preshufflev2 keep half the capacity because they are genuinely double-buffered: their GetSmemSize() returns twice the policy size, so half-of-capacity is the exact model rather than a margin. compv6 and preshufflev1 are not double-buffered; their historical halved cap is unexplained, so the ratio is preserved rather than widened without evidence. Also add gfx1250 to arch_specs.json. Its warp tables already existed in arch_specs_generated.py but not in the JSON they are generated from, so regenerating would have deleted them; they are transplanted verbatim. Both validators are updated together. Fixing only Python would emit configurations that the C++ filter then rejects at dispatch. Verified by enumerating the validator's survivor set before and after: gfx908/gfx90a/gfx942/gfx1100/gfx1200/gfx1201 byte-identical, gfx950 +2712 and gfx1250 +2885 configurations, zero lost on any target, and every gain attributable to the LDS check alone.
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
There was a problem hiding this comment.
🟡 Changes recommended
The mandatory-double-buffered comp_async pipeline incorrectly receives the full LDS budget, and C++ parity lacks automated coverage.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Makes dispatcher LDS validation architecture-aware, enabling larger GEMM tiles on gfx950 and gfx1250.
Changes:
- Adds per-architecture LDS budgets and gfx1250 specifications.
- Updates Python and C++ validators and generated artifacts.
- Adds architecture-aware LDS regression tests.
File summaries
| File | Description |
|---|---|
tests/test_lds_capacity_arch_aware.py |
Adds LDS budget regression tests. |
tests/CMakeLists.txt |
Registers the Python test suite. |
include/ck_tile/dispatcher/arch_specs_generated.hpp |
Emits architecture-specific C++ budgets. |
include/ck_tile/dispatcher/arch_filter.hpp |
Applies target architecture during validation. |
codegen/unified_gemm_codegen.py |
Displays per-architecture budgets. |
codegen/generate_arch_specs.py |
Resolves and emits LDS budgets. |
codegen/arch_specs.json |
Defines budget policy and gfx1250 data. |
codegen/arch_specs_generated.py |
Emits architecture-specific Python budgets. |
Review details
- Files reviewed: 9/9 changed files
- Comments generated: 5
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
The capacity check modelled a single staging buffer. Pipelines that ping-pong two LDS buffers allocate 2 * (A + B), so mem, compv3, compv5 and compv6 -- which make double buffering a configuration choice rather than a property of the pipeline -- could be budgeted for twice the LDS they actually use. This was harmless while the budget was capped at gfx942's 64 KB, since twice that still fit gfx950 and gfx1250. Widening the budget to the real capacity removed that accidental headroom: a double-buffered compv3 tile of 160 KB on gfx950 would have been accepted and then needed 320 KB. Thread the flag into both validators. The grouped convolution codegen already tracks it and now passes it through; the C++ filter reads algorithm.double_buffer, which registered kernels populate from SelectedKernel::DoubleSmemBuffer. Pipelines that always double already carry the halving in their per-pipeline budget, so the two are combined with a min() and compv4 and preshufflev2 are never halved twice. Defaults to single-buffered, so every existing caller is unaffected: the survivor set is unchanged on all eight architectures relative to the previous commit, and still identical to develop on the six whose LDS capacity is 64 KB.
github-actions
Bot
added
the
ck: convolution
comp_async had no entry in the pipeline budget table, so it inherited the full-capacity default. It allocates two LDS buffers unconditionally -- GemmPipelineAgBgCrCompAsync::GetSmemSize() returns num_lds_buffers * smem_size with num_lds_buffers = 2 -- so it was budgeted for twice the staging it can actually use. That was already wrong before this branch: on gfx942 the inherited 64 KB let through tiles needing 128 KB. Widening the budget per architecture made it worse rather than introducing it. Give it the half-capacity budget its buffering requires. On the 64 KB targets this removes configurations that were accepted but could never have launched; no default configuration set generates comp_async today, so nothing in tree loses a kernel. gfx1250 gains 185. Also list wavelet explicitly at full capacity. It is single-buffered, so the value is unchanged, but leaving it to fall through to "default" is exactly how comp_async went unnoticed. A test now requires every pipeline the validators can see to have a deliberate entry, and the generator maps wavelet to its enumerator so a future change to it cannot silently stop reaching C++. Fix a KeyError this branch introduced: the architecture-info listing indexed the budget table directly, so an unrecognised target raised instead of falling back to the smallest budget the way the validator does. Add Python/C++ parity tests. They compile the generated header and diff every architecture and pipeline value against the Python table, needing only a host compiler, so the two validators can no longer drift apart without CI noticing.
The C++ coverage added earlier compared budget accessor values only. One of those tests called itself an end-to-end check in its docstring, which claimed more than it did: it never reached ArchFilter::validate(), so a validator that ignored its architecture would still have passed it. Add a probe that builds a KernelKey and calls validate(), so the LDS check is reached the way a caller reaches it, and assert both that the 96 KB tile is rejected on gfx942 and accepted on gfx950 and that the rejection carries the LDS error rather than some other validator's. A second test requires the C++ and Python verdicts to agree for the same configuration. Scoped to gfx942 and gfx950 deliberately: they are the targets whose warp-tile tables agree between the two validators, so a rejection can only come from the LDS budget. Widening this to the other architectures needs the warp-tile lookup to be generated from arch_specs.json first. Verified the new tests have teeth by reverting the C++ side to an architecture-blind budget: both fail, and both pass again once restored. Renamed the accessor-level test to say what it actually checks.
|
Update For Copilot Review 1.
|
The generator formats the C++ header it emits, but passed a bare '-style=file'. clang-format resolves that by searching upward from the file being formatted, so writing the header outside the repository with --cpp-output-dir found no .clang-format, silently fell back to the built-in style, and produced a differently formatted file that does not satisfy the repository's format gate. Regenerating in place was always byte-identical, so the committed header was never stale; the output simply was not reproducible at an arbitrary destination. That is enough to make a reviewer's regenerate-and-diff check disagree with the tree, which is how this surfaced. Resolve the style file from the script's own location instead. Out-of- tree regeneration now matches the committed header exactly, where it previously differed by 613 lines and carried 414 format violations.
Review feedback: the product generation for this target is CDNA5, not RDNA4. The value was carried over verbatim from the generated module when gfx1250 was transcribed into arch_specs.json, so it predates this branch, but the transcription is where it becomes reviewable. Keep architecture as rdna and say why in the spec: arch.hpp states that gfx1250 shares the RDNA architecture with the GFX12 family while being its own standalone target family, because its MMA builtins and data-type ABI differ. The two fields record different things, the ISA lineage and the product generation, so they legitimately disagree here. The field is metadata only. Nothing branches on its value; the sole consumer checks that it is not None to reject unknown architectures in strict mode, so this cannot change which kernels are generated. Other references in the tree still describe gfx1250 as RDNA4, including a comment in arch.hpp and several test headers. Those are left alone here rather than renamed in an LDS-capacity change.
Review feedback: gfx1250 supports a 16x16x128 warp tile for fp8, and the spec listed only 16x16x64. Confirmed against the warp-gemm traits rather than taken on trust. warp_gemm_attribute_wmma_impl_8bit_traits.hpp declares dense WmmaTraits<gfx125_t, ...> specialisations for both shapes: fp8/fp8 at 16x16x64 and at 16x16x128, and likewise for bf8/bf8. The 128 form is a plain specialisation with no scale tag, so it belongs on the dense GEMM path this table governs, not only on the scaled path. Like the rest of the gfx1250 entry, the previous value was transplanted verbatim from the generated module, where it was already incomplete. Transcribing it into the spec is what made it reviewable. This widens the generated set on gfx1250: measured over a tile, wave and pipeline sweep, fp8 and bf8 each gain 123 configurations and nothing is lost. fp16 and bf16 are untouched, as are all other architectures. The mixed fp8/bf8 and bf8/fp8 pairs have traits for both shapes too, but are deliberately still absent. Adding them would newly generate kernels for dtype combinations this bridge has never exercised, which needs its own device validation rather than a spec edit. The reasoning is recorded next to the table.
|
Reviewed this on a worktree of Six findings: two I'd call blocking, four worth fixing. CriticalC1 — "zero configurations lost" is not true;
|
… the B term Review findings, all confirmed by reading the sources rather than taken on trust. kernel_key.hpp declared double_buffer with no member initializer, unlike the padding flags immediately below it. That was harmless while nothing in the validation path read the field. The LDS check now does, so a key that reached validation without assigning it would read an indeterminate value and pick a budget at random. Defaulted to false. Two generic key builders hardcoded double_buffer = true regardless of pipeline. A single-buffered pipeline built through them would be handed half the budget it is entitled to. Both now derive it from the pipeline, matching what the codegen already does. The ctypes helper also sets it true but pins the pipeline to CompV4, where true is correct, so it is left alone. The footprint model charged B at sizeof(BDataType). When the B cast policy runs before the LDS write, B is staged as ADataType instead (GetSmemSizeB), so a mixed-precision pair was under-counted. Both validators now charge B at the wider of the two, which is unchanged for equal dtypes.
|
Thank you — this is the most useful review this PR has had, and the MI350X numbers are the missing checkbox. Taking the findings in order. C1 — you are rightConfirmed and fixed in the PR body. I re-ran the enumeration with
Attribution across the six 64 KB architectures: 783 losses, all C2 — real, fixed
W1 — fixed in two of three
This also corrects something I asserted earlier in this PR. I had argued these builders were harmless because they produce query keys that never reach W4 — real, fixed
W2 — agreed in principle, and I would like your view on the remedyYour reading is correct: I have not applied a fix here because both options have a real cost. Modelling the padded size duplicates policy logic in the validator, which is its own drift risk. Dropping the fraction below 1.0 is arbitrary unless the number comes from somewhere — and it would discard headroom on exactly the tiles this PR exists to unlock. Your zero-occurrence result over 256 compiles is reassuring for fp16 but, as you say, says nothing about fp8 or gfx1250. Would you rather I take a conservative fraction now and refine it with data, or hold at 1.0 and gate it behind fp8 evidence on gfx1250 hardware when that exists? W3 — agreed, still scoping out, but the host/device split is new to meI had flagged the gfx1250 divergence and deliberately left it. The It still belongs outside an LDS-capacity change — it has three consumers in kernels this PR does not touch. But your point that this PR pushes more configurations into the divergent window is fair, so I will file it rather than leave it as a paragraph here. On the testsBoth notes are correct and both are on me. The C++ probes do On the benchmarkThank you for running it. Doubling the searched space, 128 to 256, with no compile failures and every remaining rejection accounted for by the deliberate half-capacity cap, is exactly the evidence the pending checkbox needed. The point that the sweep configs are byte-identical to the baseline is the part I would not have thought to check, and it is what makes the gain attributable to this change rather than to a different input. |
feat(ck-tile): make dispatcher LDS capacity budget
architecture-aware (#12173)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
JIRA ID : AICK-2249
## Summary
The dispatcher's codegen rejects any GEMM tile whose LDS staging
footprint exceeds a cap. The cap was keyed on the **pipeline only** and
carried no architecture term, so newer targets with substantially more
LDS than gfx942 were all held to gfx942's budget.
The effect is that the largest and deepest tiles — the ones most likely
to win on large GEMM shapes — were never generated, never benchmarked
and never selectable on those targets. On a gfx950 fp16 sweep, **160 of
288 offered configurations (56%) were rejected by this cap alone**.
This PR makes the budget architecture-aware, adds gfx1250 to the spec,
fixes both the Python and C++ validators together, and makes the check
account for ping-pong LDS staging.
**Scope:** the check is not GEMM-specific. `_validate_lds_capacity` sits
in the common `validate_kernel` chain, ungated by operator, and
`OperatorType` covers 11 operators (5 GEMM variants, 6 conv variants).
In Python it is reached by `unified_gemm_codegen.py:1595` (universal
GEMM) and `unified_grouped_conv_codegen.py:1941` (grouped convolution);
in C++ by `Registry::filter_by_arch()` (`registry.cpp:160`), which is
operator-agnostic and covers every registered kernel. The ctypes bridge
paths (quant variants, batched contraction, multi-ABD, FMHA) validate
via `validate_kernel_config` in `python/ctypes_utils.py`, which has no
LDS check at all, so they are unaffected either way.
Implements AICK-2249. Baseline evidence from AICK-2244.
## Motivation
`arch_specs.json` had `pipeline_lds_limits` as a **top-level** key — a
sibling of `architectures`, not nested inside it:
```json
"pipeline_lds_limits": {
"mem": <bytes>, "compv3": <bytes>, "compv4": <bytes>, ... "default": <bytes>
}
```
Every value was a fixed byte count, derived from gfx942 and applied to
every target.
There was no architecture axis, and nowhere to put one.
`_validate_lds_capacity` did `LDS_CAPACITY_LIMITS.get(config.pipeline,
...)` and nothing else.
Meanwhile `arch.hpp` has declared a per-architecture LDS capacity all
along, via `get_lds_size()`. On gfx942 the hardcoded cap and that
capacity coincide; on the newer targets they do not, and the shortfall
is large. This PR consumes the existing declaration rather than
introducing any new hardware figure.
This is the usual shape of this defect: a value that was correct when
there was exactly one architecture, frozen into a schema with no slot
for a second one. It survived because on gfx942 the correct answer and
the hardcoded answer are the same number, so it only ever fails
*silently, by generating less*.
The evidence that it is binding rather than theoretical: on the gfx950
sweep the cutoff falls **exactly** on the hardcoded cap for compv3/mem,
and exactly on the tighter cap for compv4, with zero exceptions in
either direction. And the winning kernels on that target were
`128x256x64` and `128x128x128` under compv3 — precisely tiles that
compv4 was forbidden from using. When the measured optimum sits on the
constraint boundary, the constraint is probably binding.
## Design note
**I did not invent a new schema.** A per-architecture `lds_capacity_kb`
field already existed in `arch_specs.json` (gfx950 already correctly
said `160`), and `ADDING_NEW_GPU.md` already documented it as a required
onboarding field. Nothing read it. Every new-GPU onboarding has been
filling in a mandatory field that no code consumed — the contract was
documented and unhonoured. So the fix is to honour it.
Pipelines now declare a **basis** rather than a byte count:
```json
"compv3": { "basis": "fraction", "value": 1.0 },
"compv4": { "basis": "fraction", "value": 0.5 }
```
resolved against each architecture's capacity at generation time. A new
GPU needs **one** number, not ten. The alternative — nesting the byte
table under each architecture — would have reintroduced the same failure
mode one level down: 8 arches x 10 pipelines of hand-maintained
constants free to drift from `arch.hpp` independently.
Resolution happens in the generator, once, so the emitted Python and C++
get literal byte counts and cannot disagree about how to read the
schema.
### On the tighter-capped pipelines — I read the sources rather than
guessing
Four pipelines sat at the tighter cap and the ticket flagged it as an
open question: genuine double-buffering that should scale with capacity,
or an independent absolute limit? On gfx942 the two readings are
numerically identical, which is why it was never forced. **The answer is
not uniform across the four:**
| pipeline | doubles LDS? | evidence |
|---|---|---|
| `compv4` | **yes** | `GemmPipelineAgBgCrCompV4::GetSmemSize()` returns
`2 * Policy::GetSmemSize<Problem>()` |
| `preshufflev2` | **yes** |
`WeightPreshufflePipelineAGmemBGmemCRegV2::GetSmemSize()` returns
`DoubleSmemBuffer ? 2 * smem_size : smem_size` |
| `compv6` | **no** | returns the policy size unmultiplied |
| `preshufflev1` | n/a | no such pipeline exists;
`preshuffle_pipelines.supported` lists only `preshufflev2` |
This is corroborated by `DOUBLE_SMEM_PIPELINES = {"compv4",
"preshufflev2", "comp_async"}`, which already exists in
`unified_batched_contraction_codegen.py` and names exactly the same two.
So for compv4 and preshufflev2, half-of-capacity is the **exact** model,
not a safety margin: the validator checks `A+B` against `capacity/2`,
which is precisely equivalent to checking the real `2*(A+B)` allocation
against full capacity.
For compv6 and preshufflev1 the halved cap is **unexplained by
buffering**. I preserved their ratio rather than widening them, because
no evidence supports a larger budget and a blind raise that regresses is
worse than the status quo. Widening those two is a separate change that
needs a measurement behind it. The reasoning is recorded in
`arch_specs.json` next to each value.
### Ping-pong staging had to be threaded in
The check modelled a *single* staging buffer. But
`unified_grouped_conv_codegen.py:148` documents that `mem`, `compv3`,
`compv5` and `compv6` make double buffering a **configuration choice**
(`--double-smem-buffer`), not a property of the pipeline — and the conv
path passed `pipeline=` into the validator without that flag, so the
budget could not see it.
That was harmless while the cap matched gfx942's capacity, because twice
that still fits the larger parts. **Widening the budget removes the
accidental headroom**, so I had to close it in the same PR:
Writing `C` for a target's LDS capacity:
| | compv3 budget | actually allocated | vs capacity | |
|---|---|---|---|---|
| develop | gfx942's cap | 2x that | below `C` | fits, by accident |
| arch-aware alone | `C` | `2C` | **over `C`** | **overflows** |
| with this fix | `C/2` | `C` | exactly `C` | fits exactly |
So the flag is now threaded into both validators. Conv passes the value
it already tracks; C++ reads `algorithm.double_buffer`, which registered
kernels populate from `SelectedKernel::DoubleSmemBuffer`
(`kernel_registration.hpp:60`).
Pipelines that *always* double already carry the halving in their
per-pipeline budget, so the two signals are combined with a `min()` —
`compv4` and `preshufflev2` are never halved twice, which is what keeps
gfx942 byte-identical. Independent confirmation that those are exactly
the two: `unified_gemm_codegen.py:1531` sets `double_buffer = pipeline
in ("compv4", "preshufflev2")`.
One deliberate non-use: `KernelConfig::build_key()`
(`kernel_config.hpp:280`) and `utils.hpp:676` hardcode `double_buffer =
true` for *any* pipeline. Those build **query** keys, which never reach
`validate_lds` — only registered-instance keys do — so the flag is
trustworthy at the one site that reads it. Worth fixing separately.
### A landmine worth calling out
`arch_specs_generated.py` **already contained gfx1250** — family, warp
configs, warp tile combos — but `arch_specs.json`, the file it is
generated from, did not. The generated C++ header did not either.
Three-way drift, in a file stamped `AUTO-GENERATED - DO NOT EDIT
DIRECTLY`.
This means following the documented workflow (edit JSON, regenerate,
commit) would have **silently deleted gfx1250's warp tables**. I
transplanted them into the JSON verbatim rather than authoring anything,
and the regeneration check below confirms every table came back
byte-identical.
## Test plan
- [x] New regression suite passes (15/15). Asserts budgets differ across
gfx942/gfx950/gfx1250, that gfx942's budget is byte-identical to the
historical table, that no budget exceeds the declared hardware capacity,
that unknown targets get the *smallest* budget rather than the largest,
and end-to-end that a large staging tile is rejected on gfx942 and
accepted on gfx950
- [x] Double-buffer cases covered: configurable pipelines halve, `2 x
budget <= capacity` on every arch/pipeline pair, always-double pipelines
are not halved twice, single-buffered remains the default, and
end-to-end a large staging tile is accepted on gfx950 single-buffered
but rejected double-buffered
- [x] Existing suites pass: `test_arch_filter_constraints`,
`test_gemm_utils`, `test_codegen_common`, `test_dispatcher_common`,
`test_tile_math`, `test_grouped_conv_codegen`, `test_grouped_conv_utils`
- [x] Generator is idempotent; regeneration leaves all non-LDS tables
**byte-identical** (`ARCH_FAMILY_MAP`, `WARP_SUPPORTED_COMBINATIONS`,
`WARP_TILE_SUPPORTED_COMBINATIONS`,
`PRESHUFFLE_WARP_TILE_SUPPORTED_COMBINATIONS`,
`TRAIT_UNSUPPORTED_COMBINATIONS`, `ELEMENT_SIZE_MAP`,
`DTYPE_COMBINATIONS`, `PRESHUFFLE_PIPELINES`)
- [x] C++ compiles standalone and returns values identical to Python for
every architecture and pipeline, single- and double-buffered, including
the unknown-architecture fallback
- [x] `clang-format-18 -style=file` clean on both headers
- [ ] **GPU validation pending** — see below
### Equivalence proof
This edits a shared table, so I enumerated the validator's survivor set
per architecture, before and after. The axes now include **every
pipeline the validators can see** — `comp_async` and `wavelet` as well
as the nine that were in the old table — and both settings of the
ping-pong staging flag.
| arch | develop | this PR | gained | lost |
|---|---:|---:|---:|---:|
| gfx908 | 4,613 | 7,386 | +3,003 | **230** |
| gfx90a | 2,701 | 4,138 | +1,595 | **158** |
| gfx942 | 5,324 | 8,360 | +3,322 | **286** |
| gfx950 | 7,116 | 19,016 | +11,900 | 0 |
| gfx1100 | 1,155 | 1,870 | +770 | **55** |
| gfx1200 | 442 | 668 | +253 | **27** |
| gfx1201 | 442 | 668 | +253 | **27** |
| gfx1250 | 2,012 | 10,926 | +8,914 | 0 |
| **total** | | | **+30,010** | **783** |
**There are losses, and an earlier version of this section claimed there
were none.** That claim was measured over the old parameter space, which
could not contain `comp_async` because the pipeline had no entry in the
old table to enumerate. It was true of what it measured and wrong as a
general statement.
**All 783 losses are `comp_async`, on the six architectures whose budget
is unchanged**, and they are intended. Attribution by pipeline across
those six:
```
comp_async 783
(no other pipeline appears)
```
`comp_async` had no entry before, so it inherited the full-capacity
default. It allocates two LDS buffers unconditionally — `GetSmemSize()`
returns `num_lds_buffers * smem_size` with `num_lds_buffers = 2` — so it
was budgeted for twice the staging it can actually use. A tile such as
conv `128x64x128` at fp16 needs 48 KB of staging and was accepted, while
the kernel would then have asked the hardware for twice that. Rejecting
it is the fix, not a regression.
The two architectures with a raised budget lose nothing, and no default
configuration set in tree generates `comp_async`, so no shipped kernel
disappears.
## Explicitly out of scope
Two other filters in the same validation chain share this defect class
(gfx942-derived constraints applied to every architecture). Neither
rejected anything in the baseline campaign, so I left them alone:
- `TRAIT_UNSUPPORTED_COMBINATIONS` — 10 tuples, no arch key
- `_cshuffle_store_ok` — docstring says "GPU-verified on gfx942",
applied to all architectures
A wider audit of the validation chain found this same pattern in several
more places, including a hand-written warp-tile lookup on the native
side that disagrees with the generated Python table on a majority of
supported architectures. That work is tracked separately in AICK-2264
and AICK-2266, under epic AICK-2265, and is deliberately not part of
this PR: it is a different table on a different axis, and unlike this
change it removes configurations, so it needs its own before/after
enumeration and device validation.
Also left alone: `get_smem_capacity()` in `arch.hpp` special-cases a
single architecture and returns a fixed value otherwise, which disagrees
with `get_lds_size()` for at least one target. It has three live
consumers outside this ticket's scope (`cshuffle_epilogue.hpp`,
`moe_sorting_kernel.hpp`, `grouped_convolution_forward_kernel.hpp`) and
may legitimately encode a per-workgroup addressable limit rather than
per-CU capacity. Flagging rather than changing it.
TheRock Submodule Bump ActivityNewest first
|
JIRA ID : AICK-2249 ## Summary gfx1250 is labelled **RDNA4** in several comments and one CI config string. The product generation is **CDNA5**. This corrects the label. Comments and one JSON `_comment` only — no code, no behaviour change. Raised in review of #12173, where the same value was corrected in `arch_specs.json`. Split out because renaming terminology across the tree does not belong in an LDS-capacity change. ## Why the RDNA4 label was there It was not arbitrary. gfx1250 has no MFMA units and multiplies through **WMMA**, which is an RDNA-family instruction, and the `gfx12xx` numbering points the same way. The observation about the *instruction path* was right; using it as a *generation* label was not. `arch.hpp` already draws exactly this distinction: > GFX1250 is its own standalone family. Although it **shares the RDNA architecture** with the GFX12 family, its MMA builtins and data-type ABI differ, so it must not be treated as a GFX12-family device. Sharing an ISA lineage is not the same as belonging to that generation. So there are two separate claims in play, and only the second one holds. ## Changes Where a generation label is meant, it now says CDNA5. Where only the instruction path matters, the generation is dropped rather than replaced — a WMMA path is a WMMA path regardless of which generation the part belongs to. | file | before | after | |---|---|---| | `dispatcher/tests/test_gemm_utils.py` | `gfx1250 (MI400 / RDNA4-WMMA) enablement` | `gfx1250 (MI400 / CDNA5, WMMA) enablement` | | `dispatcher/tests/test_batched_bridge.py` | same | same | | `dispatcher/tests/test_batched_contraction_bridge.py` | same | same | | `dispatcher/tests/test_multi_d_bridge.py` | `it runs the RDNA4 WMMA path` | `it runs the WMMA path` | | `dispatcher/codegen/grouped_conv/grouped_config_rules_full.py` | `such as rdna4/gfx1250` | `such as gfx1250` | | `tile_engine/.../bridge_default_ci_config_gfx1250.json` | `(MI400, RDNA4/WMMA)` | `(MI400, CDNA5/WMMA)` | ## Deliberately not changed - **`include/ck_tile/core/arch/arch.hpp`** and **`.../mma/scale/wmma/selector.hpp`** describe the ISA lineage and the selector strategy as RDNA. Both are accurate — the first is the quote above, the second documents an RDNA-style selector. Renaming these would replace a correct statement with a wrong one. - **`dispatcher/codegen/arch_specs_generated.py`** still reads `rdna4`. It is generated, and its source is the gfx1250 entry in `arch_specs.json` that #12173 adds; the value is corrected there. Hand-editing a generated file is the exact trap this codebase keeps hitting, so it is not done here. That last point means the tree is briefly inconsistent until #12173 lands. The affected field is metadata: nothing branches on its value, and its only consumer is a `None` check that rejects unknown architectures in strict mode. ## Test plan - [x] `test_gemm_utils`, `test_batched_bridge`, `test_batched_contraction_bridge`, `test_multi_d_bridge`, `test_grouped_conv_codegen` all pass - [x] The edited CI config still parses as JSON - [x] No `gfx1250`/`MI400` reference to RDNA4 remains outside the two files listed above as intentionally kept ## Note for the reviewer The label originated in #10921, the gfx1250 bridge enablement PR, and spread from there; AICK-2082 uses the same wording. So this is a retroactive terminology correction rather than a one-off typo. If **CDNA5 is not the right generation either**, say so before this merges — I have taken it from review feedback on #12173 and have not independently confirmed it. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
JIRA ID : AICK-2249
Summary
The dispatcher's codegen rejects any GEMM tile whose LDS staging footprint exceeds a cap. The cap was keyed on the pipeline only and carried no architecture term, so newer targets with substantially more LDS than gfx942 were all held to gfx942's budget.
The effect is that the largest and deepest tiles — the ones most likely to win on large GEMM shapes — were never generated, never benchmarked and never selectable on those targets. On a gfx950 fp16 sweep, 160 of 288 offered configurations (56%) were rejected by this cap alone.
This PR makes the budget architecture-aware, adds gfx1250 to the spec, fixes both the Python and C++ validators together, and makes the check account for ping-pong LDS staging.
Scope: the check is not GEMM-specific.
_validate_lds_capacitysits in the commonvalidate_kernelchain, ungated by operator, andOperatorTypecovers 11 operators (5 GEMM variants, 6 conv variants). In Python it is reached byunified_gemm_codegen.py:1595(universal GEMM) andunified_grouped_conv_codegen.py:1941(grouped convolution); in C++ byRegistry::filter_by_arch()(registry.cpp:160), which is operator-agnostic and covers every registered kernel. The ctypes bridge paths (quant variants, batched contraction, multi-ABD, FMHA) validate viavalidate_kernel_configinpython/ctypes_utils.py, which has no LDS check at all, so they are unaffected either way.Implements AICK-2249. Baseline evidence from AICK-2244.
Motivation
arch_specs.jsonhadpipeline_lds_limitsas a top-level key — a sibling ofarchitectures, not nested inside it:Every value was a fixed byte count, derived from gfx942 and applied to every target.
There was no architecture axis, and nowhere to put one.
_validate_lds_capacitydidLDS_CAPACITY_LIMITS.get(config.pipeline, ...)and nothing else.Meanwhile
arch.hpphas declared a per-architecture LDS capacity all along, viaget_lds_size(). On gfx942 the hardcoded cap and that capacity coincide; on the newer targets they do not, and the shortfall is large. This PR consumes the existing declaration rather than introducing any new hardware figure.This is the usual shape of this defect: a value that was correct when there was exactly one architecture, frozen into a schema with no slot for a second one. It survived because on gfx942 the correct answer and the hardcoded answer are the same number, so it only ever fails silently, by generating less.
The evidence that it is binding rather than theoretical: on the gfx950 sweep the cutoff falls exactly on the hardcoded cap for compv3/mem, and exactly on the tighter cap for compv4, with zero exceptions in either direction. And the winning kernels on that target were
128x256x64and128x128x128under compv3 — precisely tiles that compv4 was forbidden from using. When the measured optimum sits on the constraint boundary, the constraint is probably binding.Design note
I did not invent a new schema. A per-architecture
lds_capacity_kbfield already existed inarch_specs.json(gfx950 already correctly said160), andADDING_NEW_GPU.mdalready documented it as a required onboarding field. Nothing read it. Every new-GPU onboarding has been filling in a mandatory field that no code consumed — the contract was documented and unhonoured. So the fix is to honour it.Pipelines now declare a basis rather than a byte count:
resolved against each architecture's capacity at generation time. A new GPU needs one number, not ten. The alternative — nesting the byte table under each architecture — would have reintroduced the same failure mode one level down: 8 arches x 10 pipelines of hand-maintained constants free to drift from
arch.hppindependently.Resolution happens in the generator, once, so the emitted Python and C++ get literal byte counts and cannot disagree about how to read the schema.
On the tighter-capped pipelines — I read the sources rather than guessing
Four pipelines sat at the tighter cap and the ticket flagged it as an open question: genuine double-buffering that should scale with capacity, or an independent absolute limit? On gfx942 the two readings are numerically identical, which is why it was never forced. The answer is not uniform across the four:
compv4GemmPipelineAgBgCrCompV4::GetSmemSize()returns2 * Policy::GetSmemSize<Problem>()preshufflev2WeightPreshufflePipelineAGmemBGmemCRegV2::GetSmemSize()returnsDoubleSmemBuffer ? 2 * smem_size : smem_sizecompv6preshufflev1preshuffle_pipelines.supportedlists onlypreshufflev2This is corroborated by
DOUBLE_SMEM_PIPELINES = {"compv4", "preshufflev2", "comp_async"}, which already exists inunified_batched_contraction_codegen.pyand names exactly the same two.So for compv4 and preshufflev2, half-of-capacity is the exact model, not a safety margin: the validator checks
A+Bagainstcapacity/2, which is precisely equivalent to checking the real2*(A+B)allocation against full capacity.For compv6 and preshufflev1 the halved cap is unexplained by buffering. I preserved their ratio rather than widening them, because no evidence supports a larger budget and a blind raise that regresses is worse than the status quo. Widening those two is a separate change that needs a measurement behind it. The reasoning is recorded in
arch_specs.jsonnext to each value.Ping-pong staging had to be threaded in
The check modelled a single staging buffer. But
unified_grouped_conv_codegen.py:148documents thatmem,compv3,compv5andcompv6make double buffering a configuration choice (--double-smem-buffer), not a property of the pipeline — and the conv path passedpipeline=into the validator without that flag, so the budget could not see it.That was harmless while the cap matched gfx942's capacity, because twice that still fits the larger parts. Widening the budget removes the accidental headroom, so I had to close it in the same PR:
Writing
Cfor a target's LDS capacity:CC2CCC/2CCSo the flag is now threaded into both validators. Conv passes the value it already tracks; C++ reads
algorithm.double_buffer, which registered kernels populate fromSelectedKernel::DoubleSmemBuffer(kernel_registration.hpp:60).Pipelines that always double already carry the halving in their per-pipeline budget, so the two signals are combined with a
min()—compv4andpreshufflev2are never halved twice, which is what keeps gfx942 byte-identical. Independent confirmation that those are exactly the two:unified_gemm_codegen.py:1531setsdouble_buffer = pipeline in ("compv4", "preshufflev2").One deliberate non-use:
KernelConfig::build_key()(kernel_config.hpp:280) andutils.hpp:676hardcodedouble_buffer = truefor any pipeline. Those build query keys, which never reachvalidate_lds— only registered-instance keys do — so the flag is trustworthy at the one site that reads it. Worth fixing separately.A landmine worth calling out
arch_specs_generated.pyalready contained gfx1250 — family, warp configs, warp tile combos — butarch_specs.json, the file it is generated from, did not. The generated C++ header did not either. Three-way drift, in a file stampedAUTO-GENERATED - DO NOT EDIT DIRECTLY.This means following the documented workflow (edit JSON, regenerate, commit) would have silently deleted gfx1250's warp tables. I transplanted them into the JSON verbatim rather than authoring anything, and the regeneration check below confirms every table came back byte-identical.
Test plan
2 x budget <= capacityon every arch/pipeline pair, always-double pipelines are not halved twice, single-buffered remains the default, and end-to-end a large staging tile is accepted on gfx950 single-buffered but rejected double-bufferedtest_arch_filter_constraints,test_gemm_utils,test_codegen_common,test_dispatcher_common,test_tile_math,test_grouped_conv_codegen,test_grouped_conv_utilsARCH_FAMILY_MAP,WARP_SUPPORTED_COMBINATIONS,WARP_TILE_SUPPORTED_COMBINATIONS,PRESHUFFLE_WARP_TILE_SUPPORTED_COMBINATIONS,TRAIT_UNSUPPORTED_COMBINATIONS,ELEMENT_SIZE_MAP,DTYPE_COMBINATIONS,PRESHUFFLE_PIPELINES)clang-format-18 -style=fileclean on both headersEquivalence proof
This edits a shared table, so I enumerated the validator's survivor set per architecture, before and after. The axes now include every pipeline the validators can see —
comp_asyncandwaveletas well as the nine that were in the old table — and both settings of the ping-pong staging flag.There are losses, and an earlier version of this section claimed there were none. That claim was measured over the old parameter space, which could not contain
comp_asyncbecause the pipeline had no entry in the old table to enumerate. It was true of what it measured and wrong as a general statement.All 783 losses are
comp_async, on the six architectures whose budget is unchanged, and they are intended. Attribution by pipeline across those six:comp_asynchad no entry before, so it inherited the full-capacity default. It allocates two LDS buffers unconditionally —GetSmemSize()returnsnum_lds_buffers * smem_sizewithnum_lds_buffers = 2— so it was budgeted for twice the staging it can actually use. A tile such as conv128x64x128at fp16 needs 48 KB of staging and was accepted, while the kernel would then have asked the hardware for twice that. Rejecting it is the fix, not a regression.The two architectures with a raised budget lose nothing, and no default configuration set in tree generates
comp_async, so no shipped kernel disappears.Explicitly out of scope
Two other filters in the same validation chain share this defect class (gfx942-derived constraints applied to every architecture). Neither rejected anything in the baseline campaign, so I left them alone:
TRAIT_UNSUPPORTED_COMBINATIONS— 10 tuples, no arch key_cshuffle_store_ok— docstring says "GPU-verified on gfx942", applied to all architecturesA wider audit of the validation chain found this same pattern in several more places, including a hand-written warp-tile lookup on the native side that disagrees with the generated Python table on a majority of supported architectures. That work is tracked separately in AICK-2264 and AICK-2266, under epic AICK-2265, and is deliberately not part of this PR: it is a different table on a different axis, and unlike this change it removes configurations, so it needs its own before/after enumeration and device validation.
Also left alone:
get_smem_capacity()inarch.hppspecial-cases a single architecture and returns a fixed value otherwise, which disagrees withget_lds_size()for at least one target. It has three live consumers outside this ticket's scope (cshuffle_epilogue.hpp,moe_sorting_kernel.hpp,grouped_convolution_forward_kernel.hpp) and may legitimately encode a per-workgroup addressable limit rather than per-CU capacity. Flagging rather than changing it.