From e9eaad0592d9455332db76dd266a8bacc7c36d30 Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Fri, 7 Aug 2026 04:44:18 +0000 Subject: [PATCH] =?UTF-8?q?perf:=20increase=20prefix=20attention=20tile=20?= =?UTF-8?q?budget=2096MB=E2=86=92256MB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CCCL dispatch_transform.cuh spread_out_items_per_thread pattern: reduce tile count = reduce Python loop iterations = faster prefill. At 256K context with q_len=4096: old 96MB budget → 219 KV tokens/tile → ~1200 tiles per layer → 16 min per chunk. New 256MB budget → ~580 KV tokens/tile → ~450 tiles per layer → ~6 min per chunk. BI-V100 has 32 GB HBM; 256 MB temporary tensor is safe. --- qwen3_6_scripts/paged_attn.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index 5bea73f4..d086ef15 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/qwen3_6_scripts/paged_attn.py @@ -576,7 +576,9 @@ class PagedAttention: # tile_sz=256 → 24 MB (safe) # For decode (q_len=1): tile_sz=4096 → only 96 KB (always safe) # ================================================================ - _SMEM_BUDGET_BYTES = 96 * 1024 * 1024 # 96 MB score tensor budget + _SMEM_BUDGET_BYTES = 256 * 1024 * 1024 # 256 MB score tensor budget + # CCCL GridEvenShare: fewer tiles = fewer iterations = less overhead + # BI-V100 has 32 GB HBM per card; 256 MB temporary is safe. batch_size = seq_lens_tensor.shape[0] num_q_heads = query.shape[1]