a53d1a28b0
Two-kernel design following vllm's paged_attention_v2_kernel.cu: Phase 1: _paged_attn_v2_partition_kernel grid = (num_seqs, num_heads, num_partitions) Each instance: Q[head] @ K[partition]^T → softmax → @ V[partition] Status: SKELETON — paged K/V gather from indirect block_tables is complex in Triton (requires scatter/gather through block_tables). Currently falls back to PyTorch partition loop. Phase 2: _paged_attn_v2_reduce_kernel grid = (num_seqs, num_heads) Each instance: log-sum-exp reduction across partitions Status: COMPLETE — replaces Python einsum with single Triton launch. Algorithm: global_max → rescale → weighted sum (same pattern as CCCL summary_statistics binary_op for combining partial statistics). SMEM: Phase 1 needs BLOCK_N=64 × head_dim=128 × 2B × 2 = 32KB ≤ 48KB. Phase 2 needs no SMEM (partitions fit in registers). The Phase 1 paged gather is the hard part. The key_cache layout [blocks, kv_heads, head_dim/x, block_size, x] requires: 1. block_tables[seq, token // block_size] → physical_block_id 2. key_cache[physical_block_id, kv_head, :, token % block_size, :] This is indirect indexed access — possible in Triton via tl.load with computed offsets, but needs careful stride arithmetic.