Kernel: layernorm_residual (Level 2, category: norm_residual)
Input shapes: [[2048, 1024], [2048, 1024]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 1e-06)
Pallas runtime: 461.5373 ms; JAX XLA compiled runtime: 0.1136 ms; speedup vs compiled baseline: 0.0002

JAX reference implementation:
@jax.jit
def jax_layernorm_residual(x: jax.Array, residual: jax.Array) -> jax.Array:
    mean = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.var(x, axis=-1, keepdims=True)
    normed = (x - mean) / jnp.sqrt(var + 1e-5)
    return normed + residual

Original Pallas kernel (TPU-oriented): """Level 2: Fused LayerNorm + Residual Add via Pallas.

Provenance: standard transformer pre-norm pattern, MaxText attention blocks
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L2/layernorm_residual", __doc__)

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


def _layernorm_residual_kernel(x_ref, residual_ref, o_ref):
    x = x_ref[...]
    residual = residual_ref[...]
    mean = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.var(x, axis=-1, keepdims=True)
    normed = (x - mean) / jnp.sqrt(var + 1e-5)
    o_ref[...] = normed + residual


def pallas_layernorm_residual(x: jax.Array, residual: jax.Array) -> jax.Array:
    n_rows = x.shape[0]
    n_cols = x.shape[1]
    block_rows = min(128, n_rows)
    grid_size = n_rows // block_rows

    return pl.pallas_call(
        _layernorm_residual_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
        ],
        out_specs=pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
    )(x, residual)


pallas_kernel = pallas_layernorm_residual
task_name = "layernorm_residual"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "norm_residual"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused LayerNorm + Residual Add via Pallas.

Provenance: standard transformer pre-norm pattern, MaxText attention blocks
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L2/layernorm_residual", __doc__)

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


def _layernorm_residual_kernel(x_ref, residual_ref, o_ref):
    x = x_ref[...]
    residual = residual_ref[...]
    mean = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.var(x, axis=-1, keepdims=True)
    normed = (x - mean) / jnp.sqrt(var + 1e-5)
    o_ref[...] = normed + residual


def pallas_layernorm_residual(x: jax.Array, residual: jax.Array) -> jax.Array:
    n_rows = x.shape[0]
    n_cols = x.shape[1]
    block_rows = min(128, n_rows)
    while block_rows * n_cols > 16384 and block_rows > 1:
        block_rows //= 2
    grid_size = n_rows // block_rows

    return pl.pallas_call(
        _layernorm_residual_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
        ],
        out_specs=pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
    )(x, residual)


pallas_kernel = pallas_layernorm_residual
task_name = "layernorm_residual"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "norm_residual"
level = 2

Diff of GPU fix: --- a/layernorm_residual.py
+++ b/layernorm_residual.py
@@ -26,6 +26,8 @@
     n_rows = x.shape[0]
     n_cols = x.shape[1]
     block_rows = min(128, n_rows)
+    while block_rows * n_cols > 16384 and block_rows > 1:
+        block_rows //= 2
     grid_size = n_rows // block_rows
 
     return pl.pallas_call(

Jaxpr IR: { lambda ; a:f32[2048,1024] b:f32[2048,1024]. let
    c:f32[2048,1024] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(128,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024))), BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024))), BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[16,1024]} e:Ref<default>{f32[16,1024]}
          f:Ref<default>{f32[16,1024]}. let
          g:f32[16,1024] <- d[...]
          h:f32[16,1024] <- e[...]
          i:f32[16] = reduce_sum[axes=(1,) out_sharding=None] g
          j:f32[16,1] = broadcast_in_dim[broadcast_dimensions=(0,)] i
          k:f32[16,1] = div j 1024.0:f32[]
          l:f32[16,1] = jit[
            name=_var
            jaxpr={ lambda ; g:f32[16,1024] m:i32[]. let
                n:f32[16] = reduce_sum[axes=(1,) out_sharding=None] g
                o:f32[16,1] = broadcast_in_dim[broadcast_dimensions=(0,)] n
                p:f32[16,1] = div o 1024.0:f32[]
                q:f32[16,1024] = sub g p
                r:f32[16,1024] = square q
                s:f32[] = convert_element_type[
                  new_dtype=float32
                  weak_type=False
                ] m
                t:f32[] = sub 1024.0:f32[] s
                u:f32[16] = reduce_sum[axes=(1,) out_sharding=None] r
                v:f32[16,1] = broadcast_in_dim[broadcast_dimensions=(0,)] u
                w:f32[16,1] = div v t
                x:bool[] = gt t 0.0:f32[]
                l:f32[16,1] = jit[
                  name=_where
                  jaxpr={ lambda ; x:bool[] w:f32[16,1] y:f32[]. let
                      z:f32[] = convert_element_type[
                        new_dtype=float32
                        weak_type=False
                      ] y
                      ba:f32[16,1] = broadcast_in_dim z
                      l:f32[16,1] = select_n x ba w
                    in (l,) }
                ] x w nan:f32[]
              in (l,) }
          ] g 0:i32[]
          bb:f32[16,1024] = sub g k
          bc:f32[16,1] = add l 9.999999747378752e-06:f32[]
          bd:f32[16,1] = sqrt bc
          be:f32[16,1024] = div bb bd
          bf:f32[16,1024] = add be h
          bg:f32[16,1024], f[...] <- f[...], bf
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[2048,1024]),)
    ] a b
  in (c,) }