Merged
Conversation
5742bb0 to
632f364
Compare
f40daa4 to
94a0c5c
Compare
This was referenced Jan 5, 2026
drisspg
added a commit
to drisspg/flash-attention
that referenced
this pull request
Jan 5, 2026
Adds score_mod and mask_mod support to SM90 backward pass: - score_mod, score_mod_bwd, mask_mod, has_aux_tensors parameters - apply_score_mod() and apply_score_mod_bwd() methods - fastdiv_mods and aux_tensors plumbing through kernel/mma - mask_mod application in mask_fn for both block-sparse and dense paths - Score modification in mma_one_m_block before softmax stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
commented
Jan 5, 2026
drisspg
added a commit
to drisspg/flash-attention
that referenced
this pull request
Jan 5, 2026
Adds score_mod and mask_mod support to SM90 backward pass: - score_mod, score_mod_bwd, mask_mod, has_aux_tensors parameters - apply_score_mod() and apply_score_mod_bwd() methods - fastdiv_mods and aux_tensors plumbing through kernel/mma - mask_mod application in mask_fn for both block-sparse and dense paths - Score modification in mma_one_m_block before softmax stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8
drisspg
added a commit
to drisspg/flash-attention
that referenced
this pull request
Jan 5, 2026
Adds score_mod and mask_mod support to SM90 backward pass: - score_mod, score_mod_bwd, mask_mod, has_aux_tensors parameters - apply_score_mod() and apply_score_mod_bwd() methods - fastdiv_mods and aux_tensors plumbing through kernel/mma - mask_mod application in mask_fn for both block-sparse and dense paths - Score modification in mma_one_m_block before softmax stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8
94a0c5c to
759373c
Compare
stack-info: PR: #2137, branch: drisspg/stack/8
632f364 to
246cde5
Compare
759373c to
71f18b1
Compare
drisspg
added a commit
that referenced
this pull request
Jan 5, 2026
stack-info: PR: #2137, branch: drisspg/stack/8
71f18b1 to
2361da4
Compare
drisspg
added a commit
that referenced
this pull request
Jan 6, 2026
stack-info: PR: #2137, branch: drisspg/stack/8
8951af0 to
9f66e6c
Compare
9f66e6c to
a7bea6d
Compare
a7bea6d to
c555946
Compare
c555946 to
821be8b
Compare
This was referenced Jan 7, 2026
v0i0
approved these changes
Jan 9, 2026
| assert cu_seqlens_q is None and cu_seqlens_k is None, ( | ||
| "varlen + score_mod not supported in bwd yet" | ||
| ) | ||
| assert compute_capability in [10, 11], "score_mod in bwd only supported on SM100/SM110 for now" |
Collaborator
|
why is triton faster for small sizes? |
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.
Stacked PRs:
score-mod backward SM90
stack-pr keeps deleting my damn summary, sigh
FWD:

BWD:

insert pic, of ci being all green :) which I no longer have in my clipboard