[GPU] FlashAttention performance lags behind PyTorch

[neel04]

Description

I'm benchmarking naive FlashAttention in Jax vs. the Pallas's version of FA3 vs. the new dot_product_attention interface with cudnn backend.

• JAX/XLA's performance:

• Torch's performance:

Why the discrepancy? I'd have expected performance to touch 550-600 TFLOPS/s. I'm using a few XLA flags, as specified in the script below - but is there anything I'm missing? Or is this about the maximum XLA can deliver on H100s?

Steps to reproduce

1. Recreate the environment using uv. I'm assuming the drivers are installed. If not, you can use the pytorch/pytorch:2.4.0-cuda12.4.1-cudnn8-runtime image on the GPU, run the preliminary apt-get update and apt-get upgrade to set everything up.

pip3 install uv
uv venv 'main_env' --python 3.11
source main_env/bin/activate

uv pip install -U "jax[cuda12]"
uv pip install -q einops tqdm jaxtyping optax optuna equinox rich
uv pip install -q nvitop pdbpp tabulate

3. Run either script

**JAX script**

import os, sys

os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.9'

os.environ['XLA_FLAGS'] = (
    '--xla_gpu_enable_triton_softmax_fusion=true '
    '--xla_gpu_enable_cudnn_fmha=true'
)

import math
import time
from tabulate import tabulate
import jax
import jax.numpy as jnp
import numpy as np
from jax.tree_util import tree_map
from jax.experimental.pallas.ops.gpu.attention import mha as pallas_mha
from functools import partial

class Timer(object):
    def __init__(self, into=None):
        self.into = into

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, type, value, traceback):
        if self.into is not None:
            self.into.append(time.time() - self.start)

    def elapsed(self):
        return time.time() - self.start

def cartesian(*lists):
    if lists:
        xs = lists[0]
        for x in xs:
            for rest in cartesian(*lists[1:]):
                yield (x,) + rest
    else:
        yield ()

def cross_attn_flops(B,T,TK,H,C):
    HC = H*C
    # T = TK for self-attention
    flops_fwd = (
        2 * B*H*T*TK*C + # S = Q@K
        3 * B*H*T*TK +  # P=softmax(S)
        2 * B*T*TK*H*C # O = P@V
    )
    return flops_fwd

def attn_flops(B,T,H,C):
    return cross_attn_flops(B,T,T,H,C)

dtype = jnp.float16
print(f'Using dtype: {dtype}')

def convert(xs):
    return tree_map(lambda x: x.astype(dtype), xs)
    
def ref_fwd(q,k,v):
    # reference implementation
    [n, l, h, d] = q.shape
    [n, lk, hk, d] = k.shape
    softmax_scale = 1 / math.sqrt(d)
    S = jnp.einsum('nlhd,nLhd->nhlL',q,k)
    P = jax.nn.softmax(S*softmax_scale, axis=-1)
    o = jnp.einsum('nhlL,nLhd->nlhd',P,v)
    return o.astype(q.dtype)

def jax_dpa_fwd(q, k, v):
    output = jax.nn.dot_product_attention(
        query=q,
        key=k,
        value=v,
        implementation='cudnn'
    )

    return output

# ----

Bx = [8,  16] # batch size
Tx = [1024, 2048] # seqlen
Hx = [16, 32] # number of heads
Cx = [64, 128] # head dim
sx = [2] # steps

def bench_attn(mha):
    @jax.jit
    def bench(q, k, v, steps: int):
        for i in steps:
            out = mha(q, k, v)
        return out

    times = []
    table = {}

    for B,T,H,C,s in cartesian(Bx,Tx,Hx,Cx,sx):
        q = jax.random.normal(jax.random.PRNGKey(0), [B, T, H, C], dtype=dtype)
        k = jax.random.normal(jax.random.PRNGKey(1), [B, T, H, C], dtype=dtype)
        v = jax.random.normal(jax.random.PRNGKey(2), [B, T, H, C], dtype=dtype)
        steps = jnp.arange(s)

        # Warmup
        for _ in range(2):
            bench(q, k, v, steps).block_until_ready()

        out = []
        for _ in range(8):
            with Timer(out):
                bench(q, k, v, steps).block_until_ready()

        t = sum(out[2:])/len(out[2:])

        flop = attn_flops(B,T,H,C) * s
        tf = flop / t / 1e12 / s
        print(f'flops ({B} {T} {H} {C} / {s}): {tf}T', out)
        table[(B,T,H,C,s)] = tf
        times.append(t)
    
    return table, times

