from dataclasses import dataclass from typing import List, Optional, Tuple import sys import torch import traceback from vllm import _custom_ops as ops # from vllm.attention.ops.prefix_prefill import context_attention_fwd # NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT # imported here. On Iluvatar BI-V100 that kernel hangs the GPU card # permanently. Chunked-prefill / prefix-caching attention is handled by # _forward_prefix_pytorch below (pure PyTorch, no Triton dependency). # Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`. _PARTITION_SIZE = 512 @dataclass class PagedAttentionMetadata: """Metadata for PagedAttention.""" # (batch_size,). The length of sequences (entire tokens seen so far) per # sequence. seq_lens_tensor: Optional[torch.Tensor] # Maximum sequence length in the batch. 0 if it is prefill-only batch. max_decode_seq_len: int # (batch_size, max_blocks_per_seq). # Block addresses per sequence. (Seq id -> list of physical block) # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks # in the kv cache. Each block can contain up to block_size tokens. # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph # captured. block_tables: Optional[torch.Tensor] class PagedAttention: @staticmethod def get_supported_head_sizes() -> List[int]: return [64, 80, 96, 112, 120, 128, 192, 256] @staticmethod def get_kv_cache_shape( num_blocks: int, block_size: int, num_kv_heads: int, head_size: int, ) -> Tuple[int, ...]: return (2, num_blocks, block_size * num_kv_heads * head_size) @staticmethod def split_kv_cache( kv_cache: torch.Tensor, num_kv_heads: int, head_size: int, ) -> Tuple[torch.Tensor, torch.Tensor]: x = 16 // kv_cache.element_size() num_blocks = kv_cache.shape[1] key_cache = kv_cache[0] key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x, -1, x) value_cache = kv_cache[1] value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1) return key_cache, value_cache @staticmethod def write_to_paged_cache( key: torch.Tensor, value: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, slot_mapping: torch.Tensor, kv_cache_dtype: str, k_scale: float, v_scale: float, ) -> None: ops.reshape_and_cache( key, value, key_cache, value_cache, slot_mapping.flatten(), kv_cache_dtype, k_scale, v_scale, ) @staticmethod def _forward_decode_pytorch( query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, scale: float, ) -> torch.Tensor: """Pure-PyTorch decode attention for long contexts (no hardware kernel). Architecture mirrors CCCL's three-layer reduce: dispatch_reduce.cuh → kernel_reduce.cuh → agent_reduce.cuh (work distribution) (kernel entry) (tile consumption) CCCL agent_reduce.cuh has two key patterns we translate here: 1. ConsumeFullTile vectorized path: data loaded as VectorT in striped access (no BlockLoad staging → no SMEM for data, only for BlockReduce scratch). PyTorch equivalent: single reshape+view without .contiguous() when possible; fall back to one .contiguous() per K/V gather. 2. ConsumeTiles with GridEvenShare STRIP_MINE: each CTA strides across the input with stride = grid_size * tile_items. For decode (q_len=1), we tile over KV blocks with adaptive tile_sz per the same GridEvenShare formula: max_tiles = sm_count * subscription_factor. 3. summary_statistics.cu compound reduce: accumulator = {m, l, o}. unary_op: score_tile → (max, sum_exp, weighted_V). binary_op: online softmax merge with correction factor. This is the Flash Attention online softmax — identical structure. For decode, q_len=1 per sequence. The attention weight is [H, 1, seq_len] which is small (~5 MB at 50K tokens). We tile over KV blocks to control peak memory and apply online softmax (Flash Attention Algorithm 1) per tile. Shapes ------ query : [num_seqs, num_heads, head_dim] key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x] value_cache : [num_blocks, num_kv_heads, head_dim, block_size] block_tables: [num_seqs, max_blocks_per_seq] seq_lens : [num_seqs] """ num_seqs, num_heads, head_dim = query.shape num_kv_heads = key_cache.shape[1] block_size = value_cache.shape[3] gqa_ratio = num_heads // num_kv_heads orig_dtype = query.dtype dev = query.device output = torch.empty_like(query) # ================================================================ # CCCL spread_out_items_per_thread adaptive tile sizing for decode # # Ported from dispatch_transform.cuh::spread_out_items_per_thread # and dispatch_reduce.cuh::InvokePasses GridEvenShare. # # CCCL formula (dispatch_transform.cuh line 183): # items = min(max_items, # ceil_div(num_items, sm_count * threads * max_occupancy)) # items = clamp(items, min_items, max_items) # # Our translation for PyTorch decode: # "items" = KV blocks per tile (how much work per matmul call) # "num_items" = total KV blocks in the sequence # "sm_count * max_occupancy" = target number of tiles (~4-8) # Fewer tiles = fewer Python loop iterations = less launch overhead # # For decode (q_len=1), score tensor per tile is tiny: # kv_h × gqa × 1 × (tile_blocks × block_size) × 4 bytes # = 4 × 6 × 1 × 16384 × 4 = 1.5 MB (even at kv_h=4, safe) # So the constraint is NOT memory — it's minimizing loop iterations. # # CCCL grid_even_share.cuh DispatchInit logic: # total_tiles = ceil_div(num_items, tile_size) # grid_size = min(total_tiles, max_grid_size) # big_shares = total_tiles - (avg_tiles * grid_size) # Our target: ~4 tiles max (Python overhead >> kernel launch overhead) # ================================================================ # 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): seq_len = int(seq_lens[i].item()) if seq_len == 0: output[i].zero_() continue num_blocks_i = (seq_len + block_size - 1) // block_size blk_ids = block_tables[i, :num_blocks_i] # Q reshaped once: [kv_h, gqa, 1, d] fp32 — tiny for decode q_grouped = (query[i].float() .view(num_kv_heads, gqa_ratio, head_dim) .unsqueeze(2) .mul_(scale)) # Online softmax accumulators (CCCL summary_stats_data pattern) # accumulator = {m (running max), l (running sum_exp), o (running output)} m = torch.full((num_kv_heads, gqa_ratio, 1), float('-inf'), dtype=torch.float32, device=dev) l = torch.zeros_like(m) o = torch.zeros((num_kv_heads, gqa_ratio, 1, head_dim), dtype=torch.float32, device=dev) # Tile over KV blocks — CCCL spread_out_items_per_thread pattern # Adaptive: tile_blocks = ceil(num_blocks / target_tiles) # clamped to [_MIN_TILE_BLOCKS, _MAX_TILE_BLOCKS] tile_blocks = max(_MIN_TILE_BLOCKS, min(_MAX_TILE_BLOCKS, (num_blocks_i + _BI100_TARGET_TILES - 1) // _BI100_TARGET_TILES)) for tile_start in range(0, num_blocks_i, tile_blocks): tile_end = min(tile_start + tile_blocks, num_blocks_i) tile_blk_ids = blk_ids[tile_start:tile_end] # Valid tokens in this tile tile_token_start = tile_start * block_size tile_token_end = min(tile_end * block_size, seq_len) valid_tokens = tile_token_end - tile_token_start # -------------------------------------------------------- # KV gather — agent_reduce.cuh ConsumeFullTile pattern # # agent_reduce loads VectorT in striped access when possible. # PyTorch equivalent: reshape the 5D cache layout to 3D in # one permute+contiguous, avoiding the double-contiguous # pattern of the old code. # # key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x] # Target: [kv_h, d, valid_tokens] for Q@K^T # # Optimized path: permute(1,2,4,0,3) → [kv_h, d//x, x, n_blk, blk_sz] # → reshape to [kv_h, d, n_blk*blk_sz] → slice [:valid_tokens] # This is ONE contiguous() call instead of TWO. # -------------------------------------------------------- k_gathered = key_cache[tile_blk_ids] # [n, kv_h, d//x, blk_sz, x] k_t = (k_gathered .permute(1, 2, 4, 0, 3) # [kv_h, d//x, x, n, blk_sz] .contiguous() .view(num_kv_heads, head_dim, -1) # [kv_h, d, n*blk_sz] [:, :, :valid_tokens] .unsqueeze(1) # [kv_h, 1, d, valid] .float()) del k_gathered v_gathered = value_cache[tile_blk_ids] # [n, kv_h, d, blk_sz] v_t = (v_gathered .permute(1, 2, 0, 3) # [kv_h, d, n, blk_sz] .contiguous() .view(num_kv_heads, head_dim, -1) # [kv_h, d, n*blk_sz] [:, :, :valid_tokens] .transpose(1, 2) # [kv_h, valid, d] .unsqueeze(1) # [kv_h, 1, valid, d] .float()) del v_gathered # -------------------------------------------------------- # Scores + online softmax — summary_statistics.cu pattern # # unary_op: score_tile → (max, sum_exp, weighted_V) # binary_op: merge with correction factor # # CCCL summary_stats_binary_op merges: # result.mean = x.mean + delta * y.n / n # result.M2 = x.M2 + y.M2 + delta² * x.n * y.n / n # # Online softmax merge: # m_new = max(m_old, m_tile) # corr = exp(m_old - m_new) ← rescale factor # l_new = l_old * corr + l_tile # o_new = o_old * corr + tile_exp @ V # # Structurally identical: m↔max, l↔n, o↔mean×n. # -------------------------------------------------------- # [kv_h, gqa, 1, valid_tokens] s = torch.matmul(q_grouped, k_t) del k_t # Online softmax update (Flash Attention Algorithm 1) m_tile = s.amax(dim=-1, keepdim=True) # [kv_h, gqa, 1, 1] m_new = torch.maximum(m, m_tile.squeeze(-1)) corr = torch.exp(m - m_new) # rescale old accum exp_s = torch.exp(s - m_new.unsqueeze(-1)) del s m.copy_(m_new) l.mul_(corr).add_(exp_s.sum(dim=-1)) o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t)) del exp_s, v_t, corr, m_new, m_tile # Finalize: normalize o.div_(l.unsqueeze(-1)) output[i] = (o.view(num_heads, head_dim) .to(orig_dtype)) except Exception as e: print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}", file=sys.stderr, flush=True) traceback.print_exc(file=sys.stderr) raise return output # ================================================================ # CCCL Design Pattern: summary_statistics.cu transform_reduce # # CCCL packs {n, min, max, mean, M2, M3, M4} into one struct and # computes ALL statistics in a single pass via transform_reduce. # The binary_op merges two partial results (Welford parallel algo). # # Our online softmax is the same pattern: # accumulator = {m (running max), l (running sum_exp), o (running output)} # unary_op: score_tile → {max(tile), sum(exp(tile-max)), exp(tile-max) @ V} # binary_op: merge two accumulators with correction factor # # Key insight: kv_heads are INDEPENDENT — no cross-head dependency. # Current code already batches via [kv_h, gqa, q_len, tile_sz] tensor ops. # The CCCL pattern validates this is optimal: one matmul per tile across # all heads simultaneously, not per-head iteration. # # Future optimization: if we ever get Triton/CUDA access, the binary_op # merge step ({m,l,o} update) could be fused with the matmul via a # custom epilogue — this is what FlashAttention-2/3 does at the CUDA level. # ================================================================ # paged_attention_v1 on BI-V100 fails for long contexts. # Route on actual sequence length (seq_lens.max()), not the max_seq_len # parameter which is inflated to max_model_len in CUDA graph mode. _PYTORCH_DECODE_THRESHOLD = 999999 @staticmethod def forward_decode( query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, max_seq_len: int, kv_cache_dtype: str, num_kv_heads: int, scale: float, alibi_slopes: Optional[torch.Tensor], k_scale: float, v_scale: float, tp_rank: int = 0, blocksparse_local_blocks: int = 0, blocksparse_vert_stride: int = 0, blocksparse_block_size: int = 64, blocksparse_head_sliding_step: int = 0, ) -> torch.Tensor: actual_max = int(seq_lens.max().item()) if seq_lens.numel() > 0 else max_seq_len if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD: return PagedAttention._forward_decode_pytorch( query, key_cache, value_cache, block_tables, seq_lens, scale) if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1: # use blocksparse paged attention block_size = value_cache.size(-1) assert (blocksparse_block_size > 0 and blocksparse_block_size % block_size == 0), \ (f"{blocksparse_block_size=} needs to be a multiple of" f"{block_size=} used in block_tables.") output = torch.empty_like(query) block_size = value_cache.shape[3] num_seqs, num_heads, head_size = query.shape max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE) # NOTE(woosuk): We use a simple heuristic to decide whether to use # PagedAttention V1 or V2. If the number of partitions is 1, we use # V1 to avoid the overhead of reduction. Also, if the number of # sequences or heads is large, we use V1 since there is enough work # to parallelize. # TODO(woosuk): Tune this heuristic. # For context len > 8192, use V2 kernel to avoid shared memory shortage. # CCCL dispatch_reduce.cuh two-path dispatch architecture: # single-tile: num_items ≤ threads × items → one CTA, zero temp buffer # multi-tile: GridEvenShare partitions across sm_count × occupancy CTAs # # Paged attention equivalent: # V1 = single-pass: one CTA iterates ALL KV blocks (like DeviceReduceSingleTileKernel) # V2 = partitioned: KV blocks split into PARTITION_SIZE chunks across CTAs, # then a second kernel merges partition results (like InvokePasses two-phase) # # V1 is optimal when seq_len fits in one CTA's tile (small context). # V2 is optimal when seq_len >> PARTITION_SIZE (long context) — parallelism # across partitions compensates for the merge overhead. # # CCCL's GridEvenShare formula: # max_blocks = sm_occupancy × sm_count × subscription_factor # BI-V100: ~1 × 16 × 5 = 80 max blocks # V2 becomes worthwhile when max_num_partitions > 1 AND the partition # parallelism exceeds the sequence×head parallelism. # # 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_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( output, query, key_cache, value_cache, num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len, alibi_slopes, ) else: # Run PagedAttention V2. assert _PARTITION_SIZE % block_size == 0 # CCCL agent_merge_sort.cuh union _TempStorage pattern: # agent_merge_sort shares a single SMEM allocation across # load_keys, load_items, store_keys, and block_merge ops # (they don't execute concurrently, so one buffer suffices). # Our equivalent: cache V2 temp tensors across decode steps. # For max_num_seqs=1 (competition config), these shapes are # stable across all decode steps for the same sequence. _v2_key = ("v2_tmp", num_seqs, num_heads, max_num_partitions, head_size, output.dtype, output.device) _v2_cached = getattr(PagedAttention, '_v2_cache', {}).get(_v2_key) if _v2_cached is not None: tmp_output, exp_sums, max_logits = _v2_cached else: tmp_output = torch.empty( size=(num_seqs, num_heads, max_num_partitions, head_size), dtype=output.dtype, device=output.device, ) exp_sums = torch.empty( size=(num_seqs, num_heads, max_num_partitions), dtype=torch.float32, device=output.device, ) max_logits = torch.empty_like(exp_sums) if not hasattr(PagedAttention, '_v2_cache'): PagedAttention._v2_cache = {} PagedAttention._v2_cache[_v2_key] = (tmp_output, exp_sums, max_logits) ops.paged_attention_v2( output, exp_sums, max_logits, tmp_output, query, key_cache, value_cache, num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, k_scale, v_scale, tp_rank, blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size, blocksparse_head_sliding_step, ) return output @staticmethod def forward_prefix( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, kv_cache_dtype: str, key_cache: torch.Tensor, value_cache: torch.Tensor, block_tables: torch.Tensor, query_start_loc: torch.Tensor, seq_lens_tensor: torch.Tensor, context_lens: torch.Tensor, max_query_len: int, alibi_slopes: Optional[torch.Tensor], sliding_window: Optional[int], k_scale: float, v_scale: float, ) -> torch.Tensor: # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar # BI-V100 hardware (same class of issue as cudnnFlashAttnForward). # Use a pure-PyTorch fallback that reads the paged KV cache directly. return PagedAttention._forward_prefix_pytorch( query, key, value, key_cache, value_cache, block_tables, query_start_loc, seq_lens_tensor, context_lens, ) @staticmethod def _forward_prefix_pytorch( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, block_tables: torch.Tensor, query_start_loc: torch.Tensor, seq_lens_tensor: torch.Tensor, context_lens: torch.Tensor, ) -> torch.Tensor: """Pure-PyTorch prefix-attention with K-tiling (Flash-Attention online softmax). Memory complexity: O(q_len), independent of kv_len. With chunked prefill (q_len ≤ max_num_batched_tokens = 4096) peak per layer ≈ 96 MB regardless of context length. Algorithm: Flash Attention online softmax. Q is reshaped once to [kv_h, gqa, q_len, d] (24 MB) and held for all K-tiles. For each tile a running (m, l, o) accumulator is updated — the [q_len × kv_len] attention matrix is NEVER materialised in full. Tile budget (kv_h=1, gqa=6, q_len=4096, tile=256 tokens): q_seq [1, 6, 4096, 256] fp32 24 MB (held all tiles) o_acc same shape 24 MB (held all tiles) s same shape 24 MB (per tile, freed before exp_s) exp_s same shape 24 MB (per tile, brief overlap with s) Peak ≈ 96 MB (s and exp_s briefly coexist during update). Shapes ------ query : [total_q_tokens, num_q_heads, head_dim] key : [total_q_tokens, num_kv_heads, head_dim] value : [total_q_tokens, num_kv_heads, head_dim] key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x] value_cache : [num_blocks, num_kv_heads, head_dim, block_size] block_tables : [batch_size, max_blocks_per_seq] query_start_loc: [batch_size + 1] seq_lens_tensor: [batch_size] total length (context + query) context_lens : [batch_size] tokens already in KV cache """ try: # ================================================================ # Tile sizing strategy — ported from CCCL dispatch_reduce.cuh # # CCCL's GridEvenShare computes: # max_blocks = sm_occupancy × sm_count × subscription_factor # tile_size = num_items / max_blocks (evenly distributed) # # For BI-V100 (16 SMs), fixed _BLOCKS_PER_TILE=32 wastes memory # on short contexts and underutilizes on long ones. # # Key insight from kernel_reduce.cuh: # StableReductionOrder=false uses atomicAdd → single kernel pass. # For online softmax (our case), we accumulate (m, l, o) per tile # then merge — this IS a multi-pass reduce. Larger tiles = fewer # merge steps = less numerical drift + less Python loop overhead. # # CCCL subscription_factor = CUB_SUBSCRIPTION_FACTOR(0) = 5 # Effective: 16 SM × 1 CTA/SM × 5 = 80 concurrent tiles max. # But Python loop overhead dominates, so we want FEWER, LARGER tiles. # # Strategy: target ~4-8 tiles per context phase. # Fewer tiles → fewer matmul calls → less launch overhead. # SMEM constraint: score tensor [kv_h, gqa, q_len, tile_sz] fp32 # must not cause OOM. With q_len=4096, kv_h=1, gqa=6: # tile_sz=1024 → 1×6×4096×1024×4 = 96 MB (too much) # tile_sz=512 → 48 MB (borderline) # tile_sz=256 → 24 MB (safe) # For decode (q_len=1): tile_sz=4096 → only 96 KB (always safe) # ================================================================ _SMEM_BUDGET_BYTES = 96 * 1024 * 1024 # 96 MB score tensor budget batch_size = seq_lens_tensor.shape[0] num_q_heads = query.shape[1] num_kv_heads = key_cache.shape[1] head_dim = query.shape[2] gqa_ratio = num_q_heads // num_kv_heads block_size = value_cache.shape[3] scale = head_dim ** -0.5 orig_dtype = query.dtype output = torch.empty_like(query) dev = query.device for i in range(batch_size): ctx_len = int(context_lens[i].item()) q_start = int(query_start_loc[i].item()) q_end = int(query_start_loc[i + 1].item()) q_len = q_end - q_start q_i = query[q_start:q_end] # [q_len, q_h, d] k_i = key [q_start:q_end] # [q_len, kv_h, d] v_i = value[q_start:q_end] # CCCL spread_out_items_per_thread adaptive tile sizing. # # Two constraints compete: # 1. Memory: score tensor [kv_h, gqa, q_len, tile_sz] × 4 ≤ budget # 2. Iteration count: want ~4-8 tiles to minimize Python overhead # # CCCL dispatch_transform.cuh::spread_out_items_per_thread: # items = ceil_div(num_items, sm_count * threads * occupancy) # items = clamp(items, min_items, max_items) # # Our translation: tile_sz = max context tokens / target_tiles, # then clamp by memory budget. score_row_bytes = num_kv_heads * gqa_ratio * q_len * 4 if score_row_bytes > 0: mem_max_tokens = _SMEM_BUDGET_BYTES // score_row_bytes mem_max_tokens = (mem_max_tokens // block_size) * block_size else: mem_max_tokens = block_size * 256 total_kv_tokens = ctx_len + q_len # spread_out: target 4 tiles for context, 4 for current chunk spread_tile = max(block_size, (total_kv_tokens + 3) // 4) # Round to block_size spread_tile = (spread_tile // block_size) * block_size spread_tile = max(spread_tile, block_size) # Clamp by memory budget tile_sz = min(spread_tile, mem_max_tokens) tile_sz = max(tile_sz, block_size) # floor # Q reshaped and scaled once; held for all K-tiles. # [kv_h, gqa, q_len, d] fp32 — 24 MB for q_len=4096, d=256 q_seq = (q_i.permute(1, 0, 2) .float() .view(num_kv_heads, gqa_ratio, q_len, head_dim) .mul_(scale)) # Flash-Attention online-softmax accumulators. # m, l : [kv_h, gqa, q_len] fp32 — <0.1 MB # o : [kv_h, gqa, q_len, d] fp32 — 24 MB m = torch.full((num_kv_heads, gqa_ratio, q_len), float('-inf'), dtype=torch.float32, device=dev) l = torch.zeros_like(m) o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim), dtype=torch.float32, device=dev) # -------------------------------------------------------------- # Phase 1 — context tokens (positions 0 … ctx_len-1). # # Every context key has absolute position < ctx_len; every # query has position ≥ ctx_len. k_pos < q_pos is always True # → no causal mask needed for pure context tiles. # -------------------------------------------------------------- # Convert token-based tile_sz to block count for iteration blocks_per_tile = tile_sz // block_size if ctx_len > 0: num_ctx_blocks = (ctx_len + block_size - 1) // block_size if num_ctx_blocks > block_tables.shape[1]: print( f"[paged_attn WARNING] seq {i}: num_ctx_blocks={num_ctx_blocks} " f"> block_tables.shape[1]={block_tables.shape[1]}, ctx_len={ctx_len}. " "Block table is undersized (prefix_cache_hit bug). " "Capping context to available blocks — attention may be incorrect.", file=sys.stderr, flush=True) num_ctx_blocks = block_tables.shape[1] for tile_blk in range(0, num_ctx_blocks, blocks_per_tile): blk_end = min(tile_blk + blocks_per_tile, num_ctx_blocks) blk_ids = block_tables[i, tile_blk:blk_end] # Gather K/V for this tile. # key_cache [blk_ids]: [n, kv_h, d//x, blk_sz, x] # value_cache[blk_ids]: [n, kv_h, d, blk_sz] k_tile = (key_cache[blk_ids] .permute(0, 3, 1, 2, 4) .contiguous() .view(-1, num_kv_heads, head_dim)) v_tile = (value_cache[blk_ids] .permute(0, 3, 1, 2) .contiguous() .view(-1, num_kv_heads, head_dim)) # Trim padding in the last block of the tile. valid = (min(blk_end * block_size, ctx_len) - tile_blk * block_size) k_tile = k_tile[:valid] # [valid, kv_h, d] v_tile = v_tile[:valid] # k_t: [kv_h, 1, d, valid] (broadcast over gqa_ratio) # v_t: [kv_h, 1, valid, d] k_t = (k_tile.permute(1, 0, 2) .unsqueeze(1) .transpose(-1, -2) .float()) v_t = (v_tile.permute(1, 0, 2) .unsqueeze(1) .float()) del k_tile, v_tile # Scores: [kv_h, gqa, q_len, valid] s = torch.matmul(q_seq, k_t) del k_t # No causal mask: all context keys precede all queries. # Online softmax update — Flash-Attention Algorithm 1. # exp_s = s - new_max (in-place exp after del s) m_blk = s.amax(dim=-1) m_new = torch.maximum(m, m_blk) exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() corr = torch.exp(m - m_new) m.copy_(m_new) del m_blk, m_new l.mul_(corr).add_(exp_s.sum(dim=-1)) o.mul_(corr.unsqueeze(-1)).add_( torch.matmul(exp_s, v_t)) del exp_s, v_t, corr # -------------------------------------------------------------- # Phase 2 — current-chunk tokens (positions ctx_len … ctx_len+q_len-1). # # Causal mask: query at relative position j sees key at relative # position k only when k ≤ j. Tiles of tile_sz tokens each. # -------------------------------------------------------------- for kc_start in range(0, q_len, tile_sz): kc_end = min(kc_start + tile_sz, q_len) kc_len = kc_end - kc_start k_blk = k_i[kc_start:kc_end] # [kc_len, kv_h, d] v_blk = v_i[kc_start:kc_end] k_t = (k_blk.permute(1, 0, 2) .unsqueeze(1) .transpose(-1, -2) .float()) # [kv_h, 1, d, kc_len] v_t = (v_blk.permute(1, 0, 2) .unsqueeze(1) .float()) # [kv_h, 1, kc_len, d] s = torch.matmul(q_seq, k_t) # [kv_h, gqa, q_len, kc_len] del k_t # Causal mask: key at (kc_start+k) must not exceed query j. k_rel = torch.arange(kc_start, kc_end, device=dev) q_rel = torch.arange(q_len, device=dev) mask = k_rel.unsqueeze(0) > q_rel.unsqueeze(1) # [q_len, kc_len] s.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf')) del mask, k_rel, q_rel # Online softmax update (identical to context phase). m_blk = s.amax(dim=-1) m_new = torch.maximum(m, m_blk) exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() corr = torch.exp(m - m_new) m.copy_(m_new) del m_blk, m_new l.mul_(corr).add_(exp_s.sum(dim=-1)) o.mul_(corr.unsqueeze(-1)).add_( torch.matmul(exp_s, v_t)) del exp_s, v_t, corr # -------------------------------------------------------------- # Finalize: normalize running output by normalization factor. # o: [kv_h, gqa, q_len, d] → [q_len, q_h, d] # -------------------------------------------------------------- o.div_(l.unsqueeze(-1)) output[q_start:q_end] = ( o.view(num_q_heads, q_len, head_dim) .permute(1, 0, 2) .to(orig_dtype) ) except Exception as e: print(f"[paged_attn ERROR] {type(e).__name__}: {e}", file=sys.stderr, flush=True) traceback.print_exc(file=sys.stderr) raise return output @staticmethod def swap_blocks( src_kv_cache: torch.Tensor, dst_kv_cache: torch.Tensor, src_to_dst: torch.Tensor, ) -> None: src_key_cache = src_kv_cache[0] dst_key_cache = dst_kv_cache[0] ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) src_value_cache = src_kv_cache[1] dst_value_cache = dst_kv_cache[1] ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) @staticmethod def copy_blocks( kv_caches: List[torch.Tensor], src_to_dists: torch.Tensor, ) -> None: key_caches = [kv_cache[0] for kv_cache in kv_caches] value_caches = [kv_cache[1] for kv_cache in kv_caches] ops.copy_blocks(key_caches, value_caches, src_to_dists)