dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[v2] PARTITION_SIZE 512→1024 + fix import path

Browse Source 
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
This commit is contained in:
project_6
2026-08-05 03:32:23 +00:00
parent a0cddf2ddc
commit 44e4f6f947
2 changed files with 11 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
paged_attention_v2_pytorch.py
View File

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

5
vllm/_custom_ops.py
View File

@@ -150,6 +150,11 @@ def paged_attention_v2(
# Our PyTorch V2 implementation follows the same pattern: # Our PyTorch V2 implementation follows the same pattern:
# Phase 1: partition attention (each partition = one tile) # Phase 1: partition attention (each partition = one tile)
# Phase 2: cross-partition log-sum-exp reduction (summary_statistics binary_op) # 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 from paged_attention_v2_pytorch import paged_attention_v2_pytorch
paged_attention_v2_pytorch( paged_attention_v2_pytorch(
out, exp_sum, max_logits, tmp_out, out, exp_sum, max_logits, tmp_out,