[Kernel] Skip fully masked key blocks in windowed Triton prefill (#50776)
Signed-off-by: almogtavor <almogtavor@gmail.com>
Almog Tavor committed
Aug 3, 2026 at 21:48 UTC
65cf1276a25b5e7dbbed7d3febdd54a815ced8c8
1 file changed
+10
-1
vllm/v1/attention/ops/triton_prefill_attention.py
+10
-1
@@ -113,8 +113,17 @@ def _fwd_kernel(
113
# Apply causal attention pruning and sliding window attention pruning
114
end_n = tl.minimum(end_n, (start_m + 1) * BLOCK_M) if IS_CAUSAL else end_n
115
116
- # Calculate the start position for backward sliding window
116
+ # Calculate the start position for backward sliding window.
117
+ # Keys outside the window are fully masked below, so skip those blocks
118
+ # rather than loading them and discarding the result. Without this a
119
+ # windowed pass still costs O(seq_len^2).
120
start_n_limit = 0
121
+ if SLIDING_WINDOW_Q > 0:
122
+ first_needed = start_m * BLOCK_M - SLIDING_WINDOW_Q
123
+ start_n_limit = tl.maximum(0, (first_needed // BLOCK_N) * BLOCK_N)
124
+ if SLIDING_WINDOW_K > 0:
125
+ last_needed = (start_m + 1) * BLOCK_M - 1 + SLIDING_WINDOW_K
126
+ end_n = tl.minimum(end_n, last_needed + 1)
127
end_n_limit = block_mask * end_n
128
129
for start_n in range(start_n_limit, end_n_limit, BLOCK_N):