[Bugfix] Don't read KV cache past seq_len in triton paged attn kernels (#47305)

vllm/v1/attention/ops/chunked_prefill_paged_decode.py

@@ -156,6 +156,11 @@ def kernel_paged_attention_2d(
         # Supports non-contiguous mapping
         # from logical blocks to physical blocks
         abs_token_idx = start_n + offs_n
+        # Slots >= seq_len are unwritten KV cache and may hold NaN/garbage
+        # (e.g. the tail of the last partial block). They are score-masked
+        # below, but 0 * NaN = NaN would still poison the output, so exclude
+        # them from the K/V loads too.
+        kv_load_mask = abs_token_idx < seq_len
         l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE
         # Vectorized loading of physical block IDs
         p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)
@@ -181,7 +186,7 @@ def kernel_paged_attention_2d(
         # K : (HEAD_SIZE, BLOCK_SIZE)
         K_load = tl.load(
             key_cache_ptr + k_offset,
-            mask=dim_mask[:, None],
+            mask=dim_mask[:, None] & kv_load_mask[None, :],
             other=0.0,
             eviction_policy="evict_last",
         )
@@ -194,7 +199,7 @@ def kernel_paged_attention_2d(
         # V : (BLOCK_SIZE, HEAD_SIZE)
         V_load = tl.load(
             value_cache_ptr + v_offset,
-            mask=dim_mask[None, :],
+            mask=dim_mask[None, :] & kv_load_mask[:, None],
             other=0.0,
             eviction_policy="evict_last",
         )

vllm/v1/attention/ops/triton_attention_helpers.py

@@ -183,10 +183,12 @@ def compute_tile_loop_bounds(
         + 1
     )
     if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):
-        # Non-causal or mixed batches need the full sequence range.
-        # Per-element masking in compute_kv_seq_mask handles the
-        # actual causal/non-causal boundary per sequence.
-        max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
+        # Read the full sequence but never past seq_len: the causal-style
+        # formula above can overshoot for non-causal sequences, and slots
+        # >= seq_len are unwritten KV (last-block tail) that may hold NaN
+        # (0 * NaN poisons the output). Per-element masking in
+        # compute_kv_seq_mask handles the causal/non-causal boundary.
+        max_seq_prefix_len = seq_len
     else:
         max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)