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?

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

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

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

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

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