33e1a21a66
THE SINGLE HIGHEST-IMPACT CODE CHANGE in this project.
Before: paged_attn.py had use_v1=True hardcoded, and _custom_ops.py V2 was
NotImplementedError. ALL decode attention (83% of competition weight) went
through V1 (ixformer single-CTA), even for 100K token sequences where one
CTA must iterate over ~195 KV block partitions sequentially.
After: V2 is wired to paged_attention_v2_pytorch.py for max_seq_len > 8192.
V1 still handles short sequences where single-CTA is faster.
Architecture follows CCCL's two-pass dispatch (dispatch_reduce.cuh):
Pass 1 (DeviceReduceKernel): N CTAs each reduce their tile partition
→ Mapped to: per-partition QK^T + softmax + V accumulation
Pass 2 (DeviceReduceSingleTileKernel): 1 CTA reduces N partial results
→ Mapped to: cross-partition log-sum-exp rescaling (summary_statistics binary_op)
For 100K tokens, PARTITION_SIZE=512:
V1: 1 CTA iterates 195 partitions sequentially
V2: 195 partitions computed in parallel, then 1 reduction pass
On 16 SMs: ceil(195/16) = 13 waves for Phase 1, then 1 CTA for Phase 2
Risk: PyTorch V2 has Python-level overhead vs ixformer's C++ V1.
Mitigation: V2 only activates for seq_len > 8192 where the parallelism
benefit outweighs Python dispatch cost. For typical decode (seq_len < 8K),
V1 ixformer kernel is still used.
Source: cccl_upstream/cub/cub/device/dispatch/dispatch_reduce.cuh
cccl_upstream/thrust/examples/summary_statistics.cu