The single biggest performance bottleneck in the baseline:
paged_attention_v2 = raise NotImplementedError()
paged_attn.py: use_v1 = True (hardcoded to avoid calling V2)
V1 limitation: processes entire KV sequence in one kernel launch.
For seq_len=100K, this is a single massive attention computation.
V2: splits into PARTITION_SIZE=512 chunks, runs them in parallel,
then reduces with log-sum-exp. 195 parallel partitions vs 1.
Implementation (paged_attention_v2_pytorch.py):
Phase 1: Per-partition attention
- For each (seq, head, partition): compute QK^T, softmax, weighted V sum
- Store partial: tmp_output, exp_sums, max_logits (per partition)
Phase 2: Cross-partition reduction (log-sum-exp)
- global_max = max(max_logits across partitions)
- rescale = exp(partition_max - global_max) × partition_exp_sum
- output = Σ (rescale / total_sum) × partition_output
This is the same algorithm as vllm's paged_attention_v2_kernel.cu:
- The reduction pattern is identical to CCCL's block_reduce_warp_reductions
(combine partial statistics from independent segments)
- The online softmax tiling is the same as Flash Attention's partitioning
Integration:
- patch_paged_attention_v2.py patches _custom_ops.py and paged_attn.py
- Removes use_v1=True hardcode → V2 used for seq_len > 8192
- Dockerfile adds the patch step
This is a PyTorch implementation (no CUDA compilation needed).
Next step: if /usr/local/corex/ has ixcc or nvcc-compatible compiler,
replace with compiled CUDA kernel for further speedup.
