[gfx1151] flash_attn_triton_amd: enable in-thread transpose - #1
Closed
mgehre-amd wants to merge 1 commit into
Closed
mgehre-amd wants to merge 1 commit into
mgehre-amd wants to merge 1 commit into
Conversation
Triton only enables its AMD in-thread-transpose pass for gfx1151 by default on its main branch (the upcoming 3.8); the Triton version in use here still gates it to gfx942 (use_in_thread_transpose defaults to None -> arch == "gfx942"). On gfx1151 (RDNA3.5) the prefill FMHA kernels are LDS-bound and benefit from it, so turn it on at kernel import time, gated on arch and on the user not having set TRITON_HIP_USE_IN_THREAD_TRANSPOSE explicitly. Measured on gfx1151 with triton 3.6.0, dao_ai fwd_varlen, fp16, causal=False, b1 hq16 hk16 (min ms over a warm-clock do_bench window): shape (sq=sk, d) ITT off ITT on speedup 1024, d72 0.447 0.443 +0.9% 2048, d72 1.299 1.198 +7.8% 3200, d72 2.801 2.479 +11.5% 4096, d72 4.545 4.045 +11.0% 3200, d64 1.660 1.592 +4.1% 3200, d128 3.314 3.096 +6.6% Changes: - Win is largest for the ViT-style d=72, long-sequence shapes this kernel targets; no shape regressed. - Respects an explicit TRITON_HIP_USE_IN_THREAD_TRANSPOSE override so the pass can still be forced on/off for debugging.
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
1 participant
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.
Stacked on top of ROCm#3419.
Triton 3.7 enables its AMD in-thread-transpose pass only on gfx942 by default
(
knobs.amd.use_in_thread_transposedefaults toNone→arch == "gfx942").Triton main will enable it also on gfx1151 by default. For older triton versions,
this PR enabled it.
Benchmarks
gfx1151, triton 3.6.0,
dao_aifwd_varlen, fp16, causal=False, b1 hq16 hk16.Min ms over a warm-clock
do_benchwindow (min is robust to the iGPU's DVFS):Win is largest for the ViT-style d=72 long-sequence shapes this kernel targets;
no shape regressed. An explicit
TRITON_HIP_USE_IN_THREAD_TRANSPOSEoverride isstill respected so the pass can be forced on/off for debugging.