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

JAX reference implementation:
@jax.jit
def jax_swiglu(x: jax.Array, w_gate: jax.Array, w_up: jax.Array) -> jax.Array:
    gate = jax.nn.silu(x @ w_gate)
    up = x @ w_up
    return gate * up

Original Pallas kernel (TPU-oriented): """Level 2: Fused SwiGLU activation via Pallas.

SwiGLU: gate = silu(x @ W_gate), up = x @ W_up, output = gate * up.
Demonstrates: multi-input fusion, gated activation, silu transcendental.
Inspired by pallas-forge's SwiGLU kernel.
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _swiglu_kernel(x_ref, w_gate_ref, w_up_ref, o_ref):
    x = x_ref[...]
    gate = x @ w_gate_ref[...]
    gate = gate / (1.0 + jnp.exp(-gate))  # silu
    up = x @ w_up_ref[...]
    o_ref[...] = gate * up


def pallas_swiglu(
    x: jax.Array, w_gate: jax.Array, w_up: jax.Array
) -> jax.Array:
    m, k = x.shape
    _, n = w_gate.shape
    bm = min(256, m)

    return pl.pallas_call(
        _swiglu_kernel,
        out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
        grid=(m // bm,),
        in_specs=[
            pl.BlockSpec((bm, k), lambda i: (i, 0)),
            pl.BlockSpec((k, n), lambda i: (0, 0)),
            pl.BlockSpec((k, n), lambda i: (0, 0)),
        ],
        out_specs=pl.BlockSpec((bm, n), lambda i: (i, 0)),
    )(x, w_gate, w_up)


pallas_kernel = pallas_swiglu
task_name = "swiglu"
input_shapes = [(512, 1024), (1024, 2048), (1024, 2048)]
category = "mlp_fusion"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused SwiGLU activation via Pallas.

SwiGLU: gate = silu(x @ W_gate), up = x @ W_up, output = gate * up.
Demonstrates: multi-input fusion, gated activation, silu transcendental.
Inspired by pallas-forge's SwiGLU kernel.
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _swiglu_kernel(x_ref, w_gate_ref, w_up_ref, o_ref):
    x = x_ref[...]
    gate = x @ w_gate_ref[...]
    gate = gate / (1.0 + jnp.exp(-gate))  # silu
    up = x @ w_up_ref[...]
    o_ref[...] = gate * up


def pallas_swiglu(
    x: jax.Array, w_gate: jax.Array, w_up: jax.Array
) -> jax.Array:
    m, k = x.shape
    _, n = w_gate.shape
    bm = min(16, m)
    bn = min(16, n)

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


pallas_kernel = pallas_swiglu
task_name = "swiglu"
input_shapes = [(512, 1024), (1024, 2048), (1024, 2048)]
category = "mlp_fusion"
level = 2

Diff of GPU fix: --- a/swiglu.py
+++ b/swiglu.py
@@ -28,18 +28,19 @@
 ) -> jax.Array:
     m, k = x.shape
     _, n = w_gate.shape
-    bm = min(256, m)
+    bm = min(16, m)
+    bn = min(16, n)
 
     return pl.pallas_call(
         _swiglu_kernel,
         out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-        grid=(m // bm,),
+        grid=(m // bm, n // bn),
         in_specs=[
-            pl.BlockSpec((bm, k), lambda i: (i, 0)),
-            pl.BlockSpec((k, n), lambda i: (0, 0)),
-            pl.BlockSpec((k, n), lambda i: (0, 0)),
+            pl.BlockSpec((bm, k), lambda i, j: (i, 0)),
+            pl.BlockSpec((k, bn), lambda i, j: (0, j)),
+            pl.BlockSpec((k, bn), lambda i, j: (0, j)),
         ],
-        out_specs=pl.BlockSpec((bm, n), lambda i: (i, 0)),
+        out_specs=pl.BlockSpec((bm, bn), lambda i, j: (i, j)),
     )(x, w_gate, w_up)
 
 

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