[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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user