feat(rocke): dense pipe + dispatcher wiring (fp16 + bf16) by bartekxk · Pull Request #9057 · ROCm/rocm-libraries

bartekxk · 2026-07-03T12:02:16Z

Motivation

Provide mha gfx942 dense-pipe fast path,

Technical Details

dense pipe + dispatcher wiring (fp16 + bf16)

Test Plan

dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_pipe_wiring.py

prefill_bench + run_all.py

Test Result

Passed

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

JIRA ID : AICK-1491

therock-pr-bot · 2026-07-03T12:06:31Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

therock-pr-bot · 2026-07-03T12:06:32Z

🎉 All checks passed! This PR is ready for review.

Copilot

Pull request overview

This PR introduces a gfx942 “dense pipe” fast path for unified attention and adjusts gfx942 routing rules so MHA shapes can use the flash/dense pipeline across more sequence lengths (with corresponding dispatcher wiring).

Changes:

Extend gfx942 flash-eligibility logic (fp16 + bf16) to avoid carving out the small-q narrow path for MHA (while keeping it for GQA).
Add a new high-priority dense_pipe dispatcher candidate for gfx942 fp16 prefill to outrank the generic unified_2d candidate.
Update the tiled kernel cache key to include bf16-wide eligibility (but additional bf16-wide geometry/cache-key alignment issues remain).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py	Adjusts gfx942 fp16/bf16 flash routing and updates tiled cache-key/geometry selection logic.
dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py	Registers a new gfx942 fp16 `dense_pipe` kernel candidate with higher selection priority.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…arkocot/aick-1491

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

therock-pr-bot · 2026-07-07T08:56:35Z

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

python -m pip install pre-commit
pre-commit install
pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

yraparti

Move universal attention spec creation from kernel library to builders
Add the algorithm, the new knobs and the cases where the work to README.md

… ring #9057 excluded D128 bf16 from the sliced-K ring with an early-return in _enable_gfx942_flash_k_sliced_ring, on the basis that the D128 bf16 ring "produces wrong results (max_abs ~3-5, NaN/Inf for MHA)". Root-cause: that failure came from grafting the ring onto the bf16-WIDE geometry (nw=2, tile=32), which is a malformed spec on D128 -- tile=32 is not a multiple of the production block_size=64, and the wide path has no cfvst, which the ring requires. The spec builder's ring branch already uses the fp16-flash geometry (nw=4, tile=64, cfvst) when the ring is active, so removing the exclusion routes D128 bf16 through that correct geometry -- the same recipe fp16 uses, just with the gfx942-legal mfma_f32_32x32x8_bf16 atom. Also make _get_2d_launch_meta's bf16 branch ring-aware (nw=4 when the ring is active) so the launch grid/block matches the geometry the launcher builds, instead of the stale bf16-wide nw=2. Verified on MI300X (gfx942), D128 bf16 GQA (32/8) + MHA (16/16), S1024-S8192: - Numeric: max_abs 0.00049, no NaN/Inf, on both the Python and C++ engines. - Perf: 2.5-5.4x-slower-than-AOTriton -> ~1.0-1.2x (fp16 parity); bf16 D128 GQA prefill win-rate 0% -> 20%, overall 40-shape win-rate 20% -> 30%. - The bf16 D128 ring kernel is byte-structurally identical to the fp16 kernel (256 VGPR / 3 spill), differing only in the MFMA opcode. This is a Python selector/launch-meta change only (no emitter or C++ change), so per-spec emission is unchanged and the byte-identity gate is unaffected. Co-Authored-By: Claude <noreply@anthropic.com>

