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, 03:45 UTC
[Bugfix] [Kernel] Triton attention kernels: mask out V blocks that fall outside sliding window (#30887)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Dec 19, 2025 at 14:39 UTC
b5545d9d5cab2625ac04a19f552631a2034c8f47
1 file changed
+12
vllm/attention/ops/triton_unified_attention.py
+12
@@ -363,6 +363,12 @@ def kernel_unified_attention_2d(
363
L = L * alpha + l_j
364
M = m_j
365
366
+ if SLIDING_WINDOW:
367
+ qpos_lo = q_block_local_idx * BLOCK_Q
368
+ V = tl.where(
369
+ (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0
370
+ )
371
+
372
# acc : (BLOCK_M, HEAD_SIZE_PADDED)
373
acc += tl.dot(P.to(V.dtype), V)
374
@@ -678,6 +684,12 @@ def kernel_unified_attention_3d(
684
L = L * alpha + l_j
685
M = m_j
686
687
+ if SLIDING_WINDOW:
688
+ qpos_lo = q_block_local_idx * BLOCK_Q
689
+ V = tl.where(
690
+ (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0
691
+ )
692
+
693
# acc : (BLOCK_M, HEAD_SIZE_PADDED)
694
acc += tl.dot(P.to(V.dtype), V)
695