[fix] paged_attn: re-force V1 — V2 is pure PyTorch, not C++ (confirmed from _custom_ops.py)

Reading _custom_ops.py as input revealed:
  paged_attention_v1 → ixf_F.vllm_single_query_cached_kv_attention (C++ fused kernel)
  paged_attention_v2 → paged_attention_v2_pytorch (pure Python for-loop)

The previous commit incorrectly removed use_v1=True assuming V2 had a C++ backend.
V2 tensor pre-allocation kept for future C++/Triton V2 implementation.
PARTITION_SIZE=1024 change kept (benefits future V2).
This commit is contained in:
Claude
2026-08-05 06:27:17 +00:00
parent b6e7bca45a
commit 730831f267

View File

@@ -136,22 +136,15 @@ class PagedAttention:
# For context len > 8192, use V2 kernel to avoid shared memory shortage.
use_v1 = (max_seq_len <= 8192
and (max_num_partitions == 1 or num_seqs * num_heads > 512))
# V1/V2 adaptive dispatch (restored from original vllm logic):
# FORCE V1: _custom_ops.py V2 falls through to paged_attention_v2_pytorch
# which is pure PyTorch (for-loop over seqs + multiple kernel launches).
# V1 (ixf_F.vllm_single_query_cached_kv_attention) is a single fused C++ kernel.
# Until a C++ or Triton V2 implementation exists, V1 is always faster.
#
# V1: single fused C++ kernel, iterates ALL KV blocks in one CTA.
# Best when: short sequences (≤8192), or enough seqs×heads for parallelism.
#
# V2: partitioned attention with cross-partition reduce.
# ops.paged_attention_v2 IS a C++ kernel (not pure PyTorch).
# Best when: long sequences where V1's single CTA cannot saturate 16 SMs.
#
# CCCL parallel: V2 reduce pass = DeviceReduce over compound accumulator
# (max_logits, exp_sums, partial_output). The accumulator merge uses the
# same online softmax pattern as thrust/examples/summary_statistics.cu.
#
# With PARTITION_SIZE=1024 and 16 SMs, V2 is beneficial for sequences
# longer than 1024 * 16 * 2 = 32768 tokens (where V1 would have a single
# CTA iterating for too long while other SMs sit idle).
# The V2 tensor pre-allocation below is kept for when C++ V2 becomes available.
# CCCL parallel: V2 reduce = DeviceReduce over compound (max, exp_sum, output)
# using thrust/examples/summary_statistics.cu Welford merge pattern.
use_v1 = True
if use_v1:
# Run PagedAttention V1.
ops.paged_attention_v1(