@triton.jit
def alibi_attention_triton(score, batch, head, seq_len_q, seq_len_kv, num_heads
    ):
    alibi_scale = tl.math.exp2(-((head + 1) * 8.0 / num_heads))
    bias = seq_len_kv - seq_len_q
    score = score + alibi_scale * bias
    return score
