What is splash attention and how is it different from flash attention?

Splash attention is the TPU Pallas attention kernel in JAX that treats the mask as loop structure rather than as a value. Flash attention tiles the sequence and runs an online softmax so the full sequence-by-sequence score matrix is never materialised, but every tile is still computed and the mask is applied to the scores afterwards, so a causal mask costs the same arithmetic as no mask at all. Splash attention adds a compile-time pass that cuts the mask into blocks the size of one grid step and labels each one: 0 for a block that is all zeros, 1 for a block with a mix, and 2 for a block that is entirely ones. The kernel reads that label at each grid position out of TPU scalar memory. A 0 means the body does not run and the key and value blocks are not fetched. A 2 means the kernel ignores the mask block it was handed, because on a block that is entirely ones the stored mask pointer is not pointing at anything meaningful. Only a 1 both runs and masks. For a plain causal mask at 8192 tokens with 128 by 128 blocks that is 2016 of 4096 blocks never computed.

What is MaskInfo in splash attention?

MaskInfo is the compiled form of the mask, produced by process_mask before the kernel is traced and passed into the Pallas call as scalar-prefetch operands. It carries block_mask, an array of 0, 1 and 2 labels with one entry per grid position; data_next, which for each grid position holds the index of the next non-empty key and value block so the pipeline fetches the right thing and does not re-fetch across a skipped block; mask_next, which points into partial_mask_blocks, the deduplicated stack of the blocks that were mixed; and q_sequence, the query indices. The three integer arrays live in TPU scalar memory, which is scarce, so they are downcast to int16 or int8 whenever their largest value fits. There is a large exception: if the mask is a single computable mask such as CausalMask, LocalMask or ChunkedCausalMask, then mask_next and partial_mask_blocks are returned as None and the kernel recomputes the mask from q_sequence and the mask function instead of loading any of it.

Does splash attention make every sparse mask faster?

No. The sparsity has to line up with the block grid. The classification asks two questions per block, chunk.any() and chunk.all(), so a block earns a skip only when every single one of its 128 by 128 entries is zero. A mask that is 90 percent zeros but scattered uniformly at random has a zero in almost every block and a one in almost every block, so every block comes back partial: nothing is skipped, the grid does not shrink, and because a random mask is a dense NumPy array rather than a computable one, the mixed blocks are now materialised and loaded as well. At 2048 tokens with 128 by 128 blocks that is 256 of 256 blocks still visited plus four mebibytes of partial mask blocks that flash attention would never have touched. Causal, local window and chunked masks work because their zeros arrive in aligned rectangles.

What does grid shrinking do in splash attention?

Skipping a block saves the arithmetic but the grid still steps over it, so a mask with a narrow band still launches a grid as wide as the sequence. The shrink pass fixes that. For each row of the block grid it takes the indices of the non-empty columns, pads every row to the width of the longest one so the result is rectangular, and rewrites block_mask, data_next and mask_next to hold only those columns. The kernel then takes its grid width from the last dimension of data_next. For a local mask with a 1024-token left window at 8192 tokens, the grid width drops from 64 to 9. For a plain causal mask it does not drop at all, because the last query block genuinely needs every key block, and the padding to the longest row keeps the full width. That is why the shrink is the thing that makes windowed attention fast and skipping alone is what makes causal attention fast.

What does ring attention add on top of the splash kernel?

It shards the keys and values across a ring of devices so a sequence too long for one device still runs, and it reuses the splash forward pass unchanged for each shard. The forward is a scan over the ring size. At each step the device sends its current key and value shards to its neighbour with a collective permute while the local splash call runs on the shards it already has, so the transfer hides behind the compute. The MaskInfo is sliced along the key axis to match whichever shard is local this step. Because each step produces its own running maximum and normaliser, the outputs are combined with the same rescaling flash attention uses inside one kernel, one step at a time, and the whole accumulator is divided by the final normaliser once at the end. The variant in the tokamax library also masks a shard with no active blocks to zero rather than letting uninitialised values into that accumulation, and it refuses attention sinks outright.

Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop

Flash attention treats a causal mask as a value: compute the whole block of logits, then overwrite half of them with negative infinity. Splash attention treats the same mask as a schedule. Before the kernel is ever traced, a Python pass cuts the mask into blocks the size of one grid step, labels each block empty, partial or full, and turns that label grid into three small integer arrays that live in TPU scalar memory. The kernel then reads its own loop structure out of them: an empty block is not computed and its data is not fetched, a full block skips the masking entirely, and for a banded mask the grid itself is rebuilt narrower. The signature exhibit and the dial both run the real classification out of the pinned file, so the block counts on this page are the counts the compiler would produce.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop — built to be understood, not skimmed.

