From 94d2f598956dbd5ab8a6097d6eaa8176035d3220 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 25 Aug 2023 12:22:01 +0900 Subject: [PATCH] Set replacement=True in torch.multinomial (#858) --- vllm/model_executor/layers/sampler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/sampler.py b/vllm/model_executor/layers/sampler.py index 6a50ee59..9717e008 100644 --- a/vllm/model_executor/layers/sampler.py +++ b/vllm/model_executor/layers/sampler.py @@ -302,7 +302,9 @@ def _sample_from_prompt( # Random sampling. # Sample `best_of` tokens for the prompt. num_seqs = sampling_params.best_of - next_token_ids = torch.multinomial(prob, num_samples=num_seqs) + next_token_ids = torch.multinomial(prob, + num_samples=num_seqs, + replacement=True) next_token_ids = next_token_ids.tolist() return next_token_ids