Repository navigation
Conversation
✅ 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. |
5721e25 to
292ea8f
Compare
f7d4e89 to
9e21b59
Compare
6a24b2b to
3b81076
Compare
|
Location: Non-blocking — design question, not a requested change. This assembles four independently-maintained compatibility rules ( Given this pattern already needed containment once ( |
3b81076 to
a596eab
Compare
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.
a596eab to
5c7973a
Compare
|
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 ( |
…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
5c7973a to
8df9a29
Compare
…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.
8df9a29 to
86d40b3
Compare
… -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.
86d40b3 to
f4292af
Compare
…ers/yanxishi/ck/gfx1250-fmha-fwd-short-query-ping-pong
TheRock Submodule Bump ActivityNewest first
|
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://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
A 16-row short-query tile for
qr_tdm_pingpong(#13471), taken whenmax_seqlen_q <= 16on the gridsqr_tdm_pingpongtakes (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 onqr_tdm;qr_tdmstays as ondevelop.With the warps on N:
lis 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 astatic_assertat the call.The tile is owned by
qr_tdm_pingpong(F_pipeline_tag) and listed before the 64-row tile, withmax_seqlen_q <= 16added to the shared grid condition; on smaller grids short queries stay onqr_tdmas ondevelop. 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 ondevelop9951b48.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 gridmax_seqlen_q <= 16dispatches 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 onqr_tdm.Base
Built on #13471 (on
developat 9951b48); it will be rebased ontodeveloponce #13471 lands.Series
The gfx1250 FMHA forward optimization series; every PR is based on
developand benchmarked against it (bf16, d=128). #12426 and #13468 were folded into #13471, #13469 into this PR.qr_tdm_pingpongpipeline, coalesced O store and its dispatchqr_tdm_pingpong(this PR)JIRA ID : AICK-2262