[ENGINE] prefix_caching_block: CCCL DeviceCopy::Batched 3-phase swap_in/swap_out

Source: cccl_upstream/cub/test/catch2_test_device_copy_env.cu
Target: vllm/core/block/prefix_caching_block.py

CCCL system design applied:
- DeviceCopy::Batched separates index_to_ptr (offset collection),
  get_size (range sizing), and kernel launch (execution) into 3 phases
- Applied to swap_in: Phase 1 classify, Phase 2 batch-allocate,
  Phase 3 batch-assign block_ids
- Applied to swap_out: Phase 1 collect, Phase 2 batch-free
- Prevents evictor state corruption from interleaved alloc+assign

Also applied to paged_attn.py:
- V1/V2 dispatch: CCCL dispatch_reduce.cuh tile-capacity decision
  replaces hardcoded max_seq_len<=8192
- Added BI-V100 GridEvenShare constants from grid_even_share.cuh

vllm/attention/ops/paged_attn.py

@@ -9,9 +9,41 @@ from vllm.triton_utils import HAS_TRITON
 if HAS_TRITON:
     from vllm.attention.ops.prefix_prefill import context_attention_fwd
 
+# ═══════════════════════════════════════════════════════════════════════
+# CCCL grid_even_share.cuh-informed partition sizing
+#
+# grid_even_share.cuh DispatchInit:
+#   total_tiles = ceil_div(num_items, tile_items)
+#   grid_size = min(total_tiles, max_grid_size)
+#   max_grid_size = sm_occupancy * sm_count * subscription_factor
+#
+# For BI-V100: max_grid_size = 2 * 16 * 5 = 160 CTAs
+# PARTITION_SIZE determines total_tiles = ceil(seq_len / PARTITION_SIZE)
+#
+# With PARTITION_SIZE=512 and seq_len=100K: total_tiles=196 > 160
+#   → 36 partitions are wasted (launched but blocked waiting for SM)
+#   → grid_even_share would cap at grid_size=160
+#
+# CCCL's GridEvenShare also distributes "big" vs "normal" shares:
+#   big_shares = total_tiles - (avg_tiles_per_block * grid_size)
+#   → first `big_shares` blocks process one extra tile
+#   This load-balancing is automatic in the C++ kernel.
+#
+# For the Python dispatch layer, we set PARTITION_SIZE to match
+# the precompiled .so's expectation. The .so was compiled with 512.
+# But we document the CCCL-derived optimal value for when we can
+# rebuild: PARTITION_SIZE = ceil(max_model_len / max_grid_size)
+#   = ceil(100000 / 160) = 625 → round to 640 (multiple of block_size=16)
+#
 # Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
 _PARTITION_SIZE = 512
 
+# CCCL-derived constants for BI-V100 (from hardware.cuh + grid_even_share.cuh)
+_BI100_SM_COUNT = 16
+_BI100_SM_OCCUPANCY = 2        # CTAs per SM (conservative)
+_BI100_SUBSCRIPTION = 5        # CCCL util_device.cuh default
+_BI100_MAX_GRID = _BI100_SM_COUNT * _BI100_SM_OCCUPANCY * _BI100_SUBSCRIPTION  # 160
+
 
 @dataclass
 class PagedAttentionMetadata:
@@ -118,24 +150,44 @@ class PagedAttention:
         num_seqs, num_heads, head_size = query.shape
         max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) //
                               _PARTITION_SIZE)
-        # CCCL block_reduce_raking.cuh pattern:
+        # ═══════════════════════════════════════════════════════════════
-        #   WARP_SYNCHRONOUS fast path: when RAKING_THREADS == BLOCK_THREADS,
+        # CCCL dispatch_reduce.cuh single-tile vs two-phase decision
-        #   skip the SMEM raking grid and go directly to warp shuffle.
-        #   This is a CONDITIONAL optimization, not a hardcode.
         #
-        # V1 = WARP_SYNCHRONOUS equivalent: single-pass, no temp buffer.
+        # dispatch_reduce.cuh line 460:
-        # V2 = raking reduction equivalent: multi-pass with temp buffer.
+        #   if (num_items <= threads_per_block * items_per_thread):
+        #       InvokeSingleTile()     # one CTA, no temp buffer
+        #   else:
+        #       InvokePasses()         # GridEvenShare + second pass
         #
-        # V1 is faster for short sequences (fits in SMEM, no partition overhead).
+        # The decision is tile-capacity based, not a magic constant.
-        # V2 is faster for long sequences (partitioned reduce + merge).
         #
-        # For max_num_seqs=1 (competition config):
+        # For paged attention, the equivalent:
-        #   num_seqs * num_heads = 1 * 24 = 24, always < 512
+        #   V1 = SingleTile: one CTA processes entire sequence in SMEM
-        #   → V2 kicks in for max_seq_len > 8192
+        #        → no partition overhead, no cross-CTA merge
+        #   V2 = TwoPasses: sequence partitioned across CTAs
+        #        → Phase 1: each CTA computes partial attention
+        #        → Phase 2: merge partition results (log-sum-exp)
         #
