Set replacement=True in torch.multinomial (#858)

94d2f598 · Woosuk Kwon · GitHub · 75c0ca9d · 94d2f598
Unverified Commit 94d2f598 authored Aug 25, 2023 by Woosuk Kwon Committed by GitHub Aug 25, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

vllm/model_executor/layers/sampler.py vllm/model_executor/layers/sampler.py +3 -1

No files found.
--- 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