Skip to content

Fix missing seqlen_info param in softcap scoremod#2366

Open
rucnyz wants to merge 2 commits intoDao-AILab:mainfrom
rucnyz:fix/softcap-scoremod-seqlen-info
Open

Fix missing seqlen_info param in softcap scoremod#2366
rucnyz wants to merge 2 commits intoDao-AILab:mainfrom
rucnyz:fix/softcap-scoremod-seqlen-info

Conversation

@rucnyz
Copy link
Copy Markdown

@rucnyz rucnyz commented Mar 17, 2026

Summary

apply_score_mod_inner in softmax.py calls score_mod with 7 arguments including seqlen_info introduced in #2043:

https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/cute/softmax.py#L457-L465

post_mod_scores = score_mod(
    score_ssa, batch_idx_ssa, head_idx_ssa,
    q_idx=q_idx_ssa, kv_idx=kv_idx_ssa,
    seqlen_info=seqlen_info,  # here
    aux_tensors=aux_args,
)

But create_softcap_scoremod in utils.py returns a function with only 6 parameters (missing seqlen_info):

https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/cute/utils.py#L120

def scoremod_premask_fn(acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, aux_tensors):

This causes TypeError: got an unexpected keyword argument 'seqlen_info' when using softcap.

Fix

Add seqlen_info parameter to scoremod_premask_fn signature to match apply_score_mod_inner.

apply_score_mod_inner now passes seqlen_info to score_mod functions,
but create_softcap_scoremod's scoremod_premask_fn was not updated to
accept it, causing a TypeError when compiling softcap kernel variants.
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