diff --git a/paged_attention_v2_pytorch.py b/paged_attention_v2_pytorch.py index b097f49d..0ee00294 100644 --- a/paged_attention_v2_pytorch.py +++ b/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( diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 8e6092b4..67d7c4a5 100644 --- a/vllm/_custom_ops.py +++ b/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,