diff --git a/paged_attn.py b/paged_attn.py index 1be54382..68861641 100644 --- a/paged_attn.py +++ b/paged_attn.py @@ -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):