Repository navigation
perf(ck_tile/fmha): [gfx1250 opt 5] optimize qr_tdm_pingpong, extend it to every head dim and simplify the dispatch - #15122
Draft
Yanxing-Shi wants to merge 2 commits into
Draft
Yanxing-Shi wants to merge 2 commits into
Yanxing-Shi wants to merge 2 commits into
Conversation
…line and extend it to d64
- qr_tdm_pingpong serves hdim 64 as well as 128 (fp16/bf16, gfx1250).
- Second V buffer (is_qr_tdm_pingpong_double_v_v): V_{i+1} is issued a whole
trip ahead of its PV gemm on the 128-row tiles and the d64 masked 64-row
tile; Q moves onto the last V buffer.
- 128-row tiles for long keys, halving the K/V LDS reads per query row.
- SchedReadsAheadOfWmma: the B-operand LDS reads of both gemms run ahead of
their WMMAs instead of a dscnt-0 wait before each; three fragments ahead on
the d128 and masked 128-row tiles, one elsewhere.
- The next K tile is issued at the loop head together with V; the last trip
issues no K/V.
- Dispatch: d64/d128 queries in the pingpong range go to qr_tdm_pingpong,
with its 128-row tile from 128 query rows for long keys or dense 512-row
queries; d128 dense and d64 causal long-key grids stay on qr_tdm_sched.
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. |
…m, LDS read schedule and simpler dispatch qr_tdm_pingpong (fp16/bf16, gfx1250): - Head dims 32, 96, 160, 192/128, 256 and 512 besides 64 and 128. Non-power-of-2 rows take padded LDS rows; d32 stages Q/K at 32 columns instead of padding them to 64. - LDS read schedule per head dim (QrTdmPingPongLdsReadSchedule, codegen GFX125_QR_TDM_PINGPONG_LDS_READ_SCHEDULE): prefetching the next kK0/kK1 slice of K/V while the current slice's WMMAs run, and reading V's first slice before the softmax, each on the tiles where its VGPRs leave the occupancy alone (_sp<rows> / _vbs<rows> instance suffixes where they take effect). - d96/d160/d192: the 128-row tile from 128 query rows at any key length. - Logits soft cap, elementwise bias and ALiBi, dropout, sink tokens and column-major V. qr_tdm and qr_tdm_sched take the same features (soft cap, bias, dropout, sink tokens; column-major V for qr_tdm). The three TDM pipelines share one soft-cap helper and refuse a CK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT other than tanh. Dispatch on gfx1250: qr_tdm_pingpong serves every fp16/bf16 query of its head dims at any query and key length, except d128 with at least 128 query rows and fewer than two 128-row query blocks per CU, which goes to qr_tdm_sched. qr_tdm_sched and the fp16/bf16 qr_tdm instances are generated for d128 only. qr_tdm_pingpong's unpadded-V check takes the exact V row (bn1). The generated API returns from flat ifs per candidate instead of one else-if chain, which overflowed clang's parser stack on a full build. LdsShuffle2DEpilogue keeps whole, lane-contiguous vectors when a row is longer than a wavefront of them (MakeODramTileDistribution). tile_example_fmha_fwd: on gfx125, Pack-GQA is skipped for 2..16 fp16/bf16 query rows (hdim_v >= 64) that it would fold into more than 16 rows on a grid under two (batch, head) pairs per CU. README: qr_tdm_pingpong section; qr_tdm_sched and Pack-GQA notes updated.
Yanxing-Shi
changed the title
This branch has not been deployed
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.
What
Optimizes
qr_tdm_pingpong(#13471, #13484) on gfx1250 fp16/bf16 and makes it the default forward pipeline there:GFX125_QR_TDM_PINGPONG_LDS_READ_SCHEDULE).qr_tdmandqr_tdm_sched.qr_tdm_pingpongtakes every query of these head dims; only d128 on a small grid (fewer than two 128-row query blocks per CU) goes toqr_tdm_sched.qr_tdm_schedand fp16/bf16qr_tdmare therefore generated for d128 only (gfx1250 fwd instances 77005 -> 49357).ifs (theelse ifchain overflowed clang's stack on a full build);LdsShuffle2DEpiloguestores whole vectors for 512-column rows;tile_example_fmha_fwdskips Pack-GQA where it would cost parallelism on gfx125 (short queries on small grids).qr_tdm_pingpongsection.Performance
gfx1250, bf16, batch mode, bshd, TFLOPS, median of 3 runs. "causal" is
-mask=b. The three pipeline columns force one pipeline; "develop dispatch" and "this PR dispatch" are the default builds ofdevelop(c6542db) and this PR. This PR vs develop: geomean +23.8% (min +1.7%, max +55.9%) over the rows above 5 TFLOPS.Tests
tile_example_fmha_fwd -v=1(default build, gfx1250): all 8 head dims x batch/group x masks x bias x lse x sink, plus soft cap, dropout and column-major V: 1584 passed, 0 failed.test_fmha_fwd_codegen.py(45 tests) andtest_lds_shuffle_2d_epilogue.cpppass.