diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index 67d7c4a5..4ba1c1ad 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -150,13 +150,34 @@ def paged_attention_v2( # Our PyTorch V2 implementation follows the same pattern: # Phase 1: partition attention (each partition = one tile) # Phase 2: cross-partition log-sum-exp reduction (summary_statistics binary_op) + # paged_attention_v2_pytorch.py — try multiple import locations + # In docker: may be at /workspace/, next to vllm package, or in vllm/ itself import sys, os - # paged_attention_v2_pytorch.py is in the repo root, not inside vllm package - _repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - if _repo_root not in sys.path: - sys.path.insert(0, _repo_root) - from paged_attention_v2_pytorch import paged_attention_v2_pytorch - paged_attention_v2_pytorch( + _pav2 = None + # Try 1: same package (patch_ops copies it next to _custom_ops.py) + try: + from vllm.paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + except ImportError: + pass + # Try 2: /workspace/ (Dockerfile WORKDIR) + if _pav2 is None: + try: + _ws = '/workspace' + if _ws not in sys.path: + sys.path.insert(0, _ws) + from paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + except ImportError: + pass + # Try 3: repo root relative to this file + if _pav2 is None: + _repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + from paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + _pav2( out, exp_sum, max_logits, tmp_out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, seq_lens, diff --git a/qwen3_6_scripts/paged_attention_v2_pytorch.py b/qwen3_6_scripts/paged_attention_v2_pytorch.py new file mode 100644 index 00000000..e0b18a30 --- /dev/null +++ b/qwen3_6_scripts/paged_attention_v2_pytorch.py @@ -0,0 +1,325 @@ +""" +paged_attention_v2_pytorch.py — BI-V100 PagedAttention V2 (CCCL-informed) +=========================================================================== + +Fills the `raise NotImplementedError()` hole in vllm/_custom_ops.py. + +Algorithm: Partitioned attention with log-sum-exp reduction. +Architecture informed by CCCL patterns: + - summary_statistics.cu: fuse multiple statistics in a single reduction pass + - warp_reduce_shfl.cuh: accumulate (max, sum, weighted_output) as one compound type + - block_reduce_warp_reductions.cuh: reduce across partitions via shared accumulators + +Key optimization: Batched partition attention via reshaped 3D bmm. + Instead of looping over P partitions with P × torch.bmm calls, + reshape KV into [H, P*part_len, d] and Q into [H, 1, d], then + slice scores into [H, P, part_len] for partition-wise softmax. + This gives ONE bmm launch for all partitions. + + For seq_len=100K, PARTITION_SIZE=512: + Before: 195 × bmm([H,1,d] @ [H,d,512]) = 195 kernel launches + After: 1 × bmm([H,1,d] @ [H,d,100K]) + reshape = 1 kernel launch + + The partition-wise softmax is then a reshape + per-chunk operation: + scores: [H, 100K] → [H, P, 512] → max/exp/sum per partition + +Phase 2 reduction (cross-partition combine) follows CCCL's summary_statistics +binary_op pattern: combine (max_a, sum_a, out_a) with (max_b, sum_b, out_b) +using the numerically stable log-sum-exp rescaling. +""" + +import torch +from typing import Optional + +_PARTITION_SIZE = 1024 # CCCL dispatch_scan.cuh insight: tile_size balances +# parallelism (num_partitions >= SM_count * 2 to fill one wave) vs overhead +# (fewer partitions = smaller Phase 2 reduction). +# BI-V100: 16 SMs, max ~32 concurrent CTAs. +# For 100K tokens: 1024 → 98 partitions (3 waves), 512 → 195 (6 waves). +# 98 > 32 so parallelism is sufficient; halving partitions halves Phase 2 cost. + +# CCCL dispatch_reduce.cuh GridEvenShare formula (line ~180): +# max_blocks = sm_occupancy * sm_count * subscription_factor +# subscription_factor = 5 (default in cub/util_device.cuh) +# For BI-V100: sm_count=16, sm_occupancy ~= 2 (limited by registers/SMEM) +# → max_blocks = 2 * 16 * 5 = 160 +# If seq_len=100K with PARTITION_SIZE=1024 → 98 partitions < 160 → fine. +# Threshold for V1→V2 handoff: when single-tile can't hold all tokens. +# CCCL single_tile threshold = threads * items_per_thread +# = 512 * 24 = 12288 tokens → V1 handles ≤12288, V2 handles >12288. +# This aligns with BI-V100 paged_attn.py _PARTITION_SIZE=512: +# V2 triggers when seq_len > 512 * (max_blocks_per_seq_for_v1). +_BI100_SM_COUNT = 16 +_BI100_SM_OCCUPANCY = 2 # conservative: 2 CTAs per SM +_BI100_SUBSCRIPTION_FACTOR = 5 # CCCL default +_BI100_MAX_GRID = _BI100_SM_OCCUPANCY * _BI100_SM_COUNT * _BI100_SUBSCRIPTION_FACTOR # 160 + + +def paged_attention_v2_pytorch( + output: torch.Tensor, # [num_seqs, num_heads, head_size] + exp_sums: torch.Tensor, # [num_seqs, num_heads, max_num_partitions] + max_logits: torch.Tensor, # [num_seqs, num_heads, max_num_partitions] + tmp_output: torch.Tensor, # [num_seqs, num_heads, max_num_partitions, head_size] + query: torch.Tensor, # [num_seqs, num_heads, head_size] + key_cache: torch.Tensor, # [num_blocks, num_kv_heads, head_size/x, block_size, x] + value_cache: torch.Tensor, # [num_blocks, num_kv_heads, head_size, block_size] + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, # [num_seqs, max_blocks_per_seq] + seq_lens: torch.Tensor, # [num_seqs] + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str = "auto", + k_scale: float = 1.0, + v_scale: float = 1.0, + 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, +) -> None: + num_seqs, num_heads, head_size = query.shape + gqa_ratio = num_heads // num_kv_heads + max_num_partitions = tmp_output.shape[2] + + # Initialize unused slots + max_logits.fill_(float('-inf')) + exp_sums.zero_() + tmp_output.zero_() + + # CCCL kernel_reduce.cuh SingleTile fast path (line ~270): + # if (num_items <= threads_per_block * items_per_thread) + # → InvokeSingleTile() — one CTA, no temp buffer, no Phase 2 + # PyTorch translation: if seq_len fits in one partition, skip Phase 2 entirely. + # This avoids the partition/reshape/bmm overhead for short decode sequences. + # Qwen3.6 typical decode: seq_len grows from 1 to 100K over generation. + # Early tokens (seq_len < 1024) hit this fast path every step. + _SINGLE_TILE_THRESHOLD = _PARTITION_SIZE # sequences this short skip partitioning + + for seq_idx in range(num_seqs): + seq_len = int(seq_lens[seq_idx].item()) + if seq_len == 0: + output[seq_idx].zero_() + continue + + num_blocks_seq = (seq_len + block_size - 1) // block_size + num_partitions = (seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE + + # ─── CCCL SingleTile fast path ─────────────────────────── + # From kernel_reduce.cuh: when everything fits in one tile, + # do a single-pass attention without partition overhead. + # agent_reduce.cuh ConsumeRange → BlockReduce → done. + if num_partitions == 1: + blk_ids = block_tables[seq_idx, :num_blocks_seq] + q = query[seq_idx].float() # [H, d] + + # Gather KV (same as below but no partition reshape) + k_gathered = key_cache[blk_ids] + k_flat = (k_gathered + .permute(0, 3, 1, 2, 4) + .reshape(-1, num_kv_heads, head_size))[:seq_len] + v_flat = (value_cache[blk_ids] + .permute(0, 3, 1, 2) + .reshape(-1, num_kv_heads, head_size))[:seq_len] + + if k_scale != 1.0: + k_flat = k_flat.float().mul_(k_scale) + if v_scale != 1.0: + v_flat = v_flat.float().mul_(v_scale) + + if gqa_ratio > 1: + k_kv = k_flat.permute(1, 2, 0).float().contiguous() + v_kv = v_flat.permute(1, 0, 2).float().contiguous() + q_grouped = q.view(num_kv_heads, gqa_ratio, 1, head_size) + scores = torch.matmul(q_grouped, k_kv.unsqueeze(1)).squeeze(2) + scores = scores.reshape(num_heads, seq_len) * scale + else: + k_t = k_flat.permute(1, 2, 0).float().contiguous() + scores = torch.bmm(q.unsqueeze(1), k_t).squeeze(1) * scale + + if alibi_slopes is not None: + positions = torch.arange(seq_len, device=query.device, dtype=torch.float32) + scores = scores + alibi_slopes.unsqueeze(1) * positions.unsqueeze(0) + + # Direct softmax + V weighted sum — no partition overhead + weights = torch.softmax(scores, dim=-1) # [H, seq_len] + if gqa_ratio > 1: + w_grouped = weights.view(num_kv_heads, gqa_ratio, 1, seq_len) + result = torch.matmul(w_grouped, v_kv.unsqueeze(1)).squeeze(2) + output[seq_idx] = result.reshape(num_heads, head_size).to(output.dtype) + else: + v_perm = v_flat.permute(1, 0, 2).float().contiguous() + result = torch.bmm(weights.unsqueeze(1), v_perm).squeeze(1) + output[seq_idx] = result.to(output.dtype) + + # Store dummy partition values for compatibility + max_logits[seq_idx, :, 0] = scores.max(dim=-1).values + exp_sums[seq_idx, :, 0] = weights.sum(dim=-1) + tmp_output[seq_idx, :, 0, :] = output[seq_idx].float() + continue + # ─── End SingleTile fast path ──────────────────────────── + + # ============================================================= + # Batched KV gather: ONE index_select, ONE reshape + # Pattern: avoid per-block Python loop (CCCL does this via + # block-cooperative load, we do it via batched indexing) + # ============================================================= + blk_ids = block_tables[seq_idx, :num_blocks_seq] + + # Key: [nblk, kv_h, d/x, blk_sz, x] → [nblk*blk_sz, kv_h, d] + k_gathered = key_cache[blk_ids] + k_flat = (k_gathered + .permute(0, 3, 1, 2, 4) + .reshape(-1, num_kv_heads, head_size))[:seq_len] + + # Value: [nblk, kv_h, d, blk_sz] → [nblk*blk_sz, kv_h, d] + v_flat = (value_cache[blk_ids] + .permute(0, 3, 1, 2) + .reshape(-1, num_kv_heads, head_size))[:seq_len] + + if k_scale != 1.0: + k_flat = k_flat.float().mul_(k_scale) + if v_scale != 1.0: + v_flat = v_flat.float().mul_(v_scale) + + # ============================================================= + # GQA broadcast: avoid materializing the expanded KV tensor + # + # Qwen3.6: H=24, kv_h=4, gqa_ratio=6, head_dim=256 + # Old: expand kv_h→H then contiguous → allocates seq_len×H×d (1.2GB at 100K) + # New: reshape Q as [kv_h, gqa, 1, d], K as [kv_h, 1, d, seq_len] + # → bmm with broadcasting → [kv_h, gqa, 1, seq_len] + # → reshape to [H, seq_len] + # Saves: gqa_ratio × memory (6x for Qwen3.6 = 1GB per decode step) + # ============================================================= + q = query[seq_idx].float() # [H, d] + + if gqa_ratio > 1: + # K: [seq_len, kv_h, d] → [kv_h, d, seq_len] (no GQA expansion) + k_kv = k_flat.permute(1, 2, 0).float().contiguous() # [kv_h, d, seq_len] + v_kv = v_flat.permute(1, 0, 2).float().contiguous() # [kv_h, seq_len, d] + + # Q: [H, d] → [kv_h, gqa, 1, d] + q_grouped = q.view(num_kv_heads, gqa_ratio, 1, head_size) + + # Scores: [kv_h, gqa, 1, d] @ [kv_h, 1, d, seq_len] → [kv_h, gqa, 1, seq_len] + scores_all = torch.matmul(q_grouped, k_kv.unsqueeze(1)).squeeze(2) # [kv_h, gqa, seq_len] + scores_all = scores_all.reshape(num_heads, seq_len) * scale # [H, seq_len] + else: + k_t = k_flat.permute(1, 2, 0).float().contiguous() # [H, d, seq_len] + scores_all = torch.bmm(q.unsqueeze(1), k_t).squeeze(1) * scale # [H, seq_len] + + # Alibi bias (if needed) + if alibi_slopes is not None: + positions = torch.arange(seq_len, device=query.device, dtype=torch.float32) + scores_all = scores_all + alibi_slopes.unsqueeze(1) * positions.unsqueeze(0) + + # Pad to exact multiple of _PARTITION_SIZE for clean reshape + padded_len = num_partitions * _PARTITION_SIZE + if padded_len > seq_len: + pad_size = padded_len - seq_len + scores_padded = torch.full( + (num_heads, padded_len), float('-inf'), + dtype=scores_all.dtype, device=scores_all.device) + scores_padded[:, :seq_len] = scores_all + else: + scores_padded = scores_all + + # Reshape: [H, padded_len] → [H, P, part_sz] + scores_parts = scores_padded.view(num_heads, num_partitions, _PARTITION_SIZE) + + # Per-partition online softmax (vectorized over H and P simultaneously) + # Pattern from CCCL summary_statistics: compute (max, sum) in one pass + part_max = scores_parts.max(dim=-1).values # [H, P] + scores_exp = torch.exp(scores_parts - part_max.unsqueeze(-1)) # [H, P, part_sz] + part_sum = scores_exp.sum(dim=-1) # [H, P] + + # Weighted values per partition: need V reshaped the same way + # V: [seq_len, H, d] → pad → [padded_len, H, d] → [H, P, part_sz, d] + if gqa_ratio > 1: + v_perm = v_kv # already [kv_h, seq_len, d], no GQA expansion needed + # Will handle GQA in the bmm below via broadcast + else: + v_perm = v_flat.permute(1, 0, 2).float().contiguous() # [H, seq_len, d] + # Weighted V sum per partition + # NOTE: v_perm shape differs by GQA mode: + # GQA: v_perm = v_kv = [kv_h, seq_len, d] + # No GQA: v_perm = [H, seq_len, d] + # scores_exp: [H, P, part_sz] → [kv_h, gqa, P, part_sz] + # v_perm: [kv_h, seq_len, d] → [kv_h, P, part_sz, d] + if gqa_ratio > 1: + se_grouped = scores_exp.view(num_kv_heads, gqa_ratio, num_partitions, _PARTITION_SIZE) + # V: pad and reshape to [kv_h, P, part_sz, d] + if padded_len > seq_len: + v_padded_kv = torch.zeros( + (num_kv_heads, padded_len, head_size), + dtype=v_kv.dtype, device=v_kv.device) + v_padded_kv[:, :seq_len, :] = v_kv + else: + v_padded_kv = v_kv + v_parts_kv = v_padded_kv.view(num_kv_heads, num_partitions, _PARTITION_SIZE, head_size) + # Broadcast: [kv_h, gqa, P, 1, part_sz] @ [kv_h, 1, P, part_sz, d] + # → [kv_h, gqa, P, 1, d] + part_out_grouped = torch.matmul( + se_grouped.unsqueeze(3), # [kv_h, gqa, P, 1, part_sz] + v_parts_kv.unsqueeze(1) # [kv_h, 1, P, part_sz, d] + ).squeeze(3) # [kv_h, gqa, P, d] + part_out = part_out_grouped.reshape(num_heads, num_partitions, head_size) + else: + # Non-GQA: v_perm is [H, seq_len, d], pad and reshape normally + if padded_len > seq_len: + v_padded = torch.zeros( + (num_heads, padded_len, head_size), + dtype=v_perm.dtype, device=v_perm.device) + v_padded[:, :seq_len, :] = v_perm + else: + v_padded = v_perm + v_parts = v_padded.view(num_heads, num_partitions, _PARTITION_SIZE, head_size) + HP = num_heads * num_partitions + scores_exp_flat = scores_exp.reshape(HP, 1, _PARTITION_SIZE) + v_parts_flat = v_parts.reshape(HP, _PARTITION_SIZE, head_size) + part_out_flat = torch.bmm(scores_exp_flat, v_parts_flat) # [HP, 1, d] + part_out = part_out_flat.view(num_heads, num_partitions, head_size) # [H, P, d] + + # Store partition results + max_logits[seq_idx, :, :num_partitions] = part_max + exp_sums[seq_idx, :, :num_partitions] = part_sum + tmp_output[seq_idx, :, :num_partitions, :] = part_out.to(tmp_output.dtype) + + # ============================================================= + # Phase 2: Cross-partition reduction (CCCL binary_op pattern) + # + # CCCL kernel_reduce.cuh insight: when grid_size fits in a single + # tile (num_partitions <= threads * items_per_thread), the reduce + # uses SingleTile path — one CTA, no temp buffer, no pass 2 kernel. + # + # For BI-V100 with 98 partitions (100K tokens / 1024 partition_size): + # SingleTile threshold = 512 * 24 = 12288 >> 98 → always SingleTile + # This means Phase 2 is never the bottleneck. + # + # CCCL single_pass_scan_operators.cuh insight: delay() has a + # GridThreshold=500 gate. BI-V100 scan grids are always < 500 blocks, + # so ALL delay strategies (no_delay, fixed_delay, exponential_backon) + # collapse to __threadfence_block(). Delay tuning is irrelevant here. + # + # Phase 2 follows summary_statistics.cu binary_op: combine + # (max_a, sum_a, out_a) ⊕ (max_b, sum_b, out_b) via log-sum-exp. + # Fully vectorized — no loop over partitions. + # ============================================================= + pm = max_logits[seq_idx, :, :num_partitions] # [H, P] + ps = exp_sums[seq_idx, :, :num_partitions] # [H, P] + po = tmp_output[seq_idx, :, :num_partitions, :] # [H, P, d] + + global_max = pm.max(dim=-1).values # [H] + rescale = torch.exp(pm - global_max.unsqueeze(-1)) * ps # [H, P] + total = rescale.sum(dim=-1, keepdim=True) # [H, 1] + + # CCCL norm.cu principle: fuse transform with reduce to minimize traversals. + # Instead of: weights = rescale/total; final = bmm(weights, po) + # Do: final = bmm(rescale, po) / total + # Saves one element-wise division kernel launch (rescale/total → H*P elements). + # The division moves to the output (H*d elements, typically smaller than H*P). + # [H, 1, P] @ [H, P, d] → [H, 1, d] → [H, d] + final = torch.bmm(rescale.unsqueeze(1), po.float()).squeeze(1) / total # [H, d] + output[seq_idx] = final.to(output.dtype) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index f76f8156..f23f553f 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -42,6 +42,14 @@ echo "[patch_ops] _custom_ops.py → / (SMEM 32KB→48KB fix)" cp ./paged_attn.py $V/attention/ops/paged_attn.py echo "[patch_ops] paged_attn.py → attention/ops/" +# --- prefix_prefill.py: Triton-free prefix attention ------------------------- +# On BI-V100, Triton is not installed. This file provides the context_attention_fwd +# function that paged_attn.py imports. Even though our paged_attn.py comments out +# the Triton import, the base xformers.py may still try to import it. +# Deploy it so the import doesn't crash — the function itself won't be called. +cp ./prefix_prefill.py $V/attention/ops/prefix_prefill.py +echo "[patch_ops] prefix_prefill.py → attention/ops/" + # --- model_runner.py: prefix_cache_hit fix ----------------------------------- # Bug: Case 1 (prefix_cache_len <= context_len) leaves prefix_cache_hit=True, # causing undersized block_tables in chunked prefill chunk 2+. @@ -99,6 +107,14 @@ cp ./qwen3_5.py $V/model_executor/models/qwen3_5.py cp ./registry.py $V/model_executor/models/registry.py echo "[patch_ops] qwen3_5.py + registry.py deployed" +# --- paged_attention_v2_pytorch.py: PyTorch V2 attention fallback ------------ +# _custom_ops.py imports this from the parent of its own directory. +# In docker, vllm lives at $V/, so we place it one level up AND next to _custom_ops. +# Belt-and-suspenders: also copy to /workspace/ where _custom_ops.py's _repo_root points. +cp ./paged_attention_v2_pytorch.py $V/paged_attention_v2_pytorch.py +cp ./paged_attention_v2_pytorch.py /workspace/paged_attention_v2_pytorch.py +echo "[patch_ops] paged_attention_v2_pytorch.py → $V/ + /workspace/" + # --- sequence.py: fix completion_tokens inflation ---------------------------- cp ./sequence.py $V/sequence.py echo "[patch_ops] sequence.py → /"