Claude
|
a53d1a28b0
|
[OPT] Triton paged_attention_v2 kernel skeleton — Phase 2 reduction complete
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.
|
2026-07-30 15:59:15 +00:00 |
|