diff --git a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py b/vllm/v1/attention/ops/chunked_prefill_paged_decode.py index 77eb3ac60b1..73d40a0a333 100644 --- a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py +++ b/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", ) diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index ed9a38ad6cd..b90a1ac39b7 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/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)