[CuTe, SM120] Fix forward/backward bugs and add SM120 test guards - #2420
Closed
2imi9 wants to merge 2 commits into
Closed
[CuTe, SM120] Fix forward/backward bugs and add SM120 test guards#24202imi9 wants to merge 2 commits into
2imi9 wants to merge 2 commits into
Conversation
Bug fixes: - flash_fwd_sm120.py: Override self.arch after parent __init__ to prevent SM80 epilogue from enabling TMA-O store on SM120 - interface.py: Initialize dQ_single_wg and num_stages_PdS in SM120 backward config (fixes NameError in compile_key) - interface.py: Fix hdim>128 forward tile sizes for SM120 99KB SMEM - utils.py: Replace nvvm.atomicrmw with inline PTX red.global.add.f32 for newer NVVM Python bindings compatibility - __init__.py: Graceful FA2 import fallback when flash_attn_2_cuda is not built Test guards (tests/cute/test_flash_attn.py): - Add IS_SM120 flag - Skip: SplitKV, deterministic backward, learnable_sink, paged KV - Skip: GQA/MQA (pack_gqa crd2idx incompatibility on newer CUTLASS DSL) - TODO markers reference tracking PRs (#2336, #2348, #2389) Tested on RTX 5090 Laptop GPU (SM 12.0): 240 test_flash_attn_output cases passed (MHA, bf16, hdim 64/128/192/256, causal/non-causal, multiple sequence lengths, forward + backward). Builds on #2404 (bug fixes by @NJX-njx). Co-Authored-By: NJX-njx <NJX-njx@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d3334eb55f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| flash_attn_varlen_qkvpacked_func, | ||
| flash_attn_with_kvcache, | ||
| ) | ||
| except ImportError: |
There was a problem hiding this comment.
Catching a bare ImportError here suppresses all failures from flash_attn.flash_attn_interface, not just the missing flash_attn_2_cuda case this fallback is meant to handle. In scenarios like ABI mismatch, dependency regressions, or import-time bugs inside flash_attn_interface, import flash_attn will now appear to succeed but silently drop the public symbols, leading to later AttributeError/missing-function failures that are much harder to diagnose.
Useful? React with 👍 / 👎.
Only suppress the missing flash_attn_2_cuda module error, not unexpected ImportError from ABI mismatch or other failures.
blake-snc
added a commit
to blake-snc/flash-attention-dao
that referenced
this pull request
Apr 2, 2026
Override self.arch = Arch.sm_80 after parent __init__ to prevent base class code paths from seeing the runtime arch (12.x) and enabling SM90+ features. The parent __init__ overwrites the class-level arch=80 attribute with the actual GPU arch. This was found by @2imi9 in Dao-AILab#2420 for the CpAsync kernel — same bug applies here. Add can_implement() check before TMA dispatch in interface.py so that configs exceeding SM120's 99KB SMEM (e.g. hdim=192 with kv_stages=2) fall back to the CpAsync kernel instead of failing at instantiation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>
blake-snc
added a commit
to blake-snc/flash-attention-dao
that referenced
this pull request
Apr 16, 2026
Override self.arch = Arch.sm_80 after parent __init__ to prevent base class code paths from seeing the runtime arch (12.x) and enabling SM90+ features. The parent __init__ overwrites the class-level arch=80 attribute with the actual GPU arch. This was found by @2imi9 in Dao-AILab#2420 for the CpAsync kernel — same bug applies here. Add can_implement() check before TMA dispatch in interface.py so that configs exceeding SM120's 99KB SMEM (e.g. hdim=192 with kv_stages=2) fall back to the CpAsync kernel instead of failing at instantiation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>
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
1 participant
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.
Bug fixes:
Test guards (tests/cute/test_flash_attn.py):
Tested on RTX 5090 Laptop GPU (SM 12.0): 240 test_flash_attn_output cases passed (MHA, bf16, hdim 64/128/192/256, causal/non-causal, multiple sequence lengths, forward + backward).
Builds on #2404 (bug fixes by @NJX-njx).