Kernel: qk_softmax (Level 2, category: attention_component)
Input shapes: [[256, 64], [256, 64]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 4430.2044 ms; JAX XLA compiled runtime: 0.1238 ms; speedup vs compiled baseline: 0

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

Original Pallas kernel (TPU-oriented): """Level 2: Fused QK^T + Softmax via Pallas.

Provenance: jax-ml/jax flash_attention.py attention score computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _qk_softmax_kernel(q_ref, k_ref, o_ref):
    q = q_ref[...]
    k = k_ref[...]
    d_k = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / jnp.sqrt(jnp.float32(d_k))
    row_max = jnp.max(scores, axis=-1, keepdims=True)
    exp_scores = jnp.exp(scores - row_max)
    o_ref[...] = exp_scores / jnp.sum(exp_scores, axis=-1, keepdims=True)


def pallas_qk_softmax(q: jax.Array, k: 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(
        _qk_softmax_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, seq_len), 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)),
        ],
        out_specs=pl.BlockSpec((block_q, seq_len), lambda i: (i, 0)),
    )(q, k)


pallas_kernel = pallas_qk_softmax
task_name = "qk_softmax"
input_shapes = [(256, 64), (256, 64)]
category = "attention_component"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused QK^T + Softmax via Pallas.

Provenance: jax-ml/jax flash_attention.py attention score computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _qk_softmax_kernel(q_ref, k_ref, o_ref):
    q = q_ref[...]
    k = k_ref[...]
    d_k = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / jnp.sqrt(jnp.float32(d_k))
    row_max = jnp.max(scores, axis=-1, keepdims=True)
    exp_scores = jnp.exp(scores - row_max)
    o_ref[...] = exp_scores / jnp.sum(exp_scores, axis=-1, keepdims=True)


def pallas_qk_softmax(q: jax.Array, k: jax.Array) -> jax.Array:
    seq_len, d_model = q.shape
    block_q = min(128, seq_len)
    while block_q * d_model > 16384 and block_q > 1:
        block_q //= 2
    grid_size = seq_len // block_q

    return pl.pallas_call(
        _qk_softmax_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, seq_len), 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)),
        ],
        out_specs=pl.BlockSpec((block_q, seq_len), lambda i: (i, 0)),
    )(q, k)


pallas_kernel = pallas_qk_softmax
task_name = "qk_softmax"
input_shapes = [(256, 64), (256, 64)]
category = "attention_component"
level = 2

Diff of GPU fix: --- a/qk_softmax.py
+++ b/qk_softmax.py
@@ -26,6 +26,8 @@
 def pallas_qk_softmax(q: jax.Array, k: jax.Array) -> jax.Array:
     seq_len, d_model = q.shape
     block_q = min(128, seq_len)
+    while block_q * d_model > 16384 and block_q > 1:
+        block_q //= 2
     grid_size = seq_len // block_q
 
     return pl.pallas_call(

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