Kernel: flash_attention (Level 3, category: attention)
Input shapes: [[512, 64], [512, 64], [512, 64]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 1.1e-05)
Pallas runtime: 2653.4589 ms; JAX XLA compiled runtime: 0.1282 ms; speedup vs compiled baseline: 0

JAX reference implementation:
@jax.jit
def jax_flash_attention(
    q: jax.Array, k: jax.Array, v: jax.Array
) -> jax.Array:
    d_k = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / jnp.sqrt(d_k)
    weights = jax.nn.softmax(scores, axis=-1)
    return weights @ v

Original Pallas kernel (TPU-oriented): """Level 3: Tiled Flash Attention via Pallas.

Implements the core Flash Attention pattern: tiled QK^T computation with
online softmax accumulation to avoid materializing the full N x N attention
matrix.

Demonstrates: multi-dimensional grid, online accumulation with fori_loop,
memory-efficient tiling, the full Pallas repertoire.

Reference: jax/experimental/pallas/ops/tpu/flash_attention.py
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L3/flash_attention", __doc__)

from functools import partial

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _flash_attention_kernel(q_ref, k_ref, v_ref, o_ref):
    q = q_ref[...]
    k = k_ref[...]
    v = v_ref[...]
    d_k = q.shape[-1]

    scores = q @ k.swapaxes(-2, -1) / jnp.sqrt(jnp.float32(d_k))
    weights = jnp.exp(scores - jnp.max(scores, axis=-1, keepdims=True))
    weights = weights / jnp.sum(weights, axis=-1, keepdims=True)
    o_ref[...] = weights @ v


def pallas_flash_attention(
    q: jax.Array, k: jax.Array, v: jax.Array
) -> jax.Array:
    seq_len, d_model = q.shape
    block_q = min(128, seq_len)
    grid_size = seq_len // block_q

    return pl.pallas_call(
        _flash_attention_kernel,
        out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_q, d_model), lambda i: (i, 0)),
            pl.BlockSpec((seq_len, d_model), lambda i: (0, 0)),
            pl.BlockSpec((seq_len, d_model), lambda i: (0, 0)),
        ],
        out_specs=pl.BlockSpec((block_q, d_model), lambda i: (i, 0)),
    )(q, k, v)


pallas_kernel = pallas_flash_attention
task_name = "flash_attention"
input_shapes = [(512, 64), (512, 64), (512, 64)]
category = "attention"
level = 3

Fixed Pallas kernel (GPU-compatible): """Level 3: Tiled Flash Attention via Pallas.

Implements the core Flash Attention pattern: tiled QK^T computation with
online softmax accumulation to avoid materializing the full N x N attention
matrix.

Demonstrates: multi-dimensional grid, online accumulation with fori_loop,
memory-efficient tiling, the full Pallas repertoire.

Reference: jax/experimental/pallas/ops/tpu/flash_attention.py
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L3/flash_attention", __doc__)

from functools import partial

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _flash_attention_kernel(q_ref, k_ref, v_ref, o_ref):
    q = q_ref[...]
    k = k_ref[...]
    v = v_ref[...]
    d_k = q.shape[-1]

    scores = q @ k.swapaxes(-2, -1) / jnp.sqrt(jnp.float32(d_k))
    weights = jnp.exp(scores - jnp.max(scores, axis=-1, keepdims=True))
    weights = weights / jnp.sum(weights, axis=-1, keepdims=True)
    o_ref[...] = weights @ v


def pallas_flash_attention(
    q: jax.Array, k: jax.Array, v: jax.Array
) -> jax.Array:
    seq_len, d_model = q.shape
    block_q = min(32, seq_len)
    grid_size = seq_len // block_q

    return pl.pallas_call(
        _flash_attention_kernel,
        out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_q, d_model), lambda i: (i, 0)),
            pl.BlockSpec((seq_len, d_model), lambda i: (0, 0)),
            pl.BlockSpec((seq_len, d_model), lambda i: (0, 0)),
        ],
        out_specs=pl.BlockSpec((block_q, d_model), lambda i: (i, 0)),
    )(q, k, v)


pallas_kernel = pallas_flash_attention
task_name = "flash_attention"
input_shapes = [(512, 64), (512, 64), (512, 64)]
category = "attention"
level = 3

Diff of GPU fix: --- a/flash_attention.py
+++ b/flash_attention.py
@@ -38,7 +38,7 @@
     q: jax.Array, k: jax.Array, v: jax.Array
 ) -> jax.Array:
     seq_len, d_model = q.shape
-    block_q = min(128, seq_len)
+    block_q = min(32, seq_len)
     grid_size = seq_len // block_q
 
     return pl.pallas_call(

Jaxpr IR: { lambda ; a:f32[512,64] b:f32[512,64] c:f32[512,64]. let
    d:f32[512,64] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(16,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=32), Blocked(block_size=64))), BlockMapping(block_shape=(Blocked(block_size=512), Blocked(block_size=64))), BlockMapping(block_shape=(Blocked(block_size=512), Blocked(block_size=64))), BlockMapping(block_shape=(Blocked(block_size=32), Blocked(block_size=64)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; e:Ref<default>{f32[32,64]} f:Ref<default>{f32[512,64]} g:Ref<default>{f32[512,64]}
          h:Ref<default>{f32[32,64]}. let
          i:f32[32,64] <- e[...]
          j:f32[512,64] <- f[...]
          k:f32[512,64] <- g[...]
          l:f32[64,512] = transpose[permutation=(1, 0)] j
          m:f32[32,512] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=float32
          ] i l
          n:f32[] = sqrt 64.0:f32[]
          o:f32[32,512] = div m n
          p:f32[32] = reduce_max[axes=(1,)] o
          q:f32[32,1] = broadcast_in_dim[broadcast_dimensions=(0,)] p
          r:f32[32,512] = sub o q
          s:f32[32,512] = exp r
          t:f32[32] = reduce_sum[axes=(1,) out_sharding=None] s
          u:f32[32,1] = broadcast_in_dim[broadcast_dimensions=(0,)] t
          v:f32[32,512] = div s u
          w:f32[32,64] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=float32
          ] v k
          x:f32[32,64], h[...] <- h[...], w
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[512,64]),)
    ] a b c
  in (d,) }