From 730831f26790ed217d3f1cf8fc85375150a893b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:27:17 +0000 Subject: [PATCH] =?UTF-8?q?[fix]=20paged=5Fattn:=20re-force=20V1=20?= =?UTF-8?q?=E2=80=94=20V2=20is=20pure=20PyTorch,=20not=20C++=20(confirmed?= =?UTF-8?q?=20from=20=5Fcustom=5Fops.py)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- paged_attn.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/paged_attn.py b/paged_attn.py index 1c284fcb..13b8ce3c 100644 --- a/paged_attn.py +++ b/paged_attn.py @@ -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(