From eaec7be4464c45edaf881ffe438d83e91ec1fff9 Mon Sep 17 00:00:00 2001 From: czhu-cohere Date: Mon, 4 May 2026 21:29:34 -0700 Subject: [PATCH] [BugFix] Preserve max_seq_len in ubatch metadata during CUDA graph capture (#40961) Signed-off-by: root Signed-off-by: --- vllm/v1/worker/ubatch_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index 1338b46996..f4a7652902 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -214,7 +214,9 @@ def _make_metadata_with_slice( seq_lens_cpu_upper_bound[-1] -= tokens_skipped assert seq_lens_cpu_upper_bound is not None - max_seq_len = int(seq_lens_cpu_upper_bound.max()) + # Preserve the max_seq_len override set during CUDA-graph capture so + # the attention backend selects the correct kernel for SWA layers. + max_seq_len = max(int(seq_lens_cpu_upper_bound.max()), attn_metadata.max_seq_len) num_requests = request_slice.stop - request_slice.start num_actual_tokens = token_slice.stop - token_slice.start