Repository navigation
feat(rocke): D256 gfx950 bf16 prefill 32x32 + softmax-MFMA interleave fast path - #9233
Merged
Merged
Conversation
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.
… 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.
…d256-prefill # Conflicts: # dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py
…-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.
…interleave fast path Route the gfx950 D256 bf16 prefill cohort to the 32x32 transposed fast path plus the FA3-style softmax<->MFMA interleave (mode2, groups=4). Gates the three 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; the codegen flags (mfma_32x32, transposed stack, register_pv off, K single-buffer, interleave) are applied by the spec override in _tiled_spec_from_problem. Narrow + arch/feature-guarded: every other shape is byte-identical; the kernel codegen the flags trigger already ships on the base. Validated on MI355X (direct-launch sweep, fp32-ref max_abs=1.5625e-02): +9.1% @ Sq4096 (1.29x torch-flash) and +10.1% @ Sq8192 over the 32x32 base.
… lever sweep The live preset bench (_variant_flags menu) can't express default-off levers (use_softmax_mfma_interleave / use_k_single_buffer / use_register_pv), so running it does not satisfy Step 0. Point to the raw-flag sweep (probe_config_sweep.py / _select_2d_* monkey-patch) and the autotuner, with the D256 interleave miss as the concrete failure.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates ROCm Rock-E unified-attention routing to send the gfx950, bf16, head_dim=256 prefill cohort to an existing optimized 32×32 transposed kernel configuration, additionally enabling softmax↔MFMA interleave (mode2/groups=4). It also amends the optimization runbook to clarify that the shipped “preset bench” does not exhaustively sweep default-off flags.
Changes:
- Add a new gfx950/D256/bf16-prefill predicate and use it to gate the launch-critical geometry selectors (
tile_size,num_warps,block_m_per_warp). - Override multiple codegen flags in
_tiled_spec_from_problemfor that cohort (32×32 transposed stack,register_pvoff, K single-buffer, softmax↔MFMA interleave mode2/g4). - Document a Step-0 caveat in the optimization runbook about preset sweeps missing default-off levers.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md | Adds a runbook caveat explaining why preset benches are not exhaustive (and how to do raw-flag sweeps). |
| dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py | Adds D256 gfx950 prefill routing predicate, gates geometry selectors, and applies a spec override enabling a faster kernel path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
yraparti
reviewed
Jul 9, 2026
…che key The D256 gfx950 fast route applies a codegen override in _tiled_spec_from_problem (32x32 transposed stack + softmax<->MFMA interleave), but _tiled_cache_key recomputed those flags from the ungated _enable_* helpers, so the key recorded e.g. use_transposed_qk_32x32=False / use_register_pv=True and omitted the interleave flags entirely -- the key no longer described the built kernel. Add the _d256_gfx950_fast(problem) cohort bit to the key. It flips exactly when the override applies, so key === built kernel again, without duplicating the override's ~15 flags. Keying on the predicate (not the flags) deliberately keeps the override a Layer-2 hand-pinned exception ABOVE the autotuner-regenerable _enable_* helpers, rather than folding it into tables a future sweep->gate is meant to overwrite.
Break the row-aliased LDS bank conflict on the D256 gfx950 bf16 prefill
QK
K read by padding K_lds between the 2-row async-DMA slabs. The async DMA
writes one WAVE_BYTES (2-row) contiguous slab per wave/call, so padding
BETWEEN slabs keeps the write contiguous (DMA-safe) while shifting each
slab's banks; the QK read remaps (row,col) -> (row/2, (row%2)*HD + col)
so
write and read stay consistent (numerics preserved).
Gated behind use_kq_lds_pad / kq_lds_pad_halves (default off -> byte
identical; byte-identity gate GREEN). Wired pad=16 into the D256 route
in
_tiled_spec_from_problem. Measured on MI355X (direct-launch, fp32-ref
max_abs=1.5625e-2):
SQ_LDS_BANK_CONFLICT rate 897 -> 660 per 1000 (-4.5x conflict cycles)
Sq8192 ~1592 -> ~1195 us (~25% faster); vs torch-flash 0.86x -> 1.14x
Sq4096 vs torch-flash 1.27x -> ~1.6-2.0x
Occupancy-neutral: LDS 65 KB <= 80 KB budget keeps 2 WG/CU; VGPR
unchanged.
Python-engine only for now (the C++ twin NOTIMPL-rejects the host
use_q_direct_reg / use_softmax_mfma_interleave path). Also adds the
parity harness lanes used to develop this: direct-launch D256 lane,
occupancy
probe, torch-flash comparison, interleave toggle, and pad/xor-swizzle
probes.
ISSUE ID : AICK-1495
## Submission Checklist
- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Contributor
Author
|
PR with @yraparti's optimization work now beats AOTriton perf numbers. |
✅ 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. |
…256-gfx950-opt # Conflicts: # dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py # dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py # dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp
AviralGoelAMD
changed the title
AviralGoelAMD
changed the title
…lders/common; PARITY_* env->CLI Addresses two #9233 review comments on the D256 gfx950 parity harness. Comment (hoist): _run_torch_flash and _probe_print_occupancy are not D256/gfx950-specific. Move the reusable cores to builders/common so the gfx942/decode harnesses can consume one implementation instead of drifting into near-identical copies: - builders/common/torch_flash_reference.py: run_torch_flash(...), decoupled from the harness-local Scenario/data (plain tensors/params; timing injected). - builders/common/occupancy_probe.py: print_occupancy(hsaco, num_warps, arch=) with an arch arg (was hardcoded gfx950). gfx950 keeps a thin _run_torch_flash adapter and calls print_occupancy directly. gfx942 is left untouched -- it imports these if/when it needs them. Comment (PARITY_* levers): the experimental env flags are all read in one main(), so they become real argparse flags (discoverable in --help, typed): PARITY_NO_INTERLEAVE -> --no-interleave PARITY_KQ_SWIZZLE -> --kq-swizzle PARITY_KQ_PAD/_W -> --kq-pad HALVES PARITY_PROBE_OCC -> --probe-occupancy The --no-interleave lever still patches the production _tiled_spec_from_problem (the auto/production-dispatch lane builds the spec internally, so a local _replace can't reach it), but the global mutation is now restored in a main() try/finally wrapper -- no leak to later in-process callers. GPU-verified on gfx950 (MI355X), prefill_d256_gqa8_b16_sq4096: - baseline unchanged (ck-2d 394 us, max_abs 1.56e-2); - --probe-occupancy prints the hoisted occupancy line; - --no-interleave drops the interleave (kernel name loses smxil2g4; 443 us 32x32 base); - --torch-flash uses the hoisted run_torch_flash (torch-flash 626 us, backend=flash; ck-auto vs torch-flash 1.58x).
Contributor
Author
I moved both to builders/common/ |
AviralGoelAMD
added a commit
that referenced
this pull request
Jul 15, 2026
…; ceiling behind --variant hlpv Stacked on the builders/common hoist + PARITY_*-env->CLI cleanup. Addresses the #9233 review (2D-lane routing): replace the per-cohort `if _d256_gfx950_fast` branch with a variant switch. The 2D direct-launch lane now builds the PRODUCTION spec (`_tiled_spec_from_problem`) for ANY cohort by default -- a new cohort is measured with its real geometry with no per-cohort branch. The hand-tuned "aspirational ceiling" (num_warps=4 + half-local-PV) moves behind an explicit `--variant hlpv`. - Drops the `_d256_gfx950_fast` per-cohort branch from the harness. - Collapses the two duplicated build/compile/launch/time tails into one. - The shared tail consumes the CLI-param probes (--kq-swizzle / --kq-pad / --probe-occupancy) and the hoisted print_occupancy from the cleanup commit. GPU-verified on gfx950 (MI355X), 2D direct-launch: - prod default reproduces D256 (byte-identical kernel name; 391 us @ prefill_d256_gqa8_b16_sq4096) and the combo d64/b32 64x8 (54 us); - --variant hlpv reproduces the old ceiling (combo 64x8 47.7 us).
AviralGoelAMD
added a commit
that referenced
this pull request
Jul 15, 2026
The cohort predicate is the single source of truth, so the one-line dtype extension routes fp16 D256 prefill down the same 32x32-interleave + slab-pad fast path for both production (_d256_gfx950_fast) and dispatch. Addresses the #9233 review fp16-extension follow-up.
…harness-local ProbeSpec Addresses the #9233 review: use_kq_xor_swizzle produces INTENTIONALLY WRONG numerics (a read-only bank-conflict probe) but lived on the production UnifiedAttention2DTiledSpec guarded only by a comment -- one typo from silently corrupt output. Move it off the production spec so it can never be set there by accident: - drop the field from UnifiedAttention2DTiledSpec; the kernel emit now reads it via getattr(spec, "use_kq_xor_swizzle", False) -- False for every production spec, so production/normal paths are byte-identical. - the parity harness owns a frozen _ProbeTiledSpec(UnifiedAttention2DTiledSpec) subclass; --kq-swizzle builds it via _with_kq_xor_swizzle(spec). Verified (CPU): constructing a production spec with use_kq_xor_swizzle=True now raises TypeError; the ProbeSpec carries the flag, survives dataclasses.replace (kq-pad-after-swizzle path), and kernel_name() works.
…ymmetry) The K_lds slab pad's async-DMA write always pads when KQ_LDS_PAD is set, but the read only remaps for the non-fp8 K path -- an fp8+pad combo would pad the write and NOT remap the read, addressing different slots -> silent numeric corruption. The pad is a bf16 bank-conflict optimization and fp8+pad is not a supported combo (the D256 gfx950 cohort that uses the pad is bf16-only), so disable the pad under fp8 in the geometry guard, keeping write and read consistent. Byte-identical for the supported bf16 D256 path. Addresses the #9233 review (K_lds pad fp8 write/read asymmetry).
…s/tile) Closes the block_size coverage gap flagged in #9233 review: the D256 gfx950 cohort predicate has no block_size constraint and the fast route forces T=64, so bs=32 (2 blocks/tile) took the fast path while only bs=16 (4/tile) was validated. GPU-verified on gfx950/MI355X: bs=32 builds the fast spec (b32_t64...) and is correct (max_abs=0.01562, bf16 band) -> bs=32 stays enabled, gap closed by coverage rather than by narrowing the predicate.
Closes the #9233 review gap: no golden covered the D256 fast spec (use_kq_lds_pad + use_softmax_mfma_interleave). The slab pad's async-DMA write and padded read remap must agree byte-for-byte or numerics silently corrupt, but only the spec dataclass was asserted -- an addressing drift would pass every CPU test and fail only on GPU. Adds an "attention" IR-parity case (attention/gfx950/d256_pad_interleave) that builds the production fast spec and pins its lowered LLVM IR (sha256, both llvm20/llvm22 flavors) in rocke_representative_ir_sha256. _d256_gfx950_fast gates on the memoized runtime device arch via _resolve_attention_arch (not the per-case arch, which only feeds the lowerer), and the spec builder holds a bound import of it; build_attention patches both bindings wholesale to gfx950 so the fast spec is selected deterministically on any bless/CI host, and raises if the fast spec / pad+interleave is not selected (never silently pins the fallback). Bless is add-only: no existing case sha changed.
Copilot review (#9471): if every query_len is <= 0, seqs stays empty and the seqs[0] backend probe IndexErrors into the generic "no working torch sdpa backend" NotImplementedError -- a misleading failure mode. Raise an explicit ValueError naming the query_lens so the failure is deterministic and actionable.
Self-review (rocke-review.md sec.4 / one-way dependency): the D256 IR
golden added earlier lived in the platform IR-parity harness
(platform/tests/instances/rocke_ir_parity_harness.py) and imported the
library packages kernels/builders -- a platform-test-imports-library
violation. Precedent: library/tests/test_ci_static_attention.py was moved
out of platform/tests for exactly this reason ("references the library
builders package and therefore belongs in the library test layer").
Relocate it:
- Revert the platform harness + rocke_representative_ir_sha256.json to
platform-pure (byte-identical to before; the harness only builds
rocke.instances kernels again).
- Add library/tests/test_d256_attention_ir_golden.py + its golden. Same
pinning (both _resolve_attention_arch bindings -> gfx950, raises if the
fast spec / pad+interleave is not selected) and same lowered IR sha
(llvm20/llvm22), but importing only rocke.core.lower_llvm from platform
(library -> platform is the allowed direction). CPU-only, no GPU/comgr.
Resolve dispatch/attention.py: keep BOTH D256 candidates — ours _make_gfx950_d256_candidate (prefill/2D) and develop's #9462 _make_d256_decode_candidate (decode/3D). Disjoint cohorts (2D vs 3D select_path), so both register and never collide at selection.
Resolve attention_unified.py: keep BOTH function sets — ours (_d256_gfx950_cohort/ _fast/_spec_overrides, D256 prefill) and develop's #9403 _enable_softmax_mfma_interleave (d128 nw=4 interleave gate). Independent, disjoint cohorts; both defs + the d128 call site preserved. GPU-verified D256 perf unchanged; both gates fire scoped.
…esc is valid; drop redundant HLPV parity variant The d64/b32/h64kv8 bf16 single-seq long-prefill combo enables use_fast_paged_kv_desc, but _select_2d_tile_size picked T=128 for it -> the spec validator (fast_paged_kv_desc requires T=64) raised, so production crashed on these shapes on gfx950. num_warps was already 4; tile_size was the incompatibility. - _select_2d_tile_size: force T=64 for the bf16 h64kv8 no-SW single-batch d64 family so its fast_paged_kv_desc enablement is valid. Measured ~1.3-1.4x over the T=128 no-fast_paged path on MI355X; other d64 single-batch shapes keep the wider T=128. - attention_spec_builder: gate use_fast_paged_kv_desc on _select_2d_tile_size==64 (self-consistency -- never set the flag with an incompatible tile). - parity_unified_attention: the 2d lane now builds the production spec directly; removed the hand-tuned HLPV variant (use_hlpv_variant + the r4_hlpv_parity spec), now redundant because production emits the same T=64 + fast_paged kernel. Verified on MI355X (gfx950): the 4 family shapes build + run correct (bf16) with the T=64 + fast_paged kernel and match the former HLPV ceiling.
AnaghaRaoAMD
approved these changes
Jul 21, 2026
yraparti
added a commit
that referenced
this pull request
Jul 22, 2026
…gfx950 develop advanced (incl. #9233 "D256 gfx950 bf16 prefill 32x32 + softmax-MFMA interleave"), re-conflicting dispatch/attention.py. Same keep-both resolution as the prior merge: - AttentionSpec: keep BOTH the dense `kernel_name_override` and develop's `tiled_overrides` fields. - Keep BOTH candidate factories as separate functions: this branch's `_make_gfx950_attention_dense_candidate` (+ `_dense_spec`/`dense_spec_for_request`, priority 3, opt-in) and develop's new `_make_gfx950_d256_candidate` (priority 5); register both (alongside the existing d256_decode candidate). Post-merge: all 72 dispatch/attention tests pass, attention_dense golden 18/18 (no IR drift). Co-authored-by: Cursor <cursoragent@cursor.com>
AviralGoelAMD
added a commit
that referenced
this pull request
Jul 29, 2026
JIRA ID : AICK-1495 ## Summary Adds a fast attention path for **D256 (head_size 256) bf16 causal-prefill on gfx942**. gfx942 had no fast route for this shape before, so it fell back to a slow path. Stacked on #9233 (`users/avirgoel/ck/d256-gfx950-opt`). ## What changed - New D256 gfx942 kernel (`build_gfx942_4warp_gqa`), a faster 4-warp GQA design. - Dispatch now routes the D256 gfx942 bf16 prefill shape to it; every other shape is unchanged. - Removed the old, now-unused D256 builder it replaces. - Added a routing test that runs on CPU (no GPU needed). ## How the speedup works The original D256 path was stuck at ~1 workgroup per CU, bottlenecked on shared-memory traffic/bank conflicts with the matrix cores nearly idle. The new kernel keeps the hardware busy through: - **Wider tile per workgroup (4-warp, BLOCK_M=128).** 128 query tokens are handled by 4 wave64 warps at once, so the per-launch prelude is amortized and the matrix (MFMA) units stay fed. This geometry was previously unreachable on the default builder: it Python-unrolls the head_dim=256 body, and `num_warps>1` blew up the DSL→LLVM lowering pass (the multi-minute "lowerer hang"). The new kernel is **runtime-rolled** (`scf_for`), so it compiles in seconds at 4 warps — which is what makes the occupancy win buildable. - **K and Q read directly from HBM into the MFMA operands** — no `K_lds`/`Q_lds` staging (K/Q are L2-resident across the tile). Removes an LDS round-trip. - **V loaded via wide (16-byte / `dwordx4`) vectorized global reads into registers, then stored wide to `V_lds`.** This bypasses the CDNA3 async global→LDS DMA's 1-dword/lane cap (which otherwise issues many narrow loads), cutting VMEM traffic. `V_lds` remains the transpose medium — only the load path is widened. - **LDS-free softmax reduction** via `ds_bpermute` (register cross-lane butterfly), so the online-softmax row-max / row-sum need no shared memory. `V_lds` is the kernel's only LDS consumer. - **Cheaper softmax math** — the scale is folded into the log2 domain so the exponential is a single `exp2`; the accumulator rescale uses a fused `fma`; the final normalization uses a fast reciprocal. Orientation note: this kernel computes scores in the **transposed-QK** orientation (`S^T = K·Qᵀ`), the same as the previous path — it does **not** change the QK orientation, and V / output still use the transpose (`V_lds` + an `O^T` epilogue). The wins above are the load path, softmax cost, and occupancy — not transpose elimination. ## Performance **Improvement over the original (pre-optimization) D256 gfx942 kernel** — wall-clock of our own kernel, then vs now: | Shape | Before any optimization | Now | Speedup | |---|---|---|---| | Sq4096 | 22.0 ms | ≈ 1.6 ms | ~14x | | Sq8192 | 87.5 ms | ≈ 5.7 ms | ~15x | The pre-optimization baseline (~24-27x behind AOTriton, shared-memory-conflict-bound at 1 workgroup/CU) was measured on MI300X; today's numbers are on MI300A, which has fewer CUs, so the true speedup is if anything larger. **Vs AITER `unified_attention` today** (gfx942 MI300A, same-run, D256 bf16 prefill, GQA 16/2), measured as OURS / AITER: | Shape | New kernel | |---|---| | Sq4096 | ≈ 0.98x AITER (parity) | | Sq8192 | ≈ 0.94x AITER (~6% faster) | - Ratio is the headline (device-robust); confirmed with the shipped `dispatch_perf_vs_aiter.py` + a 3x-repeat, and consistent with the earlier MI300X run (0.97x / 0.94x). - Output verified against the fp32 reference and AITER within bf16 tolerance. ## TODOs / measured follow-ups (per @yraparti review) 1. **Perf vs FlyDSL v0.2.4 (gfx942):** N/A — FlyDSL v0.2.4 does not support head_dim=256 for paged-KV attention. `flydsl_flash_attn_func` raises `NotImplementedError: native paged KV supports head_dim=64 or 128, got 256` (`kernels/attention/flash_attn_interface.py:424-425` @ tag `v0.2.4`; the paged-light path is gated the same at line 472). The D256 prefill cohort is paged → FlyDSL can't run it on this arch, so there is no head-to-head number. 2. **Raw MFMA utilization + 8-wave scope (rocprofv3, MI300X):** - **MfmaUtil ≈ 8-10%** (Sq4096 8.3%, Sq8192 10.1%). - **Not pipe-throughput-bound:** `VALUBusy` ~9%, `LdsUtil` ~9%, `SALUBusy` ~0.3%, `MemUnitStalled` ~0.1% — every execution pipe idles ~90% and memory is not stalled. (`VALUUtilization=100%` is per-instruction lane efficiency, not pipe-busy.) The kernel is **latency-bound at low occupancy** (`OccupancyPercent` ~11-16%, ~3.5-5 waves/CU), register-capped at **128 arch VGPR + 128 AGPR = 256 reg/wave**. - **8-wave: limited/no scope.** A bigger tile (BLOCK_M=256) does not raise waves/CU — the per-lane accumulator is invariant to wave count, and an 8-wave CTA needs 8 co-resident waves at the same 256-reg footprint, so occupancy stays capped (and it risks not launching). The lever to lift MFMA utilization is **reducing register pressure to raise occupancy** (e.g. HD-split the accumulator), not more waves per CTA. → follow-up ticket. --------- Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
AnaghaRaoAMD
added a commit
that referenced
this pull request
Aug 17, 2026
…10508) ISSUE ID : AICK-1898 ## Problem [`attention_spec_builder.py`](dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py) bound-imported `_resolve_attention_arch`. A `from ... import sym` freezes the reference at import time, so `mock.patch.object(au, "_resolve_attention_arch", ...)` never reached the builder. It then resolved the real device arch, `_tiled_2d_impl` returned that arch's spec class, and the gfx950-only override fields hit: ``` TypeError: UnifiedAttention2DTiledSpec.__init__() got an unexpected keyword argument 'use_q_direct_reg' ``` Only fires when the builder is already in `sys.modules` before the patch, so it looked like flaky ordering. Pre-existing on develop; introduced by #9233, which added the `_kau` module handle and a comment about this exact hazard but applied it to `_d256_gfx950_fast` only. ## Fix Route both call sites through the existing `_kau` module handle. Import plumbing only — no emission or spec change. Four follow-up commits from @aosewski (`756dc0e3..55c3ecb`): `818c75f` repairs [`rocke_ir_parity_harness.py`](dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py), which read the removed attribute and took `rocke_golden_static` red with 12 drift entries; `adfc04a` contains `_RESOLVED_ATTENTION_ARCH` leaks via a conftest fixture; `7c47e74` adds [`test_arch_binding_guard.py`](dnn-providers/hip-kernel-provider/rocke/library/tests/test_arch_binding_guard.py) (verified to fail on the unfixed builder); `55c3ecb` makes the resolver test reproduce both preconditions in-test. ## Verification `pytest library/tests -m "not gpu"` on gfx942: | | result | |---|---| | develop @ `223bff09f90` | **5 failed**, 499 passed, 34 skipped | | this branch @ `55c3ecb` | **507 passed**, 34 skipped | Also at head: `platform/tests` 668 passed, golden `test_rocke_ci_static.py` 5 passed (`IR parity OK [llvm20, llvm22, llvm23]`, no re-bless), and 520 passed / 0 failed on a gfx942 MI300X node. [`test_attention_ir_cpp_parity.py`](dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py) is the dual-engine gate for these families; it skips where `rocke_engine` isn't built, so it did not execute in the runs above. ## Side effect [`prove_parity_binding.py`](dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/prove_parity_binding.py) patches the resolver at :1465 but imports `_tiled_spec_from_problem` at :1431, so `run_selector_parity(arch=...)` previously compared a Python spec for the **host** arch against a C++ spec for the **requested** arch. Fixed as a side effect, but earlier selector-parity numbers aren't comparable and masked mismatches may surface — worth a re-run across arches. Two adjacent defects are out in #10625, not here: the parity gate's skip-on-first-case behaviour, and a stale `parity_unified_attention` path missing its `.prefill` segment. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Adam Osewski <Adam.Osewski@amd.com> Co-authored-by: Adam Osewski <19374865+aosewski@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.
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:num_warps/tile_size/block_m_per_warp) so the builder spec,_tiled_cache_key, and_get_2d_launch_metaall agree;use_mfma_32x32, transposed stack,register_pvoff, K single-buffer,use_softmax_mfma_interleavemode2/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-refmax_abs=1.5625e-02— all correct)D256, bf16, GQA 16/2, causal prefill, single long sequence (batch 1) — the Qwen3-Next
gated_attnprefill regime (AICK-1495, Sq 4096/8192). Numbers are post the K_lds slab-pad (#9260), i.e. what this branch actually ships.bsNhere = paged block_size (batch = 1), not batch size.CK beats AOTriton at every ticket shape. Two stacked levers: (1) the FA3-style softmax↔MFMA interleave (mode2/g4) hides the exposed
v_expon 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_CONFLICT897→660 /1000). GPU-verified via the direct-launch lane; numerics checked against an fp32 paged reference (all PASS,max_abs=1.5625e-02).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 likeuse_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
d256-prefillLDS-budget resolver) — base isusers/avirgoel/ck/d256-prefill.