Merge branch 'main' of https://github.com/dylanyunlon/project_6

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(

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(