perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4 dispatch by AnaghaRaoAMD · Pull Request #9403 · ROCm/rocm-libraries

AnaghaRaoAMD · 2026-07-14T16:24:15Z

JIRA ID: AICK-1490

Summary

The shipped dispatcher routes the single-batch d128 prefill cohort on gfx950 to num_warps=2 with no MFMA/softmax interleave, leaving the Matrix cores starved. This enables use_softmax_mfma_interleave (iglp_opt(1)) and widens to num_warps=4 for that cohort, closing the occupancy gap and delivering a large prefill speedup with identical numerics.

Results (Llama-3.1-405B, bf16, S8192, B1, causal — MI355X)

	baseline (nw2)	nw4 (this PR)
TFLOP/s	333.9	440.7
Occupancy	11.8%	23.5%
waves/CU	4	8
VGPR (allocated)	296	250
AGPR	40	0
MFMA-busy	13.4%	17.7%

The widened kernel reaches 8 waves/CU by lowering register pressure (AGPR 40→0, VGPR 296→250, 0 spill), so it does not trip the _enable_d128_small_tile nw=2 occupancy reconciliation (LDS stays 64 KB, well under the 160 KB/CU budget — the kernel is VGPR-bound, not LDS-bound).

The win generalizes across the gap-doc shape set (Llama-3/3.1, Qwen2.5/3, GPT-3-13B MHA), both bf16 and fp16:

seqlen	speedup vs shipped
S512	1.9x
S1024	2.6x
S2048	3.2x
S4096	3.5x
S8192	3.6x

Risk

Scoped to the single-batch d128 combo cohort (_enable_softmax_mfma_interleave). Multi-seq, sliding-window, d64, and fp8 paths are unchanged — verified the selector output is unchanged off-cohort.
Occupancy-safe: the nw=4 kernel is strictly better on registers/waves than the shipped nw=2 (measured via probe_occupancy + llvm-readelf on the produced HSACO).
Escape hatch: HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0 restores prior nw=2 / no-interleave routing.
Additive change (65 lines, 2 files); no per-spec golden EMIT change for existing configs.

Testing summary

Correctness: identical to the shipped kernel (max_abs=0.0) on every cohort shape; validated against an independent torch dense causal reference (max_abs 8.2e-3, within the 4e-2 bf16 tolerance).
Perf: end-to-end through production auto-dispatch — 3.58x on D128 S8192 (437 vs 122 TF this session; 440.7 vs 333.9 TF in the summary table above, different measurement run).
Occupancy: probe_occupancy confirms 8 waves/CU for nw4 vs 4 for nw2.

Testing checklist

CI attention parity suite (gfx950)
Byte-identity gate (check_byte_identity.py --only attention) — expected GREEN (no per-spec EMIT change; routing-only)
Perf regression check on off-cohort shapes (multi-seq / sliding-window / d64 / fp8)

Technical changes

attention_unified.py: new _enable_softmax_mfma_interleave(problem) gate (single-batch d128 combo cohort, with env escape hatch); _select_2d_num_warps widens this cohort to nw=4.
attention_spec_builder.py: wire use_softmax_mfma_interleave=True / softmax_interleave_mode=1 into the gfx950 schedule fields for the cohort (field-presence guarded; mutually exclusive with use_sched_barrier, cohorts do not overlap).

🤖 Generated with Claude Code

The shipped dispatcher routes the single-batch d128 prefill cohort to num_warps=2 with no MFMA/softmax interleave, leaving the Matrix cores starved (rocprofv3 on D128 S8192: MFMA-busy ~5%, occupancy ~12%). Enable use_softmax_mfma_interleave (iglp_opt(1)) and widen to num_warps=4 for the cohort: MFMA-busy rises to ~18%, occupancy to ~24%, for a 1.9x-3.6x prefill speedup (S512->S8192) across Llama-3/3.1, Qwen2.5/3, and GPT-3-13B in both bf16 and fp16, correctness-identical to the shipped kernel (max_abs=0.0). The widened kernel is occupancy-better than the shipped nw=2 (250 vs 296 VGPR, 0 spill, 8 vs 4 waves/CU), so it does not trip the d128 small-tile nw=2 reconciliation. Scoped to the single-batch d128 combo cohort; other paths (multi-seq, sliding-window, d64, fp8) are unchanged. Escape hatch: HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0. Co-Authored-By: Claude <noreply@anthropic.com>

