Use is_family_of for sm_90 and sm_103 arch checks#2589
Merged
Johnsonms merged 2 commits intoMay 26, 2026
Conversation
Follow-up to Dao-AILab#2572 — apply the same is_family_of pattern to the two remaining range-style arch checks for consistency: - flash_fwd_sm90.py:69 (SM 9.x assert) - flash_fwd_sm100.py:195 (is_sm103 flag) Same semantic narrowing as Dao-AILab#2572: bare-base SMs (sm_90, sm_103) are excluded. These kernels rely on wgmma / UMMA / 2CTA paths that require the a/f PTX variant anyway, so bare-base targets could not compile.
Collaborator
Author
|
Cc, @ocss884 @lingolin128 helping review |
ocss884
approved these changes
May 25, 2026
| # Does S1 need to wait for S0 to finish | ||
| # self.s0_s1_barrier = self.head_dim_padded in [64, 96] and (not self.is_causal and not self.is_local) | ||
| is_sm103 = self.arch >= Arch.sm_103 and self.arch <= Arch.sm_103f | ||
| is_sm103 = self.arch.is_family_of(Arch.sm_103f) |
Contributor
There was a problem hiding this comment.
is_family_of will test (self.major == arch.major and self.minor >= arch.minor), so it is correct for now. If there is another sm10.x hardware with x > 3 in the future then it would be wrong.
Collaborator
Author
There was a problem hiding this comment.
@ocss884 @lingolin128 Good catch — added an inline comment in 5833281 to make the forward-inclusive semantics explicit. The flag gates ex2 emulation, and since sm_103 (B300) brought fast hardware ex2, any future sm_10x Blackwellvariant should inherit the same behavior. So is_family_of's minor >= 3 behavior is actually what we want here, even though the literal name is_sm103 is a bit misleading.
|
LGTM, @ocss884's consideration is completely correct in terms of code. |
is_family_of(sm_103f) also matches any future sm_10x with x > 3, not just sm_103a/f. This was raised in PR review (@ocss884) — adding an inline comment clarifying that this forward-inclusive behavior is intentional: the flag gates ex2 emulation, sm_103 (B300) has fast hardware ex2, and later Blackwell variants in the same family are assumed to inherit it. No code-behavior change.
jayhshah
approved these changes
May 26, 2026
reubenconducts
pushed a commit
to reubenconducts/flash-attention
that referenced
this pull request
Jun 2, 2026
* Use is_family_of for sm_90 and sm_103 arch checks Follow-up to Dao-AILab#2572 — apply the same is_family_of pattern to the two remaining range-style arch checks for consistency: - flash_fwd_sm90.py:69 (SM 9.x assert) - flash_fwd_sm100.py:195 (is_sm103 flag) Same semantic narrowing as Dao-AILab#2572: bare-base SMs (sm_90, sm_103) are excluded. These kernels rely on wgmma / UMMA / 2CTA paths that require the a/f PTX variant anyway, so bare-base targets could not compile. * Clarify is_sm103 forward-inclusive semantics is_family_of(sm_103f) also matches any future sm_10x with x > 3, not just sm_103a/f. This was raised in PR review (@ocss884) — adding an inline comment clarifying that this forward-inclusive behavior is intentional: the flag gates ex2 emulation, sm_103 (B300) has fast hardware ex2, and later Blackwell variants in the same family are assumed to inherit it. No code-behavior change.
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
4 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.
Follow-up to #2572 — apply the same is_family_of pattern to the two remaining range-style arch checks for consistency:
Same semantic narrowing as #2572: bare-base SMs (
sm_90,sm_103) are excluded — they were accepted by the old range check but cannot actually run these kernels, since wgmma / UMMA / 2CTA paths require thea/fPTX variant. The new assertion fails earlier with a clearer error instead of letting compilation get to ptxas before failing.