From 951afd0c02a2ce78a0570b2d528ea0bdf2ceb41a Mon Sep 17 00:00:00 2001
From: Dylan <dylan@project6.dev>
Date: Fri, 7 Aug 2026 01:19:50 +0000
Subject: [PATCH] [ENGINE] apply CCCL GridEvenShare dispatch pattern to V1/V2
 attention decision
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Source: cccl_upstream/cub/cub/device/dispatch/dispatch_reduce.cuh
        cccl_upstream/cub/cub/grid/grid_even_share.cuh

Replace ad-hoc V1/V2 heuristic with CCCL's precise work distribution:
- max_blocks = sm_occupancy × sm_count × subscription_factor (1×16×5=80)
- total_tiles = ceil_div(max_seq_len, PARTITION_SIZE)
- grid_size = min(total_tiles, max_blocks)
- V1 when grid_size==1 OR seq×head parallelism saturates GPU

CCCL kernel_reduce.cuh insight: !StableReductionOrder uses atomicAdd
for single-kernel finish. BI-V100 with 16 SMs -> max 80 CTAs ->
atomic contention negligible -> nondeterministic path is optimal.
---
 paged_attn.py | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/paged_attn.py b/paged_attn.py
index 68861641..4228e126 100644
--- a/paged_attn.py
+++ b/paged_attn.py
@@ -397,11 +397,36 @@ class PagedAttention:
         # Original heuristic (before hardcode): V1 when max_seq_len ≤ 8192 OR
         # when batch×heads already saturates the GPU (num_seqs*num_heads > 512).
         # Restored with BI-V100 SM count awareness.
+        # ──── CCCL GridEvenShare dispatch (from dispatch_reduce.cuh) ────
+        # CCCL formula: max_blocks = sm_occupancy × sm_count × subscription_factor
+        # Then: grid_size = min(total_tiles, max_blocks)
+        # If grid_size == 1 → single-tile (V1). If grid_size > 1 → multi-tile (V2).
+        #
+        # BI-V100 hardware (confirmed):
+        #   sm_count = 16, sm_occupancy ≈ 1 CTA/SM (conservative for attention),
+        #   subscription_factor = 5 (CCCL default from util_arch.cuh)
+        #
+        # Tile size = _PARTITION_SIZE (512 tokens per partition)
+        # total_tiles = ceil_div(max_seq_len, _PARTITION_SIZE)
+        # max_blocks = 1 × 16 × 5 = 80
+        #
+        # This replaces the ad-hoc "num_seqs * num_heads > 512" heuristic
+        # with CCCL's precise GridEvenShare work distribution.
         bi100_sm_count = 16
-        bi100_saturation = bi100_sm_count * 32  # ~512 concurrent warps
-        use_v1 = (max_num_partitions == 1
-                  or max_seq_len <= 8192
-                  or num_seqs * num_heads > bi100_saturation)
+        bi100_sm_occupancy = 1  # conservative: 1 attention CTA per SM
+        bi100_subscription = 5  # CCCL default subscription_factor
+        bi100_max_blocks = bi100_sm_occupancy * bi100_sm_count * bi100_subscription  # 80
+        
+        total_tiles = (max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE
+        grid_size = min(total_tiles, bi100_max_blocks)
+        
+        # CCCL single-tile vs multi-tile decision:
+        # V1 (single-tile) when problem fits in one CTA's work,
+        # OR when sequence×head parallelism already saturates the GPU
+        # (no benefit from partitioning — each sequence already has its own CTA)
+        seq_head_parallelism = num_seqs * num_heads
+        use_v1 = (grid_size == 1
+                  or seq_head_parallelism >= bi100_max_blocks)
         if use_v1:
             # Run PagedAttention V1.
             ops.paged_attention_v1(