What does a BlockSpec do in Pallas?

It answers one question per input and per output: for this invocation of the kernel, which piece of this array do I get? Its own docstring says it specifies how an array should be sliced for each invocation of a kernel. Concretely it carries a block_shape, which is the size of that piece and therefore the shape of the ref your kernel body receives, an index_map, which is a Python function from the grid coordinates to a tuple of indices, and a memory_space saying where the piece should live while the kernel reads it. A fourth field, pipeline_mode, lets you ask for a specific number of buffers. You hand one BlockSpec per input through in_specs and one per output through out_specs. Everything else about a Pallas kernel is downstream: the grid says how many times the body runs, the BlockSpecs say what it runs on, and together they are the entire memory plan, because there is no separate place to say which data moves into fast memory and when.

Does a Pallas index_map return element indices or block indices?

Block indices, in the default blocked mode, and this is the single most common first-kernel mistake. For a block of 128 rows the map returns 0, then 1, then 2, not 0, then 128, then 256. Pallas does the multiply itself: the start element on an axis is the block size times the index your map returned. The code is a match statement inside BlockMapping in jax/_src/pallas/core.py, and it has three live branches. A Blocked dimension, which is what a bare integer in your block shape becomes, multiplies. A Squeezed dimension and an Element dimension pass the index through untouched. Anything else raises. The consequence worth holding on to is that in blocked mode the only reachable start positions are multiples of the block size, so a block that starts halfway through cannot be expressed. The Element block-dimension kind exists for exactly that case, taking element indices directly plus an optional padding pair, and the documentation notes it is currently supported only on TPUs.

What happens when a Pallas block shape does not divide the array?

The kernel still runs on a full-sized block, and the elements past the end are padding. The grid extent along an axis is the ceiling division of the array dimension by the block dimension, so the partial blocks are included rather than skipped, and every invocation receives a ref of exactly the block shape, because shapes inside a compiled kernel are static and a smaller final block would be a second kernel. The documentation is explicit that the out-of-bounds elements are padded on input and discarded on output, that the padding values are unspecified, and that you should assume they are garbage. Discarded on output is what saves an elementwise kernel: whatever you write into the out-of-bounds part is thrown away, so it is correct on a ragged shape without any work from you. A kernel that reduces over its block is a different story, because summing garbage gives garbage and nothing raises. That is why masking against program_id is standard in real kernels. In interpret mode the padding is NaN on purpose, so a reduction over padding shows up rather than producing a plausible wrong number.

What is the difference between the grid and the block shape in Pallas?

The grid is how many times, the block shape is how much. The grid is a tuple of integers that behaves exactly like a nest of for loops: a grid of (n, m) runs the kernel n times m times, and inside the body pl.program_id(axis) tells you where you are on one axis while pl.num_programs(axis) gives that axis its size. Nothing in the grid mentions your data. The block shape is the size of the slice one invocation receives, and it belongs to a BlockSpec, one per array. The two are linked in only one direction and it is worth getting the direction right: you choose the block shape and the grid follows, because the grid extent along an axis is the ceiling division of the array dimension by the block dimension. The exception is a grid axis that does not correspond to data at all, such as a reduction axis you sweep while the index map holds one coordinate fixed. That is how operand reuse is expressed, and it is why a matmul kernel has three BlockSpecs with three different index maps.

Why must a Pallas block shape be a multiple of 8 and 128 on a TPU?

Because those are the dimensions of the native register tile the vector unit works on, so the last two axes of your block are the ones that have to line up with hardware. The documentation states the rule directly: on TPU only blocks with rank at least one are supported, and the last two dimensions of the block shape must either equal the corresponding dimension of the overall array or be divisible by 8 and 128 respectively. For a rank-one block there is a separate rule: the block dimension must equal the array dimension, or be a multiple of 1024, or be a power of two and at least 128 times 32 divided by the bit width of the dtype, which is the only place the element type enters the constraint. The other backends draw different lines for their own hardware reasons. Mosaic GPU leaves block sizes unrestricted except that the minormost array dimension must be a multiple of 16 bytes. Triton leaves block sizes unrestricted but requires every operation, loads and stores included, to work on a power-of-two size. The practical consequence is that a BlockSpec which runs fine in interpret mode on a laptop can still fail to lower on a chip, so check the block shape against the tile first.

BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces

A Pallas kernel never sees your array. It sees one block of it, chosen by a Python function you wrote and a loop the compiler ran, and this page reads the arithmetic that connects the two at a pinned commit of jax-ml/jax. The grid is a loop nest and the kernel is its body, so the kernel runs prod(grid) times. A BlockSpec is a block shape plus an index map, and the index map returns block indices rather than element indices: the start of the slice is the block size times the number your map returned, which is one match statement in core.py and nothing more. From there the page walks the five kinds of block dimension and what each one asks your map to return, the three defaults and why omitting a BlockSpec entirely does not mean the whole array but an unblocked ref in whichever memory the backend chose, the two memory-space vocabularies and where DEFAULT lands on a TPU, and the conversion that turns your closure into a traced jaxpr along with the three rules that conversion enforces. Then the honest parts: a block shape that does not divide the array runs anyway, on a full-sized block, reading padding whose values the documentation tells you to assume are garbage, discarded on output and not on input; the grid runs with the last axis fastest on the interpreter and in an order you may not assume on a chip; and a TPU wants the last two dimensions of your block to be multiples of 8 and 128. The signature exhibit ports the real block-index arithmetic, the ceiling division, the start indices and the padding rule, so you can dial an array shape against a block shape and read the slice the kernel would receive.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces — built to be understood, not skimmed.

