Kernel: multi_head_attention (Level 3, category: attention)
Input shapes: [[8, 256, 64], [8, 256, 64], [8, 256, 64]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 4.5e-05)
Pallas runtime: 710.0619 ms; JAX XLA compiled runtime: 0.1116 ms; speedup vs compiled baseline: 0.0002

JAX reference implementation:
@jax.jit
def jax_multi_head_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: Multi-Head Attention via Pallas.

Provenance: jax-ml/jax pallas/ops/tpu/flash_attention.py
             AI-Hypercomputer/maxtext splash attention training kernel
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _mha_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_multi_head_attention(
    q: jax.Array, k: jax.Array, v: jax.Array
) -> jax.Array:
    n_heads, seq_len, d_head = q.shape

    return pl.pallas_call(
        _mha_kernel,
        out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
        grid=(n_heads,),
        in_specs=[
            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
        ],
        out_specs=pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
    )(q, k, v)


pallas_kernel = pallas_multi_head_attention
task_name = "multi_head_attention"
input_shapes = [(8, 256, 64), (8, 256, 64), (8, 256, 64)]
category = "attention"
level = 3

Fixed Pallas kernel (GPU-compatible): """Level 3: Multi-Head Attention via Pallas.

Provenance: jax-ml/jax pallas/ops/tpu/flash_attention.py
             AI-Hypercomputer/maxtext splash attention training kernel
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _mha_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_multi_head_attention(
    q: jax.Array, k: jax.Array, v: jax.Array
) -> jax.Array:
    n_heads, seq_len, d_head = q.shape

    bq = min(32, seq_len)
    return pl.pallas_call(
        _mha_kernel,
        out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
        grid=(n_heads, seq_len // bq),
        in_specs=[
            pl.BlockSpec((1, bq, d_head), lambda h, i: (h, i, 0)),
            pl.BlockSpec((1, seq_len, d_head), lambda h, i: (h, 0, 0)),
            pl.BlockSpec((1, seq_len, d_head), lambda h, i: (h, 0, 0)),
        ],
        out_specs=pl.BlockSpec((1, bq, d_head), lambda h, i: (h, i, 0)),
    )(q, k, v)


pallas_kernel = pallas_multi_head_attention
task_name = "multi_head_attention"
input_shapes = [(8, 256, 64), (8, 256, 64), (8, 256, 64)]
category = "attention"
level = 3

Diff of GPU fix: --- a/multi_head_attention.py
+++ b/multi_head_attention.py
@@ -30,16 +30,17 @@
 ) -> jax.Array:
     n_heads, seq_len, d_head = q.shape
 
+    bq = min(32, seq_len)
     return pl.pallas_call(
         _mha_kernel,
         out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
-        grid=(n_heads,),
+        grid=(n_heads, seq_len // bq),
         in_specs=[
-            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
-            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
-            pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
+            pl.BlockSpec((1, bq, d_head), lambda h, i: (h, i, 0)),
+            pl.BlockSpec((1, seq_len, d_head), lambda h, i: (h, 0, 0)),
+            pl.BlockSpec((1, seq_len, d_head), lambda h, i: (h, 0, 0)),
         ],
-        out_specs=pl.BlockSpec((1, seq_len, d_head), lambda h: (h, 0, 0)),
+        out_specs=pl.BlockSpec((1, bq, d_head), lambda h, i: (h, i, 0)),
     )(q, k, v)
 
 

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