81972a05c6
CCCL source read: cub/device/dispatch/kernels/kernel_segmented_reduce.cuh
Three agent tiers based on segment size:
Small (≤ small_items_per_tile) → 1 thread per segment (AgentSmallReduce)
Medium (≤ medium_items_per_tile) → 1 warp per segment (AgentMediumReduce)
Large (> medium) → 1 block per segment (AgentReduce)
All three share a union __shared__ memory — only one tier active at a time.
Applied to paged_attention forward_decode:
OLD: use_v1=True forced V1 for all sequence lengths.
V2's partitioned execution was never attempted on BI-V100.
NEW: Three-tier dispatch mirroring CCCL's segmented_reduce:
Small (seq_len ≤ 8192) → V1 native (single CTA, optimal for short seqs)
Medium (8192 < seq ≤ 32K) → V2 native attempt with try/except fallback to V1
V2 partitions work across multiple CTAs, better
for 16-SM BI-V100 on medium sequences
Large (seq > 32K) → PyTorch fallback (V1 SMEM overflow)
Also added CCCL CachingDeviceAllocator buffer reuse pattern to prefix attention:
Pre-allocated _m_blk, _m_new, _corr buffers outside tile loops,
reused via torch.amax(out=), torch.maximum(out=), torch.exp(out=).