Skip to content

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

Merged
asleepzzz merged 13 commits into
developfrom
users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong
Oct 30, 2026
Merged

asleepzzz merged 13 commits into
developfrom
users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong

Conversation

@Yanxing-Shi

@Yanxing-Shi Yanxing-Shi commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

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
opt 2 #13466 NT_HT cache hint on the K/V TDM loads
opt 3 #13471 qr_tdm_pingpong pipeline, coalesced O store and its dispatch
opt 4 #13484 16-row short-query tile on qr_tdm_pingpong (this PR)

JIRA ID : AICK-2262

@Yanxing-Shi
Yanxing-Shi requested review from a team as code owners October 12, 2026 08:09
@Yanxing-Shi Yanxing-Shi added ck: attention-moe Used to tag composablekernel PRs that require approval from Attention/MOE review team. organization: ROCm project: composablekernel labels Oct 12, 2026
@therock-pr-bot

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

Copy link
Copy Markdown

✅ 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

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

Copy link
Copy Markdown

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

@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch from 5721e25 to 292ea8f Compare October 13, 2026 08:25
@Yanxing-Shi Yanxing-Shi changed the title perf(ck_tile/fmha): [gfx1250 opt 7] 16-row short-query tile on the K ping/pong body perf(ck_tile/fmha): [gfx1250 opt 5] 16-row short-query tile on the K ping/pong body Oct 13, 2026
@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch 5 times, most recently from f7d4e89 to 9e21b59 Compare October 14, 2026 01:43
@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch 2 times, most recently from 6a24b2b to 3b81076 Compare October 15, 2026 01:01
@poyenc

poyenc commented Oct 15, 2026

Copy link
Copy Markdown
Contributor

Location: KernelComponentFactoryGfx125.get_rules()

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

This assembles four independently-maintained compatibility rules (check_gemm0_k, check_qr_tdm_kv_lds_buffering, check_short_query_tile, check_qr_tdm_sched_codegen_compatibility) on top of the tile tables and runtime CppConstraints, which are also ordered separately. Each new tile, head dim, or pipeline variant has to be threaded through all of these by hand.

Given this pattern already needed containment once (_QR_TDM_SCHED_TILES), is the plan to keep adding rules this way, or would collecting supported tile/pipeline-body pairings into one inventory that these rules derive from be worth it before the next head-dim or tile expansion?

@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch from 3b81076 to a596eab Compare October 17, 2026 14:43
@Yanxing-Shi Yanxing-Shi changed the title perf(ck_tile/fmha): [gfx1250 opt 5] 16-row short-query tile on the K ping/pong body perf(ck_tile/fmha): [gfx1250 opt 4] 16-row short-query tile on qr_tdm_pingpong Oct 17, 2026
std::min on two long_index_t values lowers to v_min_i64, which has no scalar
form. The compiler parks the wave-uniform size in a VGPR, every buffer resource
built from it looks divergent, and every buffer op on it gets a waterfall loop.

Clamp only sizes of 2^31 or more (a zero-length dimension makes the size
negative, which must stay unclamped), tested on the two 32-bit halves: a 64-bit
signed compare would move the value back into VGPRs.
…mha fwd epilogue can reuse

The wave32 WMMA C layout leaves each lane holding non-contiguous O columns,
so a direct store scatters every 128-byte write over 16 cache lines.
LdsShuffle2DEpilogue stages the tile through LDS so each wave writes
contiguous rows. Unlike CShuffleEpilogue it takes only the tile shape and
warp count and reads the C layout from the incoming tile distribution. Where
a row's 16-byte vectors cannot tile a wavefront (head dims 96, 160) it keeps
the direct store. The global store itself is Default2DEpilogue's, so padded
tiles keep its store_tile_raw path.

FmhaFwdKernel declares one LDS block for the pipeline and hands it to an
epilogue that asks for it (kUsesSmem), after the pipeline is done with it,
so LDS size and occupancy stay as they are. Pipelines that declare their own
arrays (non-arena prefill, fixed segmented arena) keep the direct store.
No generated kernel selects the epilogue yet.
@Yanxing-Shi

Copy link
Copy Markdown
Contributor Author

Re this comment:

Reworked in #13471. The K ping/pong body is now its own pipeline, qr_tdm_pingpong, with its own tiles: a tile can name the pipeline that owns it (FmhaFwdTileSize::F_pipeline_tag), and one rule (check_tile_owner) pairs owned tiles with their pipeline only. That replaced check_qr_tdm_kv_lds_buffering and check_short_query_tile, and the _o2 occupancy tag. The gfx125 rules are now develop's plus check_tile_owner, and the runtime choice of qr_tdm_pingpong is a single shared condition, max_seqlen_q < 640 && batch * nhead_q >= 2 * num_cus (see the "Dispatch" section of #13471). This PR only adds the 16-row tile, owned by qr_tdm_pingpong, under the same condition plus max_seqlen_q <= 16.

…es on big grids

A qr_tdm variant with a single V LDS buffer and a whole-tile K ping/pong: trip i
reads K buffer (i & 1) while the TDM fills the other with tile i + 1, V streams
through one buffer under the QK gemm, and Q sits on V's region so K's first
transfer goes out right behind Q's. Each LDS row is padded once (16 B for Q/K,
32 B for V), expressed as a row stride. The softmax rebases o_acc and l lazily
(only when a row max moves by more than 2^8) and sums P on the matrix unit; when
the LSE is stored, an fp32 row sum is kept next to it for the LSE, so O is the
same with or without LSE.

