squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[TPU][Bugfix] Fix next_token_ids shape (#8128)

Browse Source
This commit is contained in:
Woosuk Kwon 2024-09-03 13:29:24 -07:00 committed by GitHub
parent f1575dc99f
commit 0af3abe3d3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/worker/tpu_model_runner.py
View File

@ -601,7 +601,7 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):
batch_idx += 1
else:
for seq_id in seq_ids:
next_token_id = next_token_ids[batch_idx][0]
next_token_id = next_token_ids[batch_idx]
seq_outputs.append(
SequenceOutput(seq_id, next_token_id,
{next_token_id: zero_logprob}))
@ -722,6 +722,9 @@ class ModelWrapper(TorchCompileWrapperWithCustomDispacther):
sampled_token_ids = torch.multinomial(probs,
num_samples,
replacement=True)
if num_samples == 1:
argmax_token_ids = argmax_token_ids.squeeze(dim=-1)
sampled_token_ids = sampled_token_ids.squeeze(dim=-1)
next_token_ids = torch.where(t != 0, sampled_token_ids,
argmax_token_ids)
return next_token_ids