Three changes based on reading CCCL agent_reduce.cuh + single_pass_scan_operators.cuh:
1. Restore V1/V2 adaptive dispatch (was hardcoded V1 for all cases).
ops.paged_attention_v2 IS a C++ kernel, not pure PyTorch.
For sequences > 8192 tokens, V2's partitioned parallelism better
utilizes 16 SMs than V1's single-CTA sequential iteration.
2. Pre-allocate V2 intermediate tensors (tmp_output, exp_sums, max_logits)
using module-level cache, same pattern as MoE commit d3b1108.
Eliminates 3 CUDA mallocs per decode step when V2 is active.
3. PARTITION_SIZE 512→1024. CCCL GridEvenShare insight: with 16 SMs,
fewer larger partitions (98 vs 196 for 100K tokens) produce 6.1
CTAs/SM vs 12.3, reducing inter-CTA sync overhead in V2 reduce.
CCCL sources read as input:
- agent_reduce.cuh: tile consumption loop, vectorized load, SMEM union
- single_pass_scan_operators.cuh: delay() GridThreshold=500 logic,
no_delay_constructor_t is empty on SM70+, l2w is one-time constructor
- agent_scan.cuh: SMEM = union{load, store, {prefix+scan}} not sum
- block_scan_warp_scans.cuh: warp aggregate exchange pattern
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
