def METHOD_NAME(proj_shape, dtype, dropout):
    (B, S, H, D_HEAD) = proj_shape
    D = H * D_HEAD
    q = torch.randn((B, S, H, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
    k = torch.randn((B, S, H, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
    v = torch.randn((B, S, H, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
    attn = ColoAttention(D, H, dropout=dropout)
    y = attn(q, k, v)
    assert list(y.shape) == [B, S, D]
    out_ref = attention_ref(q, k, v, None, causal=False)
    dy = torch.rand_like(y)
    grad_q, grad_k, grad_v = torch.autograd.grad(y, (q, k, v), dy)
    grad_ref_q, grad_ref_k, grad_ref_v = torch.autograd.grad(out_ref, (q, k, v), dy)
    torch.allclose(y, out_ref, atol=1e-7), f"{(y - out_ref).abs().max()}"
    torch.allclose(grad_q, grad_ref_q, atol=1e-7), f"{(grad_q - grad_ref_q).abs().max()}"
    torch.allclose(grad_k, grad_ref_k, atol=1e-7), f"{(grad_k - grad_ref_k).abs().max()}"
    torch.allclose(grad_v, grad_ref_v, atol=1e-7), f"{(grad_v - grad_ref_v).abs().max()}"