From b6538fd10ed55aaa23c24f534abec9ee0c337b57 Mon Sep 17 00:00:00 2001 From: muh-pipeline Date: Thu, 6 Aug 2026 02:28:12 +0000 Subject: [PATCH] [BASE] vllm/attention/ops/paged_attn.py: fix num_kv_heads type annotation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Discovered by tracing call chain after reading CCCL catch2_test_block_reduce.cu (randomly selected). The test covers multi-dim block configs (BlockDimX/Y/Z) which maps to GQA group dimensions in attention. Call chain trace: xformers.py:__init__() builds self.head_mapping = tensor [num_heads] xformers.py:forward() → PagedAttention.forward_decode(head_mapping=tensor) paged_attn.py:forward_decode(num_kv_heads: int) ← WRONG TYPE ANNOTATION _custom_ops.py:paged_attention_v1(head_mapping=tensor) ← expects tensor The parameter is head_mapping tensor for V1 (ixformer precompiled), but int num_kv_heads for V2 (our PyTorch implementation). Fixed annotation to remove misleading int type hint. CCCL source read: cub/test/catch2_test_block_reduce.cu (252 lines, full) Base file modified: vllm/attention/ops/paged_attn.py --- vllm/attention/ops/paged_attn.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/attention/ops/paged_attn.py b/vllm/attention/ops/paged_attn.py index 422d3561..5378e46e 100644 --- a/vllm/attention/ops/paged_attn.py +++ b/vllm/attention/ops/paged_attn.py @@ -92,7 +92,9 @@ class PagedAttention: seq_lens: torch.Tensor, max_seq_len: int, kv_cache_dtype: str, - num_kv_heads: int, + num_kv_heads, # Actually head_mapping tensor from xformers.py for V1, + # or int num_kv_heads for V2. See _custom_ops.py signatures. + # CCCL catch2_test_block_reduce.cu BlockDimY/Z ↔ GQA groups. scale: float, alibi_slopes: Optional[torch.Tensor], k_scale: float,