Repository navigation
feat(rocke): enable head-dim-256 bf16 flash prefill on gfx942 + gfx950 - #9129
Merged
Merged
Conversation
✅ 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. |
Contributor
There was a problem hiding this comment.
Pull request overview
Enables bf16 head-dim-256 (D256) tiled-attention prefill on gfx950 by adding a deterministic “LDS budget resolver” that adjusts 2D register-PV tiling parameters when the calculated LDS footprint exceeds the architecture cap, avoiding comgr compile aborts.
Changes:
- Add
_resolve_lds_budget()to model register-PV LDS usage and apply validated reductions (K single-buffer and/orT=64) to fit within the arch’s LDS cap. - Hook the resolver into
_tiled_spec_from_problem()so the 2D spec is post-processed before returning. - Provide a clearer runtime error when no validated reduction can bring LDS under budget.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-prefill
branch
from
July 7, 2026 03:33
3b82fd6 to
fd23a64
Compare
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-prefill
branch
from
July 7, 2026 07:40
fd23a64 to
8f31884
Compare
AviralGoelAMD
changed the title
AviralGoelAMD
marked this pull request as ready for review
July 7, 2026 07:42
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-prefill
branch
from
July 7, 2026 07:52
8f31884 to
24f91ae
Compare
The register-PV 2D-tiled prefill path stages only K/V tiles in LDS; at
head_dim=256 the default double-buffered K tile is 204800 B, over gfx950's
163840 B cap, so the kernel fails to compile ("local memory (204800) exceeds
limit (163840)"). Add _resolve_lds_budget, hooked on _tiled_spec_from_problem's
main return: it computes the exact register-PV LDS footprint and, when over the
arch cap, deterministically shrinks the geometry (single-buffer K, then T=64)
until it fits -- a strict no-op otherwise, so all currently-compiling configs
are byte-identical. Scoped to gfx950 register-PV; mechanism reads the cap
dynamically.
GPU-verified on MI350X (bf16, GQA 16/2): D256 prefill compiles (139264 <=
163840), runs, and matches the fp32 paged reference (max_abs 1.5625e-02) at
Sq 4096 and 8192.
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-prefill
branch
from
July 7, 2026 09:08
24f91ae to
192cc66
Compare
AviralGoelAMD
marked this pull request as draft
July 7, 2026 09:09
… dedup footprint model The Python D256 gfx950 register-PV resolver had no C++ twin: once attention is wired through the shipped provider, its spec-assembly path would select the overflowing K-double geometry and fail comgr CODEGEN. Mirror the selection-layer resolver in the C++ selectors so the provider shrinks the geometry the same way (single-buffer K, then T=64) before build. - rocke_unified_attn_tiled_2d_lds_bytes: single source of truth for the tiled-2D LDS footprint, shared by the gfx942 admission gate and the register-PV budget resolver in both engines -- removes the duplicated per-buffer arithmetic that could drift. - rocke_unified_attn_resolve_lds_budget: byte-faithful C twin of _resolve_lds_budget; strict no-op for fitting / non-register-PV / non-gfx950. - Python: _tiled_2d_lds_bytes now backs both _lds_bytes_regpv and the gfx942 gate (byte-identical numbers); the shrink levers narrow their except to ValueError and thread the reject reason into the resolver diagnostic. - Tests: footprint single-source, lever reason-threading, infeasible-raises.
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-prefill
branch
from
July 7, 2026 20:28
192cc66 to
3738343
Compare
AviralGoelAMD
marked this pull request as ready for review
July 9, 2026 01:16
…d256-prefill # Conflicts: # dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py
AviralGoelAMD
changed the title
…-bytes The shared _tiled_2d_lds_bytes footprint model counted fp8 K/V cache tiles at 2 B/elem: _kv_storage_dtype() emits "fp8e4m3" but _kv_lds_elem_bytes() (Py) and attn_kv_lds_elem_bytes() (C++) only matched "fp8"/"bf8"/"e4m3"/"e5m2". Add the full dtype strings so fp8-cache specs are modeled at 1 B/elem in both engines. bf16/fp16 (incl. D256) unaffected; byte-identity green.
yraparti
reviewed
Jul 9, 2026
yraparti
approved these changes
Jul 9, 2026
Addresses review feedback: the resolver already read the LDS cap dynamically from the arch-target API, but was gated behind a hard-coded `gfx950` literal (and a 163840 fallback). Drop both: engage purely on the register-PV path + the dynamically-read `ArchTarget.lds_capacity_bytes`, and no-op when the arch has no declared cap. Behaviour is unchanged on all current arches: register-PV over-budget specs on non-gfx950 arches are already filtered upstream by the more-conservative (Q/P-staged) supports_tiled_2d admission gate, so any spec that reaches the resolver there already fits -> strict no-op, byte-identity preserved. Mirrored in the C++ selector. New unit test asserts the resolver targets whatever cap the arch API reports (fits at a larger cap; engages+raises at a tiny one) rather than a gfx950 constant.
AnaghaRaoAMD
approved these changes
Jul 10, 2026
AviralGoelAMD
added a commit
that referenced
this pull request
Jul 22, 2026
… fast path (#9233) ## What Routes the **gfx950 D256 bf16 prefill** cohort to the optimized **32×32 transposed fast path** + **FA3-style softmax↔MFMA interleave** (mode2 / groups=4). A single routing change in `_tiled_spec_from_problem`: - **gate the 3 launch-critical geometry selectors** (`num_warps` / `tile_size` / `block_m_per_warp`) so the builder spec, `_tiled_cache_key`, and `_get_2d_launch_meta` all agree; - **spec-override the codegen flags** (`use_mfma_32x32`, transposed stack, `register_pv` off, K single-buffer, `use_softmax_mfma_interleave` mode2/g4). The kernel codegen the flags trigger already ships on the base; this is routing only. Narrow + arch/feature-guarded — **every non-D256 shape is byte-identical**. ## Perf (MI355X / gfx950, direct-launch `path="2d"`, fp32-ref `max_abs=1.5625e-02` — all correct) D256, bf16, GQA 16/2, causal prefill, **single long sequence (batch 1)** — the Qwen3-Next `gated_attn` prefill regime (AICK-1495, Sq 4096/8192). Numbers are **post the K_lds slab-pad (#9260)**, i.e. what this branch actually ships. `bsN` here = paged **block_size** (batch = 1), not batch size. | shape | CK TF/s | AOTriton TF/s | CK µs | AOTriton µs | CK vs AOTriton | |---|---|---|---|---|---| | Sq4096, block_size 64 | 440.0 | 226.7 | 312.4 | 606.2 | **1.94× (+94%)** | | Sq8192, block_size 64 | 556.5 | 403.5 | 987.9 | 1362.5 | **1.38× (+38%)** | | Sq4096, block_size 16 | 356.2 | 228.0 | 385.8 | 602.9 | **1.56× (+56%)** | | Sq8192, block_size 16 | 442.4 | 401.9 | 1242.6 | 1368.0 | **1.10× (+10%)** | **CK beats AOTriton at every ticket shape.** Two stacked levers: (1) the FA3-style softmax↔MFMA interleave (mode2/g4) hides the exposed `v_exp` on the low-occupancy (4 waves/CU) kernel; (2) the slab-granularity K_lds pad (#9260) breaks the row-aliased QK bank conflict (`SQ_LDS_BANK_CONFLICT` 897→660 /1000). GPU-verified via the direct-launch lane; numerics checked against an fp32 paged reference (all PASS, `max_abs=1.5625e-02`). > **Regime note.** These are batch-1 long-prefill (the ticket workload). At high *aggregate batch* the Sq8192 margin narrows/reverses (AOTriton's persistent work-stealing grid fills better at long context), and true batch-64 × Sq8192 currently exceeds the 2D launcher's `grid.y` limit (65600 > 65535) — a separate, pre-existing scaling limit, not on the ticket path. ## Also included A **Step 0 caveat** in the optimization runbook: the preset bench (`_variant_flags`) is *not* the exhaustive lever sweep Step 0 prescribes — it can't express default-off flags like `use_softmax_mfma_interleave`, which is exactly why this win hid for a full loop. (Docs-only; happy to split into its own PR if preferred.) ## Notes - **Stacked on #9129** (`d256-prefill` LDS-budget resolver) — base is `users/avirgoel/ck/d256-prefill`. - Relates to **AICK-1495 (S7)**. - Draft: numeric verified on-silicon via the direct-launch sweep; leaving as draft for CI + byte-identity gate + a maintainer eye on the routing. --------- Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
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.
Enables the unoptimized head-dim-256 (D256) bf16 flash prefill kernel to compile and run correctly on gfx942 + gfx950..
Problem
hipcc: local memory (204800) exceeds limit (163840).Changes (6 files, +522/−16)
_resolve_lds_budget(gfx950) — deterministically shrinks an over-budget register-PV 2D spec to fit the arch LDS cap via a validated, cheapest-first ladder: single-buffer-K → T=64 → both. No-op for anything that already fits or is not register-PV (every currently-compiling config stays byte-identical); raises a clear diagnostic if nothing fits. Hooked as a post-processor on_tiled_spec_from_problem._tiled_2d_lds_bytes(shared footprint model) — single source of truth for the tiled-2D LDS math, used by both the gfx942 admission gate and the gfx950 resolver so they can't silently drift. gfx942supports_tiled_2drefactored onto it (byte-identical for bf16/fp16; Acc term now correctly fixed 16-bit).attention_unified_selectors.{h,cpp}+ gfx942 glue, preserving the dual-backend (C++ ≡ Python) contract.test_lds_budget_resolver.py(8 cases: shrink ladder, no-op off-gfx950, non-register-PV passthrough, infeasible-raise).Verification
JIRA ID - #1495