perf(ck_tile/fmha): [gfx1250 opt 2] NT_HT cache hint on the K/V loads of short-query tiles by Yanxing-Shi · Pull Request #13466 · ROCm/rocm-libraries

Yanxing-Shi · 2026-10-11T03:00:24Z

What

Builds the K/V global tensor views of the gfx1250 qr_tdm forward kernel with the NT_HT cache hint (non-temporal in the near cache, high-priority temporal in the far cache) for every qr_tdm tile of at most 64 rows (kM0 <= 64), with or without a mask.

qr_tdm moves K/V from HBM straight into LDS through the TDM and never reads them back through the vector path, so allocating them in the near cache only evicts other data. NT_HT drops that allocation and keeps K/V high priority in the far cache, where the m-tiles of one (batch, head) share them. The TDM load takes its coherence from the tensor view (buffer_view::tdm_get -> amd_tdm_load<Coherence>), so the change is the two K/V view constructions in FmhaFwdKernel.

The bm0=128 tile, which serves long sequences, keeps the default: it relies on K/V reuse in the near cache and is 1-3% slower with the hint. The bm0=64 tiles of d=256 (taken at every length) were measured flat with the hint, 1.00x-1.01x from 8/532 up to 8192/8192, dense and causal.

Performance

gfx1250, bf16, d=128, lse=0. Baseline is develop at 46d0941 (2026-10-13); median of 3 interleaved runs, each after >= 1 s of warmup (steady-state clocks), on an idle host. "causal" is the bottom-right mask (-mask=b) unless noted; bshd unless noted.

Large batch, small seqlen

batch	nhead	s_q	s_k	mask	TFLOPS develop	TFLOPS this PR	speedup
1280	2	8	532	-	94.2	103.3	1.10x
1280	2	532	532	-	1069.2	1067.3	1.00x
1280	2	8	232	-	73.1	82.0	1.12x
1280	2	232	232	-	674.8	671.3	0.99x
1792	1	32	200	-	249.8	271.2	1.09x
1280	2	8	532	causal	94.3	102.9	1.09x
1280	2	532	532	causal	566.5	568.8	1.00x
1280	2	8	232	causal	72.3	81.0	1.12x
1280	2	232	232	causal	315.2	315.8	1.00x
1792	1	32	200	causal	226.8	246.4	1.09x

