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 4, 23:45 UTC
[ROCm] Fix cu_seqlens_q off-by-one in AITER FA speculative decode path (#39120)
Signed-off-by: Bortlesboat <bortstheboat@gmail.com>
Andrew Barnes committed
Apr 19, 2026 at 14:34 UTC
f150107efd15a873617a65b09189c945b27d61f5
1 file changed
+2
-2
vllm/v1/attention/backends/rocm_aiter_fa.py
+2
-2
@@ -1181,7 +1181,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1181
)
1182
1183
descale_shape = (
1184
- attn_metadata.query_start_loc[:num_decodes].shape[0] - 1,
1184
+ num_decodes,
1185
key_cache.shape[2],
1186
)
1187
unified_attention(
@@ -1189,7 +1189,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1189
k=key_cache,
1190
v=value_cache,
1191
out=output[:num_decode_tokens],
1192
- cu_seqlens_q=attn_metadata.query_start_loc[:num_decodes],
1192
+ cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1],
1193
max_seqlen_q=decode_max_query_len,
1194
seqused_k=attn_metadata.seq_lens[:num_decodes],
1195
max_seqlen_k=attn_metadata.max_seq_len,