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

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

Browse Source
This commit is contained in:
muh-bot
2026-08-05 03:32:38 +00:00
parent 41ecb8c9e9 44e4f6f947
commit fa8a02c642
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,