FIXED BASE FILE (not root custom file):
vllm/attention/ops/paged_attn.py — the actual vllm paged attention
Two changes from reading cub/block/specializations/block_reduce_raking.cuh:
1. V1/V2 dispatch restored (was hardcoded use_v1=True on line 119)
CCCL block_reduce_raking has WARP_SYNCHRONOUS conditional fast path:
when RAKING_THREADS == BLOCK_THREADS, skip SMEM and go to warp shuffle.
This is CONDITIONAL — not hardcoded. Our equivalent:
V1 (single-pass) is the WARP_SYNCHRONOUS fast path for short seqs.
V2 (partitioned reduce) is the raking path for long seqs.
For max_num_seqs=1: num_seqs*num_heads=24 < 512, so V2 triggers
when max_seq_len > 8192.
2. V2 temp tensor caching (agent_merge_sort union _TempStorage pattern)
Cache tmp_output/exp_sums/max_logits by shape key across decode steps.
For max_num_seqs=1, shapes are stable → zero CUDA malloc after warmup.
CCCL files: cub/block/specializations/block_reduce_raking.cuh,
cub/agent/agent_merge_sort.cuh
29f119c094