flex_attention_benchmark_causal_mask.py

# This benchmark requires a Pytorch version with FlexAttention support for XPU available
from functools import lru_cache
import os
from torch.nn.attention.flex_attention import (
    create_block_mask,
    create_mask,
    flex_attention,
)

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.nn.attention.bias import causal_lower_right
import torch._inductor  # pylint: disable=protected-access
import torch._inductor.choices  # pylint: disable=protected-access
from torch._inductor.choices import InductorChoices
from torch._inductor.template_heuristics.triton import FlexBwDConfig, FlexConfig, FlexDecodeConfig

import triton_kernels_benchmark as benchmark_suite
import triton

DEVICE = triton.runtime.driver.active.get_active_torch_device()

# Use TORCHINDUCTOR_MAX_AUTOTUNE_GEMM=1 or uncomment the following line to print the auto-tune results.
# torch._inductor.config.max_autotune_gemm = True


def get_flex_attn_fwd_configs(*args, **kwargs):  # pylint: disable=unused-argument
    configs = [
        FlexConfig(32, 16, 2, 4),
        FlexConfig(128, 64, 2, 16),
        FlexConfig(128, 64, 2, 8),
        FlexConfig(128, 32, 2, 16),
        FlexConfig(128, 32, 2, 8),
    ]
    return configs


def get_flex_attn_bwd_configs(*args, **kwargs):  # pylint: disable=unused-argument
    configs = [FlexBwDConfig(BLOCK, BLOCK, BLOCK, BLOCK, s, w) for BLOCK in [32, 64] for s in [1, 2] for w in ([4])]
    return configs


def get_flex_decode_configs(*args, **kwargs):  # pylint: disable=unused-argument
    configs = [
        FlexDecodeConfig(32, 1, 2),
        FlexDecodeConfig(32, 1, 1),
        FlexDecodeConfig(32, 2, 2),
        FlexDecodeConfig(32, 2, 1),
        FlexDecodeConfig(64, 1, 2),
        FlexDecodeConfig(64, 1, 1),
        FlexDecodeConfig(64, 2, 2),
        FlexDecodeConfig(64, 2, 1),
    ]
    return configs


# There is a auto-tuning requirement to get the best configuration for the flex attention.
# The pytorch flex attention doesn't support auto-tuning by user by default.
# Overriding the get_flex_attention_fwd_configs method to provide custom configurations for auto-tuning on XPU.
InductorChoices.get_flex_attention_fwd_configs = get_flex_attn_fwd_configs
InductorChoices.get_flex_attention_bwd_configs = get_flex_attn_bwd_configs
InductorChoices.get_flex_decode_configs = get_flex_decode_configs

torch._dynamo.config.recompile_limit = 100  # pylint: disable=protected-access

# Compile the flex_attention function
compiled_flex_attention = torch.compile(flex_attention, dynamic=False)


@lru_cache
def create_block_mask_cached(score_mod, B, H, M, N, device=DEVICE):
    block_mask = create_block_mask(score_mod, B, H, M, N, device=device)
    return block_mask


@lru_cache
def create_mask_cached(score_mod, B, H, M, N, device=DEVICE):
    mask = create_mask(score_mod, B, H, M, N, device=device)
    return mask


def create_causal_mask_fn(N_CTX_q, N_CTX_kv):
    offset = N_CTX_kv - N_CTX_q

    def causal_mask(_, __, q_idx, kv_idx):
        # Bottom-right aligned lower triangular mask
        # For square: offset=0, gives q_idx >= kv_idx (standard causal)
        # For non-square: offset>0, shifts alignment to bottom-right
        return q_idx >= (kv_idx - offset)

    return causal_mask


def get_attn_bias(provider, use_causal, N_CTX_q, N_CTX_kv, device):
    if use_causal:
        return None
    if provider == 'sycl-tla':
        return causal_lower_right(N_CTX_q, N_CTX_kv)
    causal_mask_fn = create_causal_mask_fn(N_CTX_q, N_CTX_kv)
    return create_mask_cached(causal_mask_fn, 1, 1, N_CTX_q, N_CTX_kv, device=device)