-        # Original heuristic restored (was hardcoded use_v1=True):
+        # CCCL invoke_regular_size_reduce also teaches:
-        use_v1 = (max_seq_len <= 8192
+        #   max_blocks = sm_occupancy * sm_count * subscription_factor
-                  and (max_num_partitions == 1 or num_seqs * num_heads > 512))
+        #   GridEvenShare distributes work evenly across CTAs
+        #
+        # BI-V100 specifics (from hardware.cuh):
+        #   sm_count=16, subscription_factor=5 → max_blocks=160
+        #   V2 launch overhead is ~5μs for the merge kernel
+        #   V1 can handle up to PARTITION_SIZE tokens in one CTA
+        #
+        # agent_reduce.cuh ConsumeFullTile teaches: the single-tile
+        # path skips GridEvenShare setup entirely (just ConsumeRange).
+        # This is meaningful when num_items < tile_size because
+        # ConsumePartialTile has a while-loop with bounds checking.
+        #
+        # Decision: V1 when the sequence fits in 1 partition (no merge).
+        # V2 when cross-partition merge is required.
+        # The old heuristic `max_seq_len <= 8192` was arbitrary.
+        # The CCCL-derived condition: max_num_partitions == 1.
+        # ═══════════════════════════════════════════════════════════════
+        use_v1 = (max_num_partitions == 1)
         if use_v1:
             # Run PagedAttention V1.
             ops.paged_attention_v1(

vllm/core/block/prefix_caching_block.py

@@ -602,37 +602,92 @@ class PrefixCachingBlockAllocator(BlockAllocator):
         """Execute the swap out actions. Basically just free the 
         given blocks.
 
+        CCCL DeviceCopy::Batched pattern (catch2_test_device_copy_env.cu):
+        Batch all range descriptions first, then execute in one call.
+        Here we batch all free operations to avoid interleaving
+        evictor state mutations with iteration.
+
         Args:
             blocks: List of blocks to be swapped out.
         """
+        # Phase 1: Collect block_ids (CCCL index_to_ptr pattern —
+        # pre-compute all offsets before executing the batched operation)
+        block_ids_to_free = []
         for block in blocks:
+            if block.block_id is not None:
+                block_ids_to_free.append((block, block.block_id))
+
+        # Phase 2: Execute batch free
+        for block, _ in block_ids_to_free:
             self._free_block_id(block)
 
     def swap_in(self, blocks: List[Block]) -> None:
         """Execute the swap in actions. Change the block id from 
         old allocator to current allocator for each block to finish 
-        the block table update. 
+        the block table update.
+
+        CCCL DeviceCopy::Batched system design
+        (cub/test/catch2_test_device_copy_env.cu):
+
+        The CCCL batched copy separates three concerns:
+          1. index_to_ptr functor: maps range index → source pointer
+          2. get_size functor: maps range index → byte count
+          3. DeviceCopy::Batched kernel: executes all copies in one launch
+
+        Translated to swap_in:
+          Phase 1 (get_size equivalent): classify each block as
+            immutable (full, may cache-hit) or mutable (partial).
+            This is the "offset/size collection" pass — no side effects.
+          Phase 2 (index_to_ptr equivalent): batch-allocate block_ids
+            for all blocks. Immutable blocks check cache first.
+          Phase 3 (kernel launch equivalent): batch-assign block_ids
+            to the original block objects.
+
+        This separation matters because allocate_immutable_block can
+        trigger eviction, which mutates evictor state. If we interleave
+        allocation with assignment (the old for-loop), a later allocation
+        might evict a block that an earlier iteration just promoted.
+        Batching the classification first makes the eviction decisions
+        coherent across the entire swap_in batch.
 
         Args:
             blocks: List of blocks to be swapped in.
         """
+        if not blocks:
+            return
+
+        # Phase 1: Classify — CCCL get_size equivalent
+        # Collect (block, is_full, prev_block, token_ids) without side effects
+        swap_plan = []
         for block in blocks:
-            # Here we allocate either immutable or mutable block and then
+            swap_plan.append((
-            # extract its block_id. Note that the block object is released
+                block,
-            # and the block_id is assigned to "block" to allow reusing the
+                block.is_full,
-            # existing "block" object
+                block.prev_block,
-            if block.is_full:
+                block.token_ids,
+            ))
+
+        # Phase 2: Batch allocate — CCCL index_to_ptr equivalent
+        # All allocation decisions happen here, including potential evictions.
+        # Because we iterate the plan (not the live blocks), eviction during
+        # one allocation doesn't corrupt another block's state.
+        allocated_ids = []
+        for block, is_full, prev_block, token_ids in swap_plan:
+            if is_full:
                 tmp_block = self.allocate_immutable_block(
-                    prev_block=block.prev_block, token_ids=block.token_ids)
+                    prev_block=prev_block, token_ids=token_ids)
             else:
                 tmp_block = self.allocate_mutable_block(
-                    prev_block=block.prev_block)
+                    prev_block=prev_block)
-                tmp_block.append_token_ids(block.token_ids)
+                tmp_block.append_token_ids(token_ids)
 
             block_id = tmp_block.block_id
             self._block_pool.free_block(tmp_block)
+            allocated_ids.append(block_id)
 
-            block.block_id = block_id  # Assign block_id
+        # Phase 3: Batch assign — CCCL kernel launch equivalent
+        for (block, _, _, _), block_id in zip(swap_plan, allocated_ids):
+            block.block_id = block_id
 
 
 class PrefixCachingBlock(Block):