[ROCm] [bugfix] Chunked prefill paged decode masked load perf (#50017)
Signed-off-by: Andy Friedrich <afriedri@amd.com>
afriedri committed
Aug 12, 2026 at 18:04 UTC
98f86b9c02329200a0390aecfe598e27928cbf40
1 file changed
+34
-20
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
+34
-20
@@ -156,11 +156,6 @@ 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
159
l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE
160
# Vectorized loading of physical block IDs
161
p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)
@@ -183,27 +178,46 @@ def kernel_paged_attention_2d(
178
+ internal_offsets[:, None] * stride_v_cache_3
179
)
180
186
- # K : (HEAD_SIZE, BLOCK_SIZE)
187
- K_load = tl.load(
188
- key_cache_ptr + k_offset,
189
- mask=dim_mask[:, None] & kv_load_mask[None, :],
190
- other=0.0,
191
- eviction_policy="evict_last",
192
- )
181
+ # Only the final tile can straddle seq_len. Slots >= seq_len are
182
+ # unwritten KV cache that may hold NaN/garbage; they are score-masked
183
+ # below, but 0 * NaN = NaN would still poison the output, so mask them
184
+ # out of the K/V loads too. Earlier tiles are fully written, so they
185
+ # use the cheaper token-uniform dim_mask (matching the pre-0.25.0 fast
186
+ # path) and skip the per-token predicate entirely.
187
+ # K : (HEAD_SIZE, BLOCK_SIZE), V : (BLOCK_SIZE, HEAD_SIZE)
188
+ if j == num_blocks - 1:
189
+ kv_load_mask = abs_token_idx < seq_len
190
+ K_load = tl.load(
191
+ key_cache_ptr + k_offset,
192
+ mask=dim_mask[:, None] & kv_load_mask[None, :],
193
+ other=0.0,
194
+ eviction_policy="evict_last",
195
+ )
196
+ V_load = tl.load(
197
+ value_cache_ptr + v_offset,
198
+ mask=dim_mask[None, :] & kv_load_mask[:, None],
199
+ other=0.0,
200
+ eviction_policy="evict_last",
201
+ )
202
+ else:
203
+ K_load = tl.load(
204
+ key_cache_ptr + k_offset,
205
+ mask=dim_mask[:, None],
206
+ other=0.0,
207
+ eviction_policy="evict_last",
208
+ )
209
+ V_load = tl.load(
210
+ value_cache_ptr + v_offset,
211
+ mask=dim_mask[None, :],
212
+ other=0.0,
213
+ eviction_policy="evict_last",
214
+ )
215
216
if K_load.dtype.is_fp8():
217
K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
218
else:
219
K = K_load
220
199
- # V : (BLOCK_SIZE, HEAD_SIZE)
200
- V_load = tl.load(
201
- value_cache_ptr + v_offset,
202
- mask=dim_mask[None, :] & kv_load_mask[:, None],
203
- other=0.0,
204
- eviction_policy="evict_last",
205
- )
206
-
221
if V_load.dtype.is_fp8():
222
V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
223
else: