15ef28e863
Before: 3 nested Python for-loops
for seq_idx: (1 iteration at max_num_seqs=1)
for block_idx: (6250 iterations at seq_len=100K, block_size=16)
key_cache[physical_block] + permute + reshape per block
for part_idx: (195 iterations at seq_len=100K, PARTITION=512)
torch.einsum per partition
After: 1 seq loop (trivial) + batched gather + bmm partition loop
for seq_idx: (1 iteration — same)
key_cache[blk_ids] (ONE index_select for all 6250 blocks)
.permute().reshape() (ONE reshape for entire sequence)
for part_idx: (195 iterations, each uses torch.bmm)
torch.bmm (batched over all heads simultaneously)
Key changes:
- Block gather: block-by-block Python loop → single key_cache[blk_ids]
Eliminates 6250 Python iterations for 100K sequence
- GQA: repeat_interleave (allocates) → expand (view, zero-copy)
- Partition attn: torch.einsum → torch.bmm (more efficient for batched)
- Phase 2 reduction: unchanged (already vectorized)
The block_idx loop was the real killer: 6250 Python-level tensor operations
(index + permute + reshape + slice) per decode step. Now it's one operation.