The gain is on the short-query shapes (s_q = 8 and 32), with or without the causal mask: they run the bm0=64 qr_tdm tile and are bound by the K/V stream. On this develop, d=128 problems with max_seqlen_q >= 128 take the qr_tdm_sched pipeline (#13063), so 532/532 and 232/232 do not run the qr_tdm kernels that take the hint, and are unchanged.

Other shapes (bhsd)

batch	nhead	s_q	s_k	mask	TFLOPS develop	TFLOPS this PR	speedup
1	8	4096	4096	causal (top-left)	948.1	948.2	1.00x
1	16	4096	4096	causal (top-left)	1550.9	1543.0	0.99x
2	8	2048	2048	causal (top-left)	827.8	827.2	1.00x
1	8	4096	4096	causal	951.1	952.7	1.00x
2	16	2048	2048	causal (top-left)	1070.2	1069.0	1.00x

These are the only long sequences that reach a hinted kernel: bhsd square causal grids small enough to take the causal bm0=64 kernel. They stay within run-to-run noise.

Numerics

-v=1 passes 52/52 bf16/fp16 d=128 cases, all on qr_tdm: batch/group mode, padded M/N, causal top-left/bottom-right, sliding window, elementwise bias, alibi, GQA, LSE, the 427x3 uneven grid, bshd, and batch-mode q/kv_eff_lens; and 176/176 bf16/fp16 cases for d = 64, 96, 160, 192/128 and 256 (batch/group, no mask and causal top-left/bottom-right, short and long sequences). With the hint on d=256, another 24/24 bf16 d=256 cases (bshd/bhsd, no mask and causal top-left/bottom-right, up to 2048) pass. Only the cache hint of the K/V loads changes; the data path is identical.

Base

Based on develop at 46d0941 ("fix(rocke): normalize gfx11 processors to generic catalog (#13531)"), and measured against that commit.

JIRA ID : AICK-2262

therock-pr-bot · 2026-10-11T03:04:09Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR 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.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-10-11T03:04:12Z

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

…dense short-query tiles on gfx1250 qr_tdm moves K/V from HBM straight into LDS through the TDM and never reads them back through the vector path, so allocating them in the near cache only evicts other data. Build the K/V views with NT_HT (non-temporal near, high priority far) for the unmasked kM0 <= 64 tiles, which serve the dense short-query, memory-bound shapes. The bm0=128 tile and the masked bm0=64 kernels keep the default: they also serve long sequences (square causal grids on bm0=64), where NT_HT costs K/V reuse.

…es, masks included d=256 and the fp8 types use the bm0=64 tile at every sequence length, so the kM0 <= 64 test also put NT_HT on their long sequences, where it costs K/V reuse. Limit the hint to fp16/bf16 with head dim below 256, whose bm0=64 tile is only dispatched for short queries. Apply it to masked kernels too: causal is as common as dense on these short-query shapes and gains the same (b=1280 h=2 s_q=8 s_k=532 bshd, bottom-right causal 95.0 -> 103.7 TFLOPS, 1.09x). The only long sequences a masked bm0=64 kernel serves are short bhsd square causal grids; those measure 0.98x-1.03x (median of 5) with the hint. Move the hint below kHasMask and state the codegen dispatch it relies on.

poyenc · 2026-10-15T08:57:20Z

Location: kKVCoherence

Non-blocking — process question, not a requested change.

This hint has no test or benchmark coverage anywhere in the repository — its only justification is TFLOPS numbers quoted in commit messages, not checked by CI. The condition gating it has already changed scope twice within this PR's own three commits, the final change undocumented.

Would a cheap compile-time inventory test (asserting which (dtype, head dim, pipeline, tile) combinations receive NT_HT versus the default) be worth adding, so a future change to this condition can't silently widen or narrow its scope without a reviewer noticing?

…6 below d256 The previous commit widened the hint to every qr_tdm tile of at most 64 rows, which also reached the fp8 tiles and the d256 tile; both serve every sequence length, the fp8 ones were never measured and d256 is neutral (1.00x-1.01x). Restore the fp16/bf16, head dim below 256 bound. Test the pipeline property instead of its name (kUsesLdsArena), so TDM pipelines that stage K/V in the same arena get the hint without another edit, and pin the predicate in detail::uses_kv_nt_ht_hint_v with compile-time cases for every axis (dtype, head dim, bm0, arena).

Yanxing-Shi · 2026-10-18T02:36:47Z

Re this comment:

Added. The predicate is detail::uses_kv_nt_ht_hint_v, keyed on the TDM LDS arena (kUsesLdsArena) rather than the pipeline name, and test_qr_tdm_causal_mapping.cpp pins it per axis (dtype, head dim, bm0, arena) plus two real qr_tdm pipelines.

…hint

assistant-librarian · 2026-10-28T12:42:40Z

TheRock Submodule Bump Activity

Newest first

2026-10-28 — Included in TheRock via ROCm/TheRock#9553.

…_pingpong (#13484) ## What A 16-row short-query tile for `qr_tdm_pingpong` (#13471), taken when `max_seqlen_q <= 16` on the grids `qr_tdm_pingpong` takes (`max_seqlen_q < 640 && batch * nhead_q >= 2 * num_cus`; d=128 fp16/bf16, gfx1250). The 64-row tile idles 3/4 of its M rows on such queries; this one spends its four warps on N instead (gemm_0 and gemm_1 warps 1x4x1). This folds in #13469, which put the same tile on `qr_tdm`; `qr_tdm` stays as on `develop`. With the warps on N: - each warp holds a column slice of S, so S goes through an LDS region between the K buffers and V (57600 B arena) and comes back in P's distribution, which gemm_1's A operand takes. An encoding built from the QK gemm's warp layout would permute P's columns within a row: the row sums (and LSE) stay right while each key meets the wrong value row; - the row sum of P cannot run on the matrix unit (one warp tile of columns no longer spans a row), so `l` is the fp32 sum, which the LSE shares; - `MakePForGemm1`'s in-thread (MIter, KIter) reorder is only valid with the warps on M or one MIter per warp, now a `static_assert` at the call. The tile is owned by `qr_tdm_pingpong` (`F_pipeline_tag`) and listed before the 64-row tile, with `max_seqlen_q <= 16` added to the shared grid condition; on smaller grids short queries stay on `qr_tdm` as on `develop`. No other rule changes. ## Performance gfx1250, bf16, d=128, `lse=0`, bshd; median of 3 interleaved runs, each after >= 1 s of warmup, on an idle host. "causal" is the bottom-right mask (`-mask=b`). Baseline is this PR's parent, #13471 on `develop` 9951b48. | batch | nhead | s_q | s_k | mask | TFLOPS develop | TFLOPS base | TFLOPS this PR | vs base | vs develop | |------:|------:|----:|-----:|--------|---------------:|------------:|---------------:|--------:|-----------:| | 1280 | 2 | 8 | 532 | - | 95.0 | 127.6 | 136.8 | **1.07x** | **1.44x** | | 1280 | 2 | 8 | 232 | - | 73.0 | 112.5 | 120.2 | **1.07x** | **1.65x** | | 1280 | 2 | 16 | 532 | - | 186.3 | 250.9 | 269.0 | **1.07x** | **1.44x** | | 2048 | 4 | 12 | 300 | - | 136.0 | 186.8 | 199.9 | **1.07x** | **1.47x** | | 512 | 8 | 1 | 1024 | - | 15.4 | 17.5 | 17.9 | 1.02x | **1.16x** | | 256 | 16 | 4 | 2048 | - | 65.2 | 71.0 | 73.3 | 1.03x | **1.12x** | | 64 | 32 | 8 | 4096 | - | 131.7 | 133.7 | 155.4 | **1.16x** | **1.18x** | | 1280 | 2 | 8 | 532 | causal | 94.5 | 123.4 | 135.5 | **1.10x** | **1.43x** | | 1280 | 2 | 8 | 232 | causal | 72.3 | 108.5 | 118.0 | **1.09x** | **1.63x** | | 1280 | 2 | 16 | 532 | causal | 182.3 | 240.7 | 263.5 | **1.09x** | **1.45x** | | 2048 | 4 | 12 | 300 | causal | 133.4 | 175.5 | 195.2 | **1.11x** | **1.46x** | | 512 | 8 | 1 | 1024 | causal | 15.4 | 16.6 | 17.9 | **1.08x** | **1.16x** | | 256 | 16 | 4 | 2048 | causal | 64.6 | 67.7 | 73.3 | **1.08x** | **1.13x** | | 64 | 32 | 8 | 4096 | causal | 130.6 | 138.8 | 155.3 | **1.12x** | **1.19x** | With `lse=1` (instances built without expert scheduling, as in #13471) the 16-row tile measures 1.15x/1.19x over the base on 1280 x 2 x 8 rows (s_k 232/532) and 1.42x/1.42x with the causal mask. Only queries of at most 16 rows take the tile; every other shape runs the same kernels as the base (the other five-workload rows: 532/532, 232/232 and 32/200 are within 1%). ## Tests - `test_fmha_fwd.cpp`: `QrTdmShortQuery`, 6 cases (s_q 1/8/16, seqlen_k a tile multiple and a padded tail, causal and windowed masks, group mode), each sized to reach the grid threshold on the device it runs on and checking that the 16-row kernel ran. They check O, which the column permutation above would break while leaving the LSE right. - `test_qr_tdm_lds_padding.cpp`: the 16-row tile's arena (S region at 34800, Q/V at 39168, 57600 B) and that it takes the fp32 row sum. - `test_fmha_fwd_codegen.py`: on a big grid `max_seqlen_q <= 16` dispatches the 16-row tile and 17 rows the 64-row one (including the 1792 x 1 x 32-row workload); on a small grid 8 rows stay on `qr_tdm`. ## Base Built on #13471 (on `develop` at 9951b48); it will be rebased onto `develop` once #13471 lands. ## Series The gfx1250 FMHA forward optimization series; every PR is based on `develop` and benchmarked against it (bf16, d=128). #12426 and #13468 were folded into #13471, #13469 into this PR. | # | PR | Change | Status | |-------|--------|----------------------------------------------------------------|--------| | opt 1 | #12474 | Batch-mode early exit and XCD-aware block remap | [![](https://img.shields.io/github/pulls/detail/state/ROCm/rocm-libraries/12474)](https://github.com/ROCm/rocm-libraries/pull/12474) | | opt 2 | #13466 | NT_HT cache hint on the K/V TDM loads | [![](https://img.shields.io/github/pulls/detail/state/ROCm/rocm-libraries/13466)](https://github.com/ROCm/rocm-libraries/pull/13466) | | opt 3 | #13471 | `qr_tdm_pingpong` pipeline, coalesced O store and its dispatch | [![](https://img.shields.io/github/pulls/detail/state/ROCm/rocm-libraries/13471)](https://github.com/ROCm/rocm-libraries/pull/13471) | | opt 4 | #13484 | 16-row short-query tile on `qr_tdm_pingpong` (this PR) | [![](https://img.shields.io/github/pulls/detail/state/ROCm/rocm-libraries/13484)](https://github.com/ROCm/rocm-libraries/pull/13484) | JIRA ID : AICK-2262

perf(ck_tile/fmha): [gfx1250 opt 2] NT_HT cache hint on the K/V loads of short-query tiles (#13466) ## What Builds the K/V global tensor views of the gfx1250 `qr_tdm` forward kernel with the `NT_HT` cache hint (non-temporal in the near cache, high-priority temporal in the far cache) for every `qr_tdm` tile of at most 64 rows (`kM0 <= 64`), with or without a mask. `qr_tdm` moves K/V from HBM straight into LDS through the TDM and never reads them back through the vector path, so allocating them in the near cache only evicts other data. `NT_HT` drops that allocation and keeps K/V high priority in the far cache, where the m-tiles of one (batch, head) share them. The TDM load takes its coherence from the tensor view (`buffer_view::tdm_get` -> `amd_tdm_load<Coherence>`), so the change is the two K/V view constructions in `FmhaFwdKernel`. The `bm0=128` tile, which serves long sequences, keeps the default: it relies on K/V reuse in the near cache and is 1-3% slower with the hint. The `bm0=64` tiles of d=256 (taken at every length) were measured flat with the hint, 1.00x-1.01x from 8/532 up to 8192/8192, dense and causal. ## Performance gfx1250, bf16, d=128, `lse=0`. Baseline is `develop` at 46d0941043 (2026-10-13); median of 3 interleaved runs, each after >= 1 s of warmup (steady-state clocks), on an idle host. "causal" is the bottom-right mask (`-mask=b`) unless noted; bshd unless noted. **Large batch, small seqlen** | batch | nhead | s_q | s_k | mask | TFLOPS develop | TFLOPS this PR | speedup | |------:|------:|----:|----:|--------|---------------:|---------------:|----------:| | 1280 | 2 | 8 | 532 | - | 94.2 | 103.3 | **1.10x** | | 1280 | 2 | 532 | 532 | - | 1069.2 | 1067.3 | 1.00x | | 1280 | 2 | 8 | 232 | - | 73.1 | 82.0 | **1.12x** | | 1280 | 2 | 232 | 232 | - | 674.8 | 671.3 | 0.99x | | 1792 | 1 | 32 | 200 | - | 249.8 | 271.2 | **1.09x** | | 1280 | 2 | 8 | 532 | causal | 94.3 | 102.9 | **1.09x** | | 1280 | 2 | 532 | 532 | causal | 566.5 | 568.8 | 1.00x | | 1280 | 2 | 8 | 232 | causal | 72.3 | 81.0 | **1.12x** | | 1280 | 2 | 232 | 232 | causal | 315.2 | 315.8 | 1.00x | | 1792 | 1 | 32 | 200 | causal | 226.8 | 246.4 | **1.09x** | The gain is on the short-query shapes (s_q = 8 and 32), with or without the causal mask: they run the bm0=64 `qr_tdm` tile and are bound by the K/V stream. On this `develop`, d=128 problems with `max_seqlen_q >= 128` take the `qr_tdm_sched` pipeline (#13063), so 532/532 and 232/232 do not run the `qr_tdm` kernels that take the hint, and are unchanged. **Other shapes (bhsd)** | batch | nhead | s_q | s_k | mask | TFLOPS develop | TFLOPS this PR | speedup | |------:|------:|-----:|-----:|-------------------|---------------:|---------------:|--------:| | 1 | 8 | 4096 | 4096 | causal (top-left) | 948.1 | 948.2 | 1.00x | | 1 | 16 | 4096 | 4096 | causal (top-left) | 1550.9 | 1543.0 | 0.99x | | 2 | 8 | 2048 | 2048 | causal (top-left) | 827.8 | 827.2 | 1.00x | | 1 | 8 | 4096 | 4096 | causal | 951.1 | 952.7 | 1.00x | | 2 | 16 | 2048 | 2048 | causal (top-left) | 1070.2 | 1069.0 | 1.00x | These are the only long sequences that reach a hinted kernel: bhsd square causal grids small enough to take the causal `bm0=64` kernel. They stay within run-to-run noise. ## Numerics `-v=1` passes 52/52 bf16/fp16 d=128 cases, all on `qr_tdm`: batch/group mode, padded M/N, causal top-left/bottom-right, sliding window, elementwise bias, alibi, GQA, LSE, the 427x3 uneven grid, bshd, and batch-mode `q/kv_eff_lens`; and 176/176 bf16/fp16 cases for d = 64, 96, 160, 192/128 and 256 (batch/group, no mask and causal top-left/bottom-right, short and long sequences). With the hint on d=256, another 24/24 bf16 d=256 cases (bshd/bhsd, no mask and causal top-left/bottom-right, up to 2048) pass. Only the cache hint of the K/V loads changes; the data path is identical. ## Base Based on `develop` at 46d0941043 ("fix(rocke): normalize gfx11 processors to generic catalog (#13531)"), and measured against that commit. JIRA ID : AICK-2262

Yanxing-Shi requested review from a team as code owners October 11, 2026 03:00

github-actions Bot added project: composablekernel ck: attention-moe Used to tag composablekernel PRs that require approval from Attention/MOE review team. labels Oct 11, 2026

Yanxing-Shi added the organization: ROCm label Oct 11, 2026

Yanxing-Shi changed the title ~~perf(ck_tile/fmha): [gfx1250 opt 3] NT_HT cache hint on the K/V loads of dense short-query tiles~~ perf(ck_tile/fmha): [gfx1250 opt 3] NT_HT cache hint on the K/V loads of short-query tiles Oct 11, 2026

Yanxing-Shi mentioned this pull request Oct 12, 2026

perf(ck_tile/fmha): [gfx1250 opt 4] 16-row short-query tile on qr_tdm_pingpong #13484

Merged

Yanxing-Shi added 3 commits October 13, 2026 03:54

perf(fmha/fwd): K/V NT_HT hint on every qr_tdm tile of at most 64 rows

d41e269

Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-kv-nt-hint branch from f1ff1ae to d41e269 Compare October 13, 2026 08:25

Yanxing-Shi changed the title ~~perf(ck_tile/fmha): [gfx1250 opt 3] NT_HT cache hint on the K/V loads of short-query tiles~~ perf(ck_tile/fmha): [gfx1250 opt 2] NT_HT cache hint on the K/V loads of short-query tiles Oct 13, 2026

Yanxing-Shi requested a review from Jeff-Huang October 14, 2026 02:47

poyenc reviewed Oct 15, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp Resolved

poyenc approved these changes Oct 22, 2026

View reviewed changes

Yanxing-Shi mentioned this pull request Oct 22, 2026

feat(ck_tile/fmha): gfx1250 TDM pipelines for aiter's attention APIs #13711

Draft

Merge branch 'develop' into users/yanxishi/ck/gfx1250-fmha-fwd-kv-nt-…

cef52d4

…hint

asleepzzz merged commit 8b0fd31 into develop Oct 28, 2026
103 checks passed

asleepzzz deleted the users/yanxishi/ck/gfx1250-fmha-fwd-kv-nt-hint branch October 28, 2026 06:57

Navigation Menu

Repository navigation

perf(ck_tile/fmha): [gfx1250 opt 2] NT_HT cache hint on the K/V loads of short-query tiles - #13466

perf(ck_tile/fmha): [gfx1250 opt 2] NT_HT cache hint on the K/V loads of short-query tiles#13466
asleepzzz merged 5 commits into
developfrom
users/yanxishi/ck/gfx1250-fmha-fwd-kv-nt-hint

Yanxing-Shi commented Oct 11, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Oct 11, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Oct 11, 2026

Uh oh!

poyenc commented Oct 15, 2026

Uh oh!

Uh oh!

Yanxing-Shi commented Oct 18, 2026

Uh oh!

Uh oh!

assistant-librarian Bot commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Yanxing-Shi commented Oct 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Performance

Numerics

Base

Uh oh!

therock-pr-bot Bot commented Oct 11, 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 Oct 11, 2026

Uh oh!

poyenc commented Oct 15, 2026

Uh oh!

Uh oh!

Yanxing-Shi commented Oct 18, 2026

Uh oh!

Uh oh!

assistant-librarian Bot commented Oct 28, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Yanxing-Shi commented Oct 11, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 11, 2026 •

edited

Loading