therock-pr-bot · 2026-07-14T16:28:55Z

❌ PR Check — Action Required

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	❌ Fail	Error: Source/code files changed without an accompanying unit test. Expected: add at least one test file named like `test_<name>.py` / `test_<name>.cpp` (or `<name>_test.`). Current:* code file(s) changed: `dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py`, `dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py`, `dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c`, `dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py`, `dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h` (+4 more); no test file found
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

⚠️ 1 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

❌ Unit Test

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

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

therock-pr-bot · 2026-07-14T16:28:56Z

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

❌ Unit Test

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

…+ nw=4 Mirror the Python dispatcher/emitter change so the C++ engine routes and emits identically: - selectors.cpp: add enable_softmax_mfma_interleave (byte-faithful env gate) + the nw=4 override in select_2d_num_warps for the d128 interleave cohort; export the gate in the selector header. - gfx950 glue: carry USE_SOFTMAX_INTERLEAVE / SOFTMAX_INTERLEAVE_MODE into the emit ctx, emit the smxil<mode> kernel-name suffix, and narrow the NOTIMPL rejection to mode 2 only (modes 0/1 are now ported). - kv_body_qk_softmax.cpp: emit iglp_opt(mode) after the iter-start K drain (mirrors Python 3093-3094). - parity: add idx53 (production nw=4 interleave config) to both the Python _CONFIGS table and the C make_spec() twin, in lockstep. Byte-identity gate GREEN across all 10 attention families (ll + ir canonical), including gfx950_attention_tiled_2d at 54 configs with the new interleave path. Co-Authored-By: Claude <noreply@anthropic.com>

AnaghaRaoAMD · 2026-07-14T16:49:21Z

Added the C++ engine twin for this change (commit 46f70c4), so the provider routes and emits identically to the Python path:

selectors.cpp: enable_softmax_mfma_interleave gate (byte-faithful env hatch) + the nw=4 override in select_2d_num_warps.
gfx950 glue: carry the interleave flags into the emit ctx, emit the smxil<mode> kernel-name suffix, and narrow the NOTIMPL rejection to mode 2 only (modes 0/1 now ported).
kv_body_qk_softmax.cpp: emit iglp_opt(mode) after the iter-start K drain.
parity: added the production nw=4 interleave config (idx53) to both the Python _CONFIGS table and the C make_spec() twin.

Byte-identity gate GREEN across all 10 attention families (both ll — the hard contract — and ir canonical), including gfx950_attention_tiled_2d at 54 configs with the new interleave path exercised on both engines.

Copilot

Pull request overview

This PR updates the gfx950 “single-batch d128 prefill combo” auto-routing to enable softmax↔MFMA interleaving (via iglp_opt) and prefer num_warps=4, and wires the corresponding schedule/spec fields through the Python builder, C++ selectors, and the gfx950 C-twin emitter (plus parity coverage).

Changes:

Add a new _enable_softmax_mfma_interleave / enable_softmax_mfma_interleave cohort gate (with env escape hatch) and route it to num_warps=4.
Plumb use_softmax_mfma_interleave + softmax_interleave_mode into gfx950 C-twin build ctx and emit iglp_opt(mode) for modes 0/1; extend kernel naming to include smxil….
Add a parity config (idx53) to lock C/Python byte-identity for the production-reachable interleave+nw4 configuration.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp	Accept interleave modes 0/1, reject mode 2, plumb ctx fields, and include interleave in kernel names.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp	Emit `iglp_opt(mode)` once per loop for interleave modes 0/1.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp	Add env-gated cohort enablement + select `num_warps=4` for the interleave cohort; export a helper.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h	Extend build ctx with softmax-interleave fields.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h	Expose `rocke_unified_attn_enable_softmax_mfma_interleave`.
dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py	Add idx53 parity spec for nw4 + interleave mode 1.
dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c	Add idx53 parity spec for nw4 + interleave mode 1.
dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py	Add Python `_enable_softmax_mfma_interleave` and route the cohort to `num_warps=4`.
dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py	Set gfx950 schedule fields `use_softmax_mfma_interleave=True` and `softmax_interleave_mode=1` for the cohort.

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

