f7f8113c73
Source: cccl_upstream/cub/cub/device/dispatch/dispatch_reduce.cuh
cccl_upstream/cub/cub/device/dispatch/kernels/kernel_reduce.cuh
CCCL's reduce dispatch uses GridEvenShare to compute optimal tile count:
max_blocks = sm_occupancy × sm_count × subscription_factor
tile_size = num_items / max_blocks
Applied to _forward_prefix_pytorch's KV-cache tile iteration:
- OLD: fixed _BLOCKS_PER_TILE=32 (512 tokens per tile regardless of q_len)
- NEW: adaptive tile_sz based on score tensor memory budget
- q_len=1 (decode): tile_sz grows to 2048 tokens (fewer iterations)
- q_len=4096 (prefill): tile_sz stays ~256 (fits in 96MB budget)
- Score tensor = kv_h × gqa × q_len × tile_sz × 4 bytes ≤ 96MB
CCCL kernel_reduce.cuh insight: StableReductionOrder=false uses atomic
aggregation in a single kernel launch. Our online softmax accumulator
(m, l, o) similarly benefits from fewer, larger tiles — each merge step
has Python loop overhead that dominates BI-V100's 16-SM execution.