def attention_ref_sdpa(qkv, attn_mask, dropout_p, upcast=False, causal=False):
    """
    Arguments:
        qkv: (batch_size, seqlen, 3, nheads, head_dim)
        attn_mask: (batch_size, seqlen)
        dropout_p: float
    Output:
        output: (batch_size, seqlen, nheads, head_dim)
        attention: softmax after dropout
    """
    q, k, v = (qkv.float() if upcast else qkv).unbind(dim=2)
    q = q.transpose(1, 2)  # to (batch_size, nheads, seqlen, head_dim)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    if attn_mask is not None:
        # to (batch_size, nheads, seqlen, seqlen)
        attn_mask = attn_mask.reshape(q.shape[0], 1, 1, q.shape[2])
    output = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=causal
    )
    output = output.transpose(1, 2)  # to (batch_size, seqlen, nheads, head_dim)
    return output.to(dtype=qkv.dtype)
Arguments:
qkv: (batch_size, seqlen, 3, nheads, head_dim)
attn_mask: (batch_size, seqlen)
dropout_p: float
Output:
output: (batch_size, seqlen, nheads, head_dim)
attention: softmax after dropout
tests/unittest/ops/test_attention.py