def get_sdpa_benchmark(q, k, v, attn_bias, use_causal, sm_scale, H_q, H_kv, D_HEAD_qk, D_HEAD_v, MODE, provider,
                       backwards_grad=None):
    """Get SDPA function for SYCL-TLA or OneDNN providers."""
    # SDPA backends have limitations:
    # - Require D_HEAD_qk == D_HEAD_v (no support for different head dimensions)
    # - OneDNN doesn't support backward pass for SDPA
    if D_HEAD_qk != D_HEAD_v or (provider == 'onednn' and MODE == 'bwd'):
        return None

    backend = SDPBackend.FLASH_ATTENTION if provider == 'sycl-tla' else SDPBackend.OVERRIDEABLE

    def sdpa_fn():
        with sdpa_kernel(backends=[backend]):
            return F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attn_bias,
                is_causal=use_causal,
                scale=sm_scale,
                enable_gqa=(H_q != H_kv),
            )

    if MODE == 'bwd':
        o = sdpa_fn()
        do = backwards_grad if backwards_grad is not None else torch.randn_like(o)
        bwd_fn = lambda: o.backward(do, retain_graph=True)
        return bwd_fn, do, o
    return sdpa_fn


throughput_test = os.getenv('THROUGHPUT_TEST', '0') == '1'
batch_size = int(os.getenv('BATCH_SIZE', '1'))
batch_sizes = [16, 32, 64] if throughput_test else [batch_size]
fa_kernel_mode = os.getenv('FA_KERNEL_MODE', 'fwd')

if 'B580' in torch.xpu.get_device_name():
    old_count = len(batch_sizes)
    batch_sizes = [size for size in batch_sizes if size < 16]
    if len(batch_sizes) != old_count:
        print('Skipping running batch_sizes >= 16 on b580')


# Kernel profiling for Backward mode is not working as expected:
# For details: https://github.com/pytorch/pytorch/issues/144778
@benchmark_suite.perf_report(
    benchmark_suite.Benchmark(
        x_names=['Z', 'H_q', 'H_kv', 'N_CTX_q', 'N_CTX_kv', 'D_HEAD_qk', 'D_HEAD_v', 'MODE'],
        x_vals=[[z, *params, fa_kernel_mode] for z in batch_sizes for params in [
            # Multi-head attention. H_q equals H_kv
            (32, 32, 1024, 1024, 96, 96),  # Prefill shapes of Phi3-mini-4k-instruct
            (32, 32, 1024, 1024, 128, 128),  # Prefill shapes of Qwen3-4B
            (128, 128, 1024, 1024, 192, 128),  # Prefill shapes of DeepSeek-v3
            (32, 32, 512, 1024 + 128 + 512, 96, 96),  # Append shapes of Phi3-mini-4k-instruct

            # Grouped-query attention. H_q / H_kv > 1
            (32, 8, 1024, 1024, 128, 128),  # Prefill shapes of Llama-3.1-8B
            (24, 8, 1024, 1024, 128, 128),  # Prefill shapes of meta-llama-Llama-3.2-3B
            (40, 8, 1024, 1024, 128, 128),  # Prefill shapes of Deepseek-R1-Distill-Qwen-14B
            (32, 8, 512, 1024 + 128 + 512, 128, 128),  # Append shapes of Llama-3.1-8B and Qwen3-4B
            (24, 8, 512, 1024 + 128 + 512, 128, 128),  # Append shapes of meta-llama-Llama-3.2-3B
            (40, 8, 512, 1024 + 128 + 512, 128, 128),  # Append shapes of Deepseek-R1-Distill-Qwen-14B

            # FlexDecoding configuration. N_CTX_q equals 1. N_CTX_kv < 1k
            (32, 8, 1, 1024 + 64, 128, 128),  # Decode shapes of Llama-3.1-8B and Qwen3-4B
            (24, 8, 1, 1024 + 64, 128, 128),  # Decode shapes of meta-llama-Llama-3.2-3B
            (16, 16, 1, 1024, 128, 128),  # Additional Hq=Hkv=16 PyTorch benchmark case
            (16, 2, 1, 1024, 128, 128),  # Additional Hq=16, Hkv=2 PyTorch benchmark case
            (32, 32, 1, 1024 + 64, 96, 96),  # Decode shapes of Phi3-mini-4k-instruct
            (40, 8, 1, 1024 + 64, 128, 128),  # Decode shapes of Deepseek-R1-Distill-Qwen-14B

            # Multi-query attention. H_kv equals 1
            (128, 1, 1, 1024 + 64, 576, 512),  # Decode shapes of Deepseek-v3
            (128, 1, 512, 1024 + 128 + 512, 576, 512),  # Append shapes of Deepseek-v3
        ] + ([
            # Shapes only for bwd
            [h, h, seq_len, seq_len, 128, 128]  #
            for h in [1, 2, 4, 16, 24, 32]  #
            for seq_len in [4096, 8192]  #
        ] if fa_kernel_mode == 'bwd' else [])],
        line_arg='provider',
        line_vals=['triton', 'sycl-tla', 'onednn'],
        line_names=['Triton', 'SYCL-TLA', 'OneDNN'],
        styles=[('green', '-'), ('green', '--'), ('blue', '-'), ('blue', '--')],
        ylabel=['GB/s', 'TFlops'],
        plot_name='flexAttnCausal-performance',
        args={},
    ))
