A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Quick fix for IMA with the Prefix Prefill kernel during graph capture (#25983)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Oct 3, 2025 at 04:28 UTC
5f2cacdb1e62e468c707eab168c6d6961254bc93
1 file changed
+8
vllm/v1/attention/backends/rocm_attn.py
+8
@@ -83,6 +83,14 @@ class RocmAttentionMetadataBuilder(
83
# max_model_len will cause graph capture to be extremely
84
# slow, so here we set it to 1.
85
attn_metadata.seq_lens.fill_(1)
86
+
87
+ if envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION:
88
+ # Here we set the query start locs to 0. This is to
89
+ # cover up an invalid memory access in the prefix_prefil kernel
90
+ # that we run into during graph capture (#25985)
91
+ common_attn_metadata.query_start_loc.zero_()
92
+ common_attn_metadata.query_start_loc_cpu.zero_()
93
+
94
return attn_metadata
95
96
def build(self,