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