Kernel: matmul_silu (Level 2, category: matmul_activation)
Input shapes: [[1024, 1024], [1024, 1024]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 707.243 ms; JAX XLA compiled runtime: 0.1617 ms; speedup vs compiled baseline: 0.0002

JAX reference implementation:
@jax.jit
def jax_matmul_silu(x: jax.Array, w: jax.Array) -> jax.Array:
    return jax.nn.silu(x @ w)

Original Pallas kernel (TPU-oriented): """Level 2: Fused MatMul + SiLU via Pallas.

Provenance: openxla/tokamax gated_linear_unit uses SiLU gate path
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _matmul_silu_kernel(x_ref, w_ref, o_ref):
    z = x_ref[...] @ w_ref[...]
    o_ref[...] = z / (1.0 + jnp.exp(-z))


def pallas_matmul_silu(x: jax.Array, w: jax.Array) -> jax.Array:
    m, k = x.shape
    _, n = w.shape
    bm = min(512, m)
    bn = min(512, n)
    grid = (m // bm, n // bn)

    return pl.pallas_call(
        _matmul_silu_kernel,
        out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
        grid=grid,
        in_specs=[
            pl.BlockSpec((bm, k), lambda i, j: (i, 0)),
            pl.BlockSpec((k, bn), lambda i, j: (0, j)),
        ],
        out_specs=pl.BlockSpec((bm, bn), lambda i, j: (i, j)),
    )(x, w)


pallas_kernel = pallas_matmul_silu
task_name = "matmul_silu"
input_shapes = [(1024, 1024), (1024, 1024)]
category = "matmul_activation"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused MatMul + SiLU via Pallas.

Provenance: openxla/tokamax gated_linear_unit uses SiLU gate path
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _matmul_silu_kernel(x_ref, w_ref, o_ref):
    z = x_ref[...] @ w_ref[...]
    o_ref[...] = z / (1.0 + jnp.exp(-z))


def pallas_matmul_silu(x: jax.Array, w: jax.Array) -> jax.Array:
    m, k = x.shape
    _, n = w.shape
    bm = min(16, m)
    bn = min(16, n)
    grid = (m // bm, n // bn)

    return pl.pallas_call(
        _matmul_silu_kernel,
        out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
        grid=grid,
        in_specs=[
            pl.BlockSpec((bm, k), lambda i, j: (i, 0)),
            pl.BlockSpec((k, bn), lambda i, j: (0, j)),
        ],
        out_specs=pl.BlockSpec((bm, bn), lambda i, j: (i, j)),
    )(x, w)


pallas_kernel = pallas_matmul_silu
task_name = "matmul_silu"
input_shapes = [(1024, 1024), (1024, 1024)]
category = "matmul_activation"
level = 2

Diff of GPU fix: --- a/matmul_silu.py
+++ b/matmul_silu.py
@@ -21,8 +21,8 @@
 def pallas_matmul_silu(x: jax.Array, w: jax.Array) -> jax.Array:
     m, k = x.shape
     _, n = w.shape
-    bm = min(512, m)
-    bn = min(512, n)
+    bm = min(16, m)
+    bn = min(16, n)
     grid = (m // bm, n // bn)
 
     return pl.pallas_call(

Jaxpr IR: { lambda ; a:f32[1024,1024] b:f32[1024,1024]. let
    c:f32[1024,1024] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(64, 64), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024))), BlockMapping(block_shape=(Blocked(block_size=1024), Blocked(block_size=16))), BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=16)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[16,1024]} e:Ref<default>{f32[1024,16]}
          f:Ref<default>{f32[16,16]}. let
          g:f32[16,1024] <- d[...]
          h:f32[1024,16] <- e[...]
          i:f32[16,16] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=float32
          ] g h
          j:f32[16,16] = neg i
          k:f32[16,16] = exp j
          l:f32[16,16] = add 1.0:f32[] k
          m:f32[16,16] = div i l
          n:f32[16,16], f[...] <- f[...], m
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[1024,1024]),)
    ] a b
  in (c,) }