Align vLLM's beam search implementation with HF generate (#857)

2023-09-04 17:29:42 -07:00
parent e15932bb60
commit 002800f081
24 changed files with 596 additions and 260 deletions
--- a/vllm/outputs.py
+++ b/vllm/outputs.py
@@ -75,10 +75,12 @@ class RequestOutput:
        # Get the top-n sequences.
        n = seq_group.sampling_params.n
        seqs = seq_group.get_seqs()
-        assert n <= len(seqs)
-        sorted_seqs = sorted(seqs,
-                             key=lambda seq: seq.get_cumulative_logprob(),
-                             reverse=True)
+        if seq_group.sampling_params.use_beam_search:
+            sorting_key = lambda seq: seq.get_beam_search_score(
+                seq_group.sampling_params.length_penalty)
+        else:
+            sorting_key = lambda seq: seq.get_cumulative_logprob()
+        sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)
        top_n_seqs = sorted_seqs[:n]

        # Create the outputs.