…liced-K ring (#9662) ## Summary The correctness fix (#9647) put fp16 D128 prefill on the non-ring flash path — correct but ~0.55–0.99× AOTriton flash. This restores the sliced-K ring for **fp16 D128** using a **depth-2** pipeline that is both correct and faster, recovering (and exceeding) the ring's performance without the bug. - **fp16 D128** → depth-2 ring: correct at magnitude + **1.18–1.56× faster** than the non-ring path (0.65–1.55× flash; beats flash at S4096). - **bf16 D128** → stays non-ring (its bs=64 optimum; unchanged from #9647). - **D64** (both dtypes) → depth-3 ring, unchanged. ## How it fixes the ring The original bug: depth-3 ring, `kg % 3` slot map, k_groups=4 (D128) → slice 3 reuses slot 0, and the reusing DMA was **unfenced** → corrupted QK accumulation. Two changes: 1. **Drain-on-reuse fence (now unconditional).** Before a prefetch DMA that reuses the slot slice(kg-1) was just read from, drain LDS reads + barrier. Previously only fired under `K_SLICED_LDSSEQ`. 2. **`ring_depth` spec field.** Depth-2 (`k%2`, 2 slots) keeps the live set `{kg, kg+1}` with no slot reuse in the k_groups=4 case and uses less LDS → lower occupancy pressure. Measured best for fp16 D128. ## Risk assessment - **Python-only**, consistent with how the ring already ships: the sliced-K ring schedule lives only in the Python engine (C++ has ring primitives but no schedule loop), and attention is **not in the byte-identity gate**. #9198 enabled the D128 ring with no C++ change for the same reason — no C++ mirror or gate re-bless required. - bf16 D128 and D64 paths are unchanged; the fence refactor is byte-equivalent for k_groups=2 (D64 verified). ## Perf (fp16 D128, magnitude randn, MI300X, vs AOTriton flash) | shape | non-ring (#9647) | depth-2 ring (this PR) | speedup | |---|---|---|---| | 32/8 S1024 | 0.68× | 0.88× | 1.29× | | 32/8 S4096 | 0.92× | **1.42×** | 1.54× | | 128/8 S1024 | 0.70× | 1.06× | 1.51× | | 128/8 S4096 | 0.99× | **1.55×** | 1.56× | | 16/16 S4096 | 0.80× | **1.17×** | 1.46× | ## Testing summary - **Numeric:** fp16 D128 depth-2 ring PASS vs fp32 oracle at magnitude (max_abs ~0.002) — GQA 32/8, 128/8 + MHA 16/16, S1024–4096. - **No regression:** bf16 D128 (non-ring) and D64 depth-3 ring (both dtypes) still PASS at magnitude. - **Suite:** 177 attention tests pass. Guardrail test updated for per-dtype routing; on-GPU numeric check gates both the bf16 non-ring and fp16 depth-2 ring paths. ## Testing checklist - [x] fp16 D128 depth-2 ring numeric vs fp32 oracle at magnitude - [x] fp16 D128 perf recovered (1.18–1.56× vs non-ring; beats flash at S4096) - [x] bf16 D128 non-ring unchanged; D64 depth-3 ring unregressed - [x] launch-meta matches depth-2 spec (nw=4) - [x] 177 attention tests pass - [ ] Reviewer: confirm Python-only ring authoring is acceptable (matches #9198 precedent); C++ ring + gate coverage tracked separately ## Technical changes - `attention_tiled_2d.py` — `ring_depth` spec field (+ kernel_name tag); unconditional drain-on-reuse fence; depth-2 (`k%2`) schedule in the ring loop. - `attention_unified.py` — `_select_gfx942_flash_ring_depth`; ring selector re-enables fp16 D128 (depth-2), keeps bf16 D128 non-ring. - `attention_spec_builder.py` — wire `ring_depth` into both flash spec builders. - `test_attn_bf16_d128_ring.py` — per-dtype routing assertions + on-GPU numeric gate for both paths. ## Note: pre-existing gap (not introduced here) The sliced-K ring (shipping since #9057/#9198) has no C++ engine implementation and no byte-identity gate coverage for attention. This PR follows that existing Python-only pattern. Closing that gap (C++ ring + attention in the gate) is separate follow-up work. 🤖 Generated with [Claude Code](https://claude.com/claude-code) [AICK-1538]: https://amd-hub.atlassian.net/browse/AICK-1538?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ [AICK-1489]: https://amd-hub.atlassian.net/browse/AICK-1489?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ Co-authored-by: Claude <noreply@anthropic.com>

feat(rocke): dense pipe + dispatcher wiring (fp16 + bf16)

1f77f31

bartekxk self-assigned this Jul 3, 2026

github-actions Bot added the project: hip-kernel-provider label Jul 3, 2026

therock-pr-bot Bot added the Not ready to Review label Jul 3, 2026

assistant-librarian Bot added the organization: ROCm label Jul 3, 2026

bartekxk requested a review from Copilot July 5, 2026 11:53

Copilot started reviewing on behalf of bartekxk July 5, 2026 11:53

View session

Copilot AI reviewed Jul 5, 2026

View reviewed changes

bartekxk added 3 commits July 6, 2026 22:29

improvements

83dfeb7

Merge branch 'develop' of github.com:ROCm/rocm-libraries into users/b…

f8e5a49

…arkocot/aick-1491

update bench

755ebf8

therock-pr-bot Bot removed the Not ready to Review label Jul 6, 2026

gfx950 measurements

076eb65

bartekxk requested a review from Copilot July 7, 2026 08:41

Copilot started reviewing on behalf of bartekxk July 7, 2026 08:42

View session

bartekxk marked this pull request as ready for review July 7, 2026 08:42

bartekxk requested review from a team as code owners July 7, 2026 08:42

Merge branch 'develop' into users/barkocot/aick-1491

77686ab

Copilot AI reviewed Jul 7, 2026

View reviewed changes

bartekxk and others added 4 commits July 7, 2026 10:51

Potential fix for pull request finding

eb6757f

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

703868a

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

c6fa666

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

fixes

9f19504

tmp

0e96156

bartekxk requested review from AnaghaRaoAMD and arai713 July 7, 2026 09:03

bartekxk requested a review from AviralGoelAMD July 7, 2026 09:03

yraparti requested changes Jul 7, 2026

View reviewed changes

fixes

ed0ea25

github-actions Bot added the documentation label Jul 8, 2026

yraparti approved these changes Jul 8, 2026

View reviewed changes

bartekxk enabled auto-merge (squash) July 8, 2026 14:39

AnaghaRaoAMD mentioned this pull request Jul 8, 2026

feat(rocke/attention): enable gfx942 bf16 D128 prefill ring for fp16 parity #9198

Merged

4 of 5 tasks

bartekxk merged commit 7e104fc into develop Jul 8, 2026
39 of 43 checks passed

bartekxk deleted the users/barkocot/aick-1491 branch July 8, 2026 18:10

AnaghaRaoAMD mentioned this pull request Jul 21, 2026

perf(rocke/attention): recover fp16 gfx942 D128 prefill via depth-2 sliced-K ring #9662

Merged

5 of 6 tasks

AnaghaRaoAMD mentioned this pull request Aug 7, 2026

fix(rocke): reach _resolve_attention_arch through the module handle #10508

Merged

Navigation Menu

Repository navigation

feat(rocke): dense pipe + dispatcher wiring (fp16 + bf16) - #9057

feat(rocke): dense pipe + dispatcher wiring (fp16 + bf16)#9057
bartekxk merged 12 commits into
developfrom
users/barkocot/aick-1491

bartekxk commented Jul 3, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Jul 3, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Jul 3, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

therock-pr-bot Bot commented Jul 7, 2026

Uh oh!

yraparti left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

bartekxk commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

therock-pr-bot Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

therock-pr-bot Bot commented Jul 7, 2026

Pre-commit check failed

Uh oh!

yraparti left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bartekxk commented Jul 3, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 3, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 3, 2026 •

edited

Loading