Kernel: batched_matmul (Level 1, category: matmul)
Input shapes: [[8, 256, 256], [8, 256, 256]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 241.3267 ms; JAX XLA compiled runtime: 0.1036 ms; speedup vs compiled baseline: 0.0004

JAX reference implementation:
@jax.jit
def jax_batched_matmul(x: jax.Array, y: jax.Array) -> jax.Array:
    return x @ y

Original Pallas kernel (TPU-oriented): """Level 1: Batched matrix multiplication via Pallas.

Provenance: jnp.matmul with batch dims, used in multi-head attention
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/batched_matmul", __doc__)

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


def _batched_matmul_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] @ y_ref[...]


def pallas_batched_matmul(x: jax.Array, y: jax.Array) -> jax.Array:
    batch, m, k = x.shape
    _, _, n = y.shape

    return pl.pallas_call(
        _batched_matmul_kernel,
        out_shape=jax.ShapeDtypeStruct((batch, m, n), x.dtype),
        grid=(batch,),
        in_specs=[
            pl.BlockSpec((1, m, k), lambda b: (b, 0, 0)),
            pl.BlockSpec((1, k, n), lambda b: (b, 0, 0)),
        ],
        out_specs=pl.BlockSpec((1, m, n), lambda b: (b, 0, 0)),
    )(x, y)


pallas_kernel = pallas_batched_matmul
task_name = "batched_matmul"
input_shapes = [(8, 256, 256), (8, 256, 256)]
category = "matmul"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Batched matrix multiplication via Pallas.

Provenance: jnp.matmul with batch dims, used in multi-head attention
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/batched_matmul", __doc__)

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


def _batched_matmul_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] @ y_ref[...]


def pallas_batched_matmul(x: jax.Array, y: jax.Array) -> jax.Array:
    batch, m, k = x.shape
    _, _, n = y.shape

    bm_ = min(32, m)
    bn_ = min(32, n)
    return pl.pallas_call(
        _batched_matmul_kernel,
        out_shape=jax.ShapeDtypeStruct((batch, m, n), x.dtype),
        grid=(batch, m // bm_, n // bn_),
        in_specs=[
            pl.BlockSpec((1, bm_, k), lambda b, i, j: (b, i, 0)),
            pl.BlockSpec((1, k, bn_), lambda b, i, j: (b, 0, j)),
        ],
        out_specs=pl.BlockSpec((1, bm_, bn_), lambda b, i, j: (b, i, j)),
    )(x, y)


pallas_kernel = pallas_batched_matmul
task_name = "batched_matmul"
input_shapes = [(8, 256, 256), (8, 256, 256)]
category = "matmul"
level = 1

Diff of GPU fix: --- a/batched_matmul.py
+++ b/batched_matmul.py
@@ -20,15 +20,17 @@
     batch, m, k = x.shape
     _, _, n = y.shape
 
+    bm_ = min(32, m)
+    bn_ = min(32, n)
     return pl.pallas_call(
         _batched_matmul_kernel,
         out_shape=jax.ShapeDtypeStruct((batch, m, n), x.dtype),
-        grid=(batch,),
+        grid=(batch, m // bm_, n // bn_),
         in_specs=[
-            pl.BlockSpec((1, m, k), lambda b: (b, 0, 0)),
-            pl.BlockSpec((1, k, n), lambda b: (b, 0, 0)),
+            pl.BlockSpec((1, bm_, k), lambda b, i, j: (b, i, 0)),
+            pl.BlockSpec((1, k, bn_), lambda b, i, j: (b, 0, j)),
         ],
-        out_specs=pl.BlockSpec((1, m, n), lambda b: (b, 0, 0)),
+        out_specs=pl.BlockSpec((1, bm_, bn_), lambda b, i, j: (b, i, j)),
     )(x, y)
 
 

Jaxpr IR: { lambda ; a:f32[8,256,256] b:f32[8,256,256]. let
    c:f32[8,256,256] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(8, 8, 8), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=1), Blocked(block_size=32), Blocked(block_size=256))), BlockMapping(block_shape=(Blocked(block_size=1), Blocked(block_size=256), Blocked(block_size=32))), BlockMapping(block_shape=(Blocked(block_size=1), Blocked(block_size=32), Blocked(block_size=32)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[1,32,256]} e:Ref<default>{f32[1,256,32]}
          f:Ref<default>{f32[1,32,32]}. let
          g:f32[1,32,256] <- d[...]
          h:f32[1,256,32] <- e[...]
          i:f32[32,256] = squeeze[dimensions=(0,)] g
          j:f32[256,32] = squeeze[dimensions=(0,)] h
          k:f32[32,32] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=float32
          ] i j
          l:f32[1,32,32] = broadcast_in_dim[broadcast_dimensions=(1, 2)] k
          m:f32[1,32,32], f[...] <- f[...], l
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[8,256,256]),)
    ] a b
  in (c,) }