def METHOD_NAME(query, key, value, num_heads, causal=False):
    query = query.view(query.shape[0], query.shape[1], num_heads, -1).permute(
        0, 2, 1, 3
    )
    key = key.view(key.shape[0], key.shape[1], num_heads, -1).permute(0, 2, 1, 3)
    value = value.view(value.shape[0], value.shape[1], num_heads, -1).permute(
        0, 2, 1, 3
    )
    query = query.reshape(-1, query.shape[2], query.shape[3])
    key = key.reshape(-1, key.shape[2], key.shape[3]).permute(0, 2, 1)
    value = value.reshape(-1, value.shape[2], value.shape[3])
    scale = 1 / math.sqrt(query.shape[-1])
    scores = flow.baddbmm(
        flow.empty(
            query.shape[0],
            query.shape[1],
            key.shape[1],
            dtype=query.dtype,
            device=query.device,
        ),
        query,
        key,
        beta=0,
        alpha=scale,
    )
    if causal:
        causal_mask = flow.triu(
            flow.ones(
                scores.shape[-2], scores.shape[-1], dtype=flow.bool, device="cuda"
            ),
            1,
        )
        scores = flow.masked_fill(scores, causal_mask, float("-inf"))
    attn = flow.softmax(scores, dim=-1)
    out = flow.matmul(attn, value)
    out = out.reshape(-1, num_heads, out.shape[1], out.shape[2])
    out = out.permute(0, 2, 1, 3)
    out = out.reshape(out.shape[0], out.shape[1], -1)
    return out