The score_mod function passed to flex_attention should operate on the pre-softmax attention scores, but the snippet below appears to be applying the attention biases (s_aux) and computing the post-softmax scores.

I don't think it is possible to apply (gpt-oss-style) attention sinks using the score_mod alone, but you can do it by passing return_lse=True to flex_attention and renormalising using the extra return value. If someone can point me to where unit tests for this code should live I'm happy to PR a fix.