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

paged_attention_v2_pytorch.py

@@ -31,7 +31,12 @@ using the numerically stable log-sum-exp rescaling.
 import torch
 from typing import Optional
 
-_PARTITION_SIZE = 512
+_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.
 
 
 def paged_attention_v2_pytorch(

vllm/_custom_ops.py

@@ -150,6 +150,11 @@ 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)
+    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(
         out, exp_sum, max_logits, tmp_out,