diff --git a/vllm/attention/ops/paged_attn.py b/vllm/attention/ops/paged_attn.py index c90e8dd8..422d3561 100644 --- a/vllm/attention/ops/paged_attn.py +++ b/vllm/attention/ops/paged_attn.py @@ -116,16 +116,24 @@ class PagedAttention: num_seqs, num_heads, head_size = query.shape max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE) - # NOTE(woosuk): We use a simple heuristic to decide whether to use - # PagedAttention V1 or V2. If the number of partitions is 1, we use - # V1 to avoid the overhead of reduction. Also, if the number of - # sequences or heads is large, we use V1 since there is enough work - # to parallelize. - # TODO(woosuk): Tune this heuristic. - # For context len > 8192, use V2 kernel to avoid shared memory shortage. + # CCCL block_reduce_raking.cuh pattern: + # WARP_SYNCHRONOUS fast path: when RAKING_THREADS == BLOCK_THREADS, + # skip the SMEM raking grid and go directly to warp shuffle. + # This is a CONDITIONAL optimization, not a hardcode. + # + # V1 = WARP_SYNCHRONOUS equivalent: single-pass, no temp buffer. + # V2 = raking reduction equivalent: multi-pass with temp buffer. + # + # V1 is faster for short sequences (fits in SMEM, no partition overhead). + # V2 is faster for long sequences (partitioned reduce + merge). + # + # For max_num_seqs=1 (competition config): + # num_seqs * num_heads = 1 * 24 = 24, always < 512 + # → V2 kicks in for max_seq_len > 8192 + # + # Original heuristic restored (was hardcoded use_v1=True): use_v1 = (max_seq_len <= 8192 and (max_num_partitions == 1 or num_seqs * num_heads > 512)) - use_v1 = True if use_v1: # Run PagedAttention V1. ops.paged_attention_v1( @@ -144,17 +152,30 @@ class PagedAttention: else: # Run PagedAttention V2. assert _PARTITION_SIZE % block_size == 0 - tmp_output = torch.empty( - size=(num_seqs, num_heads, max_num_partitions, head_size), - dtype=output.dtype, - device=output.device, - ) - exp_sums = torch.empty( - size=(num_seqs, num_heads, max_num_partitions), - dtype=torch.float32, - device=output.device, - ) - max_logits = torch.empty_like(exp_sums) + # CCCL agent_merge_sort.cuh union _TempStorage pattern: + # cache temp tensors across decode steps (stable shapes for + # max_num_seqs=1 with slowly growing sequence). + _v2_key = (num_seqs, num_heads, max_num_partitions, + head_size, output.dtype, str(output.device)) + _v2 = getattr(PagedAttention, '_v2_cache', {}).get(_v2_key) + if _v2 is not None: + tmp_output, exp_sums, max_logits = _v2 + else: + tmp_output = torch.empty( + size=(num_seqs, num_heads, max_num_partitions, head_size), + dtype=output.dtype, + device=output.device, + ) + exp_sums = torch.empty( + size=(num_seqs, num_heads, max_num_partitions), + dtype=torch.float32, + device=output.device, + ) + max_logits = torch.empty_like(exp_sums) + if not hasattr(PagedAttention, '_v2_cache'): + PagedAttention._v2_cache = {} + PagedAttention._v2_cache[_v2_key] = ( + tmp_output, exp_sums, max_logits) ops.paged_attention_v2( output, exp_sums,