The scalar world

pl.dsA dynamic slice of a ref inside the kernel: reads only the chunk, at an offset that can depend on runtime values.taught in /l/pallas →(start, size) slices a ref dynamically from inside the kernel body, at a start position computed at run time rather than fixed when the kernel is traced. Call it on a ref and you read only that chunk into a value, not the whole block. The compiler does not need to know the start position ahead of time, only that the size stays fixed and the memory space supports it, which is what separates it from ordinary Python slicing: a Python slice needs its bounds known when the kernel is traced, and pl.ds is exactly the escape hatch for when they are not.

That start position often comes from the same place a scalar comes from anywhere else in a kernel: an SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → ref holding a value read earlier in the kernel body, a per-sequence length or a computed row offset. Combine the two ideas and a kernel can walk a ragged sequence by reading its true length out of SMEM once, then using pl.dsA dynamic slice of a ref inside the kernel: reads only the chunk, at an offset that can depend on runtime values.taught in /l/pallas → with that length to read only the valid rows, never touching the padded tail at all inside the loop that matters.

This matters most at the edges of the grid itself, not just inside a single dynamic read. A grid built for a block size that does not evenly divide the array leaves a partial block on the last step: the block is still shaped for the full block size, but only part of it holds real data, and the rest holds whatever was left over from stale 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 → on a previous step. The fix is a mask, and it starts with an iota, generating a per-lane index inside the block and comparing it against the array's true bound to decide which lanes are real and which are leftover.

Multiply the loaded values by that mask before you use them, or select against a padding value, and the partial block behaves correctly no matter how far past the true bound the physical block extends. Do the masking after a reduction has already mixed real and stale values together, rather than before, and the mask no longer has anything valid left to select between: the damage is already folded into the sum. Get the comparison direction wrong, or compare the iota against the block's size instead of the array's true size, and the kernel reads or writes past the array's real edge without raising an error to tell you so.

None of this masking machinery is needed when the array happens to divide evenly by the block size, which is exactly the common path section one's divisibility rule already covers. The mask matters specifically for the boundary case: an array whose last dimension does not divide cleanly by the block, where the grid still has to walk in fixed-size steps, and the very last step is the one that runs past the edge.

pl.dsA dynamic slice of a ref inside the kernel: reads only the chunk, at an offset that can depend on runtime values.taught in /l/pallas → works on any ref, not only the array carrying the main algebra of the kernel. It applies just as well to a 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 → block the pipeline staged, to scratch memory you allocated yourself, or to a start position computed from a value you read out of SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → a few lines earlier. The pattern is the same regardless of which memory space the ref sits in; only the cost of the read or write it performs changes.

def sum_ragged_kernel(x_ref, o_ref, *, n_valid, rows):
    i = pl.program_id(0)
    x = x_ref[...]
    idx = i * rows + jax.lax.broadcasted_iota(jnp.int32, x.shape, 0)
    x = jnp.where(idx < n_valid, x, 0.0)
    @pl.when(i == 0)
    def _():
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.sum(x, axis=0, keepdims=True)

Every 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 → so far has an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → that depends only on the grid coordinate: something the compiler can evaluate before it ever sees real data. PrefetchScalarGridSpec breaks that constraint. It copies scalar arrays into SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → before the pipeline starts running, and makes those SMEM values available to the index map itself, so which block gets loaded on step 5 can depend on what a specific input value actually is, not just on which step number 5 happens to be.

That ordering is the whole reason this works rather than deadlocking on itself. An index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → that reads a value out of a ref the pipeline has not staged yet would have nothing to read; PrefetchScalarGridSpec solves that by moving the scalar arrays into SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → before the pipeline's normal staging even begins, as a separate step that happens first. By the time step 0 of the main pipeline runs, every scalar the schedule depends on is already sitting in SMEM, not still in flight, so the index map's read of it is never racing the copy that produced it.

That is the mechanism, not a special case bolted on top of it. It is exactly what splash attention and ragged paged attention use to route each query to the right block of keys and values, where the routing decision is data (which page a token belongs to, how long a sequence really runs) rather than a static grid pattern known ahead of time. A fixed grid walk has no way to express "load whichever block this particular query needs"; a data-dependent index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → is the only way to say that at all.

The scalar arrays involved stay small on purpose. Page tables, sequence lengths, and block indices are the kind of thing that fits in SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → cheaply; the actual keys and values, the large data the routing decision points at, still move through the ordinary 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 → pipeline once the index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → has decided which block of them to fetch. Scalar prefetch decides where to look. It is not a replacement for the pipeline that does the looking.

The tradeoff is set-up complexity against schedule flexibility. An ordinary 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 →'s index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → is pure and needs no extra wiring around it, while PrefetchScalarGridSpec needs the scalar arrays carved out as a separate argument group before the grid spec is even built. That extra step is exactly what buys back the ability to route on data instead of on a fixed pattern known ahead of time.

None of this is free at the framework level, either. The prefetch step runs before the main pipeline begins, so scalar arrays large enough to matter still cost some setup time before the first block of the real computation starts. The pattern earns its keep specifically when the alternative is a schedule that cannot be expressed at all, not when the alternative is a marginally simpler static index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → that would have worked just as well.

def permute_kernel(perm_ref, x_ref, o_ref):
    o_ref[...] = x_ref[...]

def permute_blocks(x, perm, bm=64):
    m, d = x.shape
    grid_spec = pltpu.PrefetchScalarGridSpec(
        num_scalar_prefetch=1,
        grid=(m // bm,),
        in_specs=[pl.BlockSpec((bm, d), lambda i, perm: (perm[i], 0))],
        out_specs=pl.BlockSpec((bm, d), lambda i, perm: (i, 0)),
    )
    return pl.pallas_call(
        permute_kernel,
        grid_spec=grid_spec,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        interpret=True,
    )(perm, x)