[perf/decode] paged_attn: CCCL GridEvenShare-informed tile sizing

CCCL dispatch_reduce.cuh uses:
  max_blocks = sm_occupancy * sm_count * subscription_factor
  BI-V100: 1 * 16 * 5 = 80 max CTAs

But paged_attn._forward_decode_pytorch runs in Python (torch.matmul),
not as CUDA CTA launches. Python loop overhead >> kernel launch overhead.
Each iteration = torch.matmul + online softmax update (2-3 CUDA launches).

Change: TARGET_TILES 4→2, MIN_TILE_BLOCKS 64→128, MAX_TILE_BLOCKS 4096→8192

Effect: For seq_len=100K (6250 blocks), tile_blocks goes from
  ceil(6250/4)=1563 → ceil(6250/2)=3125 blocks per tile
  = 2 Python iterations instead of 4
  = 50% fewer torch.matmul launches for long contexts

Memory check: 3125 blocks × 16 tokens/block = 50K tokens per tile
  Score: 4 kv_heads × 6 gqa × 50K × 4B = 4.8 MB ✓ (fits in 48KB SMEM for the
  matmul kernel; actual memory is HBM-allocated by PyTorch)

Source: CCCL grid_even_share.cuh DispatchInit + subscription_factor=5
This commit is contained in:
muh
2026-08-06 06:00:52 +00:00
parent 2ee9571575
commit 7552365c7f

View File

@@ -166,9 +166,17 @@ class PagedAttention:
# big_shares = total_tiles - (avg_tiles * grid_size)
# Our target: ~4 tiles max (Python overhead >> kernel launch overhead)
# ================================================================
_BI100_TARGET_TILES = 4 # minimize Python loop iterations
_MIN_TILE_BLOCKS = 64 # floor: avoid tiny matmuls
_MAX_TILE_BLOCKS = 4096 # ceiling: avoid single huge allocation
# CCCL GridEvenShare: max_blocks = sm_occupancy * sm_count * subscription_factor
# BI-V100: 1 * 16 * 5 = 80 max CTAs for CUDA kernels.
# But this is Python (PyTorch ops), not CUDA launches — Python loop
# overhead dominates. Each iteration = 1 torch.matmul launch + online
# softmax update. Target 2 iterations (not 4): the matmul itself is
# already parallelized across SMs, so fewer Python loops = less overhead.
# For seq_len=100K with block_size=16: 6250 blocks / 2 = 3125 blocks/tile.
# Score tensor: 4 kv_heads × 6 gqa × 1 × 50000 × 4B = 4.8 MB — fits.
_BI100_TARGET_TILES = 2 # 2 iterations: minimize Python loop overhead
_MIN_TILE_BLOCKS = 128 # floor: ensure matmul is large enough to saturate 16 SMs
_MAX_TILE_BLOCKS = 8192 # ceiling: 8192 × 16 = 128K tokens per tile — fits in memory
try:
for i in range(num_seqs):