naive_flops, custom_time = bench_attn(ref_fwd)
pallas_flops, pallas_time = bench_attn(partial(pallas_mha, segment_ids=None))
jax_dpa_flops, dpa_time = bench_attn(jax_dpa_fwd)

table = []

for idx, (B,T,H,C,s) in enumerate(cartesian(Bx,Tx,Hx,Cx,sx)):
    n_flops, n_time = naive_flops[(B,T,H,C,s)], custom_time[idx]
    p_flops, p_time = pallas_flops[(B,T,H,C,s)], pallas_time[idx]
    j_flops, j_time = jax_dpa_flops[(B,T,H,C,s)], dpa_time[idx]
    
    table.append((B,T,H,C, n_flops, p_flops, j_flops, attn_flops(B,T,H,C), n_time, p_time, j_time))

print(tabulate(table, headers=['B','T','H','C','TFlop/s (naive)','TFlop/s (pallas)','TFlop/s (jax_dpa)', 'FLOPs', 'Naive time', 'Pallas time', 'DPA time'], floatfmt='.5f'))

**PyTorch Benchmark script**

import torch
torch._dynamo.config.cache_size_limit = 10000  # Increase cache size to 10,000

import time
from tabulate import tabulate
from typing import Optional
from itertools import product
from torch.utils.flop_counter import FlopCounterMode

class Timer:
    def __init__(self, into=None):
        self.into = into

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, type, value, traceback):
        if self.into is not None:
            self.into.append(time.time() - self.start)

