Function _test_benchmark_attn_func

tests/test_interface.py:90–138

(D, dtype, device, is_causal, is_fp8=False)

@torch.no_grad()
def _test_benchmark_attn_func(D, dtype, device, is_causal, is_fp8=False):
    import triton

    torch.manual_seed(0)

    B = 16
    H = 16
    S_Q = 8192
    S_KV = 8192

    query = torch.randn(B, H, S_Q, D, dtype=dtype, device=device)
    key = torch.randn(B, H, S_KV, D, dtype=dtype, device=device)
    value = torch.randn(B, H, S_KV, D, dtype=dtype, device=device)

    def attention_fn():
        if is_fp8:
            fp8_attention(query, key, value, is_causal)
        else:
            vanilla_attention(query, key, value, is_causal)

    try:
        attention_fn()
    except ValueError as e:
        pytest.skip(str(e))

    def fa_fn():
        flash_attention(query, key, value, is_causal)

    def cudnn_sdpa_fn():
        cudnn_sdpa(query, key, value, is_causal)

    flops_per_matmul = 2 * B * H * S_Q * S_KV * D
    total_flops = 2 * flops_per_matmul

    if is_causal:
        total_flops //= 2

    ms_fa = triton.testing.do_bench(fa_fn)
    tflops_fa = total_flops * 1e-12 / (ms_fa * 1e-3)
    print(f"TFLOPS (Flash Attention): {tflops_fa:.2f}")

    if D <= 128:
        ms_cudnn_sdpa = triton.testing.do_bench(cudnn_sdpa_fn)
        tflops_cudnn_sdpa = total_flops * 1e-12 / (ms_cudnn_sdpa * 1e-3)
        print(f"TFLOPS (CUDNN SDPA): {tflops_cudnn_sdpa:.2f}")

    ms_quantum_attention = triton.testing.do_bench(attention_fn)
    tflops_quantum_attention = total_flops * 1e-12 / (ms_quantum_attention * 1e-3)
    print(f"TFLOPS (Quantum Attention): {tflops_quantum_attention:.2f}")


@pytest.mark.parametrize("D", [64, 128, 256])