[ROCm] [Bugfix] Preserve CPU query offsets during capture (#51585)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
Aug 20, 2026 at 07:33 UTC
cd5035379c4c565fc9265cf121bc1c8364bc5eca
1 file changed
+2
-4
vllm/v1/attention/backends/rocm_attn.py
+2
-4
@@ -103,11 +103,9 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
103
# slow, so here we set it to 1.
104
attn_metadata.seq_lens.fill_(1)
105
106
- # Here we set the query start locs to 0. This is to
107
- # cover up an invalid memory access in the prefix_prefil kernel
108
- # that we run into during graph capture (#25985)
106
+ # Zero device query start locations to avoid invalid memory access in
107
+ # the prefix prefill kernel during graph capture (#25985).
108
common_attn_metadata.query_start_loc.zero_()
110
- common_attn_metadata.query_start_loc_cpu.zero_()
109
110
return attn_metadata
111