class OptimizedMHA(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        B, T, H, C = q.shape
        scale = 1.0 / (C ** 0.5)

        # Compute attention scores
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale  # [B, H, T, T]
        attn = torch.softmax(attn, dim=-1)

        # Apply attention to values
        out = torch.matmul(attn, v)  # [B, H, T, C]
        out = out.transpose(1, 2)  # [B, T, H, C]

        return out

class TorchMHA(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.scaled_dot_product_attention(
            q, k, v,
            dropout_p=0.0,
            is_causal=False
        )

def get_flops(model, sample_inputs):
    """Get FLOPs using FlopCounterMode"""
    with FlopCounterMode(model) as flop_counter:
        _ = model(*sample_inputs)
    return flop_counter.get_total_flops()
def bench_implementation(model, q, k, v, warmup=10, steps=100):
    """Benchmark a specific implementation"""
    # Warmup
    for _ in range(warmup):
        with torch.no_grad():
            _ = model(q, k, v)
    torch.cuda.synchronize()

    # Benchmark timing
    times = []
    for _ in range(steps):
        with torch.no_grad(), Timer(times):
            _ = model(q, k, v)
            torch.cuda.synchronize()

    # Calculate statistics
    times = times[10:]  # Discard first 10 runs
    avg_time = sum(times) / len(times)
    return avg_time

def bench_attention():
    # Configuration
    device = torch.device("cuda")
    dtype = torch.float16

    # Test parameters
    Bx = [8, 16]               # batch size
    Tx = [1024, 2048]            # sequence length
    Hx = [16, 32]          # number of heads
    Cx = [64, 128]     # head dimension
    sx = [4]               # steps to run

    # Initialize models
    custom_model = OptimizedMHA().to(device)
    torch_model = TorchMHA().to(device)

    # Compile both models
    compiled_custom = torch.compile(
        custom_model,
        mode="max-autotune-no-cudagraphs",
        fullgraph=True,
    )
    compiled_torch = torch.compile(
        torch_model,
        mode="max-autotune-no-cudagraphs",
        fullgraph=True,
    )

    results = []

    # Run benchmarks for each configuration
    for B, T, H, C, s in product(Bx, Tx, Hx, Cx, sx):
        # Create input tensors
        q = torch.randn(B, T, H, C, device=device, dtype=dtype)
        k = torch.randn(B, T, H, C, device=device, dtype=dtype)
        v = torch.randn(B, T, H, C, device=device, dtype=dtype)

        q = q.transpose(1, 2)  # [B, H, T, C]
        k = k.transpose(1, 2)  # [B, H, T, C]
        v = v.transpose(1, 2)  # [B, H, T, C]

        # Get FLOPs using FlopCounterMode (on CPU with float32)
        model_cpu = OptimizedMHA()
        q_cpu = q.cpu().float()
        k_cpu = k.cpu().float()
        v_cpu = v.cpu().float()
        flops = get_flops(model_cpu, (q_cpu, k_cpu, v_cpu))

        # Benchmark both implementations
        custom_time = bench_implementation(compiled_custom, q, k, v)
        torch_time = bench_implementation(compiled_torch, q, k, v)

        # Calculate TFLOPs/s for both
        custom_tflops = flops / custom_time / 1e12
        torch_tflops = flops / torch_time / 1e12

        # Calculate speedup
        speedup = custom_time / torch_time  # >1 means torch is faster

        print(f"\nConfig (B={B}, T={T}, H={H}, C={C}):")
        print(f"Custom impl: {custom_tflops:.2f} TFlop/s")
        print(f"Torch impl: {torch_tflops:.2f} TFlop/s")
        print(f"Speedup (Torch vs Custom): {speedup:.2f}x")
        print(f"Measured FLOPs: {flops:,}")

        results.append((
            B, T, H, C,
            round(custom_tflops, 2),
            round(torch_tflops, 2),
            round(speedup, 2),
            flops,
            custom_time,
            torch_time
        ))

    # Print results table
    headers = [
        'Batch', 'SeqLen', 'Heads', 'HeadDim',
        'Naive MHA TFlop/s', 'SDPA TFlop/s',
        'Advantage', 'FLOPs', 'Custom Time', 'SDPA Time'
    ]
    print("\nResults:")
    print(tabulate(results, headers=headers, floatfmt='.5f'))

if __name__ == "__main__":
    bench_attention()

---

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
jax:    0.4.35
jaxlib: 0.4.34
numpy:  2.1.3
python: 3.11.10 (main, Oct 16 2024, 04:38:48) [Clang 18.1.8 ]
device info: NVIDIA H100 PCIe-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='nifty-orthodox-whale', release='6.8.0-40-generic', version='#40~22.04.3-Ubuntu SMP PREEMPT_DYNAMIC Tue Jul 30 17:30:19 UTC 2', machine='x86_64')


$ nvidia-smi
Sun Nov 17 02:04:44 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.06             Driver Version: 535.183.06   CUDA Version: 12.4     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA H100 PCIe               On  | 00000000:00:07.0 Off |                    0 |
| N/A   31C    P0              54W / 350W |    467MiB / 81559MiB |      2%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+

[neel04]

Update: I changed the torch script to use FlopCounterMode. Now the results are more realistic/accurate but JAX still lags behind despite explicitly being forced to use CuDNN.

cc @kaixih @sbodenstein @dfm

[Rick0827]

i think your torch script might not work as expected, since the inputs format of torch.nn.functional.scaled_dot_product_attention is [B, H, T, C] [https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html]

[sbodenstein]

This is a complicated benchmarking setup, with many things potentially going wrong. Can you simplify this to just measuring milliseconds, and also have a correctness test (that PyTorch and JAX give the same output for the same input).

[neel04]

@Rick0827 Thank you for pointing that out.

@sbodenstein I have updated both scripts to now report times as well. However, I opted to skip correctness tests because reproducibility requires us to sacrifice performance which I'm afraid of touching

The variance however is very low between runs plus we can average over multiple steps (sx) so this should be a non-issue.

On A100:

• JAX:

• Torch: