[CuTe, Flex] Allow score mod use in varlen backward#2547
[CuTe, Flex] Allow score mod use in varlen backward#2547reubenconducts wants to merge 12 commits into
Conversation
| score_mod_bwd = utils.create_softcap_scoremod_bwd(softcap) | ||
| if score_mod is not None: | ||
| assert score_mod_bwd is not None, "score_mod_bwd is required when score_mod is provided" | ||
| assert cu_seqlens_q is None and cu_seqlens_k is None, ( |
There was a problem hiding this comment.
we have full varlen bwd coverage right? I havnt thought about deterministic for instance
There was a problem hiding this comment.
We don't have blocksparse varlen bwd yet (i.e. it isn't plumbed up at all), but I didn't guard against this in the interface. I'll add an assertion.
| ), | ||
| mCuBlockIdxOffsets=( | ||
| blocksparse_tensors.cu_block_idx_offsets if blocksparse_tensors is not None else None | ||
| blocksparse_tensors.cu_block_idx_offsets |
There was a problem hiding this comment.
I think the SM90 bwd score_mod partials should be rebuilt after the per-batch fastdiv_mods recompute. Right now score_mod_fn / score_mod_bwd_fn capture (via partial) the original fastdiv_mods before the work-tile loop, so reassigning fastdiv_mods inside the loop only reaches mask_fn, not the score_mod paths. SM90 fwd constructs its score_mod partial after the recompute; can we mirror that here?
The code section is above but we shoudl probs test on hopper before land
There was a problem hiding this comment.
Done, and tested on Hopper. Two failures due to small numerical error one seqlen x score mod config.
FAILED tests/cute/test_score_mod_varlen.py::test_varlen_score_mod_backward_with_global[True-score_mod_tuple0-seqlens_q23-seqlens_k23-128-dtype0] - AssertionError: dK CuTE err 9.96e-02 exceeds 3*PT err 3.12e-02 + 1.00e-03
FAILED tests/cute/test_score_mod_varlen.py::test_varlen_score_mod_backward_with_global[False-score_mod_tuple0-seqlens_q23-seqlens_k23-128-dtype0] - AssertionError: dK CuTE err 9.96e-02 exceeds 3*PT err 3.12e-02 + 1.00e-03
============================================ 2 failed, 11746 passed, 9372 skipped, 8712 xfailed, 1540 warnings in 289.50s (0:04:49)
reubenconducts
force-pushed
the
rstern/varlen-score-mod-bwd
branch
from
342d620 to
b311c04
Compare
|
Any sense for prio of full varlen bwd coverage ? This may be critical for one of my use cases. If it's not a prioritized follow-up, I may work on a PR to add that functionality once this lands. |
reubenconducts
force-pushed
the
rstern/varlen-score-mod-bwd
branch
from
b311c04 to
cafa779
Compare
3 participants
Lifts the current restriction on
score_modbeing 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.
cc @drisspg @v0i0