Kernel: fused_softmax_cross_entropy (Level 2, category: loss_fusion)
Input shapes: [[1024, 512], [1024, 512]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 1235.7604 ms; JAX XLA compiled runtime: 0.0956 ms; speedup vs compiled baseline: 0.0001

JAX reference implementation:
@jax.jit
def jax_fused_softmax_cross_entropy(logits: jax.Array, labels: jax.Array) -> jax.Array:
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    return -jnp.sum(labels * log_probs, axis=-1)

Original Pallas kernel (TPU-oriented): """Level 2: Fused Softmax + Cross-Entropy Loss via Pallas.

Provenance: openxla/tokamax linear_softmax_cross_entropy_loss pattern
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _fused_ce_kernel(logits_ref, labels_ref, o_ref):
    logits = logits_ref[...]
    labels = labels_ref[...]
    row_max = jnp.max(logits, axis=-1, keepdims=True)
    shifted = logits - row_max
    log_sum_exp = jnp.log(jnp.sum(jnp.exp(shifted), axis=-1, keepdims=True))
    log_probs = shifted - log_sum_exp
    o_ref[...] = -jnp.sum(labels * log_probs, axis=-1)


def pallas_fused_softmax_cross_entropy(
    logits: jax.Array, labels: jax.Array
) -> jax.Array:
    n_rows = logits.shape[0]
    n_cols = logits.shape[1]
    block_rows = min(128, n_rows)
    grid_size = n_rows // block_rows

    return pl.pallas_call(
        _fused_ce_kernel,
        out_shape=jax.ShapeDtypeStruct((n_rows,), logits.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,), lambda i: (i,)),
    )(logits, labels)


pallas_kernel = pallas_fused_softmax_cross_entropy
task_name = "fused_softmax_cross_entropy"
input_shapes = [(1024, 512), (1024, 512)]
category = "loss_fusion"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused Softmax + Cross-Entropy Loss via Pallas.

Provenance: openxla/tokamax linear_softmax_cross_entropy_loss pattern
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _fused_ce_kernel(logits_ref, labels_ref, o_ref):
    logits = logits_ref[...]
    labels = labels_ref[...]
    row_max = jnp.max(logits, axis=-1, keepdims=True)
    shifted = logits - row_max
    log_sum_exp = jnp.log(jnp.sum(jnp.exp(shifted), axis=-1, keepdims=True))
    log_probs = shifted - log_sum_exp
    o_ref[...] = -jnp.sum(labels * log_probs, axis=-1)


def pallas_fused_softmax_cross_entropy(
    logits: jax.Array, labels: jax.Array
) -> jax.Array:
    n_rows = logits.shape[0]
    n_cols = logits.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(
        _fused_ce_kernel,
        out_shape=jax.ShapeDtypeStruct((n_rows,), logits.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,), lambda i: (i,)),
    )(logits, labels)


pallas_kernel = pallas_fused_softmax_cross_entropy
task_name = "fused_softmax_cross_entropy"
input_shapes = [(1024, 512), (1024, 512)]
category = "loss_fusion"
level = 2

Diff of GPU fix: --- a/fused_softmax_cross_entropy.py
+++ b/fused_softmax_cross_entropy.py
@@ -29,6 +29,8 @@
     n_rows = logits.shape[0]
     n_cols = logits.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[1024,512] b:f32[1024,512]. let
    c:f32[1024] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=32), Blocked(block_size=512))), BlockMapping(block_shape=(Blocked(block_size=32), Blocked(block_size=512))), BlockMapping(block_shape=(Blocked(block_size=32),))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[32,512]} e:Ref<default>{f32[32,512]} f:Ref<default>{f32[32]}. let
          g:f32[32,512] <- d[...]
          h:f32[32,512] <- e[...]
          i:f32[32] = reduce_max[axes=(1,)] g
          j:f32[32,1] = broadcast_in_dim[broadcast_dimensions=(0,)] i
          k:f32[32,512] = sub g j
          l:f32[32,512] = exp k
          m:f32[32] = reduce_sum[axes=(1,) out_sharding=None] l
          n:f32[32,1] = broadcast_in_dim[broadcast_dimensions=(0,)] m
          o:f32[32,1] = log n
          p:f32[32,512] = sub k o
          q:f32[32,512] = mul h p
          r:f32[32] = reduce_sum[axes=(1,) out_sharding=None] q
          s:f32[32] = neg r
          t:f32[32], f[...] <- f[...], s
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[1024]),)
    ] a b
  in (c,) }