diff --git a/paged_attn.py b/paged_attn.py index 988f9032..4df75a46 100644 --- a/paged_attn.py +++ b/paged_attn.py @@ -123,7 +123,10 @@ class PagedAttention: # For context len > 8192, use V2 kernel to avoid shared memory shortage. use_v1 = (max_seq_len <= 8192 and (max_num_partitions == 1 or num_seqs * num_heads > 512)) - use_v1 = True + # V2 is now implemented via paged_attention_v2_pytorch.py (CCCL two-pass pattern). + # For short sequences (<=8192), V1 (ixformer pre-compiled) is faster. + # For long sequences (>8192), V2 partitions work across CTAs. + # On BI-V100 (16 SMs), V2's partition reduction fits in L2 (6MB). if use_v1: # Run PagedAttention V1. ops.paged_attention_v1( diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index ca9bceef..8e6092b4 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -144,7 +144,22 @@ def paged_attention_v2( blocksparse_block_size: int = 64, blocksparse_head_sliding_step: int = 0, ) -> None: - raise NotImplementedError() + # CCCL two-pass dispatch pattern (dispatch_reduce.cuh): + # Pass 1: N CTAs each reduce their tile → d_block_reductions[N] + # Pass 2: 1 CTA reduces d_block_reductions[N] → d_out + # 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) + from paged_attention_v2_pytorch import paged_attention_v2_pytorch + paged_attention_v2_pytorch( + out, exp_sum, max_logits, tmp_out, + 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, + ) def paged_attention_rocm(