[Pallas] Attention perf: further reduce spillage from pre-loading Q, by loading Q in-loop and not pipelining it#2397
Draft
AmesingFlank wants to merge 1 commit into
Draft
[Pallas] Attention perf: further reduce spillage from pre-loading Q, by loading Q in-loop and not pipelining it#2397AmesingFlank wants to merge 1 commit into
AmesingFlank wants to merge 1 commit into
Conversation
…by loading Q in-loop and not pipelining it stack-info: PR: #2397, branch: AmesingFlank/stack/51
AmesingFlank
force-pushed
the
AmesingFlank/stack/51
branch
from
caff8b8 to
18e5642
Compare
meta-cla
Bot
added
the
CLA Signed
Contributor
Author
|
CI failure, investigating |
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.
Optimization found by claude, by comparing the current Helion kernel with this reference impl
Similar to #2373, but takes it one step further
Previously, the Helion attention kernel pre-loads
qfromq_view[tile_b, tile_m, :]before the device loop. This requires registers which persists across the entire loop, which results in more spillage that damages performance.This PR modifies this by loading
q_view[tile_b, tile_m, :]within the loop body (even though neithertile_bnortile_mis part of the loop iteration). A small compiler change is needed here: when we decide whether or not to pipeline a tensor within a device loop, if the accessed block ids do not match any of the loop-itered block ids, then no need to pipeline it.On TPU, this change improves TFLOPs from 653 TFLOPs to 660 TFLOPs.