From b5545d9d5cab2625ac04a19f552631a2034c8f47 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Fri, 19 Dec 2025 14:39:54 +0100 Subject: [PATCH] [Bugfix] [Kernel] Triton attention kernels: mask out V blocks that fall outside sliding window (#30887) Signed-off-by: Thomas Parnell --- vllm/attention/ops/triton_unified_attention.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/attention/ops/triton_unified_attention.py b/vllm/attention/ops/triton_unified_attention.py index e60af7434d25c..c946dbd8a2c4e 100644 --- a/vllm/attention/ops/triton_unified_attention.py +++ b/vllm/attention/ops/triton_unified_attention.py @@ -363,6 +363,12 @@ def kernel_unified_attention_2d( L = L * alpha + l_j M = m_j + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0 + ) + # acc : (BLOCK_M, HEAD_SIZE_PADDED) acc += tl.dot(P.to(V.dtype), V) @@ -678,6 +684,12 @@ def kernel_unified_attention_3d( L = L * alpha + l_j M = m_j + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0 + ) + # acc : (BLOCK_M, HEAD_SIZE_PADDED) acc += tl.dot(P.to(V.dtype), V)