Questions

What does a BlockSpec do in Pallas?: It answers one question per input and per output: for this invocation of the kernel, which piece of this array do I get? Its own docstring says it specifies how an array should be sliced for each invocation of a kernel. Concretely it carries a block_shape, which is the size of that piece and therefore the shape of the ref your kernel body receives, an index_map, which is a Python function from the grid coordinates to a tuple of indices, and a memory_space saying where the piece should live while the kernel reads it. A fourth field, pipeline_mode, lets you ask for a specific number of buffers. You hand one BlockSpec per input through in_specs and one per output through out_specs. Everything else about a Pallas kernel is downstream: the grid says how many times the body runs, the BlockSpecs say what it runs on, and together they are the entire memory plan, because there is no separate place to say which data moves into fast memory and when.
Does a Pallas index_map return element indices or block indices?: Block indices, in the default blocked mode, and this is the single most common first-kernel mistake. For a block of 128 rows the map returns 0, then 1, then 2, not 0, then 128, then 256. Pallas does the multiply itself: the start element on an axis is the block size times the index your map returned. The code is a match statement inside BlockMapping in jax/_src/pallas/core.py, and it has three live branches. A Blocked dimension, which is what a bare integer in your block shape becomes, multiplies. A Squeezed dimension and an Element dimension pass the index through untouched. Anything else raises. The consequence worth holding on to is that in blocked mode the only reachable start positions are multiples of the block size, so a block that starts halfway through cannot be expressed. The Element block-dimension kind exists for exactly that case, taking element indices directly plus an optional padding pair, and the documentation notes it is currently supported only on TPUs.
What happens when a Pallas block shape does not divide the array?: The kernel still runs on a full-sized block, and the elements past the end are padding. The grid extent along an axis is the ceiling division of the array dimension by the block dimension, so the partial blocks are included rather than skipped, and every invocation receives a ref of exactly the block shape, because shapes inside a compiled kernel are static and a smaller final block would be a second kernel. The documentation is explicit that the out-of-bounds elements are padded on input and discarded on output, that the padding values are unspecified, and that you should assume they are garbage. Discarded on output is what saves an elementwise kernel: whatever you write into the out-of-bounds part is thrown away, so it is correct on a ragged shape without any work from you. A kernel that reduces over its block is a different story, because summing garbage gives garbage and nothing raises. That is why masking against program_id is standard in real kernels. In interpret mode the padding is NaN on purpose, so a reduction over padding shows up rather than producing a plausible wrong number.
What is the difference between the grid and the block shape in Pallas?: The grid is how many times, the block shape is how much. The grid is a tuple of integers that behaves exactly like a nest of for loops: a grid of (n, m) runs the kernel n times m times, and inside the body pl.program_id(axis) tells you where you are on one axis while pl.num_programs(axis) gives that axis its size. Nothing in the grid mentions your data. The block shape is the size of the slice one invocation receives, and it belongs to a BlockSpec, one per array. The two are linked in only one direction and it is worth getting the direction right: you choose the block shape and the grid follows, because the grid extent along an axis is the ceiling division of the array dimension by the block dimension. The exception is a grid axis that does not correspond to data at all, such as a reduction axis you sweep while the index map holds one coordinate fixed. That is how operand reuse is expressed, and it is why a matmul kernel has three BlockSpecs with three different index maps.
Why must a Pallas block shape be a multiple of 8 and 128 on a TPU?: Because those are the dimensions of the native register tile the vector unit works on, so the last two axes of your block are the ones that have to line up with hardware. The documentation states the rule directly: on TPU only blocks with rank at least one are supported, and the last two dimensions of the block shape must either equal the corresponding dimension of the overall array or be divisible by 8 and 128 respectively. For a rank-one block there is a separate rule: the block dimension must equal the array dimension, or be a multiple of 1024, or be a power of two and at least 128 times 32 divided by the bit width of the dtype, which is the only place the element type enters the constraint. The other backends draw different lines for their own hardware reasons. Mosaic GPU leaves block sizes unrestricted except that the minormost array dimension must be a multiple of 16 bytes. Triton leaves block sizes unrestricted but requires every operation, loads and stores included, to work on a power-of-two size. The practical consequence is that a BlockSpec which runs fine in interpret mode on a laptop can still fail to lower on a chip, so check the block shape against the tile first.

BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces

Questions

Related explainers