Questions

What is splash attention and how is it different from flash attention?: Splash attention is the TPU Pallas attention kernel in JAX that treats the mask as loop structure rather than as a value. Flash attention tiles the sequence and runs an online softmax so the full sequence-by-sequence score matrix is never materialised, but every tile is still computed and the mask is applied to the scores afterwards, so a causal mask costs the same arithmetic as no mask at all. Splash attention adds a compile-time pass that cuts the mask into blocks the size of one grid step and labels each one: 0 for a block that is all zeros, 1 for a block with a mix, and 2 for a block that is entirely ones. The kernel reads that label at each grid position out of TPU scalar memory. A 0 means the body does not run and the key and value blocks are not fetched. A 2 means the kernel ignores the mask block it was handed, because on a block that is entirely ones the stored mask pointer is not pointing at anything meaningful. Only a 1 both runs and masks. For a plain causal mask at 8192 tokens with 128 by 128 blocks that is 2016 of 4096 blocks never computed.
What is MaskInfo in splash attention?: MaskInfo is the compiled form of the mask, produced by process_mask before the kernel is traced and passed into the Pallas call as scalar-prefetch operands. It carries block_mask, an array of 0, 1 and 2 labels with one entry per grid position; data_next, which for each grid position holds the index of the next non-empty key and value block so the pipeline fetches the right thing and does not re-fetch across a skipped block; mask_next, which points into partial_mask_blocks, the deduplicated stack of the blocks that were mixed; and q_sequence, the query indices. The three integer arrays live in TPU scalar memory, which is scarce, so they are downcast to int16 or int8 whenever their largest value fits. There is a large exception: if the mask is a single computable mask such as CausalMask, LocalMask or ChunkedCausalMask, then mask_next and partial_mask_blocks are returned as None and the kernel recomputes the mask from q_sequence and the mask function instead of loading any of it.
Does splash attention make every sparse mask faster?: No. The sparsity has to line up with the block grid. The classification asks two questions per block, chunk.any() and chunk.all(), so a block earns a skip only when every single one of its 128 by 128 entries is zero. A mask that is 90 percent zeros but scattered uniformly at random has a zero in almost every block and a one in almost every block, so every block comes back partial: nothing is skipped, the grid does not shrink, and because a random mask is a dense NumPy array rather than a computable one, the mixed blocks are now materialised and loaded as well. At 2048 tokens with 128 by 128 blocks that is 256 of 256 blocks still visited plus four mebibytes of partial mask blocks that flash attention would never have touched. Causal, local window and chunked masks work because their zeros arrive in aligned rectangles.
What does grid shrinking do in splash attention?: Skipping a block saves the arithmetic but the grid still steps over it, so a mask with a narrow band still launches a grid as wide as the sequence. The shrink pass fixes that. For each row of the block grid it takes the indices of the non-empty columns, pads every row to the width of the longest one so the result is rectangular, and rewrites block_mask, data_next and mask_next to hold only those columns. The kernel then takes its grid width from the last dimension of data_next. For a local mask with a 1024-token left window at 8192 tokens, the grid width drops from 64 to 9. For a plain causal mask it does not drop at all, because the last query block genuinely needs every key block, and the padding to the longest row keeps the full width. That is why the shrink is the thing that makes windowed attention fast and skipping alone is what makes causal attention fast.
What does ring attention add on top of the splash kernel?: It shards the keys and values across a ring of devices so a sequence too long for one device still runs, and it reuses the splash forward pass unchanged for each shard. The forward is a scan over the ring size. At each step the device sends its current key and value shards to its neighbour with a collective permute while the local splash call runs on the shards it already has, so the transfer hides behind the compute. The MaskInfo is sliced along the key axis to match whichever shard is local this step. Because each step produces its own running maximum and normaliser, the outputs are combined with the same rescaling flash attention uses inside one kernel, one step at a time, and the whole accumulator is divided by the final normaliser once at the end. The variant in the tokamax library also masks a shard with no active blocks to zero rather than letting uninitialised values into that accumulation, and it refuses attention sinks outright.

Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop

Questions

Related explainers