Two changes based on CCCL source reading:
1. PARTITION_SIZE 512→1024 in paged_attention_v2_pytorch.py
From dispatch_scan.cuh: grid_size = num_tiles = ceil(N / tile_size).
Optimal tile_size balances parallelism vs overhead:
- BI-V100: 16 SMs, max ~32 concurrent CTAs
- Need num_partitions >= 32 to fill one wave
- 100K tokens / 1024 = 98 partitions (3 waves) ✓
- 100K tokens / 512 = 195 partitions (6 waves) — twice the Phase 2 cost
Note: only affects V2 (PyTorch path). V1 (ixformer) has its own partition size.
2. Fix V2 import path in _custom_ops.py
paged_attention_v2_pytorch.py is in repo root, not vllm package.
Added sys.path manipulation to find it at runtime.
Also read: cccl_upstream/thrust/examples/expand.cu (variable-length
replication pattern — maps to GQA expansion, but our broadcast approach
is already more efficient than physical replication).
Source: cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh lines 350-380
cccl_upstream/thrust/examples/expand.cu
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
