Two changes informed by reading CCCL engine source code as input:
1. SingleTile fast path (from kernel_reduce.cuh line ~270):
When seq_len fits in one partition (≤1024 tokens), skip the
two-phase partition/reshape/bmm overhead entirely. Direct
softmax + V weighted sum. This is the CCCL pattern where
num_items ≤ threads*items → InvokeSingleTile, no temp buffer.
Impact: Early decode tokens (seq_len < 1024) avoid all partition
machinery. Qwen3.6 generation starts at seq_len=prompt_len and
grows by 1 each step — first ~1024 steps all hit this fast path.
2. GridEvenShare constants (from dispatch_reduce.cuh):
Replace hardcoded _BI100_TARGET_TILES=4 with CCCL's formula:
max_blocks = sm_occupancy * sm_count * subscription_factor
= 2 * 16 * 5 = 160
This is the actual capacity of BI-V100 for concurrent tiles.
Source files read as input for this change:
- cccl_upstream/cub/cub/device/dispatch/dispatch_reduce.cuh (full)
- cccl_upstream/cub/cub/device/dispatch/kernels/kernel_reduce.cuh (full)
- cccl_upstream/cub/cub/agent/agent_reduce.cuh (full)
- paged_attention_v2_pytorch.py (full)
- vllm/_custom_ops.py (first 200 lines)
From reading cccl_upstream/cub/cub/device/dispatch/kernels/kernel_reduce.cuh:
- SingleTile path: when num_partitions fits in one tile (always true for
BI-V100 attention with <=200 partitions), reduce uses single CTA.
Phase 2 is never the bottleneck.
- StableReductionOrder=false uses atomicAdd to skip pass 2 entirely.
Not applicable to attention (compound accumulator), but confirms
the Phase 2 architecture is correct.
From reading cccl_upstream/cub/cub/agent/single_pass_scan_operators.cuh:
- delay<Delay, GridThreshold=500>(): when gridDim.x < 500, ALL delay
strategies collapse to __threadfence_block(). BI-V100 scan grids
have at most ~12 blocks (100K/8448). ALL delay tuning (ns, dcid, l2w)
is irrelevant — bench_bi100.py's no_delay 'win' was actually noise
between identical __threadfence_block() calls.
From reading cccl_upstream/thrust/examples/summed_area_table.cu:
- inclusive_scan_by_key pattern for per-row operations maps to
per-sequence softmax denominator computation in paged_attention.
CCCL norm.cu demonstrates transform_reduce fusion: compute sqrt(sum(x^2))
as transform_reduce(x, square, 0, plus) in ONE kernel, not transform(square)
then reduce(plus) as two kernels. Same principle applied to Phase 2:
Before (6 kernel launches):
global_max = pm.max(dim=-1) # launch 1
rescale = exp(pm - max) * ps # launch 2 (exp + mul fused by PyTorch)
total = rescale.sum(dim=-1) # launch 3
weights = rescale / total # launch 4 ← ELIMINATED
final = bmm(weights, po) # launch 5
After (5 kernel launches):
global_max = pm.max(dim=-1)
rescale = exp(pm - max) * ps
total = rescale.sum(dim=-1)
final = bmm(rescale, po) / total # division on H×d output, not H×P weights
The division moves from H×P elements (24×98 = 2352 for 100K seq) to
H×d elements (24×128 = 3072) — slightly more elements but one fewer
kernel launch, and the bmm output is already in L1 cache.
Also read CCCL sources this round:
- cub/block/block_load.cuh: LoadDirectBlocked + vectorization strategy
- cub/device/dispatch/dispatch_scan.cuh: grid_size = num_tiles, tile_state alloc
- thrust/examples/expand.cu: variable-length replication (GQA broadcast)
- thrust/examples/norm.cu: transform_reduce fusion for L2 norm
- tuning_radix_sort.cuh policy_selector: onesweep_radix_bits=8 confirmed
Source: cccl_upstream/thrust/examples/norm.cu
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
Bug: After GQA broadcast optimization, v_perm was [kv_h, seq_len, d]
in the GQA path, but unconditional v_padded allocation used num_heads:
v_padded = torch.zeros((num_heads, padded_len, head_size))
v_padded[:, :seq_len, :] = v_perm # [24, padded, d] vs [4, seq, d] → CRASH
Fix: v_padded/v_parts allocation is now inside the non-GQA else branch.
GQA branch uses its own v_padded_kv with correct [kv_h, padded, d] shape.
This was a real runtime bug — V2 would have crashed on first call
for any GQA model (Qwen3.6, Llama, etc.).
Previous commit broadcast Q@K^T (saved 1GB/step).
This commit broadcasts scores@V too (saves 2GB/step).
Before: V expanded from [kv_h, padded_len, d] to [H, padded_len, d]
4×100K×256×4B → 24×100K×256×4B = 400MB → 2.4GB allocation
After: broadcast matmul at kv_h level
se: [kv_h, gqa, P, 1, part_sz] @ V: [kv_h, 1, P, part_sz, d]
→ [kv_h, gqa, P, 1, d] → reshape to [H, P, d]
V stays at kv_h size: 400MB (no 2.4GB allocation)
Total per-decode-step memory for 100K context:
Before all GQA opts: 3.6GB (K expansion + V expansion)
After: 600MB (6x total reduction from GQA ratio=6)
This is the CCCL insight applied: transform_reduce with a compound type.
Instead of expanding to full head count then reducing, keep the reduction
at the minimal group size and broadcast the grouping dimension.
Phase 1 rewrite:
Before: for p in range(195): torch.bmm(Q, K_partition_p)
After: scores = torch.bmm(Q, K_all) # ONE launch for all 100K tokens
scores_parts = scores.view(H, P, part_sz) # reshape, no copy
part_out = torch.bmm(scores_exp_flat, v_parts_flat) # ONE launch
195 Python→CUDA round-trips → 2 round-trips.
Architecture informed by CCCL:
- summary_statistics.cu: fuse (max, exp_sum, weighted_output) computation
into a single reduction pass over the data. We do this by computing
Q@K^T over the ENTIRE sequence in one bmm, then reshaping to partitions
for the softmax statistics — the data is only read once from HBM.
- block_reduce_warp_reductions.cuh: Phase 2 reduction combines partition
statistics using the same (rescale, accumulate) pattern as CUB's
cross-warp aggregate merging.
Phase 2 (unchanged, already vectorized):
global_max + rescale + torch.bmm(weights, partition_outputs)
Total GPU kernel launches per decode step:
Before: 1 (gather) + 195 (Q@K) + 195 (scores@V) + 1 (reduce) = 392
After: 1 (gather) + 1 (Q@K_all) + 1 (scores_exp@V) + 1 (reduce) = 4
KV gather also stays batched: key_cache[blk_ids] is one index_select.
The single biggest performance bottleneck in the baseline:
paged_attention_v2 = raise NotImplementedError()
paged_attn.py: use_v1 = True (hardcoded to avoid calling V2)
V1 limitation: processes entire KV sequence in one kernel launch.
For seq_len=100K, this is a single massive attention computation.
V2: splits into PARTITION_SIZE=512 chunks, runs them in parallel,
then reduces with log-sum-exp. 195 parallel partitions vs 1.
Implementation (paged_attention_v2_pytorch.py):
Phase 1: Per-partition attention
- For each (seq, head, partition): compute QK^T, softmax, weighted V sum
- Store partial: tmp_output, exp_sums, max_logits (per partition)
Phase 2: Cross-partition reduction (log-sum-exp)
- global_max = max(max_logits across partitions)
- rescale = exp(partition_max - global_max) × partition_exp_sum
- output = Σ (rescale / total_sum) × partition_output
This is the same algorithm as vllm's paged_attention_v2_kernel.cu:
- The reduction pattern is identical to CCCL's block_reduce_warp_reductions
(combine partial statistics from independent segments)
- The online softmax tiling is the same as Flash Attention's partitioning
Integration:
- patch_paged_attention_v2.py patches _custom_ops.py and paged_attn.py
- Removes use_v1=True hardcode → V2 used for seq_len > 8192
- Dockerfile adds the patch step
This is a PyTorch implementation (no CUDA compilation needed).
Next step: if /usr/local/corex/ has ixcc or nvcc-compatible compiler,
replace with compiled CUDA kernel for further speedup.