Runbook arch/gfx950.md §21.4a described the abandoned BLOCK_M=128 experiment geometry with pre-nw4 register numbers (213->240 VGPR). Correct it to the shipped BLOCK_M=32 cohort and the iglp_opt + num_warps=4 paired lever (296->250 VGPR, 4->8 waves/CU), which is now the production routing. Trim the attention harness README to own the routing decision + full eligible-space sweep (why the lever is gated to this cohort), and link the lever mechanism (occupancy table, rocprofv3 limiter read) to the runbook so each fact lives in exactly one place.

…-1490) - glue.cpp: guard use_softmax_mfma_interleave + use_sched_barrier as mutually exclusive before the ctx copy, mirroring the Python __post_init__ ValueError so C emission cannot diverge from the contract. - selectors.cpp: extract rocke_env_forces_off() and strip leading/trailing ASCII whitespace before the OFF-keyword compare, matching Python's .strip().lower() gate (e.g. "0 " / " off"); dedupes the two env kill-switches. - selectors.cpp: keep the wide T=2*block_size tile for the interleave cohort (guard before the small-tile check) so select_2d_tile_size mirrors Python. - attention_unified.py: document why the lever is gated to the single-batch d128 combo (full eligible-space sweep; d256/SWA covered by AICK-1495/1492).

Satisfy the pre-commit clang-format hook (v18.1.x): align declaration whitespace and re-wrap rocke_g950_fail argument lists in the two files touched by the Copilot review fixes. No functional change.

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.

AnaghaRaoAMD requested review from a team as code owners July 14, 2026 16:24

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

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

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

AnaghaRaoAMD changed the title ~~perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4~~ perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4 dispatch Jul 14, 2026

AnaghaRaoAMD requested review from AviralGoelAMD and yraparti July 14, 2026 16:56

bartekxk requested a review from Copilot July 16, 2026 10:20

Copilot started reviewing on behalf of bartekxk July 16, 2026 10:21

View session

bartekxk reviewed Jul 16, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py Resolved

bartekxk reviewed Jul 16, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py Resolved

Copilot AI reviewed Jul 16, 2026

View reviewed changes

Rao added 2 commits July 20, 2026 19:34

AnaghaRaoAMD requested a review from bartekxk July 20, 2026 23:38

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

style(rocke): clang-format the gfx950 nw=4 dispatch review fixes

05bfee6

Satisfy the pre-commit clang-format hook (v18.1.x): align declaration whitespace and re-wrap rocke_g950_fail argument lists in the two files touched by the Copilot review fixes. No functional change.

bartekxk approved these changes Jul 21, 2026

View reviewed changes

AnaghaRaoAMD merged commit c1c0c41 into develop Jul 21, 2026
24 of 25 checks passed

AnaghaRaoAMD deleted the users/anarao/gfx950-nw4-dispatch branch July 21, 2026 15:57

Navigation Menu

Repository navigation

perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4 dispatch - #9403

perf(rocke): gfx950 d128 prefill softmax-MFMA interleave + nw=4 dispatch#9403
AnaghaRaoAMD merged 5 commits into
developfrom
users/anarao/gfx950-nw4-dispatch

AnaghaRaoAMD commented Jul 14, 2026 •

edited

Loading

Uh oh!

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

edited

Loading

Uh oh!

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

edited

Loading

Uh oh!

AnaghaRaoAMD commented Jul 14, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Summary

Results (Llama-3.1-405B, bf16, S8192, B1, causal — MI355X)

Risk

Testing summary

Testing checklist

Technical changes

Uh oh!

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

Uh oh!

❌ PR Check — Action Required

Uh oh!

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

Uh oh!

Uh oh!

AnaghaRaoAMD commented Jul 14, 2026

Uh oh!

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!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AnaghaRaoAMD commented Jul 14, 2026 •

edited

Loading

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

edited

Loading

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

edited

Loading