Repository navigation
Add scheduled TDM FMHA for FP16 and BF16 - #13063
Merged
asleepzzz merged 12 commits intoOct 13, 2026
Merged
Conversation
github-actions
Bot
added
documentation
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. |
asleepzzz
approved these changes
Oct 12, 2026
asleepzzz
enabled auto-merge (squash)
October 12, 2026 05:26
Jeff-Huang
force-pushed
the
users/chiahuan/ck/qr-tdm-sched-develop-20261005
branch
2 times, most recently
from
October 12, 2026 10:15
dec5ece to
86c97e6
Compare
Run clang-format-18 over the two files flagged by the CK Static checks stage. Both changes are whitespace-only; the token streams and string literals are byte-identical to the previous revision. - tile_window.hpp: align the tdm_load_to_lds template parameter list and join the gfx125 s_max_i32 inline asm onto a single permitted line. - fmha_tdm_sched_config.hpp: narrow the kVPadAmount / kVPadInterval assignment alignment group.
Jeff-Huang
force-pushed
the
users/chiahuan/ck/qr-tdm-sched-develop-20261005
branch
from
October 12, 2026 23:12
86c97e6 to
e37c194
Compare
This was referenced Oct 13, 2026
Merged
Merged
Closed
Closed
Contributor
TheRock Submodule Bump ActivityNewest first
|
1 task done
asleepzzz
pushed a commit
that referenced
this pull request
Oct 28, 2026
…ap for gfx1250 fwd (#12474) ## What - **Batch-mode early exit.** With `cu_seqlen_q_ptr` (per-batch `q_eff_lens`), the grid still covers the padded `seqlen_q`, so every block past a batch's effective length ran the whole K/V loop and stored nothing. Those blocks now return early, as group mode already does. This is the main gain of the PR: **1.21x-1.72x** on batch-mode varlen. - **XCD remap.** gfx1250 dispatches block `L` to XCD `L % 8`, so the m-tiles of one (batch, head), which read the same K/V, end up on all eight L2s. The block ids are remapped into contiguous per-XCD segments (the bijection `GemmSpatiallyLocalTilePartitioner::RemapXCD` also writes, in branch-free form, now `ck_tile::remap_xcd` in `ck_tile/core/utility/xcd_remap.hpp`) so those tiles share one L2. Both layouts are covered: bhsd in the head-major decode, and bshd in the head-fastest decode, which now decodes the m-tile fastest so a (batch, head) stays contiguous inside its segment. The remap applies only on `__gfx125__`, only in batch mode, and only without `cu_seqlen_q/k`: with per-batch lengths it would leave the XCDs unbalanced, so varlen keeps the round-robin. In the bshd path it needs more than one m-tile per (batch, head). Masked bshd grids reverse the m-tile order inside each (batch, head) (long tiles first); they are remapped only when every XCD gets at least as many (batch, head) groups as a group has m-tiles, because with fewer groups the round-robin's long-tiles-first order across all heads is what keeps the tail short (1x4x16384 causal lost 16% without that condition). The bhsd decode needs no such condition: its masked grids that skip the batch-local causal order (non-square causal, sliding window) measure 1.01x-1.12x with the remap, 1x4 heads included (1x4x8192/16384 bottom-right causal 1.02x, 1x4x16384 window 1024 1.07x, 8x16x2048/4096 causal 1.12x), and square causal and dense bhsd stay within 1%. ## 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. **Batch-mode varlen (early exit)** | layout | batch | nhead | s_q = s_k | mask | q/kv_eff_lens | ms develop | ms this PR | speedup | |--------|------:|------:|----------:|--------|--------------------------------------|-----------:|-----------:|----------:| | bhsd | 5 | 16 | 8192 | - | 8192,4096,2048,1024,512 | 0.459 | 0.335 | **1.37x** | | bhsd | 5 | 16 | 8192 | causal | 8192,4096,2048,1024,512 | 0.245 | 0.189 | **1.30x** | | bhsd | 8 | 8 | 4096 | - | 4096,3072,2048,1536,1024,512,256,128 | 0.097 | 0.076 | **1.28x** | | bshd | 8 | 8 | 4096 | - | 4096,3072,2048,1536,1024,512,256,128 | 0.095 | 0.070 | **1.36x** | | bhsd | 1280 | 2 | 532 | - | 8..532 (1) | 0.196 | 0.142 | **1.38x** | | bshd | 1280 | 2 | 532 | - | 8..532 (1) | 0.211 | 0.160 | **1.32x** | (1) 1280 lengths spread over 8..532 (`8 + (389 * i) % 525`), the same for q and kv. The gain is the padded part of the grid that no longer runs: the more the effective lengths fall short of the padded `seqlen_q`, the larger it is. Times are used here because the example's TFLOPS counts the padded problem. **Large batch, small seqlen (XCD remap, bshd)** | batch | nhead | s_q | s_k | mask | TFLOPS develop | TFLOPS this PR | speedup | |------:|------:|----:|----:|--------|---------------:|---------------:|----------:| | 1280 | 2 | 8 | 532 | - | 95.4 | 95.3 | 1.00x | | 1280 | 2 | 532 | 532 | - | 1074.3 | 1099.8 | 1.02x | | 1280 | 2 | 8 | 232 | - | 72.5 | 72.5 | 1.00x | | 1280 | 2 | 232 | 232 | - | 673.2 | 705.6 | **1.05x** | | 1792 | 1 | 32 | 200 | - | 252.5 | 250.1 | 0.99x | | 1280 | 2 | 8 | 532 | causal | 94.3 | 94.4 | 1.00x | | 1280 | 2 | 532 | 532 | causal | 571.1 | 585.0 | 1.02x | | 1280 | 2 | 8 | 232 | causal | 72.1 | 71.3 | 0.99x | | 1280 | 2 | 232 | 232 | causal | 317.0 | 332.4 | **1.05x** | | 1792 | 1 | 32 | 200 | causal | 228.4 | 228.0 | 1.00x | The `s_q = 8/32` shapes have one m-tile per (batch, head), so there is no K/V to share; they are the control and stay flat. 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; the remap sits in the kernel's tile decode and applies to that pipeline too (232/232 1.05x and 532/532 1.02x, with and without the causal mask). The gain is small here because `qr_tdm_sched` runs bm0=128, so 232 and 532 have only 2 and 5 m-tiles per (batch, head) to share K/V. With opt 3 (#13471), which runs these shapes on big grids on the bm0=64 `qr_tdm_pingpong` tile (4 and 9 m-tiles), the remap of the causal grids makes 232/232 causal 1.32x faster and 532/532 causal 1.04x (measured on the merged series). **Other shapes (bshd, XCD remap)** | batch | nhead | s_q | s_k | mask | TFLOPS develop | TFLOPS this PR | speedup | |------:|------:|------:|------:|-------------------|---------------:|---------------:|----------:| | 1280 | 2 | 232 | 232 | causal (top-left) | 316.7 | 330.7 | **1.04x** | | 1 | 4 | 16384 | 16384 | causal (top-left) | 1715.2 | 1703.4 | 0.99x | | 4 | 16 | 2048 | 2048 | causal (top-left) | 1431.8 | 1415.6 | 0.99x | | 427 | 3 | 232 | 232 | - | 626.1 | 653.8 | **1.04x** | | 2 | 16 | 4096 | 4096 | - | 2187.1 | 2170.3 | 0.99x | | 1 | 2 | 16384 | 16384 | causal (top-left) | 1070.4 | 1051.6 | 0.98x | 427x3 does not divide by 8 XCDs; the balanced segments keep the map a bijection there. Masked grids are remapped only when every XCD gets at least as many (batch, head) groups as a group has m-tiles: 1280x2x232 causal qualifies (1.04x); the long causal grids (1x4x16384, 4x16x2048, 1x2x16384) do not and keep `develop`'s long-to-short order across all heads (within 2%; remapped, 1x4x16384 lost 16% and 4x16x2048 11%). ## Numerics `-v=1` passes 52/52 bf16/fp16 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` with and without a causal mask. The remap was also checked by enumeration to be a bijection for 1760 (nhead, m-tiles, batch) grids, and to put each (batch, head) on a single XCD for the benchmark grids. New test: `BatchXCDRemap` in `test_fmha_fwd.cpp` (6 cases: bhsd and bshd, with and without per-batch `q/kv_eff_lens`, and causal) runs batch-mode grids with several m-tiles per (batch, head) and a block count that is not a multiple of the XCD count. The existing suites only reach grids where the remap is the identity, and uniform effective lengths in bhsd only. `test_ck_tile_xcd_remap` checks `ck_tile::remap_xcd` against the branching form for 1-200 ids on 1-16 XCDs, that it is a bijection, and that `remap_xcd_inverse` inverts it. The GEMM, StreamK and fmha bwd copies of the bijection are left as they are. ## Base Based on `develop` at 46d0941 ("fix(rocke): normalize gfx11 processors to generic catalog (#13531)"), and measured against that commit. JIRA ID : AICK-2262
asleepzzz
pushed a commit
that referenced
this pull request
Oct 28, 2026
… 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 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
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Nov 4, 2026
Add scheduled TDM FMHA for FP16 and BF16
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
JIRA ID: AICK-2310
# Add scheduled TDM FMHA for FP16 and BF16
## Summary
- Add the `qr_tdm_sched` forward FMHA pipeline for FP16/BF16 on gfx125:
Q/K–V dimensions 64/64, 128/128, and 192/128. The implementation uses
geometry-derived QK/PV scheduling, packed LDS, and scheduled
WMMA/LDS-read placement.
- Generate and try `qr_tdm_sched` first for its supported row-major V
cases with `max_seqlen_q >= 128`, `sink_ptr == nullptr`, and no bias,
dropout, soft cap, or quantization. Dense and supported causal masks are
included; LSE is optional.
- Select N64 for D64 and N128 for D128. D192 selects N128 when the
average K length is at least 512, K-length pointers are absent, the N128
workgroup count fits the CU count, and the mask is dense or unrestricted
square causal; other supported D192 cases select N64.
- Fall back to the existing `qr_tdm` candidate when a scheduled trait is
ineligible, including short Q or a non-null virtual-sink pointer. Shapes
or features outside `qr_tdm` continue through the existing `qr`
fallback.
- Separate the no-sink and finite-sink kernel specializations, and add
codegen, dispatch, layout, schedule, sink, and numerical regression
tests under `test/ck_tile/fmha`.
## BF16 performance
All performance numbers in this PR description are taken from the
unmodified `tile_example_fmha_fwd` output. The two binaries use the same
example main, inputs, GPU timer, and command-line options; one links
`qr_tdm_sched` from this PR and the other links upstream `qr_tdm`.
Kernel names were printed and checked on every run.
Measured on gfx1250 **B0 stepping** (256 CUs). SCLK was neither locked
nor logged by the benchmark, so no fixed SCLK value is claimed. The
historical 110 workloads per dimension use BSHD layout, `Sq = Sk` in
{256, 2048, 8192, 16384, 32768}, dense or causal masking, LSE, no sink,
and the same scale and input seed for both kernels. The example uses
GPU-event timing, five warmups, each workload’s historical repeat count,
and two adjacent alternating rounds. For each kernel, the table reports
the median of the two values printed by the example. Speedup divides the
two printed TFLOP/s medians; above 1× favors `qr_tdm_sched`. The example
prints latency to 0.001 ms and throughput to 0.01 TFLOP/s.
| Q/K–V dimension | Workloads faster with qr_tdm_sched (Introduced by
this PR) | Median speedup vs qr_tdm | Speedup range |
| --- | ---: | ---: | ---: |
| 64/64 | 110/110 | 1.317× | 1.054–1.819× |
| 128/128 | 110/110 | 1.319× | 1.145–1.882× |
| 192/128 | 110/110 | 1.179× | 1.035–1.882× |
The full per-workload results follow. Each kernel cell shows **execution
time in µs / effective TFLOP/s**. `B/H/Hkv` denotes batch size, query
heads, and key/value heads. Causal TFLOP/s uses triangular attention
work.
<details><summary>BF16 Q/K=64, V=64: all 110 workloads</summary>
| B/H/Hkv | Sq = Sk | Mask | qr_tdm_sched (Introduced by this PR)<br>µs
/ TFLOP/s | qr_tdm<br>µs / TFLOP/s | qr_tdm_sched / qr_tdm |
| --- | ---: | --- | ---: | ---: | ---: |
| 1/8/8 | 256 | dense | 6.5 / 20.60 | 8 / 16.43 | 1.254× |
| 1/8/8 | 256 | causal | 9 / 7.87 | 9.5 / 7.07 | 1.112× |
| 1/8/8 | 2048 | dense | 20 / 428.70 | 32 / 268.56 | 1.596× |
| 1/8/8 | 2048 | causal | 22 / 198.80 | 29 / 148.99 | 1.334× |
| 1/8/8 | 8192 | dense | 105 / 1306.69 | 143 / 961.35 | 1.359× |
| 1/8/8 | 8192 | causal | 76.5 / 899.52 | 104 / 658.12 | 1.367× |
| 1/8/8 | 16384 | dense | 352.5 / 1558.99 | 516 / 1065.49 | 1.463× |
| 1/8/8 | 16384 | causal | 205.5 / 1336.91 | 231 / 1191.24 | 1.122× |
| 1/8/8 | 32768 | dense | 1316 / 1670.73 | 1783.5 / 1233.18 | 1.355× |
| 1/8/8 | 32768 | causal | 683.5 / 1608.76 | 831 / 1323.27 | 1.216× |
| 1/64/8 | 256 | dense | 6 / 167.20 | 8 / 131.09 | 1.275× |
| 1/64/8 | 256 | causal | 9 / 61.97 | 10 / 55.78 | 1.111× |
| 1/64/8 | 2048 | dense | 54 / 1272.37 | 77 / 896.91 | 1.419× |
| 1/64/8 | 2048 | causal | 35 / 972.00 | 44 / 784.50 | 1.239× |
| 1/64/8 | 8192 | dense | 670.5 / 1640.03 | 889.5 / 1235.90 | 1.327× |
| 1/64/8 | 8192 | causal | 366.5 / 1499.62 | 443 / 1241.95 | 1.207× |
| 1/64/8 | 16384 | dense | 2584.5 / 1701.43 | 3442.5 / 1277.50 | 1.332×
|
| 1/64/8 | 16384 | causal | 1362 / 1614.49 | 1637.5 / 1343.09 | 1.202× |
| 1/64/8 | 32768 | dense | 10242.5 / 1717.59 | 13515.5 / 1301.62 |
1.320× |
| 1/64/8 | 32768 | causal | 5295.5 / 1661.06 | 6376.5 / 1379.45 | 1.204×
|
| 1/64/64 | 256 | dense | 6 / 166.81 | 8 / 129.20 | 1.291× |
| 1/64/64 | 256 | causal | 9 / 61.73 | 10 / 55.41 | 1.114× |
| 1/64/64 | 2048 | dense | 54.5 / 1260.94 | 77 / 895.85 | 1.408× |
| 1/64/64 | 2048 | causal | 36 / 961.11 | 44.5 / 776.22 | 1.238× |
| 1/64/64 | 8192 | dense | 700.5 / 1569.05 | 906 / 1214.18 | 1.292× |
| 1/64/64 | 8192 | causal | 375.5 / 1464.12 | 450 / 1222.10 | 1.198× |
| 1/64/64 | 16384 | dense | 2859.5 / 1538.04 | 3619.5 / 1214.97 | 1.266×
|
| 1/64/64 | 16384 | causal | 1461 / 1505.70 | 1733 / 1269.13 | 1.186× |
| 1/64/64 | 32768 | dense | 11411.5 / 1541.58 | 14538 / 1210.10 | 1.274×
|
| 1/64/64 | 32768 | causal | 5834.5 / 1507.61 | 6931.5 / 1269.02 |
1.188× |
| 1/12/12 | 256 | dense | 8 / 26.42 | 9 / 21.56 | 1.225× |
| 1/12/12 | 256 | causal | 10 / 10.31 | 11 / 9.39 | 1.098× |
| 1/12/12 | 2048 | dense | 23 / 565.14 | 35 / 368.93 | 1.532× |
| 1/12/12 | 2048 | causal | 23 / 279.37 | 30.5 / 211.40 | 1.321× |
| 1/12/12 | 8192 | dense | 147.5 / 1399.60 | 181.5 / 1137.03 | 1.231× |
| 1/12/12 | 8192 | causal | 92 / 1122.19 | 115 / 898.60 | 1.249× |
| 1/12/12 | 16384 | dense | 516.5 / 1596.95 | 670.5 / 1229.97 | 1.298× |
| 1/12/12 | 16384 | causal | 280 / 1471.83 | 349 / 1181.07 | 1.246× |
| 1/12/12 | 32768 | dense | 1977 / 1668.20 | 2623.5 / 1257.21 | 1.327× |
| 1/12/12 | 32768 | causal | 1021 / 1614.88 | 1219 / 1353.24 | 1.193× |
| 1/32/32 | 256 | dense | 6 / 85.71 | 8 / 67.95 | 1.261× |
| 1/32/32 | 256 | causal | 8 / 32.12 | 9.5 / 28.52 | 1.126× |
| 1/32/32 | 2048 | dense | 30 / 1143.37 | 41 / 831.19 | 1.376× |
| 1/32/32 | 2048 | causal | 25 / 696.09 | 32 / 541.97 | 1.284× |
| 1/32/32 | 8192 | dense | 349.5 / 1574.21 | 460 / 1194.55 | 1.318× |
| 1/32/32 | 8192 | causal | 192 / 1430.57 | 235 / 1171.82 | 1.221× |
| 1/32/32 | 16384 | dense | 1334 / 1648.75 | 1763 / 1247.51 | 1.322× |
| 1/32/32 | 16384 | causal | 696 / 1580.20 | 832.5 / 1321.26 | 1.196× |
| 1/32/32 | 32768 | dense | 5428.5 / 1620.33 | 7043.5 / 1248.83 | 1.297×
|
| 1/32/32 | 32768 | causal | 2760.5 / 1593.35 | 3289.5 / 1336.97 |
1.192× |
| 32/8/8 | 256 | dense | 13 / 341.98 | 15 / 295.50 | 1.157× |
| 32/8/8 | 256 | causal | 15 / 146.99 | 15 / 139.49 | 1.054× |
| 32/8/8 | 2048 | dense | 213 / 1290.18 | 268.5 / 1022.58 | 1.262× |
| 32/8/8 | 2048 | causal | 148.5 / 927.29 | 175.5 / 782.73 | 1.185× |
| 32/8/8 | 8192 | dense | 2637 / 1668.03 | 3477 / 1264.95 | 1.319× |
| 32/8/8 | 8192 | causal | 1474 / 1492.24 | 1785.5 / 1231.93 | 1.211× |
| 32/8/8 | 16384 | dense | 10304 / 1707.32 | 13556 / 1297.74 | 1.316× |
| 32/8/8 | 16384 | causal | 5525 / 1592.06 | 6652.5 / 1322.32 | 1.204× |
| 32/8/8 | 32768 | dense | 40970 / 1717.57 | 53682 / 1310.85 | 1.310× |
| 32/8/8 | 32768 | causal | 21349 / 1648.11 | 25695 / 1369.36 | 1.204× |
| 1/4/4 | 256 | dense | 6 / 11.62 | 8 / 8.68 | 1.339× |
| 1/4/4 | 256 | causal | 8 / 4.21 | 9 / 3.75 | 1.120× |
| 1/4/4 | 2048 | dense | 18 / 233.25 | 30 / 141.31 | 1.651× |
| 1/4/4 | 2048 | causal | 21 / 104.25 | 28 / 76.84 | 1.357× |
| 1/4/4 | 8192 | dense | 68.5 / 1005.61 | 114 / 604.32 | 1.664× |
| 1/4/4 | 8192 | causal | 67.5 / 508.64 | 97 / 353.75 | 1.438× |
| 1/4/4 | 16384 | dense | 197 / 1394.20 | 272.5 / 1008.87 | 1.382× |
| 1/4/4 | 16384 | causal | 144 / 952.23 | 200.5 / 685.69 | 1.389× |
| 1/4/4 | 32768 | dense | 681.5 / 1613.03 | 1019.5 / 1078.21 | 1.496× |
| 1/4/4 | 32768 | causal | 390 / 1411.11 | 442.5 / 1242.26 | 1.136× |
| 1/3/3 | 256 | dense | 6 / 8.69 | 8 / 6.55 | 1.327× |
| 1/3/3 | 256 | causal | 8 / 3.17 | 9 / 2.82 | 1.122× |
| 1/3/3 | 2048 | dense | 18 / 174.82 | 30 / 106.75 | 1.638× |
| 1/3/3 | 2048 | causal | 21 / 77.13 | 28 / 57.34 | 1.345× |
| 1/3/3 | 8192 | dense | 64 / 802.98 | 112 / 461.57 | 1.740× |
| 1/3/3 | 8192 | causal | 66.5 / 388.10 | 96 / 267.89 | 1.449× |
| 1/3/3 | 16384 | dense | 178.5 / 1155.61 | 263.5 / 783.27 | 1.475× |
| 1/3/3 | 16384 | causal | 137.5 / 749.85 | 194 / 531.33 | 1.411× |
| 1/3/3 | 32768 | dense | 528.5 / 1560.32 | 667 / 1236.47 | 1.262× |
| 1/3/3 | 32768 | causal | 316 / 1305.88 | 412 / 1000.44 | 1.305× |
| 1/2/2 | 256 | dense | 6 / 5.72 | 8 / 4.40 | 1.301× |
| 1/2/2 | 256 | causal | 8 / 2.21 | 9 / 1.97 | 1.124× |
| 1/2/2 | 2048 | dense | 17.5 / 123.09 | 29 / 73.66 | 1.671× |
| 1/2/2 | 2048 | causal | 20 / 54.11 | 27 / 39.39 | 1.374× |
| 1/2/2 | 8192 | dense | 61 / 561.33 | 109 / 316.07 | 1.776× |
| 1/2/2 | 8192 | causal | 65 / 264.22 | 95 / 181.44 | 1.456× |
| 1/2/2 | 16384 | dense | 127 / 1079.66 | 220 / 625.33 | 1.727× |
| 1/2/2 | 16384 | causal | 128 / 537.47 | 186.5 / 368.42 | 1.459× |
| 1/2/2 | 32768 | dense | 375.5 / 1463.87 | 533.5 / 1030.62 | 1.420× |
| 1/2/2 | 32768 | causal | 275 / 999.47 | 386 / 712.56 | 1.403× |
| 1/1/1 | 256 | dense | 5 / 3.51 | 6 / 2.59 | 1.355× |
| 1/1/1 | 256 | causal | 7 / 1.19 | 8 / 1.06 | 1.128× |
| 1/1/1 | 2048 | dense | 16 / 66.02 | 28 / 37.91 | 1.741× |
| 1/1/1 | 2048 | causal | 19 / 27.58 | 27 / 20.00 | 1.379× |
| 1/1/1 | 8192 | dense | 58 / 295.06 | 106 / 162.21 | 1.819× |
| 1/1/1 | 8192 | causal | 63 / 136.24 | 92.5 / 92.93 | 1.466× |
| 1/1/1 | 16384 | dense | 117 / 585.65 | 211 / 325.26 | 1.801× |
| 1/1/1 | 16384 | causal | 124 / 278.17 | 182 / 188.73 | 1.474× |
| 1/1/1 | 32768 | dense | 244 / 1128.26 | 427.5 / 642.97 | 1.755× |
| 1/1/1 | 32768 | causal | 247.5 / 556.17 | 365 / 376.88 | 1.476× |
| 32/1/1 | 256 | dense | 5 / 109.78 | 7 / 82.30 | 1.334× |
| 32/1/1 | 256 | causal | 7 / 37.70 | 8 / 33.35 | 1.130× |
| 32/1/1 | 2048 | dense | 27.5 / 1238.43 | 39 / 876.60 | 1.413× |
| 32/1/1 | 2048 | causal | 28 / 609.23 | 36 / 474.27 | 1.285× |
| 32/1/1 | 8192 | dense | 349 / 1575.91 | 460 / 1194.72 | 1.319× |
| 32/1/1 | 8192 | causal | 233 / 1180.41 | 299 / 920.61 | 1.282× |
| 32/1/1 | 16384 | dense | 1308.5 / 1680.83 | 1748.5 / 1258.01 | 1.336×
|
| 32/1/1 | 16384 | causal | 771 / 1426.37 | 953.5 / 1153.07 | 1.237× |
| 32/1/1 | 32768 | dense | 5088.5 / 1728.62 | 6778 / 1297.69 | 1.332× |
| 32/1/1 | 32768 | causal | 2802.5 / 1569.31 | 3492.5 / 1259.32 | 1.246×
|
</details>
<details><summary>BF16 Q/K=128, V=128: all 110 workloads</summary>
| B/H/Hkv | Sq = Sk | Mask | qr_tdm_sched (Introduced by this PR)<br>µs
/ TFLOP/s | qr_tdm<br>µs / TFLOP/s | qr_tdm_sched / qr_tdm |
| --- | ---: | --- | ---: | ---: | ---: |
| 1/8/8 | 256 | dense | 5 / 49.45 | 7 / 40.56 | 1.219× |
| 1/8/8 | 256 | causal | 7 / 19.76 | 10 / 14.05 | 1.406× |
| 1/8/8 | 2048 | dense | 19.5 / 883.62 | 26 / 652.13 | 1.355× |
| 1/8/8 | 2048 | causal | 21 / 411.26 | 36 / 238.12 | 1.727× |
| 1/8/8 | 8192 | dense | 139.5 / 1969.99 | 162.5 / 1694.40 | 1.163× |
| 1/8/8 | 8192 | causal | 91.5 / 1500.30 | 141 / 974.51 | 1.540× |
| 1/8/8 | 16384 | dense | 510 / 2154.63 | 589 / 1865.78 | 1.155× |
| 1/8/8 | 16384 | causal | 281.5 / 1951.12 | 383.5 / 1434.64 | 1.360× |
| 1/8/8 | 32768 | dense | 1956 / 2248.86 | 2269.5 / 1938.01 | 1.160× |
| 1/8/8 | 32768 | causal | 1020.5 / 2154.98 | 1377.5 / 1596.22 | 1.350×
|
| 1/64/8 | 256 | dense | 6 / 372.57 | 7 / 306.49 | 1.216× |
| 1/64/8 | 256 | causal | 7 / 153.30 | 10 / 108.67 | 1.411× |
| 1/64/8 | 2048 | dense | 78.5 / 1752.28 | 91 / 1513.26 | 1.158× |
| 1/64/8 | 2048 | causal | 50 / 1369.23 | 66 / 1035.78 | 1.322× |
| 1/64/8 | 8192 | dense | 1009 / 2180.05 | 1164 / 1888.82 | 1.154× |
| 1/64/8 | 8192 | causal | 560 / 1964.26 | 736 / 1493.53 | 1.315× |
| 1/64/8 | 16384 | dense | 3924 / 2241.60 | 4520.5 / 1945.76 | 1.152× |
| 1/64/8 | 16384 | causal | 2078 / 2116.28 | 2770 / 1587.78 | 1.333× |
| 1/64/8 | 32768 | dense | 16023 / 2195.86 | 18352 / 1917.18 | 1.145× |
| 1/64/8 | 32768 | causal | 8186.5 / 2148.93 | 10847 / 1621.95 | 1.325×
|
| 1/64/64 | 256 | dense | 6 / 368.88 | 7 / 306.31 | 1.204× |
| 1/64/64 | 256 | causal | 7 / 153.55 | 10 / 108.72 | 1.412× |
| 1/64/64 | 2048 | dense | 80.5 / 1708.78 | 92.5 / 1486.65 | 1.149× |
| 1/64/64 | 2048 | causal | 51 / 1351.80 | 67.5 / 1016.80 | 1.329× |
| 1/64/64 | 8192 | dense | 1096.5 / 2005.53 | 1267.5 / 1734.80 | 1.156×
|
| 1/64/64 | 8192 | causal | 584 / 1883.05 | 751 / 1464.62 | 1.286× |
| 1/64/64 | 16384 | dense | 4300 / 2045.51 | 5004.5 / 1757.77 | 1.164× |
| 1/64/64 | 16384 | causal | 2245.5 / 1958.67 | 2935 / 1498.59 | 1.307×
|
| 1/64/64 | 32768 | dense | 17102 / 2057.36 | 19826.5 / 1774.61 | 1.159×
|
| 1/64/64 | 32768 | causal | 8833 / 1991.76 | 11570.5 / 1520.44 | 1.310×
|
| 1/12/12 | 256 | dense | 6 / 68.67 | 7 / 57.87 | 1.187× |
| 1/12/12 | 256 | causal | 7 / 28.33 | 10 / 20.48 | 1.383× |
| 1/12/12 | 2048 | dense | 21 / 1222.07 | 28 / 923.80 | 1.323× |
| 1/12/12 | 2048 | causal | 22 / 592.16 | 37 / 348.87 | 1.697× |
| 1/12/12 | 8192 | dense | 210.5 / 1955.64 | 247.5 / 1664.91 | 1.175× |
| 1/12/12 | 8192 | causal | 125.5 / 1637.74 | 157.5 / 1311.42 | 1.249× |
| 1/12/12 | 16384 | dense | 775.5 / 2125.59 | 894.5 / 1844.38 | 1.152× |
| 1/12/12 | 16384 | causal | 420 / 1964.32 | 551.5 / 1495.67 | 1.313× |
| 1/12/12 | 32768 | dense | 3062 / 2154.51 | 3586 / 1839.65 | 1.171× |
| 1/12/12 | 32768 | causal | 1585.5 / 2079.95 | 2073 / 1591.16 | 1.307×
|
| 1/32/32 | 256 | dense | 6 / 192.74 | 7 / 157.75 | 1.222× |
| 1/32/32 | 256 | causal | 7 / 78.35 | 10 / 55.27 | 1.418× |
| 1/32/32 | 2048 | dense | 39.5 / 1733.26 | 46.5 / 1474.77 | 1.175× |
| 1/32/32 | 2048 | causal | 28 / 1225.74 | 41 / 838.59 | 1.462× |
| 1/32/32 | 8192 | dense | 520 / 2115.86 | 597 / 1842.40 | 1.148× |
| 1/32/32 | 8192 | causal | 287.5 / 1912.56 | 380 / 1445.95 | 1.323× |
| 1/32/32 | 16384 | dense | 2052 / 2143.53 | 2386 / 1843.25 | 1.163× |
| 1/32/32 | 16384 | causal | 1068.5 / 2058.40 | 1401.5 / 1569.14 |
1.312× |
| 1/32/32 | 32768 | dense | 8160 / 2155.86 | 9495 / 1852.79 | 1.164× |
| 1/32/32 | 32768 | causal | 4212.5 / 2088.30 | 5575 / 1577.80 | 1.324×
|
| 32/8/8 | 256 | dense | 9 / 948.69 | 11 / 807.17 | 1.175× |
| 32/8/8 | 256 | causal | 10 / 425.35 | 13 / 329.72 | 1.290× |
| 32/8/8 | 2048 | dense | 301.5 / 1822.08 | 354.5 / 1550.15 | 1.175× |
| 32/8/8 | 2048 | causal | 205 / 1339.68 | 264 / 1040.47 | 1.288× |
| 32/8/8 | 8192 | dense | 4004.5 / 2196.34 | 4650 / 1891.64 | 1.161× |
| 32/8/8 | 8192 | causal | 2218 / 1983.39 | 2948 / 1492.08 | 1.329× |
| 32/8/8 | 16384 | dense | 15650 / 2248.20 | 18080 / 1946.01 | 1.155× |
| 32/8/8 | 16384 | causal | 8324 / 2113.67 | 11086.5 / 1586.88 | 1.332×
|
| 32/8/8 | 32768 | dense | 62354 / 2257.07 | 71921.5 / 1956.83 | 1.153×
|
| 32/8/8 | 32768 | causal | 32432.5 / 2169.79 | 43190 / 1629.34 | 1.332×
|
| 1/4/4 | 256 | dense | 6 / 23.94 | 7 / 19.89 | 1.203× |
| 1/4/4 | 256 | causal | 7 / 9.84 | 10 / 6.97 | 1.410× |
| 1/4/4 | 2048 | dense | 18.5 / 464.11 | 25 / 338.91 | 1.369× |
| 1/4/4 | 2048 | causal | 20 / 211.88 | 36 / 120.81 | 1.754× |
| 1/4/4 | 8192 | dense | 80 / 1713.53 | 102 / 1345.16 | 1.274× |
| 1/4/4 | 8192 | causal | 71 / 961.78 | 129 / 531.40 | 1.810× |
| 1/4/4 | 16384 | dense | 268.5 / 2047.22 | 310.5 / 1771.69 | 1.156× |
| 1/4/4 | 16384 | causal | 172 / 1599.62 | 271.5 / 1013.66 | 1.578× |
| 1/4/4 | 32768 | dense | 996.5 / 2206.86 | 1157 / 1900.56 | 1.161× |
| 1/4/4 | 32768 | causal | 539.5 / 2037.87 | 746.5 / 1472.63 | 1.384× |
| 1/3/3 | 256 | dense | 6 / 17.51 | 7 / 14.69 | 1.192× |
| 1/3/3 | 256 | causal | 7 / 7.22 | 10 / 5.21 | 1.385× |
| 1/3/3 | 2048 | dense | 18 / 349.55 | 25 / 257.00 | 1.360× |
| 1/3/3 | 2048 | causal | 20 / 159.11 | 35 / 91.00 | 1.748× |
| 1/3/3 | 8192 | dense | 72 / 1433.12 | 96.5 / 1068.22 | 1.342× |
| 1/3/3 | 8192 | causal | 70 / 737.01 | 127.5 / 404.37 | 1.823× |
| 1/3/3 | 16384 | dense | 235 / 1756.28 | 277 / 1488.26 | 1.180× |
| 1/3/3 | 16384 | causal | 155.5 / 1327.24 | 262.5 / 785.17 | 1.690× |
| 1/3/3 | 32768 | dense | 761 / 2167.19 | 902.5 / 1826.72 | 1.186× |
| 1/3/3 | 32768 | causal | 423 / 1947.91 | 569.5 / 1448.76 | 1.345× |
| 1/2/2 | 256 | dense | 5 / 12.27 | 7 / 10.12 | 1.212× |
| 1/2/2 | 256 | causal | 7 / 4.96 | 10 / 3.50 | 1.414× |
| 1/2/2 | 2048 | dense | 18 / 238.93 | 25 / 174.05 | 1.373× |
| 1/2/2 | 2048 | causal | 20 / 108.42 | 35 / 61.23 | 1.771× |
| 1/2/2 | 8192 | dense | 66 / 1041.96 | 92 / 748.55 | 1.392× |
| 1/2/2 | 8192 | causal | 67.5 / 509.00 | 125.5 / 273.87 | 1.859× |
| 1/2/2 | 16384 | dense | 151.5 / 1813.10 | 193.5 / 1419.58 | 1.277× |
| 1/2/2 | 16384 | causal | 136 / 1010.78 | 250 / 549.94 | 1.838× |
| 1/2/2 | 32768 | dense | 512 / 2148.24 | 595 / 1848.09 | 1.162× |
| 1/2/2 | 32768 | causal | 326.5 / 1685.10 | 527.5 / 1041.99 | 1.617× |
| 1/1/1 | 256 | dense | 5 / 6.31 | 7 / 5.12 | 1.231× |
| 1/1/1 | 256 | causal | 7 / 2.51 | 10 / 1.77 | 1.421× |
| 1/1/1 | 2048 | dense | 17 / 123.37 | 24 / 89.28 | 1.382× |
| 1/1/1 | 2048 | causal | 20 / 54.98 | 35 / 30.91 | 1.778× |
| 1/1/1 | 8192 | dense | 62 / 553.37 | 88 / 389.71 | 1.420× |
| 1/1/1 | 8192 | causal | 65 / 262.91 | 123 / 139.94 | 1.879× |
| 1/1/1 | 16384 | dense | 126 / 1086.95 | 178 / 773.91 | 1.404× |
| 1/1/1 | 16384 | causal | 129 / 531.38 | 243 / 282.38 | 1.882× |
| 1/1/1 | 32768 | dense | 287.5 / 1911.71 | 372.5 / 1475.77 | 1.295× |
| 1/1/1 | 32768 | causal | 264.5 / 1038.78 | 487.5 / 563.97 | 1.842× |
| 32/1/1 | 256 | dense | 6 / 193.55 | 7 / 159.84 | 1.211× |
| 32/1/1 | 256 | causal | 7 / 78.92 | 10 / 55.59 | 1.420× |
| 32/1/1 | 2048 | dense | 39 / 1755.37 | 47 / 1456.70 | 1.205× |
| 32/1/1 | 2048 | causal | 35 / 980.40 | 51 / 675.58 | 1.451× |
| 32/1/1 | 8192 | dense | 524 / 2098.01 | 611 / 1798.82 | 1.166× |
| 32/1/1 | 8192 | causal | 334.5 / 1643.07 | 470 / 1170.80 | 1.403× |
| 32/1/1 | 16384 | dense | 1975 / 2226.80 | 2305.5 / 1907.85 | 1.167× |
| 32/1/1 | 16384 | causal | 1132 / 1942.68 | 1580 / 1391.70 | 1.396× |
| 32/1/1 | 32768 | dense | 7753.5 / 2268.94 | 9014.5 / 1951.53 | 1.163×
|
| 32/1/1 | 32768 | causal | 4220.5 / 2084.35 | 5883.5 / 1495.04 | 1.394×
|
</details>
<details><summary>BF16 Q/K=192, V=128: all 110 workloads</summary>
| B/H/Hkv | Sq = Sk | Mask | qr_tdm_sched (Introduced by this PR)<br>µs
/ TFLOP/s | qr_tdm<br>µs / TFLOP/s | qr_tdm_sched / qr_tdm |
| --- | ---: | --- | ---: | ---: | ---: |
| 1/8/8 | 256 | dense | 7 / 48.02 | 8 / 43.73 | 1.098× |
| 1/8/8 | 256 | causal | 9 / 19.61 | 11 / 16.02 | 1.224× |
| 1/8/8 | 2048 | dense | 23 / 925.25 | 35 / 617.25 | 1.499× |
| 1/8/8 | 2048 | causal | 26 / 416.77 | 44 / 242.61 | 1.718× |
| 1/8/8 | 8192 | dense | 182 / 1888.17 | 201.5 / 1707.81 | 1.106× |
| 1/8/8 | 8192 | causal | 133 / 1292.51 | 177 / 971.64 | 1.330× |
| 1/8/8 | 16384 | dense | 665.5 / 2064.98 | 745 / 1845.91 | 1.119× |
| 1/8/8 | 16384 | causal | 392 / 1752.87 | 452.5 / 1518.08 | 1.155× |
| 1/8/8 | 32768 | dense | 2587 / 2125.11 | 2912 / 1888.20 | 1.125× |
| 1/8/8 | 32768 | causal | 1408 / 1952.06 | 1712 / 1605.57 | 1.216× |
| 1/64/8 | 256 | dense | 7 / 369.48 | 8 / 333.00 | 1.110× |
| 1/64/8 | 256 | causal | 9 / 151.57 | 11 / 124.22 | 1.220× |
| 1/64/8 | 2048 | dense | 102 / 1678.51 | 112.5 / 1521.56 | 1.103× |
| 1/64/8 | 2048 | causal | 69 / 1250.12 | 77 / 1121.22 | 1.115× |
| 1/64/8 | 8192 | dense | 1322 / 2078.82 | 1481 / 1856.32 | 1.120× |
| 1/64/8 | 8192 | causal | 757 / 1816.07 | 891.5 / 1541.56 | 1.178× |
| 1/64/8 | 16384 | dense | 5171 / 2126.24 | 5815.5 / 1890.62 | 1.125× |
| 1/64/8 | 16384 | causal | 2859 / 1923.03 | 3381 / 1626.04 | 1.183× |
| 1/64/8 | 32768 | dense | 20964 / 2097.89 | 23605.5 / 1863.16 | 1.126×
|
| 1/64/8 | 32768 | causal | 11295.5 / 1946.87 | 13384.5 / 1642.97 |
1.185× |
| 1/64/64 | 256 | dense | 7 / 369.46 | 8 / 333.52 | 1.108× |
| 1/64/64 | 256 | causal | 9 / 152.22 | 11 / 124.11 | 1.226× |
| 1/64/64 | 2048 | dense | 105.5 / 1627.51 | 114.5 / 1499.87 | 1.085× |
| 1/64/64 | 2048 | causal | 70.5 / 1212.61 | 78.5 / 1094.97 | 1.107× |
| 1/64/64 | 8192 | dense | 1458 / 1885.32 | 1605.5 / 1712.16 | 1.101× |
| 1/64/64 | 8192 | causal | 798 / 1721.78 | 918 / 1496.87 | 1.150× |
| 1/64/64 | 16384 | dense | 5716.5 / 1923.35 | 6349 / 1731.78 | 1.111× |
| 1/64/64 | 16384 | causal | 3107.5 / 1769.16 | 3614 / 1521.28 | 1.163×
|
| 1/64/64 | 32768 | dense | 22649 / 1941.85 | 25245 / 1742.13 | 1.115× |
| 1/64/64 | 32768 | causal | 12231.5 / 1797.93 | 14312.5 / 1536.50 |
1.170× |
| 1/12/12 | 256 | dense | 7 / 68.65 | 8 / 63.35 | 1.084× |
| 1/12/12 | 256 | causal | 9 / 28.48 | 11 / 23.48 | 1.213× |
| 1/12/12 | 2048 | dense | 25.5 / 1273.39 | 37 / 877.98 | 1.450× |
| 1/12/12 | 2048 | causal | 27 / 594.66 | 46 / 353.45 | 1.682× |
| 1/12/12 | 8192 | dense | 275 / 1872.67 | 318 / 1619.62 | 1.156× |
| 1/12/12 | 8192 | causal | 161 / 1602.28 | 195 / 1324.34 | 1.210× |
| 1/12/12 | 16384 | dense | 1020 / 2021.97 | 1133.5 / 1819.31 | 1.111× |
| 1/12/12 | 16384 | causal | 567 / 1817.85 | 676 / 1525.85 | 1.191× |
| 1/12/12 | 32768 | dense | 4026.5 / 2047.83 | 4507.5 / 1829.59 | 1.119×
|
| 1/12/12 | 32768 | causal | 2155.5 / 1913.13 | 2552.5 / 1615.36 |
1.184× |
| 1/32/32 | 256 | dense | 7 / 187.38 | 8 / 168.97 | 1.109× |
| 1/32/32 | 256 | causal | 9 / 77.41 | 11 / 63.06 | 1.228× |
| 1/32/32 | 2048 | dense | 53 / 1621.17 | 58 / 1485.64 | 1.091× |
| 1/32/32 | 2048 | causal | 40 / 1071.55 | 51 / 844.96 | 1.268× |
| 1/32/32 | 8192 | dense | 687.5 / 1999.30 | 762 / 1803.20 | 1.109× |
| 1/32/32 | 8192 | causal | 391.5 / 1755.41 | 466.5 / 1473.47 | 1.191× |
| 1/32/32 | 16384 | dense | 2731 / 2013.10 | 3041 / 1807.96 | 1.113× |
| 1/32/32 | 16384 | causal | 1467.5 / 1873.78 | 1720.5 / 1597.70 |
1.173× |
| 1/32/32 | 32768 | dense | 10799.5 / 2036.24 | 12074 / 1821.28 | 1.118×
|
| 1/32/32 | 32768 | causal | 5818.5 / 1889.83 | 6846.5 / 1606.04 |
1.177× |
| 32/8/8 | 256 | dense | 12 / 910.68 | 12 / 880.00 | 1.035× |
| 32/8/8 | 256 | causal | 13 / 412.21 | 15 / 369.06 | 1.117× |
| 32/8/8 | 2048 | dense | 395.5 / 1737.20 | 434.5 / 1580.08 | 1.099× |
| 32/8/8 | 2048 | causal | 273 / 1260.08 | 313 / 1098.72 | 1.147× |
| 32/8/8 | 8192 | dense | 5265 / 2088.28 | 5900.5 / 1863.57 | 1.121× |
| 32/8/8 | 8192 | causal | 3038.5 / 1809.44 | 3584 / 1534.17 | 1.179× |
| 32/8/8 | 16384 | dense | 20551.5 / 2139.96 | 23140 / 1900.63 | 1.126×
|
| 32/8/8 | 16384 | causal | 11441.5 / 1922.14 | 13574 / 1620.11 | 1.186×
|
| 32/8/8 | 32768 | dense | 82542.5 / 2131.29 | 92859.5 / 1894.49 |
1.125× |
| 32/8/8 | 32768 | causal | 44795 / 1963.68 | 53240 / 1652.22 | 1.189× |
| 1/4/4 | 256 | dense | 7 / 23.82 | 8 / 21.95 | 1.085× |
| 1/4/4 | 256 | causal | 9 / 9.81 | 10.5 / 8.02 | 1.224× |
| 1/4/4 | 2048 | dense | 22 / 487.02 | 34 / 320.13 | 1.521× |
| 1/4/4 | 2048 | causal | 25 / 213.91 | 44 / 123.07 | 1.738× |
| 1/4/4 | 8192 | dense | 96 / 1791.19 | 136 / 1266.14 | 1.415× |
| 1/4/4 | 8192 | causal | 90 / 958.32 | 163 / 525.59 | 1.823× |
| 1/4/4 | 16384 | dense | 346 / 1984.35 | 386.5 / 1779.49 | 1.115× |
| 1/4/4 | 16384 | causal | 254.5 / 1350.47 | 342.5 / 1003.10 | 1.346× |
| 1/4/4 | 32768 | dense | 1310 / 2098.09 | 1470 / 1869.88 | 1.122× |
| 1/4/4 | 32768 | causal | 754 / 1823.20 | 878 / 1564.96 | 1.165× |
| 1/3/3 | 256 | dense | 7 / 17.62 | 8 / 16.11 | 1.094× |
| 1/3/3 | 256 | causal | 9 / 7.25 | 11 / 5.95 | 1.219× |
| 1/3/3 | 2048 | dense | 22 / 366.83 | 33 / 240.81 | 1.523× |
| 1/3/3 | 2048 | causal | 25 / 160.38 | 43 / 92.67 | 1.731× |
| 1/3/3 | 8192 | dense | 86 / 1496.30 | 130 / 988.91 | 1.513× |
| 1/3/3 | 8192 | causal | 87 / 739.11 | 161 / 400.90 | 1.844× |
| 1/3/3 | 16384 | dense | 307.5 / 1674.18 | 347.5 / 1482.52 | 1.129× |
| 1/3/3 | 16384 | causal | 238 / 1082.90 | 331 / 778.76 | 1.391× |
| 1/3/3 | 32768 | dense | 1007 / 2047.31 | 1181 / 1745.87 | 1.173× |
| 1/3/3 | 32768 | causal | 567.5 / 1816.11 | 715 / 1441.41 | 1.260× |
| 1/2/2 | 256 | dense | 7 / 12.18 | 8 / 11.12 | 1.094× |
| 1/2/2 | 256 | causal | 9 / 4.94 | 10 / 4.03 | 1.225× |
| 1/2/2 | 2048 | dense | 21 / 252.14 | 33 / 163.94 | 1.538× |
| 1/2/2 | 2048 | causal | 25 / 108.95 | 43 / 62.28 | 1.749× |
| 1/2/2 | 8192 | dense | 78 / 1099.51 | 125 / 684.67 | 1.606× |
| 1/2/2 | 8192 | causal | 85 / 504.75 | 159 / 269.97 | 1.870× |
| 1/2/2 | 16384 | dense | 179.5 / 1914.47 | 260 / 1321.57 | 1.449× |
| 1/2/2 | 16384 | causal | 171.5 / 1001.16 | 317.5 / 541.13 | 1.850× |
| 1/2/2 | 32768 | dense | 667.5 / 2059.55 | 744.5 / 1846.57 | 1.115× |
| 1/2/2 | 32768 | causal | 492.5 / 1395.38 | 670.5 / 1024.41 | 1.362× |
| 1/1/1 | 256 | dense | 7 / 6.20 | 7.5 / 5.61 | 1.105× |
| 1/1/1 | 256 | causal | 8 / 2.50 | 10 / 2.03 | 1.232× |
| 1/1/1 | 2048 | dense | 21 / 129.12 | 32 / 83.85 | 1.540× |
| 1/1/1 | 2048 | causal | 24 / 55.22 | 43 / 31.50 | 1.753× |
| 1/1/1 | 8192 | dense | 73 / 584.67 | 121 / 355.15 | 1.646× |
| 1/1/1 | 8192 | causal | 83 / 259.93 | 155 / 138.75 | 1.873× |
| 1/1/1 | 16384 | dense | 149.5 / 1148.66 | 243.5 / 705.19 | 1.629× |
| 1/1/1 | 16384 | causal | 163 / 526.39 | 307.5 / 279.66 | 1.882× |
| 1/1/1 | 32768 | dense | 341.5 / 2012.12 | 507.5 / 1354.29 | 1.486× |
| 1/1/1 | 32768 | causal | 331.5 / 1036.67 | 622.5 / 551.94 | 1.878× |
| 32/1/1 | 256 | dense | 7 / 191.09 | 8 / 175.12 | 1.091× |
| 32/1/1 | 256 | causal | 9 / 78.16 | 11 / 63.94 | 1.222× |
| 32/1/1 | 2048 | dense | 53 / 1614.35 | 58.5 / 1474.45 | 1.095× |
| 32/1/1 | 2048 | causal | 50.5 / 845.21 | 58.5 / 734.68 | 1.150× |
| 32/1/1 | 8192 | dense | 681 / 2018.51 | 760.5 / 1808.03 | 1.116× |
| 32/1/1 | 8192 | causal | 459 / 1496.40 | 565.5 / 1215.26 | 1.231× |
| 32/1/1 | 16384 | dense | 2588.5 / 2124.03 | 2905 / 1892.45 | 1.122× |
| 32/1/1 | 16384 | causal | 1559.5 / 1762.61 | 1917 / 1434.18 | 1.229× |
| 32/1/1 | 32768 | dense | 10155 / 2165.40 | 11427 / 1924.38 | 1.125× |
| 32/1/1 | 32768 | causal | 5901.5 / 1863.09 | 7211 / 1524.86 | 1.222× |
</details>
## Clang 24 performance disclosure
A separate matched Clang 24 study found slower `qr_tdm_sched`
performance in some D64/D128 causal configurations. That study used a
different timing harness, so its quantitative results are excluded from
this description to keep every reported performance number sourced from
`tile_example_fmha_fwd`. A matched Clang 24 example run is still needed
before assigning a case count or slowdown magnitude to this benchmark.
The dispatch intentionally keeps `qr_tdm_sched` priority for its
supported configurations; the Clang 23 results above should not be
extrapolated to Clang 24.
## Toolchain and build environment
| Component | Clang 23 example comparison | Clang 24 matched validation
/ separate diagnostic |
| --- | --- | --- |
| Device compiler | AMD Clang `23.0.0git`; LLVM
`2a67d831a09cc538194bd595279ee91258570c10`; `clang++` SHA256
`b552c9fa9e9158688cd4944a4aa899dfedc7426b700097ef9338f4476dc6f6f6` | AMD
Clang `24.0.0git`; LLVM `4f43f4746ede4cc49ab4129f473dbd54e0d9db4f`;
`clang++` SHA256
`1cfe274fcb8f92e5ccb423173e207dda6036e4a9efcfcdf29193049849d33979` |
| `hipcc -v` / HIP | Container `hipcc` reports Clang 23 and HIP 7.15.0;
`hipconfig --version` reports `7.15.0-0000000` | The pinned Clang 24
build invokes `clang++` directly, with `--hip-path=/opt/rocm`; container
HIP headers/runtime are 7.15.0. Container `hipcc` was not used for this
build. |
| LLVM linker and device libraries | AMD LLD 23.0.0;
`/opt/rocm/amdgcn/bitcode` | AMD LLD 24.0.0; pinned Clang 24
device-library bundle |
| Target and common build options | `gfx1250`, `-O3`, `-DNDEBUG`,
`CK_TILE_FMHA_FWD_FAST_EXP2=1`; scheduled objects use `-mllvm
-amdgpu-expert-scheduling-mode` | Same target and relevant
optimization/math settings; scheduled objects use expert scheduling |
| Source | This PR at `4688449fa9b` vs upstream `qr_tdm` at
`9e3337709130` (includes PR #12347) | Matched final-candidate source and
upstream `qr_tdm`; numerical results omitted here because they were not
produced by `tile_example_fmha_fwd` |
## Validation
- Current-candidate FP16/BF16 generated-dispatch correctness passed
924/924 configurations across D64/D128/D192; sequences through 4096
checked all rows, while longer sequences used deterministic sampled rows
plus whole-output finite and LSE scans.
- The generated-dispatch GTests passed 480/480; finite virtual-sink
fallback passed 96/96 and selected `qr_tdm`; direct layout, softmax,
fragment, and schedule GTests passed 79/79. Python codegen/dispatcher
tests passed 40/40.
- All 72 generated scheduled objects compiled under the pinned Clang 24
toolchain. Local benchmark scripts and codegen filters are excluded from
the PR.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JIRA ID: AICK-2310
Add scheduled TDM FMHA for FP16 and BF16
Summary
qr_tdm_schedforward FMHA pipeline for FP16/BF16 on gfx125: Q/K–V dimensions 64/64, 128/128, and 192/128. The implementation uses geometry-derived QK/PV scheduling, packed LDS, and scheduled WMMA/LDS-read placement.qr_tdm_schedfirst for its supported row-major V cases withmax_seqlen_q >= 128,sink_ptr == nullptr, and no bias, dropout, soft cap, or quantization. Dense and supported causal masks are included; LSE is optional.qr_tdmcandidate when a scheduled trait is ineligible, including short Q or a non-null virtual-sink pointer. Shapes or features outsideqr_tdmcontinue through the existingqrfallback.test/ck_tile/fmha.BF16 performance
All performance numbers in this PR description are taken from the unmodified
tile_example_fmha_fwdoutput. The two binaries use the same example main, inputs, GPU timer, and command-line options; one linksqr_tdm_schedfrom this PR and the other links upstreamqr_tdm. Kernel names were printed and checked on every run.Measured on gfx1250 B0 stepping (256 CUs). SCLK was neither locked nor logged by the benchmark, so no fixed SCLK value is claimed. The historical 110 workloads per dimension use BSHD layout,
Sq = Skin {256, 2048, 8192, 16384, 32768}, dense or causal masking, LSE, no sink, and the same scale and input seed for both kernels. The example uses GPU-event timing, five warmups, each workload’s historical repeat count, and two adjacent alternating rounds. For each kernel, the table reports the median of the two values printed by the example. Speedup divides the two printed TFLOP/s medians; above 1× favorsqr_tdm_sched. The example prints latency to 0.001 ms and throughput to 0.01 TFLOP/s.The full per-workload results follow. Each kernel cell shows execution time in µs / effective TFLOP/s.
B/H/Hkvdenotes batch size, query heads, and key/value heads. Causal TFLOP/s uses triangular attention work.BF16 Q/K=64, V=64: all 110 workloads
µs / TFLOP/s
µs / TFLOP/s
BF16 Q/K=128, V=128: all 110 workloads
µs / TFLOP/s
µs / TFLOP/s
BF16 Q/K=192, V=128: all 110 workloads
µs / TFLOP/s
µs / TFLOP/s
Clang 24 performance disclosure
A separate matched Clang 24 study found slower
qr_tdm_schedperformance in some D64/D128 causal configurations. That study used a different timing harness, so its quantitative results are excluded from this description to keep every reported performance number sourced fromtile_example_fmha_fwd. A matched Clang 24 example run is still needed before assigning a case count or slowdown magnitude to this benchmark. The dispatch intentionally keepsqr_tdm_schedpriority for its supported configurations; the Clang 23 results above should not be extrapolated to Clang 24.Toolchain and build environment
23.0.0git; LLVM2a67d831a09cc538194bd595279ee91258570c10;clang++SHA256b552c9fa9e9158688cd4944a4aa899dfedc7426b700097ef9338f4476dc6f6f624.0.0git; LLVM4f43f4746ede4cc49ab4129f473dbd54e0d9db4f;clang++SHA2561cfe274fcb8f92e5ccb423173e207dda6036e4a9efcfcdf29193049849d33979hipcc -v/ HIPhipccreports Clang 23 and HIP 7.15.0;hipconfig --versionreports7.15.0-0000000clang++directly, with--hip-path=/opt/rocm; container HIP headers/runtime are 7.15.0. Containerhipccwas not used for this build./opt/rocm/amdgcn/bitcodegfx1250,-O3,-DNDEBUG,CK_TILE_FMHA_FWD_FAST_EXP2=1; scheduled objects use-mllvm -amdgpu-expert-scheduling-mode4688449fa9bvs upstreamqr_tdmat9e3337709130(includes PR #12347)qr_tdm; numerical results omitted here because they were not produced bytile_example_fmha_fwdValidation
qr_tdm; direct layout, softmax, fragment, and schedule GTests passed 79/79. Python codegen/dispatcher tests passed 40/40.