test_flash_attn_triton.py

import torch
import triton
import torch.nn.functional as F
import os
import pytest

from flash_attn import triton_flash_attn_func
from flash_attn.flash_attn_triton_mqa_gqa import attention as attention_mqa_gqa
from flash_attn.flash_attn_triton_mqa_gqa import input_helper_layout



@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD, dtype',
[ (*shape, dtype)
    for shape in [(4, 48, 1024, 64),
                  (1, 48, 2048, 64),
                  (1, 48, 4096, 64),
                  (1, 48, 1024, 128),
                  (1, 48, 2048, 128),
                  (1, 48, 4096, 128)]
    for dtype in [
        torch.float16,
        torch.bfloat16
    ]
])
@pytest.mark.parametrize('causal', [ True])
def test_op_fwd(Z, H, N_CTX, D_HEAD, causal, dtype):
    torch.manual_seed(20)
    q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5)
    k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5)
    v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5)

    sm_scale = q.shape[-1]**-0.5
    dropout_p = 0.0
    # dout = torch.randn_like(q, dtype=torch.float16)
    # triton implementation
    tri_out = triton_flash_attn_func(q, k, v, dropout_p,  causal=causal, softmax_scale=sm_scale)

    # reference implementation
    M = torch.tril(torch.ones((Z, H, N_CTX, N_CTX), device="cuda"))
    p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
    if causal:
        p[M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).to(dtype)
#    p = F.dropout(p_nodrop, 0.5)

    ref_out = torch.matmul(p, v)

    # compare
    atol = 1.4e-1 if dtype == 'fp8' else 1e-2
    rtol = 1e-2 if dtype == 'fp8' else 0
    torch.testing.assert_close(ref_out, tri_out, atol=atol, rtol=rtol)


@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD',
                         [(4, 48, 1024, 64),
                          (1,16,4096,128),
                          (1,32,4096,128),
                          (1,20,4096,128),
                          (1,24,4096,128),
                          (1, 48, 2048, 64),
                          (1, 48, 4096, 64),
                          (1, 16, 8192, 128),
                          ])
