From f9e1cbf89736d072b3f702c73914003888bfbdf0 Mon Sep 17 00:00:00 2001 From: forest-david <839894616@qq.com> Date: Wed, 9 Sep 2026 15:35:32 +0800 Subject: [PATCH] Fix computed tokens for DBO request continuations Assisted-by: OpenAI Codex Signed-off-by: forest-david <839894616@qq.com> --- tests/v1/attention/test_attention_splitting.py | 14 ++++++++++++++ vllm/v1/worker/ubatch_utils.py | 5 +++++ 2 files changed, 19 insertions(+) diff --git a/tests/v1/attention/test_attention_splitting.py b/tests/v1/attention/test_attention_splitting.py index acbc22a02793..3ec3d175d2bc 100644 --- a/tests/v1/attention/test_attention_splitting.py +++ b/tests/v1/attention/test_attention_splitting.py @@ -390,6 +390,20 @@ def test_prefill_split_across_ubatches( # Second ubatch: first request (continuation) seq_len should be full # original assert int(second_meta.seq_lens[0]) == seq_lens[split_req_idx] + + # Tokens processed for the first chunk become computed context for the + # continuation. The original metadata must remain unchanged. + assert first_meta._num_computed_tokens_cpu is not None + assert second_meta._num_computed_tokens_cpu is not None + assert common._num_computed_tokens_cpu is not None + assert int(first_meta._num_computed_tokens_cpu[-1]) == context_lens[split_req_idx] + assert int(second_meta._num_computed_tokens_cpu[0]) == ( + context_lens[split_req_idx] + tokens_in_first_chunk + ) + assert int(common._num_computed_tokens_cpu[split_req_idx]) == context_lens[ + split_req_idx + ] + # Any following full requests in second ubatch should match originals for j in range(1, second_meta.num_reqs): # Map to original request index diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index f9612b1e7f82..98c291c3d829 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -296,6 +296,11 @@ def _make_metadata_with_slice( if attn_metadata._num_computed_tokens_cpu is not None else None ) + if splits_first_request and num_computed_tokens_cpu is not None: + # Tokens from the preceding microbatch are now part of this request's + # computed context. Clone to avoid mutating the parent metadata. + num_computed_tokens_cpu = num_computed_tokens_cpu.clone() + num_computed_tokens_cpu[0] += first_tok - start_locs[first_req] if splits_last_request: # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate