Add tests for flash_attention_backward perf model (PR #512)#580
Open
Add tests for flash_attention_backward perf model (PR #512)#580
Conversation
Agent-Logs-Url: https://github.com/AMD-AGI/TraceLens/sessions/57c6aab9-d187-4fb9-90bb-cacfb503a156 Co-authored-by: gabeweisz <162640284+gabeweisz@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Add tests related to PR #512
Add tests for flash_attention_backward perf model (PR #512)
Apr 3, 2026
gabeweisz
approved these changes
Apr 3, 2026
Collaborator
Collaborator
|
@spandoesai #582 has been merged. |
Collaborator
|
Thanks, @gabeweisz could you rebase and try the tests again? |
Collaborator
|
Merged main. This particular PR does not use reference reports, it just adds additional tests. All tests pass after merging main |
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.
PR #512 overhauled
flash_attention_backward— reparenting it fromflash_attentiontoSDPA, adding a properget_param_details, and registeringflash_attn::_flash_attn_backwardin the op mapping — but shipped without tests.New test file:
tests/test_flash_attention_backward.pyflash_attn::_flash_attn_backwardresolves toflash_attention_backwardand categorizes asSDPA_bwdSDPA, notflash_attentionget_param_detailsargument indexing —doutat index 0; q/k/v correctly read from indices 1/2/3get_param_detailsconcrete inputs —dropoutfromconcrete[8],causalfromconcrete[10]; both default safely whenConcrete Inputsis absent__init__attribute —self.d_his set toself.d_h_qk(required byget_simulation_time_bwd_func)flops()— delegates toflops_bwd(); backward > forward for identical shapes; non-causal > causalbytes()— zero-arg form resolvesbytes_per_elementfromparam_detailsdtype; explicitbpedelegates tobytes_bwd()faflag —type(self).__name__check inget_simulation_time_bwdcorrectly identifiesflash_attention_backwardas a flash-attention op