I expected Mosaic can canonicalize 2 same strided load but it did not. So I manually unify kv strided loads to one and this also make the code cleaner.

Perf gain: 20%

• with the same 128 kv_page_per_blk and 32 q_per_blk, we can see 20% speedup in both kernel and vllm benchmarking results which shows throughput was boosted from 6.15 to 7.49 reqs/second on TPU v6e-1.

Tested:

python test/test_pallas.py -v -k PallasTest.test_ragged_paged_attention_wrapper