Memory spaces and scratch

A block staged into VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → is the default, and the one you have already been using: the compiler copies the block described by your BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → into on-chip vector memory before your kernel body runs. Two other memory spaces exist for cases a plain block does not cover, and scratch memory sits alongside all three as workspace that belongs to none of your kernel's actual arguments.

SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu →, scalar memory, holds small scalar values: a sequence length, a flag, an index used to pick a dynamic slice later in the kernel. Pass these as SMEM refs instead of VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → refs and the compiler keeps them out of the vector datapath entirely, which is where a single integer does not belong in the first place; it never needed an (8, 128) tile to sit in. Reading a scalar out of SMEM inside the kernel body costs nothing like reading a full vector block, because there is no tile to assemble and no lane structure to respect.

ANYThe memory space that tells the compiler not to place a ref at all: a promise that the kernel will move the data itself with a manual DMA.taught in /l/pallas → tells the compiler not to place the ref at all: it is a promise that you will move the data yourself with a manual DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas →, the subject of the section after next. You reach for ANY when the access pattern the kernel needs cannot be expressed as a BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → grid walk, only as an explicit copy issued from inside the kernel body. Declaring a ref ANY is also how you tell the compiler that the automatic pipeline should not even try to stage it, which matters because staging it automatically is exactly what you are about to do by hand instead.

Scratch memory is different from all three: it is not tied to any of the kernel's inputs or outputs. You request it through scratch_shapesRequests per-invocation workspace (VMEM buffers, semaphores) that belongs to the kernel itself, not to any input or output, and persists across grid steps.taught in /l/pallas →, and it lives for the duration of the whole kernel invocation, persisting across grid steps rather than being re-allocated on each one. The pattern worth knowing by name is the f32 accumulator. Allocate scratch in f32 even when the kernel's actual output ref is bf16, accumulate every partial sum into that f32 scratch across the grid's reduction axis, and cast down to bf16 only once, on the final write.

The reason is ordinary floating point behavior, not a Pallas quirk. Accumulate directly in bf16 instead and every partial sum rounds to bf16 precision before the next one is added, and that rounding error compounds with every step of the reduction. A long enough reduction axis turns a rounding error too small to notice on step one into a visible error by the last step. Keeping the running total in f32 and only rounding down once removes every one of those intermediate rounding events, leaving exactly one.

def matmul_scratch_kernel(a_ref, b_ref, o_ref, acc_ref):
    k = pl.program_id(2)
    @pl.when(k == 0)
    def _():
        acc_ref[...] = jnp.zeros_like(acc_ref)
    acc_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
    @pl.when(k == pl.num_programs(2) - 1)
    def _():
        o_ref[...] = acc_ref[...].astype(o_ref.dtype)

def matmul_scratch(a, b, bm=128, bn=128, bk=128):
    m, k = a.shape
    _, n = b.shape
    return pl.pallas_call(
        matmul_scratch_kernel,
        grid=(m // bm, n // bn, k // bk),
        in_specs=[pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
                  pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j))],
        out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
        out_shape=jax.ShapeDtypeStruct((m, n), a.dtype),
        scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],
        interpret=True,
    )(a, b)

The same SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → mechanism is what makes a ragged batch kernel possible. Instead of padding every sequence in a batch to a common length and spending compute on the padding, the kernel reads each sequence's real length from SMEM and uses it to bound the work inside the kernel body, the same scalar-in-SMEM idea, just driving a loop bound instead of a mask. Measured on a ragged batch shaped 8x4096 where sequences average 41% of the padded length, that approach runs in 397.6 microseconds against 788.8 microseconds for the padded XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → baseline: 1.98x faster, on v6e-1.