Quickstart: TPU

This quickstart covers the basics of writing Pallas kernels for TPU: memory spaces, pl.kernel, and pipelining.

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

TPU memory spaces

Pallas kernels interact with memory through Refs – JAX’s mutable array references. On TPU, each Ref lives in a specific memory space:

• HBM: large but slow. Kernel inputs and outputs live here.

• VMEM: small but fast. Computation happens here.

You can’t compute directly on HBM Refs. Data has to be copied to VMEM first, and results copied back.

Your first kernel

Here’s a kernel that fills an output array with a constant. We allocate a temporary VMEM buffer, write the constant there, and copy it to the output HBM Ref.

@pl.kernel(
    out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
    mesh=pltpu.TensorCoreMesh(axis_name='core'),
    scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),
)
def fill_42(o_ref, o_vmem):
  # Compute in VMEM
  o_vmem[...] = jnp.full_like(o_vmem, 42.0)

  # VMEM → HBM (blocks until the transfer completes)
  pltpu.sync_copy(o_vmem, o_ref)

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

pl.kernel wraps a kernel function to make it callable with standard JAX arrays. It handles allocating the necessary Refs and executing the kernel on-device:

• out_type allocates an output Ref in HBM for the result.

• scratch_types allocates the VMEM buffers (like o_vmem) that get passed as extra arguments to the kernel.

Note

If your chip has multiple TensorCores (e.g. TPU v5p has 2), the kernel runs on all of them. As written, both cores will redundantly perform the exact same computation. In a real kernel, you would distribute work across cores either manually or by using pipelining.

Distributing work across TensorCores

To process a larger array, we can distribute work across multiple TensorCores manually. Each core runs a separate invocation of the kernel.

def iota() -> jax.Array:
  tpu_info = pltpu.get_tpu_info()

  @pl.kernel(
      out_type=jax.ShapeDtypeStruct((128 * tpu_info.num_cores,), jnp.float32),
      mesh=pltpu.TensorCoreMesh(axis_name='core'),
      scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),
  )
  def kernel(o_ref, o_vmem):
    i = jax.lax.axis_index('core')

    # Compute our chunk in VMEM
    o_vmem[...] = jnp.arange(128, dtype=jnp.float32) + i * 128

    # Copy back to our slice of HBM
    pltpu.sync_copy(o_vmem, o_ref.at[pl.ds(i * 128, 128)])

  return kernel()

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

jax.lax.axis_index('core') tells you which TensorCore you’re on – use it to determine your slice of the output.

Making it fast: pipelining

The examples above work, but are inefficient. They copy data in, compute, then copy out – each sync_copy blocks until the transfer completes. The TensorCores sit idle while memory is moving, and vice versa.

Pipelining fixes this by breaking the work into blocks and overlapping memory transfers with compute. While the TensorCores process block N, the DMA engine is already fetching block N+1.

pltpu.emit_pipeline handles the double-buffering machinery for you. It takes:

• A grid: how many pipeline blocks to process.

• BlockSpecs: how to slice the arrays into blocks.

• A body function: the computation to run for each block.

Here’s an elementwise addition kernel using pipelining:

def add_matrices_pipelined(x: jax.Array, y: jax.Array) -> jax.Array:
  @pl.kernel(
      out_type=jax.ShapeDtypeStruct.like(x),
      mesh=pltpu.TensorCoreMesh(axis_name='core'),
  )
  def kernel(x_hbm, y_hbm, o_hbm):
    def add_body(x_vmem, y_vmem, o_vmem):
      o_vmem[...] = x_vmem[...] + y_vmem[...]

    pltpu.emit_pipeline(
        add_body,
        grid=(x_hbm.shape[0] // 128, x_hbm.shape[1] // 128),
        in_specs=[
            pl.BlockSpec((128, 128), lambda i, j: (i, j)),
            pl.BlockSpec((128, 128), lambda i, j: (i, j)),
        ],
        out_specs=pl.BlockSpec((128, 128), lambda i, j: (i, j)),
        core_axis_name='core',
        dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),
    )(x_hbm, y_hbm, o_hbm)
  return kernel(x, y)

x = 2 * jnp.ones((256, 256))
y = 3 * jnp.ones((256, 256))
add_matrices_pipelined(x, y)
# Array([[5., 5., 5., ..., 5., 5., 5.],
#        [5., 5., 5., ..., 5., 5., 5.],
#        [5., 5., 5., ..., 5., 5., 5.],
#        ...,
#        [5., 5., 5., ..., 5., 5., 5.],
#        [5., 5., 5., ..., 5., 5., 5.],
#        [5., 5., 5., ..., 5., 5., 5.]], dtype=float32)

Notice the body is just the compute – no sync_copy, no scratch allocation. emit_pipeline takes care of the memory movement and buffer management.

We also passed core_axis_name='core' and dimension_semantics. This tells the pipeline how to map the grid to the hardware:

• pltpu.PARALLEL means the first dimension of the grid is distributed across the TensorCores, splitting the work automatically.

• pltpu.ARBITRARY indicates that Pallas cannot make assumptions about data independence along this dimension; it therefore cannot be parallelized across cores and is instead executed sequentially on a single core.

What’s next

• TPU Details – the full picture of TPU memory spaces and supported ops

• TPU Pipelining – deeper dive into pipelining, reductions, and accumulation

• Matrix Multiplication – putting it all together in a full kernel