Repository navigation
perf(ck_tile/fmha): [gfx1250 opt 3] 16-row short-query tile for the gfx1250 qr_tdm forward - #13469
Yanxing-Shi wants to merge 10 commits into
Conversation
github-actions
Bot
added
project: composablekernel
ck: attention-moe
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
7195ad7 to
d0a387f
Compare
At max_seqlen_q <= 16 the 64-row tile idles at least 3/4 of its M rows. Add a d=128 fp16/bf16 tile of 16 rows that spends its four warps on N (1x4x1 for both gemms), dispatched first for max_seqlen_q <= 16 and only paired with qr_tdm: qr's gemm_1 slices P along M and cannot form a slice narrower than one warp tile. Warps on N route S through LDS, a path no shipped qr_tdm tile exercised, and it needs three fixes: - MakeSRegTileDistribution built S's layout from the QK gemm's warp config, but S is consumed by gemm_1 as its A operand. The layouts differ, so P's columns reached gemm_1 permuted within each row: the row sum and LSE stay right while O is off by ~1%. Use the P distribution, which embeds the PV warp gemm's A encoding. - On the double K/V LDS buffer path S aliased K0, which the K prefetch two iterations ahead overwrites. Give S its own slot after V1 when it is in LDS at all (kSBytes > 0); tiles with the warps on M are unchanged. - The progressive K reload walks four N-subtiles per warp and needs the gemm_0 warps all on M, so it is only enabled for kNWarp == 1. The LDS padding selection also accepts the 16-row, warps-on-N tile, which otherwise ran with unpadded K/V rows. test_qr_tdm_lds_padding checks the padding selection and the arena (S at 71680, 76032 bytes) for that tile. gfx1250, bf16, d=128, bshd, b=1280 h=2, TFLOPS develop -> this: s_q=8 s_k=532 96.6 -> 111.5 (1.15x), causal 95.9 -> 110.3 (1.15x) s_q=8 s_k=232 73.7 -> 87.0 (1.18x), causal 73.3 -> 85.3 (1.16x) Shapes with s_q > 16 keep their tiles and are unchanged.
With the gemm_0 warps on N, S has its own slot in the double K/V buffer arena, but the policy's GetSmemSize still returned the K/V buffers only. The kernel sizes its LDS from the arena layout, so nothing was under- allocated; keep the policy's estimate consistent with it anyway.
Run s_q=8 against s_k=509 with LSE and assert the b16x64 qr_tdm instance is selected. A mis-distributed S on its way through LDS permutes P within a row, which keeps the row sums (and LSE) right while O is wrong, so this case needs the O comparison rather than LSE alone.
…tQueryProblem -Wshadow (an error in the CK build) flags the function-local alias in the LDS arena checks, which reused the name of the file-scope short-query problem.
…SmemSize The hand-summed region sizes dropped the 256-byte region alignment (and the double-buffer S slot is placed after V1), so they disagreed with the arena the pipeline actually allocates. The arena layout owns the offsets; use it.
…e Gfx125 rules Its codegen tests address it as get_rules()[-1].
… tile max_seqlen_q <= 16 takes the 16-row tile, not the bm0=64 one.
636f499 to
71530c7
Compare
Yanxing-Shi
changed the title
…ths and masks QrTdmShortQuery ran one batch-mode s_q=8 dense case with LSE off, while its comment asked for LSE on. It now runs six representative cases with LSE on: s_q = 1, 8 and 16, seqlen_k on and off a tile multiple, no mask, both causal masks and a sliding window, batch and group mode, and checks the tile in batch mode. MHA keeps pack-GQA from folding the heads into seqlen_q. Also: ruff-format the rule list, put the progressive-K comment back on kStagedKPairs, and word the S distribution comment for the code as it is.
…om the problem's The pipeline's kProgressiveDsLoadK (the problem's flag, off for tiles with the warps on N) shadowed Problem::kProgressiveDsLoadK, which the static_assert above and the policy still read; call it kUseProgressiveDsLoadK. No functional change.
…le_fmha target When a PR touches example/ck_tile/01_fmha, smart_build_ci.sh builds the test_ck_tile_fmha umbrella and runs every ctest matching test_ck_tile_fmha. That includes test_ck_tile_fmha_reference, which the umbrella did not build, so CI reported it as "Not Run" on gfx1201 and the other non-gfx125 targets.
|
Location: This repack unconditionally transposes |
| ) -> bool: | ||
| # The 16-row tile spends its warps on N, which only qr_tdm handles: | ||
| # qr's gemm_1 slices P along M and cannot form a slice narrower than | ||
| # one warp tile, so pairing the tile with qr fails to compile. |
There was a problem hiding this comment.
This check keys on the tile's F_bm0 dimension rather than the geometry invariant it's actually protecting. The comment above explains the real constraint: qr's gemm_1 slices P along M and can't form a slice narrower than one warp tile, which only matters when gemm_1's warps are split along N instead of M. That's exactly what F_rn1 encodes (the tile dataclass defines it as "number of warps for gemm1 along head dim v"). Today's single F_bm0==16 tile happens to be the only N-split tile, so the check works, but it only works by coincidence: a future N-split tile at a different F_bm0 would skip this guard and still fail to compile with qr, while an unrelated M-split tile that happens to use F_bm0==16 would be needlessly barred from qr.
Check the actual invariant instead:
if kernel_ctx.tile.F_rn1 <= 1:
return True
return kernel_ctx.pipeline.tag == "qr_tdm"
| return True | ||
| return kernel_ctx.pipeline.tag == "qr_tdm" | ||
|
|
||
| # The sched rule stays last (its tests address it as get_rules()[-1]). |
There was a problem hiding this comment.
This comment documents that check_qr_tdm_sched_codegen_compatibility must stay last in the list because two tests reach it positionally through get_rules()[-1]. Since get_rules()'s output is only ever consumed as a conjunction over all rules, list order has no effect on which kernels get generated — the ordering constraint exists purely to satisfy those two tests' indexing, and this comment is the production code accommodating that rather than the tests being fixed. The next engineer adding a compatibility rule has to remember to append before this one, with only a comment enforcing it; a reorder or an unrelated append placed after it will silently break both tests again, the same way this PR's own first commit did when it introduced check_short_query_tile.
Have those two tests select the rule by name or function identity instead of position, e.g. next(r for r in get_rules() if r.__name__ == "check_qr_tdm_sched_codegen_compatibility"), then this comment and the ordering constraint on rules.extend(...) here can both be removed.
|
Location: This comment documents LDS offsets only for the pre-existing warps-on-M d128 layout (double buffer ending at 71680, single buffer ending at 22784). It doesn't mention the new warps-on-N short-query tile this PR adds, where the double-buffer layout now places S after V1 and ends at 76032 instead of 71680, per the arena-offset static_asserts in Add the warps-on-N offsets to this comment, or replace the concrete numbers with a pointer to the compile-time layout tests that now cover both cases. |
|
Closing: the 16-row tile now lives only on the new qr_tdm_pingpong pipeline in #13484, and qr_tdm stays as on develop. The review points carried over: the MakePForGemm1 static_assert is in the qr_tdm_pingpong body, the tile/pipeline pairing is a tile owner field instead of a bm0 check, and the arena layouts are pinned by test_qr_tdm_lds_padding.cpp. |
…_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://github.com/ROCm/rocm-libraries/pull/12474) | | opt 2 | #13466 | NT_HT cache hint on the K/V TDM loads | [](https://github.com/ROCm/rocm-libraries/pull/13466) | | opt 3 | #13471 | `qr_tdm_pingpong` pipeline, coalesced O store and its dispatch | [](https://github.com/ROCm/rocm-libraries/pull/13471) | | opt 4 | #13484 | 16-row short-query tile on `qr_tdm_pingpong` (this PR) | [](https://github.com/ROCm/rocm-libraries/pull/13484) | JIRA ID : AICK-2262
What
Adds a 16-row short-query tile for the gfx1250
qr_tdmforward (d=128, fp16/bf16), dispatched first formax_seqlen_q <= 16.At
max_seqlen_q <= 16the 64-row tile idles at least 3/4 of its M rows: its four warps sit on M, and three of them compute nothing but padding. The new tile has 16 rows and spends its four warps on N instead (1x4x1 for both gemms), so every warp works on the real query rows and owns a quarter of the K/V columns. It is only paired withqr_tdm(a codegen rule):qr's gemm_1 slices P along M and cannot form a slice narrower than one warp tile.Warps on N route S through LDS, a path no shipped
qr_tdmtile exercised, and it needs three fixes:MakeSRegTileDistributionbuilt S's layout from the QK gemm's warp config, but S is consumed by gemm_1 as its A operand. The layouts differ, so P's columns reached gemm_1 permuted within each row: the row sum and LSE stay right while O is off by ~1%. It now returns the P distribution, which embeds the PV warp gemm's A encoding.kSBytes > 0), and the pipeline takes its address from the arena layout.kNWarp == 1.The LDS padding selection also accepts the 16-row, warps-on-N tile (it otherwise ran with unpadded K/V rows). Tiles with the warps on M generate the same code as before.
Resources of the new kernel: 152 VGPRs, no scratch, 76032 B of LDS (the d=128 padded arena plus 4352 B for S), i.e. 4 workgroups per CU like the 64-row tile. Cost: the d=128 fp16/bf16
qr_tdminstance count goes from 601 to 901. The new instances carry thekvlp_plktraits of the d=128 pipelines, but progressive K is off inside them.Performance
gfx1250, bf16, d=128,
lse=0. Baseline isdevelopat 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); bshd.Large batch, small seqlen
The s_q=8 shapes take the new tile. 32/200 (
s_q > 16) keeps its tile; on thisdevelop, d=128 problems withmax_seqlen_q >= 128take theqr_tdm_schedpipeline (#13063), so 532/532 and 232/232 do not run theqr_tdmkernels at all. Both are unchanged.Other shapes
Every shape with
s_q <= 16gains 1.13x-1.22x, with or without the causal mask.Numerics
-v=1passes 52/52 bf16/fp16 d=128 cases, all onqr_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-modeq/kv_eff_lens; and 384/384 short-query cases (bf16/fp16, batch/group, s_q = 1, 8, 12, 16 x s_k = 64, 232, 517, with no mask, causal top-left/bottom-right, sliding window, LSE, alibi, elementwise bias and MQA), 222 of them on the new tile.test_qr_tdm_lds_paddingnow also checks the padding selection and the LDS arena of the 16-row tile; those checks fail ondevelop.QrTdmShortQueryintest_fmha_fwd.cppnow runs 6 cases instead of one (s_q = 1, 8, 16; seqlen_k on and off a tile multiple; no mask, both causal masks and a sliding window; batch and group; LSE on) and checks the 16-row tile in batch mode.The PR also adds
test_ck_tile_fmha_referenceto thetest_ck_tile_fmhabuild target. CI builds that target and runs everytest_ck_tile_fmha*test when a PR touchesexample/ck_tile/01_fmha, and reported the reference test as "Not Run" on the non-gfx125 targets.Base
Based on
developat 46d0941 ("fix(rocke): normalize gfx11 processors to generic catalog (#13531)"), and measured against that commit.JIRA ID : AICK-2262