Skip to content

[CuTe, Flex] Allow score mod use in varlen backward#2547

Open
reubenconducts wants to merge 6 commits into
Dao-AILab:mainfrom
reubenconducts:rstern/varlen-score-mod-bwd
Open

[CuTe, Flex] Allow score mod use in varlen backward#2547
reubenconducts wants to merge 6 commits into
Dao-AILab:mainfrom
reubenconducts:rstern/varlen-score-mod-bwd

Conversation

@reubenconducts
Copy link
Copy Markdown
Contributor

@reubenconducts reubenconducts commented May 7, 2026

Lifts the current restriction on score_mod being used in FA-4 backward with variable sequence length. Supersedes PR #2544.

I added some examples of backwards score mods that use globally-indexed aux tensors to showcase the API.

Note: score mod is still not compatible with 2cta instrs bwd.

cc @drisspg @v0i0

@reubenconducts reubenconducts marked this pull request as draft May 8, 2026 17:00
@reubenconducts reubenconducts marked this pull request as ready for review May 8, 2026 18:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant