Skip to content

GitLab

"docs/vscode:/vscode.git/clone" did not exist on "4e95ec111cd179f2ab0f6931bf57663f828a51ec"
Unverified Commit 87d9a261 authored by jiangkuaixue123's avatar jiangkuaixue123 Committed by GitHub
Browse files

[Bugfix] Fix ubatch wrapper num_tokens calculate (#33694) 


Signed-off-by: default avatarjiangkuaixue123 <jiangxiaozhou111@163.com>
parent 80f921ba
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 3 deletions
vllm/v1/worker/gpu_ubatch_wrapper.py
View file @ 87d9a261
...@@ -412,9 +412,7 @@ class UBatchWrapper: ...@@ -412,9 +412,7 @@ class UBatchWrapper:
attn_metadata = forward_context.attn_metadata attn_metadata = forward_context.attn_metadata
slot_mapping = forward_context.slot_mapping slot_mapping = forward_context.slot_mapping
num_tokens = ( num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices)
ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start
) * 2
input_ids = kwargs["input_ids"] input_ids = kwargs["input_ids"]
positions = kwargs["positions"] positions = kwargs["positions"]
intermediate_tensors = kwargs["intermediate_tensors"] intermediate_tensors = kwargs["intermediate_tensors"]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment