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 11, 05:00 UTC
[Bugfix] Fix incorrect tiles creation for mm prefix triton attention (#30974)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Dec 19, 2025 at 16:00 UTC
ac1c934276167085dc7d9e5adb8c2089de0460f9
1 file changed
+10
-4
vllm/attention/ops/triton_unified_attention.py
+10
-4
@@ -189,9 +189,14 @@ def kernel_unified_attention_2d(
189
+ 1
190
)
191
192
- # adjust for potential padding in the last q_block by considering the
193
- # actual sequence length
194
- max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
192
+ if USE_MM_PREFIX:
193
+ # image bidirectional attention ranges require a full range
194
+ # including q_block padding to make sure doc mask is correct
195
+ max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
196
+ else:
197
+ # adjust for potential padding in the last q_block by considering the
198
+ # actual sequence length
199
+ max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
200
201
# calculate the number of tiles that need to be processed to
202
# cover the longest sequence prefix (due to causal masking, tiles beyond
@@ -202,7 +207,8 @@ def kernel_unified_attention_2d(
207
# Default: keep previous global behavior
208
tile_start = 0
209
tile_end = num_tiles
205
- if SLIDING_WINDOW > 0:
210
+ # TODO(Isotr0py): sliding window pruning with image bidirectional mask
211
+ if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:
212
# Query rows covered by this Q-block
213
qpos_lo = q_block_local_idx * BLOCK_Q
214
qpos_hi = tl.minimum(