def METHOD_NAME(q, k, v, attn_mask=None, causal=False):
    """
    attention output of the control group
    """
    dtype_og = q.dtype
    seqlen_q, seqlen_k = q.shape[1], k.shape[1]
    d = q.shape[-1]
    scale = 1.0 / math.sqrt(d)
    scores = torch.einsum('bthd,bshd->bhts', q * scale, k)
    if attn_mask is not None:
        scores.masked_fill_(rearrange(~attn_mask, 'b s -> b 1 1 s'), float('-inf'))
    if causal:
        causal_mask = torch.triu(torch.ones(seqlen_q, seqlen_k, dtype=torch.bool, device=q.device), 1)
        scores.masked_fill_(causal_mask, float('-inf'))
    attention = torch.softmax(scores, dim=-1)
    output = torch.einsum('bhts,bshd->bthd', attention, v)
    output = rearrange(output, "b s h d -> b s (h d)")
    # Modify the data at the positions of the mask to 0
    if attn_mask is not None:
        output.masked_fill_(rearrange(~attn_mask, 'b s -> b s 1'), 0.0)
    return output.to(dtype=dtype_og)