From 7552365c7fbc2711987d5b544551d4f482f5ced3 Mon Sep 17 00:00:00 2001
From: muh <muh@project6.dev>
Date: Thu, 6 Aug 2026 06:00:52 +0000
Subject: [PATCH] [perf/decode] paged_attn: CCCL GridEvenShare-informed tile
 sizing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 paged_attn.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

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):