xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-06-07 06:35:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Minor

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-25 03:28:53 +00:00
parent 5323969fcf
commit e2c7dedb3a
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
vllm/worker/tpu_model_runner.py
View File

@ -193,7 +193,7 @@ class TPUModelRunner:
kv_caches, kv_caches,
logits_indices, logits_indices,
) )
# TODO # TODO(woosuk): Support sampling with temperature and top_p.
next_token_ids = jnp.argmax(logits, axis=-1) next_token_ids = jnp.argmax(logits, axis=-1)
return next_token_ids, new_kv_caches return next_token_ids, new_kv_caches