# pylint: disable=too-many-branches
def benchmark(Z, H_q, H_kv, N_CTX_q, N_CTX_kv, D_HEAD_qk, D_HEAD_v, MODE, provider):
    print(
        f'Running case: {Z=}, {H_q=}, {H_kv=}, {N_CTX_q=}, {N_CTX_kv=}, {D_HEAD_qk=}, {D_HEAD_v=}, {MODE=}, {provider=}'
    )
    torch.xpu.empty_cache()
    torch.manual_seed(42)

    # Maximum across torch=200, triton=600
    do_bench = benchmark_suite.get_do_bench(n_warmup=600, n_repeat=10, quantiles=[0.5, 0.0, 1.0])
    if MODE not in ('fwd', 'bwd'):
        raise ValueError(f"Invalid MODE: {MODE}. Expected 'fwd' or 'bwd'.")
    dtype = torch.float16
    q = torch.randn((Z, H_q, N_CTX_q, D_HEAD_qk), device=DEVICE, dtype=dtype, requires_grad=MODE == 'bwd')
    k = torch.randn((Z, H_kv, N_CTX_kv, D_HEAD_qk), device=DEVICE, dtype=dtype, requires_grad=MODE == 'bwd')
    v = torch.randn((Z, H_kv, N_CTX_kv, D_HEAD_v), device=DEVICE, dtype=dtype, requires_grad=MODE == 'bwd')
    sm_scale = 0.125

    block_mask = create_block_mask_cached(create_causal_mask_fn(N_CTX_q, N_CTX_kv), 1, 1, N_CTX_q, N_CTX_kv,
                                          device=DEVICE)

    if provider in ('sycl-tla', 'onednn'):
        use_causal = N_CTX_q == N_CTX_kv
        attn_bias = get_attn_bias(provider, use_causal, N_CTX_q, N_CTX_kv, DEVICE)
        sdpa_result = get_sdpa_benchmark(q, k, v, attn_bias, use_causal, sm_scale, H_q, H_kv, D_HEAD_qk, D_HEAD_v, MODE,
                                         provider)
        if sdpa_result is None:
            return None, float('nan'), float('nan'), float('nan'), float('nan')

        if MODE == 'bwd':
            sdpa_fn, _, _ = sdpa_result
        else:
            sdpa_fn = sdpa_result
        _, min_ms, max_ms, mean, cv = do_bench(
            sdpa_fn, device=DEVICE,
            benchmark_label='ScaledDotProductFlashAttentionBackward0' if MODE == 'bwd' else None)

    elif provider == 'triton':
        ref_results_provider = 'sycl-tla'
        kernel_options = {'BLOCKS_ARE_CONTIGUOUS': True, 'USE_TMA': True}

        # pylint: disable=too-many-boolean-expressions
        if H_q == 128 and H_kv == 1 and N_CTX_q == 1 and N_CTX_kv == 1088 and D_HEAD_qk == 576 and D_HEAD_v == 512:
            # Workaround for DeepSeek-v3 decode shape
            # Force to use prefill kernel because decode exceeds available Per Thread Scratch Space (PTSS)
            # Due to that we cannot compile
            kernel_options['FORCE_USE_FLEX_ATTENTION'] = True

        triton_fn = lambda: compiled_flex_attention(q, k, v, block_mask=block_mask, scale=sm_scale, enable_gqa=(
            not H_q == H_kv), kernel_options=kernel_options)

        is_reference_available = True
        if D_HEAD_qk != D_HEAD_v:
            print(f'Skipping correctness check: {ref_results_provider} requires D_HEAD_qk == D_HEAD_v')
            is_reference_available = False

        if MODE == 'bwd':
            backwards_grad = torch.randn(Z, H_q, N_CTX_q, D_HEAD_v, dtype=dtype, device=DEVICE,
                                         requires_grad=MODE == 'bwd')

            if is_reference_available:
                use_causal = N_CTX_q == N_CTX_kv
                attn_bias_ref = get_attn_bias(ref_results_provider, use_causal, N_CTX_q, N_CTX_kv, DEVICE)

                sdpa_result = get_sdpa_benchmark(q, k, v, attn_bias_ref, use_causal, sm_scale, H_q, H_kv, D_HEAD_qk,
                                                 D_HEAD_v, MODE, ref_results_provider, backwards_grad=backwards_grad)
                _, _, sycl_o = sdpa_result
                sycl_grads = torch.autograd.grad((sycl_o, ), (q, k, v), backwards_grad, retain_graph=True)
                sycl_tensors = (sycl_o, *sycl_grads)

            triton_o = triton_fn()
            triton_grads = torch.autograd.grad((triton_o, ), (q, k, v), backwards_grad, retain_graph=True)
            compiled_tensors = (triton_o, *triton_grads)

            if is_reference_available:
                tensor_names = ['out', 'grad_query', 'grad_key', 'grad_value']
                for sycl, compiled, name in zip(sycl_tensors, compiled_tensors, tensor_names):  # pylint: disable=used-before-assignment
                    benchmark_suite.assert_close(
                        lambda: sycl, lambda: compiled, atol=6e-2, rtol=1e-3,  # pylint: disable=cell-var-from-loop
                        err_msg=f'Error comparing {name} between triton and {ref_results_provider}')

            triton_fn = lambda: torch.autograd.grad((triton_o, ), (q, k, v), backwards_grad, retain_graph=True)
        else:
            if is_reference_available:
                use_causal = N_CTX_q == N_CTX_kv
                attn_bias_ref = get_attn_bias(ref_results_provider, use_causal, N_CTX_q, N_CTX_kv, DEVICE)

                sycl_fn = get_sdpa_benchmark(q, k, v, attn_bias_ref, use_causal, sm_scale, H_q, H_kv, D_HEAD_qk,
                                             D_HEAD_v, MODE, ref_results_provider)

                if sycl_fn is not None:
                    benchmark_suite.assert_close(triton_fn, sycl_fn, atol=1e-2, rtol=1e-3,
                                                 err_msg=f'triton to {ref_results_provider}')

        _, min_ms, max_ms, mean, cv = do_bench(triton_fn, device=DEVICE, grad_to_none=(q, k, v),
                                               benchmark_label=None if MODE == 'fwd' else 'CompiledFunctionBackward')

    else:
        raise NotImplementedError(f'Unsupported provider {provider}')

    def flops_triangle(length):
        # the triangle with diagonal.
        return ((length + 1) * length // 2) * 2  # mul + add

    def flops_rectangle(m, n, k):
        return m * n * k * 2  # mul + add

    if N_CTX_q == 1:
        # decoding ignore the causal mask since only one query is involved.
        qk_flops = H_q * N_CTX_q * N_CTX_kv * D_HEAD_qk * 2  # mul + add.
        pv_flops = H_q * N_CTX_q * D_HEAD_v * N_CTX_kv * 2  # mul + add.
    elif N_CTX_q == N_CTX_kv:
        # Square matrix - standard causal mask (lower triangle)
        qk_flops = H_q * flops_triangle(N_CTX_q) * D_HEAD_qk
        pv_flops = H_q * flops_triangle(N_CTX_q) * D_HEAD_v
    elif N_CTX_q < N_CTX_kv:
        # More keys than queries - bottom-right aligned: rectangle + triangle
        # Each query q can attend to keys [0, q + offset] where offset = N_CTX_kv - N_CTX_q
        # This forms: all queries attend to first (N_CTX_kv - N_CTX_q) keys (rectangle)
        #             plus a lower triangle for the remaining N_CTX_q x N_CTX_q region
        rect_width = N_CTX_kv - N_CTX_q
        qk_flops = H_q * (flops_rectangle(N_CTX_q, rect_width, D_HEAD_qk) + flops_triangle(N_CTX_q) * D_HEAD_qk)
        pv_flops = H_q * (flops_rectangle(N_CTX_q, rect_width, D_HEAD_v) + flops_triangle(N_CTX_q) * D_HEAD_v)
    else:  # N_CTX_q > N_CTX_kv
        # More queries than keys - bottom-right aligned: only last N_CTX_kv queries have attention
        # First (N_CTX_q - N_CTX_kv) queries are fully masked out
        qk_flops = H_q * flops_triangle(N_CTX_kv) * D_HEAD_qk
        pv_flops = H_q * flops_triangle(N_CTX_kv) * D_HEAD_v

    tflops = lambda mean: Z * (qk_flops + pv_flops) * (1e-12) / (mean * 1e-3)

    q_elems = H_q * N_CTX_q * D_HEAD_qk
    k_elems = H_kv * N_CTX_kv * D_HEAD_qk
    v_elems = H_kv * N_CTX_kv * D_HEAD_v
    gbps = lambda mean: Z * (q_elems + k_elems + v_elems) * 2 * (1e-9) / (mean * 1e-3)  # float16 2 bytes

    if MODE == 'bwd':
        # The tflops and gbps are aligned to the one in flash_attention_benchmark.
        tflops = lambda mean: 2.5 * Z * (qk_flops + pv_flops) * (1e-12) / (mean * 1e-3)
        gbps = lambda mean: 2.5 * Z * (q_elems + k_elems + v_elems) * 2 * (1e-9) / (mean * 1e-3)

    return (gbps(mean), gbps(max_ms), gbps(min_ms)), (tflops(mean), tflops(max_ms), tflops(min_ms)), cv


def get_benchmark(providers_filter=None, fa_kernel_mode='fwd', batch_size=1):  # pylint: disable=W0613,W0621
    local_batch_sizes = [16, 32, 64] if os.getenv('THROUGHPUT_TEST', '0') == '1' else [batch_size]
    if 'B580' in torch.xpu.get_device_name():
        local_batch_sizes = [size for size in local_batch_sizes if size < 16]
    base = benchmark.benchmarks
    base_shapes = [x[1:-1] for x in base.x_vals]
    x_vals = [[z, *params, fa_kernel_mode] for z in local_batch_sizes for params in base_shapes]
    if fa_kernel_mode == 'bwd':
        x_vals += [[z, h, h, seq_len, seq_len, 128, 128, fa_kernel_mode]
                   for z in local_batch_sizes
                   for h in [1, 2, 4, 16, 24, 32]
                   for seq_len in [4096, 8192]]

    @benchmark_suite.perf_report(
        benchmark_suite.Benchmark(x_names=base.x_names, x_vals=x_vals, line_arg=base.line_arg, line_vals=base.line_vals,
                                  line_names=base.line_names, styles=base.styles, ylabel=base.ylabel,
                                  plot_name=base.plot_name, args=base.args))
    def _benchmark(Z, H_q, H_kv, N_CTX_q, N_CTX_kv, D_HEAD_qk, D_HEAD_v, MODE, provider):
        return benchmark.fn(Z, H_q, H_kv, N_CTX_q, N_CTX_kv, D_HEAD_qk, D_HEAD_v, MODE, provider)

    return _benchmark


if __name__ == '__main__':
    benchmark.run(show_plots=False, print_data=True)