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 8, 22:45 UTC
[Bugfix] Fix ubatch wrapper num_tokens calculate (#33694)
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
jiangkuaixue123 committed
Feb 5, 2026 at 00:41 UTC
87d9a261664705e0c9635014b4e2d49eddc8a056
1 file changed
+1
-3
vllm/v1/worker/gpu_ubatch_wrapper.py
+1
-3
@@ -412,9 +412,7 @@ class UBatchWrapper:
412
413
attn_metadata = forward_context.attn_metadata
414
slot_mapping = forward_context.slot_mapping
415
- num_tokens = (
416
- ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start
417
- ) * 2
415
+ num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices)
416
input_ids = kwargs["input_ids"]
417
positions = kwargs["positions"]
418
intermediate_tensors = kwargs["intermediate_tensors"]