Repository navigation
feat(rocke): gfx942 D256 bf16 prefill 4-warp GQA fast path - #9310
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.
…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.
…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
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a dedicated fast-path implementation and routing for gfx942 D256 bf16 causal-prefill (GQA) so this cohort no longer falls back to the slow tiled builder on gfx942.
Changes:
- Introduces a new gfx942 D256 4-warp GQA kernel builder and related spec/support plumbing.
- Updates unified-attention dispatch to route the targeted D256 gfx942 bf16 prefill cohort to the new builder and launch geometry.
- Adds a CPU-only routing/contract unit test to lock in the dispatch gate and builder surface.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx942_fast_routing.py | Adds CPU unit tests to validate the new gfx942 D256 fast-path gate and builder contract. |
| dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py | Adds fast-path-related flags/plumbing and introduces the new build_gfx942_4warp_gqa kernel builder. |
| dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py | Adds _d256_gfx942_fast predicate and dispatch/launch-meta routing to the new gfx942 D256 fast kernel. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Move the D256 gfx950 bf16-prefill codegen constellation from a post-build replace() override in kernels.common _tiled_spec_from_problem into the builder (builders.common.attention_spec_builder), per review: the winning spec is created in the builder, not baked into the dispatch layer. The _d256_gfx950_fast cache-key discriminator is retained so the key stays faithful to the built kernel. The predicate is called via the module (_kau._d256_gfx950_fast) so tests patching attention_unified._d256_gfx950_fast still steer the builder branch. Verified byte-identical: built spec + _tiled_cache_key unchanged for the D256 fast path and every fallback/non-D256 shape; fast vs feature-gated fallback keys stay distinct; byte-identity gate GREEN (10/10 attention families, Python==C++); attention routing + lds-budget unit tests pass. ## Motivation <!-- Explain the purpose of this PR and the goals it aims to achieve. --> ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
…es (#9343) JIRA ID : AICK-1495 ## What Hooks the D256 gfx950 bf16 prefill fast path (`_d256_gfx950_fast`: 32x32 transposed + FA3-style softmax<->MFMA interleave) into the two repeatable harnesses, per review request to wire new kernels into the dispatcher **and** the benchmark harnesses. The dispatcher hookup already shipped in #9233; this closes the harness side. - **Parity harness** (`builders/gfx950/attention/prefill/parity_unified_attention.py`): adds `prefill_d256_gqa8_b16_sq4096` / `prefill_d256_gqa8_b16_sq8192` to `default_scenarios()` (Qwen3-Next-80B-A3B cohort: GQA 16/2, head_size 256, bf16 - the two measured points). The pre-existing only head_size==256 scenario was decode-class and does not trip the prefill fast-path gate. - **Benchmark shapes** (`benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json`): committed aiter-trace-schema shape set (Sq 2048/4096/8192) for the shape-driven `benchmark_prefill2d_live.py` sweep. No committed head_size==256 prefill shapes existed before. ## Verification (gfx950 / MI350X, kreb rocm7.13, real GPU run) - **Parity: PASS.** Both scenarios route to the D256 fast path (`dispatch_path=2d`, kernel `rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_..._mphsplit_smxil2g4_qdreg`) and match the fp32 `ref_paged_attn` reference at `max_abs = 1.5625e-2` (= the documented parity value; within bf16 `atol=1e-1`). Identical `max_abs` across the auto / 2d / 3d paths. - **vs AOTriton (the ticket baseline): CK wins.** `--torch-flash` (torch SDPA FLASH = AOTriton on ROCm): Sq4096 = **1.59x**, Sq8192 = **1.03x** (ck-auto 431 / 1429 us vs torch-flash 685 / 1470 us). Consistent with "beats AOTriton at every ticket shape". - **Benchmark sweep: runs all 3 shapes, `correct=3 incorrect=0`,** and the `prod` variant is confirmed to dispatch the D256 fast kernel (`dispatch_path=2d`, `_d256_gfx950_fast=True`, same kernel name as above). - Note on the sweep's speedup column: `benchmark_prefill2d_live.py` on this branch compares against **AITER-Triton** (`aiter.ops.triton.unified_attention`), not AOTriton, so its "speedup vs Triton" is a different comparison than the AOTriton number above. It also times CK via the provider path (`run_unified_attention_torch`) while timing the Triton lane as a raw kernel call; the same CK kernel measures ~431 us direct (parity) vs ~1221 us via the provider path in the sweep - a measurement-path difference (mechanism not yet isolated). The AOTriton comparison above (direct GPU-time) is the representative metric. (An AOTriton baseline lane for this harness lives on `attn-bench-framework-toggles` @ f8f05cc, not in this line.) Stacked on the #9233 branch (`users/avirgoel/ck/d256-gfx950-opt`).
Blank line after the import block flagged by black --check in CI pre-commit.
- attention_tiled_2d.py: rewrite the `use_kq_lds_pad` field comment. It described the experimental measurement-only pad (intentionally-wrong numerics); production wires the DMA-safe slab-granularity pad (numerics preserved, enabled by the D256 fast route). Comment now matches behavior. - parity_unified_attention.py: the `pr9233_d256_*_bs64` scenarios set num_kv_heads=16 (MHA) but the PR cohort is GQA 16/2. Set num_kv_heads=2 so they match the reported cohort (num_queries_per_kv=8). - benchmark_prefill2d_live.py: `CkVariantBench.build()` compiled every variant via LLVM-direct, so --compile-backend only affected the prod lane. Route build() through the backend (self.compile_backend or _select_2d_compile_backend); the flag now applies to all lanes and the variant lanes match production compile by default.
Contributor
Author
|
I have added our shapes to benchmark scripts. |
…tcher (#9405) JIRA ID : AICK-1495 This PR adds an optional **`AttentionSpec.tiled_overrides`** field to the attention dispatcher (`library/dispatch/attention.py`). That is an **intentional extension of the dispatcher's deliberately-coarse parity identity** — the module docstring states the CTA geometry/knobs are *"intentionally OUT of the parity identity"* (which is `(path, head_size, block_size)`). ## Commit 2 — knob handoff (MVP) - `_d256_gfx950_spec_overrides()` = single source of truth for the 32×32-transposed + softmax↔MFMA-interleave (mode2/g4) + slab-pad constellation. Both the builder override (`attention_spec_builder`) and the dispatch candidate consume it → dispatched spec and built kernel cannot drift. **Byte-identical** to the previous inline `.replace(...)`. - `_tiled_spec_from_problem(problem, *, overrides=None)` applies dispatcher-supplied overrides on top of the heuristic base and **still runs `_resolve_lds_budget` in all cases**. The no-arg (production) path is unchanged. - `AttentionSpec.tiled_overrides` carries the pinned knobs (see the reviewer-attention note above). ## Honest ROI / scope **Behavior-neutral today.** Production `run_unified_attention_torch` still **bypasses the dispatcher** (it uses `select_path` + `_tiled_spec_from_problem`), and for the D256 cohort the dispatcher's overrides are **idempotent** with the builder's existing `_impl` override. So nothing routes or runs differently today. The value is **de-risking the eventual production reroute**: it proves the "dispatcher declares knobs → orchestrator accepts" path so the later flip (route production through dispatch + delete the `_impl` override) is a small, safe change. ## Verification - 13/13 gfx950 dispatch tests (8 routing + 5 knob-handoff: SSOT constellation, dispatched spec carries the overrides, generic candidate carries none, orchestrator accepts overrides [SSOT reproduces the spec], invalid override rejected by the spec `__post_init__`); existing dispatch suite 12/12 (no regression). <img width="721" height="834" alt="rocKE_dsl" src="https://github.com/user-attachments/assets/ab62e8a6-50d2-4c9e-b393-e2c6d8d8d5ab" />
…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).
…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.
Contributor
|
-Thanks for adding fast path to the d256 gqa. Can you just post the following metrics in TODOs?
|
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.
Add a fast dispatch path for D256 (head_size=256) bf16 causal-prefill
paged attention on gfx942 (CDNA3 / MI300X). gfx942 had no fast D256 route
on the default builder: the wide-flash gate excludes head_size==256 and the
narrow 16x16x16 fallback overflows LDS at a competitive tile.
Kernel (kernels/gfx942/attention_tiled_2d.py::build_gfx942_4warp_gqa):
BLOCK_M=128 / 4-wave64 CTA natural-QK (S = Q @ K^T) attention on the
32x32x8 bf16 MFMA atom. Reads paged K/V direct HBM->register (i32
element offset, i64-overflow fallback via the default builder for caches
> 2 GiB), varlen token-major Q via TensorDescriptor.naive +
in-kernel binary_search_seq_idx on cu_seqlens_q, bottom-right causal
(chunked-prefill / decode), GQA head->kv_head, softmax scale*log2e fold,
LDS 5-stage swizzle reduction, conflict-free V.
Dispatch (kernels/common/attention_unified.py): _d256_gfx942_fast gates
the cohort (bf16, head_size=256, prefill, block_size in {16,32}, no
fp8/sliding-window/softcap/sinks/alibi/qq_bias, cache <= 2 GiB); it sets
the launch-critical geometry selectors so the spec, _tiled_cache_key and
_get_2d_launch_meta agree, selects the builder in _get_2d_launcher, and
keys the cache so its HSACO never shares a slot with the default builder.
Every other shape is byte-identical.
supports_tiled_2d exposes use_d256_gfx942_fast, which bypasses the generic
staged-tile LDS-budget model (the fast-path builder manages its own LDS).
Test: tests/dispatch/attention/test_d256_gfx942_fast_routing.py covers the
gate accept/reject matrix (arch, dtype, head_size, block_size, decode,
feature flags) and the builder contract, arch-pinned so it runs host-side
without a gfx942 GPU.
The _d256_gfx942_fast docstring cited the removed std-QK research kernel's dense TF/s (115.4/124.9). Replace with the shipped 4-warp GQA builder's same-run measurement vs AITER unified_attention (D256 bf16 causal prefill, GQA 16/2, block_size=16): ~0.98x @ Sq4096 (parity), ~0.94x @ Sq8192 (OURS ~6% faster) on MI300A, matching the MI300X prototype run (0.97x/0.94x). Output verified against the fp32 reference and AITER within bf16 tol.
…nch) The _d256_gfx942_fast docstring claimed spec/_tiled_cache_key/_get_2d_launch_meta all agree on num_warps=1/block_m_per_warp=32. They do not: those values are the cohort discriminator (distinct HSACO cache slot, byte-identity for other shapes), while build_gfx942_4warp_gqa is a 4-wave64/BLOCK_M=128 kernel that _get_2d_launch_meta deliberately launches with block=(256,1,1)/BLOCK_M=128. Reword so a '4-warp' builder isn't documented as '1-warp ... all agree'.
- attention_tiled_2d.py: drop unused `mfma_k_loop` import. - supports_tiled_2d: validate the d256_gfx942_fast fast-path preconditions (bf16 / head_size==256) before returning supported, instead of trusting the caller's flag; fix the comment (K is direct HBM->reg, V is staged via V_lds). - 4-warp kernel header: correct the K/V + addressing claim -- K direct HBM->register, V staged through V_lds, i32 paged addressing with the `_d256_gfx942_fast` gate excluding > 2 GiB caches that need i64. - _select_2d_num_warps / _select_2d_block_m_per_warp: note the d256_gfx942_fast returns are cache-key discriminators only; the dedicated 4-warp kernel launches block=(256,1,1) / BLOCK_M=128 via _get_2d_launch_meta.
Adds d256_gfx942_prefill_shapes.json (GQA 16/2, head_size 256, bf16, Sq4096/8192) so the dispatch-driven benchmark_prefill2d_live.py exercises the 4-warp kernel through backend="auto" against AITER unified_attention (paged, 2D-forced) -- apples-to-apples paged-vs-paged. Fixes the stale final_shapes_check.py note that claimed no tiled d256 path on gfx942; build_gfx942_4warp_gqa now routes via _d256_gfx942_fast. Measured (MI300X, D256 bf16 causal prefill GQA 16/2, paged, vs AITER unified_attention 2D): OURS/AITER ~0.97 @SQ4096, ~0.98 @sq8192; pure-fp32-ref max_abs 0.05/0.08 (within bf16 tolerance).
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/d256-gfx942-4wgqa
branch
from
July 21, 2026 21:20
8b1b96a to
9c4fcef
Compare
… library golden Extends the D256 IR golden with build_gfx942_4warp_gqa (the byte-sensitive V wide-load -> V_lds transpose path), mirroring the gfx950 fast-spec case. CPU-only (pure-python lowerer, both llvm20/llvm22 flavors); guards the codegen/addressing drift the routing test can't catch. Kernel numerics independently verified on MI300X: bf16 worst-element max_abs is seed-variance 0.02-0.16, consistent with the documented 0.05/0.08.
Contributor
Author
|
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
Resolve 3 D256 conflicts (develop already carries the gfx950 D256 work): - attention_unified.py: union - keep both _d256_gfx950_fast (develop) and _d256_gfx942_fast (ours) selectors at all 4 sites; develop's other changes auto-merged untouched. - golden json + test: take ours (clean superset - develop's gfx950 case is byte-identical; ours adds the gfx942 4-warp case). Verified CPU-side: golden IR sha256 matches for both gfx942+gfx950 cohorts; gfx942 routing (12), gfx950 routing (13), lds-budget (9) all pass.
|
🎉 All checks passed! This PR is ready for review. |
yraparti
approved these changes
Jul 22, 2026
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.
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
build_gfx942_4warp_gqa), a faster 4-warp GQA design.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:
num_warps>1blew 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_lds/Q_ldsstaging (K/Q are L2-resident across the tile). Removes an LDS round-trip.dwordx4) vectorized global reads into registers, then stored wide toV_lds. This bypasses the CDNA3 async global→LDS DMA's 1-dword/lane cap (which otherwise issues many narrow loads), cutting VMEM traffic.V_ldsremains the transpose medium — only the load path is widened.ds_bpermute(register cross-lane butterfly), so the online-softmax row-max / row-sum need no shared memory.V_ldsis the kernel's only LDS consumer.exp2; the accumulator rescale uses a fusedfma; 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+ anO^Tepilogue). 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:
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_attentiontoday (gfx942 MI300A, same-run, D256 bf16 prefill, GQA 16/2), measured as OURS / AITER:dispatch_perf_vs_aiter.py+ a 3x-repeat, and consistent with the earlier MI300X run (0.97x / 0.94x).TODOs / measured follow-ups (per @yraparti review)
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_funcraisesNotImplementedError: native paged KV supports head_dim=64 or 128, got 256(kernels/attention/flash_attn_interface.py:424-425@ tagv0.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.Raw MFMA utilization + 8-wave scope (rocprofv3, MI300X):
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.