Quickstart: GPU

Quickstart: GPU#

This quickstart shows you how to write Pallas kernels for NVIDIA GPUs using Mosaic GPU. The examples target Hopper (H100) GPUs, but the core concepts - memory spaces, grids, pipelining - apply to all supported GPU generations.

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

The programming model#

In Mosaic GPU, each Pallas “thread” corresponds to a warpgroup (4 warps = 128 CUDA threads). You write straight-line code that operates on arrays; the warpgroup executes it in lockstep, so there are no individual CUDA threads to manage.

GPU memory spaces#

Pallas kernels access memory through Refs, JAX’s mutable array references. On GPU, each Ref lives in a specific memory space:

  • GMEM (global memory / HBM): large, slow. Used for kernel inputs and outputs.

  • SMEM (shared memory): small, fast, per-SM. Shared between threads within a block, used to stage data for Tensor Core ops.

  • TMEM (tensor memory): fast, per-SM. Used for Tensor Core ops (only available on Blackwell GPUs).

The typical data flow for Tensor Core work on Hopper is: GMEM → SMEM → Tensor Cores → registers → SMEM → GMEM (on Blackwell, TMEM replaces SMEM for staging Tensor Core inputs and outputs). Simpler ALU work (like elementwise operations) bypasses SMEM and Tensor Cores entirely: GMEM → registers → GMEM. Shared memory is primarily needed if threads within a block must exchange or reuse data.

Your first kernel#

Here’s a kernel that fills an output array with a constant:

@plgpu.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
def fill_42(o_ref):
  o_ref[...] = jnp.full_like(o_ref, 42.0)

result = fill_42()  # [42.0, 42.0, ...]

plgpu.kernel allocates the output buffer, runs the computation on-device, and returns a JAX array.

Parallel work with a grid#

To process a larger array, add a grid. Each grid point runs as a separate CUDA block in parallel across SMs:

@plgpu.kernel(
    out_type=jax.ShapeDtypeStruct((1024,), jnp.float32),
    grid=(8,),
    grid_names=('i',),
)
def iota(o_ref):
  i = jax.lax.axis_index('i')
  o_ref[pl.ds(i * 128, 128)] = jnp.arange(128, dtype=jnp.float32) + i * 128

result = iota()  # [0.0, 1.0, ..., 1023.0]

jax.lax.axis_index(name) tells you which grid block you’re in - use it to determine your slice of the output. pl.ds(start, size) creates a dynamic slice of the given size, and is equivalent to start:start+size.

Pipelining: a matmul kernel#

The examples above don’t need pipelining, but anything that hits the Tensor Cores - matmuls, attention, etc. - should overlap GMEM↔SMEM transfers with compute. Otherwise, the Tensor Cores sit idle waiting for data.

plgpu.emit_pipeline handles this. It takes:

  • A sequential grid: how many pipeline steps to run (typically over the contraction dimension).

  • BlockSpecs: how to slice your inputs per step.

  • A body function: the computation to run at each step.

The outer plgpu.kernel grid handles the parallel work, and emit_pipeline handles the sequential reduction within each block.

Here’s a matmul for Hopper GPUs:

Note

This example uses wgmma, which is specific to Hopper GPUs. For Blackwell, see Blackwell Matrix Multiplication, which uses tcgen05 instead.

def matmul(a, b, tile_m=128, tile_n=128, tile_k=64, out_dtype=jnp.float16):
  m, k = a.shape
  _, n = b.shape

  @plgpu.kernel(
      out_type=jax.ShapeDtypeStruct((m, n), out_dtype),
      scratch_types=dict(
          o_smem=plgpu.SMEM((tile_m, tile_n), out_dtype),
          acc=plgpu.ACC((tile_m, tile_n), jnp.float32),
      ),
      grid=(m // tile_m, n // tile_n),
      grid_names=('m', 'n'),
  )
  def kernel(a_gmem, b_gmem, o_gmem, o_smem, acc):
    pid_m = jax.lax.axis_index('m')
    pid_n = jax.lax.axis_index('n')

    def body(_, a_smem, b_smem):
      plgpu.wgmma(acc, a_smem, b_smem)
      plgpu.wgmma_wait(1)  # Keep one wgmma in flight.

    plgpu.emit_pipeline(
        body,
        grid=(k // tile_k,),
        in_specs=[
            plgpu.BlockSpec(
                (tile_m, tile_k), lambda ki: (pid_m, ki), delay_release=1
            ),
            plgpu.BlockSpec(
                (tile_k, tile_n), lambda ki: (ki, pid_n), delay_release=1
            ),
        ],
        max_concurrent_steps=2,
    )(a_gmem, b_gmem)

    # Drain: move the accumulated result to GMEM via SMEM.
    o_smem[...] = acc[...].astype(out_dtype)
    plgpu.commit_smem()  # Make the SMEM write visible to the TMA engine.
    plgpu.copy_smem_to_gmem(
        o_smem,
        o_gmem.at[pl.ds(pid_m * tile_m, tile_m),
                  pl.ds(pid_n * tile_n, tile_n)],
    )
    plgpu.wait_smem_to_gmem(0)  # Wait for all copies to finish.

  return kernel(a, b)

A few things to note:

  • The parallel grid (plgpu.kernel(..., grid=...)) maps one CUDA block per output tile.

  • The sequential grid (emit_pipeline(..., grid=...)) is the pipeline loop over the K dimension.

  • scratch_types specifies temporary memory allocations needed for each parallel grid point, such as shared memory buffers (SMEM) or Tensor Core accumulators (ACC). Each key in the dictionary is passed as a keyword argument to the kernel function.

  • plgpu.ACC is the Tensor Core accumulator. wgmma accumulates into it asynchronously.

  • delay_release=1 tells the pipeline to keep an extra buffer alive. Without this, the pipeline releases the corresponding buffer immediately, allowing the next iteration to overwrite input data while wgmma is still reading it.

What’s next#