def vanilla_attention(
    q: Tensor, k: Tensor, v: Tensor, scale: float = None, attn_mask: Tensor = None
) -> Tensor:
    """Vanilla attention in the most basic form.
    q,k,v: batch, seqlen, num_heads, head_dim
        Either batch or sequence dimension could be variable (but not both)
    attn_mask: attention mask is *added* to the attention,
        use 0 and -inf to mask a sequence index
    """
    batch_name, seq_name = [it._attrs["name"] for it in q._attrs["shape"]][0:2]
    B, N, G, D = _get_shape(q)
    B, M, _, _ = _get_shape(k)
    BG = B * G
    if BG < 0:
        BG = -1
    C = G * D
    if scale is None:
        scale = D ** (-0.5)
    q = q * scale

    q = ops.permute()(q, [0, 2, 1, 3])  # BxGxNxD
    q = ops.reshape()(q, (BG, N, D))  # BGxNxD

    k = ops.reshape()(k, (B, M, C))  # BxMxGD
    k = ops.permute021()(k)  # BxGDxM
    k = ops.reshape()(k, (BG, D, M))  # BGxDxM

    attention = ops.bmm_rrr()(q, k)  # BGxNxM
    if attn_mask is not None:
        attention = attention + attn_mask
    attention = ops.softmax()(attention, -1)  # BGxNxM

    v = ops.reshape()(v, (B, M, C))  # BxMxGD
    v = ops.permute021()(v)  # BxGDxM
    v = ops.reshape()(v, (BG, D, M))  # BGxDxM

    out = ops.bmm_rcr()(v, attention)  # BGxDxN
    out = ops.reshape()(out, (B, C, N))  # BxGDxN == BxCxN
    out = ops.permute021()(out)  # BxNxC
    out._attrs["shape"][0]._attrs["name"] = batch_name
    out._attrs["shape"][1]._attrs["name"] = seq_name
    return out
Vanilla attention in the most basic form.
q,k,v: batch, seqlen, num_heads, head_dim
Either batch or sequence dimension could be variable (but not both)
attn_mask: attention mask is *added* to the attention,
use 0 and -inf to mask a sequence index
python/aitemplate/frontend/nn/vanilla_attention.py