From 44e4f6f9472f977b4803229dc048fde60baa057e Mon Sep 17 00:00:00 2001 From: project_6 Date: Wed, 5 Aug 2026 03:32:23 +0000 Subject: [PATCH] =?UTF-8?q?[v2]=20PARTITION=5FSIZE=20512=E2=86=921024=20+?= =?UTF-8?q?=20fix=20import=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 7 ++++++- vllm/_custom_ops.py | 5 +++++ 2 files changed, 11 insertions(+), 1 deletion(-) 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,