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, 07:30 UTC
[BugFix] Preserve max_seq_len in ubatch metadata during CUDA graph capture (#40961)
Signed-off-by: root <conway.zhu@cohere.com> Signed-off-by: <conway.zhu@cohere.com>
czhu-cohere committed
May 4, 2026 at 21:29 UTC
eaec7be4464c45edaf881ffe438d83e91ec1fff9
1 file changed
+3
-1
vllm/v1/worker/ubatch_utils.py
+3
-1
@@ -214,7 +214,9 @@ def _make_metadata_with_slice(
214
seq_lens_cpu_upper_bound[-1] -= tokens_skipped
215
216
assert seq_lens_cpu_upper_bound is not None
217
- max_seq_len = int(seq_lens_cpu_upper_bound.max())
217
+ # Preserve the max_seq_len override set during CUDA-graph capture so
218
+ # the attention backend selects the correct kernel for SWA layers.
219
+ max_seq_len = max(int(seq_lens_cpu_upper_bound.max()), attn_metadata.max_seq_len)
220
221
num_requests = request_slice.stop - request_slice.start
222
num_actual_tokens = token_slice.stop - token_slice.start