score-mod backward SM90 #2137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

drisspg merged 1 commit into main from drisspg/stack/8

Jan 10, 2026

+293 −31

Collaborator

drisspg commented Jan 5, 2026 •

edited

Loading

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

drisspg force-pushed the drisspg/stack/7 branch from 5742bb0 to 632f364 Compare

January 5, 2026 01:02

drisspg force-pushed the drisspg/stack/8 branch from f40daa4 to 94a0c5c Compare

January 5, 2026 01:02

This was referenced Jan 5, 2026

Misc tests that should be xfailed for now #2127

Merged

block-sparse backward SM90 #2136

Merged

score-mod backward SM100 cleanup #2130

Closed

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

94a0c5c

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

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Outdated Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Outdated Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Outdated Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Outdated Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Outdated Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/flash_bwd_sm90.py Show resolved Hide resolved

drisspg commented

View reviewed changes

flash_attn/cute/interface.py Outdated Show resolved Hide resolved

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

5b689b5

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


          score-mod backward SM90

759373c

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 marked this pull request as draft

January 5, 2026 05:42

drisspg changed the base branch from drisspg/stack/7 to main

January 5, 2026 05:42

drisspg force-pushed the drisspg/stack/8 branch from 94a0c5c to 759373c Compare

January 5, 2026 05:42

drisspg changed the base branch from main to drisspg/stack/7

January 5, 2026 05:42

drisspg marked this pull request as ready for review

January 5, 2026 05:43


          score-mod backward SM90

c8aa2cb

stack-info: PR: #2137, branch: drisspg/stack/8

drisspg force-pushed the drisspg/stack/7 branch from 632f364 to 246cde5 Compare

January 5, 2026 17:01

drisspg force-pushed the drisspg/stack/8 branch from 759373c to 71f18b1 Compare

January 5, 2026 17:01

drisspg marked this pull request as draft

January 5, 2026 17:04

drisspg changed the base branch from drisspg/stack/7 to main

January 5, 2026 17:04

drisspg added a commit that referenced this pull request


          score-mod backward SM90

2361da4

stack-info: PR: #2137, branch: drisspg/stack/8

drisspg force-pushed the drisspg/stack/8 branch from 71f18b1 to 2361da4 Compare

January 5, 2026 17:04

drisspg changed the base branch from main to drisspg/stack/7

January 5, 2026 17:04

drisspg marked this pull request as ready for review

January 5, 2026 17:04

drisspg marked this pull request as draft

January 5, 2026 19:08

drisspg changed the base branch from drisspg/stack/7 to main

January 9, 2026 23:19

drisspg force-pushed the drisspg/stack/8 branch from 821be8b to fc27f96 Compare

January 9, 2026 23:19

drisspg changed the base branch from main to drisspg/stack/7

January 9, 2026 23:19

drisspg marked this pull request as ready for review

January 9, 2026 23:20

drisspg marked this pull request as draft

January 9, 2026 23:24

drisspg changed the base branch from drisspg/stack/7 to main

January 9, 2026 23:24

drisspg force-pushed the drisspg/stack/8 branch from fc27f96 to 3374930 Compare

January 9, 2026 23:24

drisspg changed the base branch from main to drisspg/stack/7

January 9, 2026 23:24

drisspg marked this pull request as ready for review

January 9, 2026 23:24

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8

drisspg marked this pull request as draft

January 9, 2026 23:39

drisspg changed the base branch from drisspg/stack/7 to main

January 9, 2026 23:39

drisspg changed the base branch from main to drisspg/stack/7

January 9, 2026 23:39

drisspg marked this pull request as ready for review

January 9, 2026 23:39

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

1b61260

stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8

drisspg marked this pull request as draft

January 10, 2026 00:29

drisspg changed the base branch from drisspg/stack/7 to main

January 10, 2026 00:29

drisspg force-pushed the drisspg/stack/8 branch from 3374930 to 1b61260 Compare

January 10, 2026 00:29

drisspg changed the base branch from main to drisspg/stack/7

January 10, 2026 00:29

drisspg marked this pull request as ready for review

January 10, 2026 00:29

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

1f46876

stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8

drisspg marked this pull request as draft

January 10, 2026 00:50

drisspg changed the base branch from drisspg/stack/7 to main

January 10, 2026 00:50

drisspg force-pushed the drisspg/stack/8 branch from 1b61260 to 1f46876 Compare

January 10, 2026 00:50

drisspg changed the base branch from main to drisspg/stack/7

January 10, 2026 00:50

drisspg marked this pull request as ready for review

January 10, 2026 00:50

drisspg force-pushed the drisspg/stack/8 branch from 1f46876 to c8aa2cb Compare

January 10, 2026 00:54

drisspg changed the base branch from drisspg/stack/7 to main

January 10, 2026 00:54

drisspg added a commit to drisspg/flash-attention that referenced this pull request


          score-mod backward SM90

fa53486

stack-info: PR: Dao-AILab#2137, branch: drisspg/stack/8

drisspg merged commit 844b10f into main

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet