11 / 11 Kernels (CUDA / Triton / Pallas)

Pallas Skips the Refetch

In a Pallas TPU flash-attention kernel with a grid over (query_block, key_block), the K/V BlockSpec's index_map returns the same key_block index across several consecutive inner-loop grid steps before advancing to the next key block. What does the Pallas pipeline emitter do because of this, and why is it worth structuring the grid this way?

A

It automatically grows the block_shape for that operand for the duration of the repeated indices, trading extra VMEM for fewer grid steps

B

It inserts a full pipeline drain and synchronization barrier whenever index_map repeats an output, to guarantee VMEM consistency before continuing

C

It evaluates each BlockSpec's index_map one grid step ahead; when the next step maps to the same block index as the current one, it skips the redundant HBM->VMEM copy and keeps the block resident, so revisiting the same block on consecutive steps directly cuts HBM traffic

D

It reruns the kernel body twice on that block as a numerical-stability check, since a repeated index_map output is treated as a signal of a potential race condition

E

It has no special-cased effect: a fresh async copy is issued on every single grid step regardless of what index_map returns, so operand ordering never changes HBM traffic