[Cute,Fwd] Extend score_mod to variable sequence length#2043

Merged

jayhshah merged 35 commits intoDao-AILab:mainfrom

reubenconducts:rstern/varlen-score-mod-rebase

Dec 15, 2025

Commits on Dec 15, 2025

rebase to main

reubenconducts
authored and
jayhshah
committed
varlen support for score mod

reubenconducts
authored and
jayhshah
committed
interface change for varlen score mod

reubenconducts
authored and
jayhshah
committed
implement varlen support for score mod

reubenconducts
authored and
jayhshah
committed
varlen score mod working; updated tests

reubenconducts
authored and
jayhshah
committed
modify varlen score mod to use fastdiv_mods updated per sequence

reubenconducts
authored and
jayhshah
committed
updated test suite

reubenconducts
authored and
jayhshah
committed
current working state of varlen score mod

reubenconducts
authored and
jayhshah
committed
refactor varlen score mod tests

reubenconducts
authored and
jayhshah
committed
fix to transpose

reubenconducts
authored and
jayhshah
committed
refactor varlen score mod tests; fix bug; clean up varlen score mod application in kernel

reubenconducts
authored and
jayhshah
committed
refactor test_score_mod.py to use external score mod definition file

reubenconducts
authored and
jayhshah
committed
update flash_fwd.py for varlen score mod

reubenconducts
authored and
jayhshah
committed
sm90 varlen score mod working; test revisions

reubenconducts
authored and
jayhshah
committed
enable packgqa for varlen score mod; set up fastdiv_mod recomputation

reubenconducts
authored and
jayhshah
committed
update flash_fwd_sm100.py for recomputing fastdiv_mods & format varlen score mod test

reubenconducts
authored and
jayhshah
committed
Overwrite pack_gqa.py, tile_scheduler.py, and test_flash_attn.py with origin/main versions

reubenconducts
authored and
jayhshah
committed
rebase to main

reubenconducts
authored and
jayhshah
committed
fix test rebase artifacts

reubenconducts
authored and
jayhshah
committed
fix floor_if_packed redundancy

reubenconducts
authored and
jayhshah
committed
correct sm90 divmods mismatch

reubenconducts
authored and
jayhshah
committed
revert test_flash_attn to main

reubenconducts
authored and
jayhshah
committed
add varlen score mod benchmark script

reubenconducts
authored and
jayhshah
committed
packgqa for varlen (independent of score mod)

reubenconducts
authored and
jayhshah
committed
rm benchmark from PR

reubenconducts
authored and
jayhshah
committed
move score mod arg wrapping to utils.py

reubenconducts
authored and
jayhshah
committed
format with ruff

reubenconducts
authored and
jayhshah
committed
major refactor: change score_mod signature to accept seqlen_info and update all tests accordingly

reubenconducts
authored and
jayhshah
committed
reinstate varlen packgqa exclusion checks

reubenconducts
authored and
jayhshah
committed
move fastdiv_mods recomputation out of apply_score_mod in prep for varlen mask_mod support

reubenconducts
authored and
jayhshah
committed
remove duplicate fastdiv_mod recomputation

reubenconducts
authored and
jayhshah
committed
[Fix] fastdiv_mods for paged attn and seqused_*

reubenconducts
authored and
jayhshah
committed
clean up PR; fix paged_kv varlen for sm90

reubenconducts
authored and
jayhshah
committed
update to varlen score mod test script (paged kv)

reubenconducts
authored and
jayhshah
committed
remove premature seqlen arguments from sm90 apply_mask_mod

reubenconducts
authored and
jayhshah
committed