[Bugfix] Don't read KV cache past `seq_len` in triton paged attn kernels (#47305)
Nick Hill committed
Jul 1, 2026 at 20:43 UTC
4787f2dd1b5705b92c095885f2f07f7253f5aed8
2 files changed
+13
-6
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
+7
-2
@@ -156,6 +156,11 @@ def kernel_paged_attention_2d(
156
# Supports non-contiguous mapping
157
# from logical blocks to physical blocks
158
abs_token_idx = start_n + offs_n
159
+ # Slots >= seq_len are unwritten KV cache and may hold NaN/garbage
160
+ # (e.g. the tail of the last partial block). They are score-masked
161
+ # below, but 0 * NaN = NaN would still poison the output, so exclude
162
+ # them from the K/V loads too.
163
+ kv_load_mask = abs_token_idx < seq_len
164
l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE
165
# Vectorized loading of physical block IDs
166
p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)
@@ -181,7 +186,7 @@ def kernel_paged_attention_2d(
186
# K : (HEAD_SIZE, BLOCK_SIZE)
187
K_load = tl.load(
188
key_cache_ptr + k_offset,
184
- mask=dim_mask[:, None],
189
+ mask=dim_mask[:, None] & kv_load_mask[None, :],
190
other=0.0,
191
eviction_policy="evict_last",
192
)
@@ -194,7 +199,7 @@ def kernel_paged_attention_2d(
199
# V : (BLOCK_SIZE, HEAD_SIZE)
200
V_load = tl.load(
201
value_cache_ptr + v_offset,
197
- mask=dim_mask[None, :],
202
+ mask=dim_mask[None, :] & kv_load_mask[:, None],
203
other=0.0,
204
eviction_policy="evict_last",
205
)
vllm/v1/attention/ops/triton_attention_helpers.py
+6
-4
@@ -183,10 +183,12 @@ def compute_tile_loop_bounds(
183
+ 1
184
)
185
if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):
186
- # Non-causal or mixed batches need the full sequence range.
187
- # Per-element masking in compute_kv_seq_mask handles the
188
- # actual causal/non-causal boundary per sequence.
189
- max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
186
+ # Read the full sequence but never past seq_len: the causal-style
187
+ # formula above can overshoot for non-causal sequences, and slots
188
+ # >= seq_len are unwritten KV (last-block tail) that may hold NaN
189
+ # (0 * NaN poisons the output). Per-element masking in
190
+ # compute_kv_seq_mask handles the causal/non-causal boundary.
191
+ max_seq_prefix_len = seq_len
192
else:
193
max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
194