From 87d9a261664705e0c9635014b4e2d49eddc8a056 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Thu, 5 Feb 2026 00:41:45 +0800 Subject: [PATCH] [Bugfix] Fix ubatch wrapper num_tokens calculate (#33694) Signed-off-by: jiangkuaixue123 --- vllm/v1/worker/gpu_ubatch_wrapper.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 97f8b92ce..765427683 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -412,9 +412,7 @@ class UBatchWrapper: attn_metadata = forward_context.attn_metadata slot_mapping = forward_context.slot_mapping - num_tokens = ( - ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start - ) * 2 + num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) input_ids = kwargs["input_ids"] positions = kwargs["positions"] intermediate_tensors = kwargs["intermediate_tensors"]