It is its own pipeline tag (qr_tdm_pingpong, BlockFmhaPipelineQRKSVSTdmPingPong
deriving from the qr_tdm pipeline and policy) with its own tile: tiles can now
name the pipeline that owns them (FmhaFwdTileSize::F_pipeline_tag,
check_tile_owner). qr_tdm and its generated instances are unchanged. The
LDS-shuffle epilogue applies to the qr_tdm_pingpong instances, and so does the
expert scheduling mode, except for the instances that store the LSE: built in
that mode they produce intermittently wrong O and LSE values on big grids.

qr_tdm_pingpong needs several resident blocks to overlap the softmax with the
gemms, while qr_tdm_sched and qr_tdm's double K/V buffer prefetch a tile ahead.
d128 fp16/bf16 takes qr_tdm_pingpong below 640 rows with at least 2 (batch,
head) pairs per num_cus (HIP's multiProcessorCount, 256 on gfx1250;
GFX125_QR_TDM_PINGPONG_RANGE); elsewhere qr_tdm_sched and qr_tdm keep their
develop ranges. On fewer pairs or longer aligned queries qr_tdm_pingpong
measured up to 21% slower than qr_tdm_sched.

gfx1250, bf16, d128, bshd, lse off, TFLOPS against develop:

  b     h  s_q  s_k  mask       develop  this    ratio
  1280  2    8  532  -            95.0  127.6   1.34x
  1280  2    8  532  causal b     94.5  123.4   1.31x
  1280  2  532  532  -          1078.7 1326.2   1.23x
  1280  2  532  532  causal b    571.2  892.5   1.56x
  1280  2    8  232  -            73.0  112.5   1.54x
  1280  2    8  232  causal b     72.3  108.5   1.50x
  1280  2  232  232  -           671.1  963.4   1.44x
  1280  2  232  232  causal b    316.4  487.5   1.54x
  1792  1   32  200  -           249.9  384.0   1.54x
  1792  1   32  200  causal b    228.7  339.7   1.49x
@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch from 5c7973a to 8df9a29 Compare October 18, 2026 09:07
…n steps as functions

The qr_tdm_pingpong body was a third copy of qr_tdm's sink pre-seed of m/l, the
sink-aware K tile range, the LSE of a block with nothing to do, the TDM padding
configs, scale_s plus elementwise/ALiBi bias, the uneven-split and mask -inf, the
LSE store and the O normalisation. Write them once as static members of
BlockFmhaPipelineQRKSVSTdmPingPong (InitRowMaxSum, GetKTileRange,
StoreLseOfEmptyRows, MakeTdmConfig, ApplyScaleAndBias, MaskScores, StoreLse,
NormalizeO); the main loop keeps only its TDM scheduling and softmax rescaling.

qr_tdm itself is unchanged. Moving its two bodies onto the same calls computes the
same values (the five fwd test data types pass) but changes register allocation of
its kernels at the VGPR limit: d256 ALiBi went from 0 to 80 spilled VGPRs and ran
at 0.74x, d512 ALiBi causal at 0.90x.
The QrTdmPingPong cases all run d128 without a sink, so the 288 sink and the
300 dpad/dvpad qr_tdm_pingpong instances (the latter storing O through the
LDS-shuffle epilogue with a padded N) had no runtime coverage: the other suites
use grids too small to take qr_tdm_pingpong. QrTdmPingPongEdges runs them on the
same grids: a sink without a mask and with a sliding window, d120 causal, and
d128 x 120 with an elementwise bias over an unaligned seqlen_k.
@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch from 8df9a29 to 86d40b3 Compare October 18, 2026 10:08
… -Werror

The compiled-dispatch tests build the generated fmha_fwd_api.cpp without warnings, so
a dispatch condition comparing ck_tile::index_t with the unsigned num_cus passed here
and failed only in the CK build, which compiles it with -Werror. Turn that warning on
in the test; with the int/unsigned comparison put back it now fails.
With max_seqlen_q <= 16 the 64-row tile idles 3/4 of its M rows. Add a 16-row
qr_tdm_pingpong tile that spends its four warps on N instead (gemm_0 and gemm_1
warps 1x4x1). S then goes through an LDS region between the K buffers and V and
comes back in P's distribution, which gemm_1's A operand takes; with the gemm_1
warps on N the row sum of P cannot run on the matrix unit, so l is the fp32 sum,
which the LSE shares.

The tile takes the same grid condition as the 64-row pingpong tile
(GFX125_QR_TDM_PINGPONG_RANGE, at least 2 (batch, head) pairs per num_cus); on
smaller grids short queries stay on qr_tdm as on develop.
@Yanxing-Shi
Yanxing-Shi force-pushed the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch from 86d40b3 to f4292af Compare October 18, 2026 10:57
@asleepzzz
asleepzzz enabled auto-merge (squash) October 30, 2026 06:36
@asleepzzz
asleepzzz merged commit 1e9db89 into develop Oct 30, 2026
175 of 177 checks passed
@asleepzzz
asleepzzz deleted the users/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong branch October 30, 2026 23:29
@assistant-librarian

Copy link
Copy Markdown
Contributor

TheRock Submodule Bump Activity

Newest first

shumway pushed a commit to ROCm/composable_kernel that referenced this pull request Nov 4, 2026
perf(ck_tile/fmha): [gfx1250 opt 4] 16-row short-query tile
 on qr_tdm_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`
9951b48ee3.

| 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 9951b48ee3); 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ck: attention-moe Used to tag composablekernel PRs that require approval from Attention/MOE review team. organization: ROCm project: composablekernel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants