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

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

Browse Source 
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
This commit is contained in:
Claude
2026-08-06 04:12:19 +00:00
parent 5aba296eba
commit dd59ec95c2
2 changed files with 131 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
vllm/attention/ops/paged_attn.py
View File

@@ -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(

75
vllm/core/block/prefix_caching_block.py
View File

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