[v2] PARTITION_SIZE 512→1024 + fix import path

Two changes based on CCCL source reading:

1. PARTITION_SIZE 512→1024 in paged_attention_v2_pytorch.py
   From dispatch_scan.cuh: grid_size = num_tiles = ceil(N / tile_size).
   Optimal tile_size balances parallelism vs overhead:
   - BI-V100: 16 SMs, max ~32 concurrent CTAs
   - Need num_partitions >= 32 to fill one wave
   - 100K tokens / 1024 = 98 partitions (3 waves) ✓
   - 100K tokens / 512 = 195 partitions (6 waves) — twice the Phase 2 cost
   Note: only affects V2 (PyTorch path). V1 (ixformer) has its own partition size.

2. Fix V2 import path in _custom_ops.py
   paged_attention_v2_pytorch.py is in repo root, not vllm package.
   Added sys.path manipulation to find it at runtime.

Also read: cccl_upstream/thrust/examples/expand.cu (variable-length
replication pattern — maps to GQA expansion, but our broadcast approach
is already more efficient than physical replication).

Source: cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh lines 350-380
        cccl_upstream/thrust/examples/expand.cu

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,