cbe6066257
Phase 1 rewrite:
Before: for p in range(195): torch.bmm(Q, K_partition_p)
After: scores = torch.bmm(Q, K_all) # ONE launch for all 100K tokens
scores_parts = scores.view(H, P, part_sz) # reshape, no copy
part_out = torch.bmm(scores_exp_flat, v_parts_flat) # ONE launch
195 Python→CUDA round-trips → 2 round-trips.
Architecture informed by CCCL:
- summary_statistics.cu: fuse (max, exp_sum, weighted_output) computation
into a single reduction pass over the data. We do this by computing
Q@K^T over the ENTIRE sequence in one bmm, then reshaping to partitions
for the softmax statistics — the data is only read once from HBM.
- block_reduce_warp_reductions.cuh: Phase 2 reduction combines partition
statistics using the same (rescale, accumulate) pattern as CUB's
cross-warp aggregate merging.
Phase 2 (unchanged, already vectorized):
global_max + rescale + torch.bmm(weights, partition_outputs)
Total GPU kernel launches per decode step:
Before: 1 (gather) + 195 (Q@K) + 195 (scores@V) + 1 (reduce) = 392
After: 1 (gather) + 1 (Q@K_all) + 1 (scores_exp@V) + 1 (reduce) = 4
KV gather also stays batched: key_cache[blk_ids] is one index_select.