dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[CCCL-PORT] Adaptive tile sizing from dispatch_reduce.cuh GridEvenShare

Browse Source 
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.
This commit is contained in:
dylanyunlon
2026-08-05 08:10:29 +00:00
parent 15753cb68d
commit f7f8113c73
1 changed files with 48 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
paged_attn.py
View File

@@ -340,11 +340,36 @@ class PagedAttention:
context_lens : [batch_size] tokens already in KV cache
"""
try:
# Paged-block tiles for context phase.
# tile_sz = _BLOCKS_PER_TILE × block_size (e.g. 16×16 = 256 tokens).
# Score tensor [kv_h, gqa, q_len, tile_sz] fp32 = 24 MB per tile.
# Same tile size reused for the current-chunk phase.
_BLOCKS_PER_TILE = 32
# ================================================================
# Tile sizing strategy — ported from CCCL dispatch_reduce.cuh
#
# CCCL's GridEvenShare computes:
# max_blocks = sm_occupancy × sm_count × subscription_factor
# tile_size = num_items / max_blocks (evenly distributed)
#
# For BI-V100 (16 SMs), fixed _BLOCKS_PER_TILE=32 wastes memory
# on short contexts and underutilizes on long ones.
#
# Key insight from kernel_reduce.cuh:
# StableReductionOrder=false uses atomicAdd → single kernel pass.
# For online softmax (our case), we accumulate (m, l, o) per tile
# then merge — this IS a multi-pass reduce. Larger tiles = fewer
# merge steps = less numerical drift + less Python loop overhead.
#
# CCCL subscription_factor = CUB_SUBSCRIPTION_FACTOR(0) = 5
# Effective: 16 SM × 1 CTA/SM × 5 = 80 concurrent tiles max.
# But Python loop overhead dominates, so we want FEWER, LARGER tiles.
#
# Strategy: target ~4-8 tiles per context phase.
# Fewer tiles → fewer matmul calls → less launch overhead.
# SMEM constraint: score tensor [kv_h, gqa, q_len, tile_sz] fp32
# must not cause OOM. With q_len=4096, kv_h=1, gqa=6:
# tile_sz=1024 → 1×6×4096×1024×4 = 96 MB (too much)
# tile_sz=512 → 48 MB (borderline)
# tile_sz=256 → 24 MB (safe)
# For decode (q_len=1): tile_sz=4096 → only 96 KB (always safe)
# ================================================================
_SMEM_BUDGET_BYTES = 96 * 1024 * 1024 # 96 MB score tensor budget
batch_size = seq_lens_tensor.shape[0]
num_q_heads = query.shape[1]
@@ -352,7 +377,6 @@ class PagedAttention:
head_dim = query.shape[2]
gqa_ratio = num_q_heads // num_kv_heads
block_size = value_cache.shape[3]
tile_sz = _BLOCKS_PER_TILE * block_size
scale = head_dim ** -0.5
orig_dtype = query.dtype
output = torch.empty_like(query)
@@ -368,6 +392,19 @@ class PagedAttention:
k_i = key [q_start:q_end] # [q_len, kv_h, d]
v_i = value[q_start:q_end]
# CCCL-style adaptive tile sizing per sequence.
# Score tensor = [kv_h, gqa, q_len, tile_sz] × 4 bytes
# Solve: kv_h × gqa × q_len × tile_sz × 4 ≤ budget
score_row_bytes = num_kv_heads * gqa_ratio * q_len * 4
if score_row_bytes > 0:
max_tile_tokens = _SMEM_BUDGET_BYTES // score_row_bytes
# Round down to block_size boundary
max_tile_tokens = (max_tile_tokens // block_size) * block_size
# Clamp: at least 1 block, at most what context needs
tile_sz = max(block_size, min(max_tile_tokens, 2048))
else:
tile_sz = block_size * 32 # fallback
# Q reshaped and scaled once; held for all K-tiles.
# [kv_h, gqa, q_len, d] fp32 — 24 MB for q_len=4096, d=256
q_seq = (q_i.permute(1, 0, 2)
@@ -391,14 +428,11 @@ class PagedAttention:
# query has position ≥ ctx_len. k_pos < q_pos is always True
# → no causal mask needed for pure context tiles.
# --------------------------------------------------------------
# Convert token-based tile_sz to block count for iteration
blocks_per_tile = tile_sz // block_size
if ctx_len > 0:
num_ctx_blocks = (ctx_len + block_size - 1) // block_size
# Safety: if block_tables is too narrow this indicates a
# prefix_cache_hit + chunked-prefill bug in model_runner.py
# (Case 1 leaves prefix_cache_hit=True but block_table is
# only computed_block_nums, not the full context blocks).
# patch_model_runner.py fixes the root cause; this guard
# prevents a zero-dim amax() crash if it still slips through.
if num_ctx_blocks > block_tables.shape[1]:
print(
f"[paged_attn WARNING] seq {i}: num_ctx_blocks={num_ctx_blocks} "
@@ -407,8 +441,8 @@ class PagedAttention:
"Capping context to available blocks — attention may be incorrect.",
file=sys.stderr, flush=True)
num_ctx_blocks = block_tables.shape[1]
for tile_blk in range(0, num_ctx_blocks, _BLOCKS_PER_TILE):
blk_end = min(tile_blk + _BLOCKS_PER_TILE, num_ctx_blocks)
for tile_blk in range(0, num_ctx_blocks, blocks_per_tile):
blk_end = min(tile_blk + blocks_per_tile, num_ctx_blocks)
blk_ids = block_tables[i, tile_blk:blk_end]
# Gather K/V for this tile.