39e32343eb
Architecture document: docs/paged_attention_kernel_architecture.md Defines every module from CCCL algorithm patterns before code. Three-level decomposition from CCCL: Level 1 (warp_reduce_shfl): shfl.down butterfly for per-thread QK scores Level 2 (block_reduce_warp_reductions): warp partials → SMEM → block aggregate Level 3 (agent_scan decoupled lookback): cross-partition combine Compound type (from summary_statistics.cu): attention_partial = (max_score, exp_sum, weighted_v[256]) combine(a, b) = online softmax rescaling (same math as Flash Attention) Key design change: Grid on num_kv_heads, not num_heads. Before: grid = (1, 24, 200) = 4800 blocks, KV loaded 6x redundantly After: grid = (1, 4, 200) = 800 blocks, KV loaded once per kv_head Each block computes GQA_RATIO=6 query heads with shared KV loads. Reduces KV cache bandwidth by 6x (the GQA ratio). SMEM budget verified: K tile [32, 256] fp16 = 16KB V tile [32, 256] fp16 = 16KB Total = 32KB ≤ 48KB ✓ Phase 1 kernel: _partition_attn_kernel Processes query heads sequentially within the GQA group to minimize register pressure (6 × 256 = 1536 registers too many if all loaded simultaneously). Phase 2 kernel: _reduce_partitions_kernel Also gridded on kv_heads, reduces all partitions for GQA_RATIO heads per block. This replaces the previous Triton V2 which was gridded on num_heads and had no GQA awareness at the kernel level.