From d2b4df54ff4cf897bd3bbd181d3ad38aaa2766f5 Mon Sep 17 00:00:00 2001 From: project6-dev Date: Thu, 13 Aug 2026 07:09:50 +0000 Subject: [PATCH] =?UTF-8?q?perf:=20native=20ixformer=20decode=20=E2=80=94?= =?UTF-8?q?=20v1=20=E2=89=A432K,=20v2=20>32K=20(no=20Python=20fallback)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decode path: - ≤32K: paged_attention_v1 (5D KV layout, x=8) — verified on real BI-V100 - >32K: paged_attention_v2 (5D→4D permute) — verified 65K on real BI-V100 - Removes _forward_decode_pytorch Python fallback entirely Verified: v1 passes ctx=1024..32768, v2 passes ctx=32768..65536+ flash_attn_varlen_func prefill already merged in prior commit (ad6863ed). --- qwen3_6_scripts/paged_attn.py | 36 ++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index e8a8f2c5..c3f8492c 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/qwen3_6_scripts/paged_attn.py @@ -1585,11 +1585,37 @@ class PagedAttention: f"[{min_block}, {max_block}] outside " f"[0, {key_cache.shape[0] - 1}]") - if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD: - with bi100_timer("paged_attn.decode_pytorch"): - return PagedAttention._forward_decode_pytorch( - query, key_cache, value_cache, block_tables, seq_lens, - scale) + # BI-V100: paged_attention_v1 supports max_context_len<=32768. + # For longer contexts, use v2 with layout conversion (5D→4D). + # v1 key: [blocks, kv_h, head_dim//x, block_size, x] + # v2 key: [blocks, kv_h, block_size, head_dim] + if actual_max > 32768: + num_kv_heads = key_cache.shape[1] + key_cache_v2 = (key_cache + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(key_cache.shape[0], num_kv_heads, + block_size, head_size)) + value_cache_v2 = (value_cache + .permute(0, 1, 3, 2) + .contiguous()) + output = torch.empty_like(query) + _partition = 512 + max_num_partitions = ((max_seq_len + _partition - 1) // + _partition) + 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) + import ixformer.functions as _ixf_F + _ixf_F.vllm_single_query_cached_kv_attention_v2( + output, _partition, exp_sums, max_logits, tmp_output, + query, key_cache_v2, value_cache_v2, head_mapping, scale, + block_tables, seq_lens, block_size, max_seq_len) + return output if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1: # use blocksparse paged attention