A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[Kernel] Optimize Prefill Attention in Unified Triton Attention Kernel (#20308)
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
jvlunteren committed
Jul 7, 2025 at 21:08 UTC
22dd9c2730dc1124b9d0ac15fff223d0b8d9020b
1 file changed
+13
-1
vllm/attention/ops/triton_unified_attention.py
+13
-1
@@ -145,7 +145,19 @@ def kernel_unified_attention_2d(
145
mask=query_mask_1,
146
other=0.0)
147
148
- num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
148
+ # compute the length of the longest sequence prefix spanned by any
149
+ # query token in the current q_block (q_block_local_idx)
150
+ max_seq_prefix_len = context_len + q_block_local_idx * BLOCK_Q + (
151
+ BLOCK_M - 1) // num_queries_per_kv + 1
152
+
153
+ # adjust for potential padding in the last q_block by considering the
154
+ # actual sequence length
155
+ max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
156
+
157
+ # calculate the number of tiles (blocks) that need to be processed to
158
+ # cover the longest sequence prefix (due to causal masking, blocks beyond
159
+ # this prefix can be skipped)
160
+ num_blocks = cdiv_fn(max_seq_prefix_len, BLOCK_SIZE)
161
162
# iterate through tiles
163
for j in range(0, num_blocks):