[Cute] Block sparse support Sm100 #1985
Open
+785
−163
Conversation
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
drisspg
commented
Nov 5, 2025
drisspg
commented
Nov 5, 2025
drisspg
changed the title
drisspg
force-pushed
the
sm100-block-sparsity
branch
6 times, most recently
from
06588af to
547cf51
Compare
drisspg
force-pushed
the
sm100-block-sparsity
branch
from
547cf51 to
66cbbdb
Compare
drisspg
commented
Nov 12, 2025
|
|
||
|
|
||
| @cute.jit | ||
| def handle_block_sparse_empty_tile_correction_sm100( |
Contributor
Author
There was a problem hiding this comment.
had to dupe alot here but still think its better, having very large IF Else indents makes it harder to rebase / iterate on all the constexpr tree
drisspg
force-pushed
the
sm100-block-sparsity
branch
from
66cbbdb to
d1ece5e
Compare
tridao
approved these changes
Nov 12, 2025
drisspg
force-pushed
the
sm100-block-sparsity
branch
from
d1ece5e to
1db7911
Compare
- Implement block-sparse attention in flash_fwd_sm100.py - Update interface.py to handle SM100 block size calculations (2x multiplier for m_block_size since 1 CTA handles 2*tile_m rows) - Add mask_mod parameter support in mask.py for block-sparse masking - Add SM100 test fixtures and tile size handling in test_mask_mod.py This enables block-sparsity on SM 10.0 architecture, including mask_mod support and proper block size accounting.
drisspg
force-pushed
the
sm100-block-sparsity
branch
from
1db7911 to
b01ba0c
Compare
Contributor
Author
|
@tridao Okay, finally rebased, perf looks good |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
Summary
(2x multiplier for m_block_size since 1 CTA handles 2*tile_m rows)
Fast follow
Do the aux_tensor fastdivmod wrapping to avoid OOB reads
Also we should land:
#1984
Before and rebase so its easier to review
Perf
Alot of perf wins (not universal for document mask ) but the delta from sol is much higher than what was found on hopper impl
Not autotuning the flex blocksparse impl gives this:
And autotuning the triton impl:
Possible problems
Looking at the Pm samples we can see a long tail:
For causal_mask with the default StaticPersistentSchedule. (We need to build a generic version of this) but we already have a better schedule for causal. If hard code the LPT schedule ![Uploading Screenshot 2025-11-04 at 5.32.58 PM.png…]()we go from :
to:
Tests