def METHOD_NAME(
    test_case,
    N=512,
    S=128,
    D=128,
    h=8,
    d=32,
    mode="row",
    ensemble_batch=8,
    inplace=False,
):
    x = flow.randn(N, S, D, requires_grad=True).cuda()  # N, S, D
    w3 = [flow.randn(D, h * d, requires_grad=True).cuda() for _ in range(3)]  # D, h*d*3
    mask = flow.randn(N, S, requires_grad=False).cuda()  # N, S
    bias = None
    scale = 1 / (d ** 0.5)
    if mode in ["row", "triangular_start", "triangular_end"]:
        bias = flow.randn(1, h, S, S, requires_grad=True).cuda()  # 1, h, S, S
        bias.retain_grad()
        mask = mask[:, None, None, :]
    if mode == "ensemble":
        x = flow.randn(ensemble_batch, N, S, D, requires_grad=True).cuda()  # N, S, D
        bias = flow.randn(
            ensemble_batch, 1, h, S, S, requires_grad=True
        ).cuda()  # E, 1, h, S, S
        bias.retain_grad()
        mask = flow.randn(ensemble_batch, N, 1, 1, S, requires_grad=False).cuda()
    if mode == "col" or mode == "global_col":
        N, S = S, N
        x = x.transpose(-2, -3)  # S, N, D
        mask = mask.transpose(-1, -2)
        if mode == "col":
            mask = mask[..., None, None, :]  # S, 1, 1, N
    q, k, v = [flow.matmul(x, w) for w in w3]  # N, S, h * d
    if mode == "template":
        n_templ = 4
        x = flow.randn(S, S, 1, D, requires_grad=True).cuda()
        k = v = flow.randn(S, S, n_templ, D, requires_grad=True).cuda()  # N, S, D
        mask = flow.randn(1, 1, 1, 1, n_templ).cuda()
        q, k, v = [flow.matmul(x_, w) for x_, w in zip([x, k, v], w3)]
    q, k, v = [
        permute_final_dims(a.view(*a.shape[:-1], h, d), (0, 2, 1, 3)) for a in [q, k, v]
    ]  # N, h, S, d
    if mode == "global_col":
        w_q = flow.randn(D, h * d, requires_grad=True).cuda()  # D, h*d
        w_kv = flow.randn(D, d * 2, requires_grad=True).cuda()  # D, h*d*2
        q = flow.sum(x * mask.unsqueeze(-1), dim=-2) / (
            flow.sum(mask, dim=-1)[..., None] + 1e-9
        )  # [N, D]
        mask = mask[..., :, None, :]  # N,1,S
        q = flow.matmul(q, w_q).view(*q.shape[:-1], h, d)  # N, h, d
        k, v = flow.matmul(x, w_kv).chunk(2, dim=-1)  # N, S, d
    qk = flow.matmul(q, k.transpose(-1, -2))
    # general op
    x.retain_grad()
    out1 = _ref_op(qk, v, scale, mask, bias, inplace)
    out1.sum().backward(retain_graph=True)
    grad_x1 = x.grad
    grad_bias1 = bias.grad if bias is not None else None
    # fused op
    out2 = _fused_op(qk, v, scale, mask, bias, inplace)
    out2.sum().backward()
    grad_x2 = x.grad
    grad_bias2 = bias.grad if bias is not None else None
    test_case.assertTrue(np.allclose(out1, out2, atol=2e-3, rtol=1e-5))
    test_case.assertTrue(np.allclose(grad_x1, grad_x2, atol=5e-3, rtol=1e-5))
    if bias is not None:
        test_case.assertTrue(np.allclose(grad_bias1, grad_bias2, atol=5e-4, rtol=1e-5))