b6e7bca45a
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