From dd59ec95c2dd6067f99fe2b1991ef485344e6b4e Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 6 Aug 2026 04:12:19 +0000 Subject: [PATCH] [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 | 80 ++++++++++++++++++++----- vllm/core/block/prefix_caching_block.py | 75 +++++++++++++++++++---- 2 files changed, 131 insertions(+), 24 deletions(-) diff --git a/vllm/attention/ops/paged_attn.py b/vllm/attention/ops/paged_attn.py index 5378e46e..bc760918 100644 --- a/vllm/attention/ops/paged_attn.py +++ b/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, - # skip the SMEM raking grid and go directly to warp shuffle. - # This is a CONDITIONAL optimization, not a hardcode. + # ═══════════════════════════════════════════════════════════════ + # CCCL dispatch_reduce.cuh single-tile vs two-phase decision # - # V1 = WARP_SYNCHRONOUS equivalent: single-pass, no temp buffer. - # V2 = raking reduction equivalent: multi-pass with temp buffer. + # dispatch_reduce.cuh line 460: + # 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). - # V2 is faster for long sequences (partitioned reduce + merge). + # The decision is tile-capacity based, not a magic constant. # - # For max_num_seqs=1 (competition config): - # num_seqs * num_heads = 1 * 24 = 24, always < 512 - # → V2 kicks in for max_seq_len > 8192 + # For paged attention, the equivalent: + # V1 = SingleTile: one CTA processes entire sequence in SMEM + # → 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): - use_v1 = (max_seq_len <= 8192 - and (max_num_partitions == 1 or num_seqs * num_heads > 512)) + # CCCL invoke_regular_size_reduce also teaches: + # max_blocks = sm_occupancy * sm_count * subscription_factor + # 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( diff --git a/vllm/core/block/prefix_caching_block.py b/vllm/core/block/prefix_caching_block.py index 7c8a2bc4..b93ff1d2 100644 --- a/vllm/core/block/prefix_caching_block.py +++ b/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 - # extract its block_id. Note that the block object is released - # and the block_id is assigned to "block" to allow reusing the - # existing "block" object - if block.is_full: + swap_plan.append(( + block, + block.is_full, + block.prev_block, + 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) - tmp_block.append_token_ids(block.token_ids) + prev_block=prev_block) + 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):