@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16])
def test_op_bwd(Z, H, N_CTX, D_HEAD, dtype):
    torch.manual_seed(20)
    causal = True
    q = (torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    k = (torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    v = (torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    q, k, v = (x.to(dtype) for x in (q, k, v))

    sm_scale= q.shape[-1]**-0.5
    dout = torch.randn_like(q)
    # reference implementation
    M = torch.tril(torch.ones((Z, H, N_CTX, N_CTX), device="cuda"))
    p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
    if causal:
        p[M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).to(dtype)
#    p = F.dropout(p_nodrop, 0.5)

    ref_out = torch.matmul(p, v)
    ref_out.backward(dout)
    ref_dv, v.grad = v.grad.clone(), None
    ref_dk, k.grad = k.grad.clone(), None
    ref_dq, q.grad = q.grad.clone(), None
    # triton implementation
    tri_out = triton_flash_attn_func(q, k, v, causal=causal, softmax_scale=sm_scale)
    tri_out.backward(dout)
    tri_dv, v.grad = v.grad.clone(), None
    tri_dk, k.grad = k.grad.clone(), None
    tri_dq, q.grad = q.grad.clone(), None

    # compare
    torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=0)
    if torch.version.hip is None:
        torch.testing.assert_close(ref_dv, tri_dv, atol=1e-2, rtol=0)
    # The current block size for MI200 series is 64x64. This results in
    # larger differences in float results due to rounding.
    else:
        torch.testing.assert_close(ref_dv, tri_dv, atol=5e-2, rtol=1e-2)
    torch.testing.assert_close(ref_dk, tri_dk, atol=5e-2, rtol=1e-2)
    torch.testing.assert_close(ref_dq, tri_dq, atol=5e-2, rtol=1e-2)

@pytest.mark.parametrize('Z, H, N_CTX_Q, N_CTX_K, D_HEAD, dtype',
[ (*shape, dtype)
    for shape in [(4, 48, 4096, 77, 64),
                  (2, 8, 4096, 4096, 40),
                  (2, 8, 1024, 1024, 80),
                  (2, 8, 1024, 77, 80),
                  (2, 8, 256, 256, 160),
                  (2, 8, 256, 77, 160)]
    for dtype in [
       torch.float16,
       torch.bfloat16
    ]
])
@pytest.mark.parametrize('causal', [False, True])
def test_op_neq_fwd(Z, H, N_CTX_Q, N_CTX_K , D_HEAD, causal, dtype):
    torch.manual_seed(20)
    q = torch.empty((Z, H, N_CTX_Q, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    k = torch.empty((Z, H, N_CTX_K, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    v = torch.empty((Z, H, N_CTX_K, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    softmax_scale= q.shape[-1]**-0.5
    # dout = torch.randn_like(q, dtype=torch.float16)
    # triton implementation
    tri_out = triton_flash_attn_func(q, k, v, causal=causal)
    if isinstance(tri_out, tuple):  # 检查 o 是否为 tuple，如果是则取第一个值
        tri_out = tri_out[0]

    # reference implementation
    M = torch.tril(torch.ones((Z, H, N_CTX_Q, N_CTX_K), device="cuda"))
    p = torch.matmul(q, k.transpose(2, 3)) * softmax_scale
    if causal:
        p[M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).to(dtype)

    ref_out = torch.matmul(p, v)

   # compare
    atol = 1.4e-1 if dtype == 'fp8' else 1e-2
    rtol = 1e-2 if dtype == 'fp8' else 0
    torch.testing.assert_close(ref_out, tri_out, atol=atol, rtol=rtol)

@pytest.mark.parametrize('Z, H, N_CTX_Q, N_CTX_K, D_HEAD',
                         [(4, 48, 2048, 77, 64),
                           (1,8, 2048,16,128),
                           (1,20,4096,4096,128),
                          (1,24,4096,4096,128),
                           (2, 10, 1024, 1024, 80),
                           (2, 20, 1024, 1024, 80),
                           (2, 8, 4096, 4096, 40),
                           (2, 8, 1024, 1024, 80),
                           (2, 8, 1024, 77, 80),
                           (2, 8, 256, 256, 160),
                          ])
@pytest.mark.parametrize('dtype', [torch.float16,torch.bfloat16])
def test_op_neq_bwd(Z, H, N_CTX_Q,N_CTX_K, D_HEAD, dtype):
    torch.manual_seed(20)
    causal = True
    q = (torch.empty((Z, H, N_CTX_Q, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    k = (torch.empty((Z, H, N_CTX_K, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    v = (torch.empty((Z, H, N_CTX_K, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
    q, k, v = (x.to(dtype) for x in (q, k, v))

    softmax_scale= q.shape[-1]**-0.5
    print("softmax_scale", softmax_scale)

    dout = torch.randn_like(q)
    # reference implementation
    M = torch.tril(torch.ones((Z, H, N_CTX_Q, N_CTX_K), device="cuda"))
    p = torch.matmul(q.float(), k.float().transpose(2, 3)) * softmax_scale
    if causal:
        p[M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).to(dtype)
#    p = F.dropout(p_nodrop, 0.5)

    ref_out = torch.matmul(p.float(), v.float())
    ref_out.backward(dout)
    ref_dv, v.grad = v.grad.clone(), None
    ref_dk, k.grad = k.grad.clone(), None
    ref_dq, q.grad = q.grad.clone(), None
    # triton implementation
    tri_out = triton_flash_attn_func(q, k, v, causal=causal)
    if isinstance(tri_out, tuple):  # 检查 o 是否为 tuple，如果是则取第一个值
        tri_out = tri_out[0]

    tri_out.backward(dout)
    tri_dv, v.grad = v.grad.clone(), None
    tri_dk, k.grad = k.grad.clone(), None
    tri_dq, q.grad = q.grad.clone(), None
    # compare
    torch.testing.assert_close(ref_out, tri_out.float(), atol=1e-2, rtol=0)
    if torch.version.hip is None:
        torch.testing.assert_close(ref_dv, tri_dv, atol=1e-2, rtol=0)
    # The current block size for MI200 series is 64x64. This results in
    # larger differences in float results due to rounding.
    else:
        torch.testing.assert_close(ref_dv, tri_dv, atol=5e-2, rtol=1e-2)
    torch.testing.assert_close(ref_dk, tri_dk, atol=5e-2, rtol=1e-2)
    torch.testing.assert_close(ref_dq, tri_dq, atol=5e-2, rtol=1e-2)

@pytest.mark.parametrize('Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype',
[ (*shape, dtype)
    for shape in [(1, 16, 16,4096, 4096,128),
        (1, 32, 32, 4096, 4096,128),
        (1, 32, 4, 4096, 4096,128),
        (1, 52, 4, 4096, 4096,128),
        (1, 16, 2, 4096, 4096,128),
        (1, 26, 2, 4096, 4096,128),
        (1, 8, 1, 4096, 4096,128),
        (1, 13, 1, 4096, 4096,128),
    ]
    for dtype in [
        torch.float16,
    ]
])
@pytest.mark.parametrize('causal', [True])
@pytest.mark.parametrize('layout', ['bhsd'])
def test_op_varlen_fwd(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dtype, layout):
    torch.manual_seed(20)
    q, k, v, input_metadata = input_helper_layout(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout)
    if causal:
        input_metadata.need_causal()

    o = torch.empty_like(q)
    print(q.size(),k.size(),v.size())
    # triton implementation
    tri_out = attention_mqa_gqa(q, k, v, None, input_metadata)

    # Transpose here if layout is bshd so we have same reference code for all layouts
    if layout == 'bshd':
        q = q.transpose(1, 2).clone()
        k = k.transpose(1, 2).clone()
        v = v.transpose(1, 2).clone()
    # Replicate K and V if using MQA/GQA
    if HQ != HK:
        k = k.view(k.shape[0], k.shape[1], -1, k.shape[2],
                   k.shape[3]).expand(-1, -1, HQ // HK, -1, -1).reshape(k.shape[0], -1, k.shape[2], k.shape[3])
        v = v.view(v.shape[0], v.shape[1], -1, v.shape[2],
                   v.shape[3]).expand(-1, -1, HQ // HK, -1, -1).reshape(v.shape[0], -1, v.shape[2], v.shape[3])
    scores = torch.einsum('bhqd,bhkd->bhqk', q, k).float() * input_metadata.sm_scale
    if causal:
        mask = torch.tril(torch.ones(N_CTX_Q, N_CTX_K, device="cuda"), diagonal=N_CTX_K - N_CTX_Q)
        scores[:, :, mask == 0] = float("-inf")

    p = torch.softmax(scores, dim=-1)
    if causal:
        # If N_CTX_Q > N_CTX_K, there is at least one row of all -infs going into
        # the softmax. This produces a row of NaNs as -inf - -inf == NaN. So we fix
        # this by converting the NaNs to 0s, which is what they should be out of the softmax.
        nan_mask = torch.isnan(p)
        p[nan_mask == 1] = 0
    ref_out = torch.einsum('bhqk,bhkd->bhqd', p.half(), v)
    # compare
    if layout == 'bshd':
        ref_out = ref_out.transpose(1, 2).clone()
    torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)


# vary seq length for fixed head and batch=4
configs = []
for mode in ['fwd', 'bwd']:
    for D_HEAD in [128, 64]:
        for causal in [False, True]:
            if mode == 'bwd' and causal == False:
                continue
            configs.append(triton.testing.Benchmark(
                x_names=['BATCH', 'H', 'N_CTX'],
                x_vals=[(4, 16, 1024),
                        (8, 16, 2048),
                        (4, 16, 4096),
                        (2, 16, 8192),
                        (1, 16, 16384),
                        (4, 48, 1024),
                        (4, 48, 2048),
                        (4, 48, 4096),
                        (4, 48, 8192),
                        (4, 48, 16384),
                        ],
                line_arg='provider',
                line_vals=['triton'],
                line_names=['Triton'],
                styles=[('red', '-'), ('blue', '-')],
                ylabel='ms',
                plot_name=f'fused-attention-{mode}-d{D_HEAD}-causal={causal}',
                args={
                    'D_HEAD': D_HEAD,
                    'dtype': torch.float16,
                    'mode': mode,
                    'causal': causal,
                },
            ))


@triton.testing.perf_report(configs)
def bench_flash_attention(BATCH, H, N_CTX, D_HEAD, causal, mode, provider, dtype=torch.float16, device="cuda"):
    assert mode in ["fwd", "bwd"]
    warmup = 25
    rep = 10
    # Bwd pass only supports causal=True right now
    if mode == 'bwd':
        causal = True
    if provider == "triton":
        q = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
        k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
        v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
        sm_scale = D_HEAD ** -0.5
        fn = lambda: triton_flash_attn_func(q, k, v, causal=causal, softmax_scale=sm_scale)
        if mode == 'bwd':
            o = fn()
            do = torch.randn_like(o)
            fn = lambda: o.backward(do, retain_graph=True)
        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
    if provider == "flash":
        qkv = torch.randn((BATCH, N_CTX, 3, H, D_HEAD), dtype=dtype, device=device, requires_grad=True)
        fn = lambda: triton_flash_attn_func(qkv, causal=causal)
        if mode == "bwd":
            o = fn()
            do = torch.randn_like(o)
            fn = lambda: o.backward(do, retain_graph=True)
        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
    flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD
    total_flops = 2 * flops_per_matmul
    if causal:
        total_flops *= 0.5
    if mode == "bwd":
        total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
    return total_flops / ms * 1e-9

# only works on post-Ampere GPUs right now
# bench_flash_attention.run(save_path=".", print_data=True)