[Bugfix] Take the VRAM usage of prompt_logprobs into account

tests/samplers/test_logprobs.py

@@ -132,3 +132,15 @@ def test_max_logprobs():
     bad_sampling_params = SamplingParams(logprobs=2)
     with pytest.raises(ValueError):
         runner.generate(["Hello world"], sampling_params=bad_sampling_params)
+
+
+@pytest.mark.parametrize("chunked_prefill", [True, False])
+def test_batched_logprobs(chunked_prefill: bool):
+    runner = VllmRunner("facebook/opt-125m",
+                        max_logprobs=1,
+                        max_num_batched_logprobs=16,
+                        enable_chunked_prefill=chunked_prefill)
+    vllm_sampling_params = SamplingParams(logprobs=1, prompt_logprobs=1)
+    num_seqs = 1024
+    runner.generate(["Hello world" for _ in range(num_seqs)],
+                    sampling_params=vllm_sampling_params)

vllm/config.py

@@ -657,6 +657,8 @@ class SchedulerConfig:
             swapping. However, when the sequence group has multiple sequences
             (e.g., beam search), recomputation is not currently supported. In
             such a case, we use swapping instead.
+        max_num_batched_logprobs: Maximum number of logprobs to be processed in
+            a single iteration. This should be carefully chosen to avoid OOM.
     """
 
     def __init__(self,
@@ -668,7 +670,8 @@ def __init__(self,
                  delay_factor: float = 0.0,
                  enable_chunked_prefill: bool = False,
                  embedding_mode: Optional[bool] = False,
-                 preemption_mode: Optional[str] = None) -> None:
+                 preemption_mode: Optional[str] = None,
+                 max_num_batched_logprobs: Optional[int] = None) -> None:
         if max_num_batched_tokens is not None:
             self.max_num_batched_tokens = max_num_batched_tokens
         else:
@@ -695,6 +698,9 @@ def __init__(self,
         self.chunked_prefill_enabled = enable_chunked_prefill
         self.embedding_mode = embedding_mode
         self.preemption_mode = preemption_mode
+        self.max_num_batched_logprobs = (
+            max_num_batched_logprobs is not None and max_num_batched_logprobs
+            or min(self.max_num_batched_tokens, 256))
 
         self._verify_args()
 
@@ -715,6 +721,12 @@ def _verify_args(self) -> None:
                 "be greater than or equal to max_num_seqs "
                 f"({self.max_num_seqs}).")
 
+        if self.max_num_batched_logprobs > self.max_num_batched_tokens:
+            raise ValueError(
+                f"max_num_batched_logprobs ({self.max_num_batched_logprobs}) "
+                f"must be less than or equal to max_num_batched_tokens "
+                f"({self.max_num_batched_tokens}).")
+
         if self.num_lookahead_slots < 0:
             raise ValueError(
                 "num_lookahead_slots "

vllm/core/scheduler.py

@@ -12,7 +12,8 @@
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.sequence import (Sequence, SequenceData, SequenceGroup,
-                           SequenceGroupMetadata, SequenceStatus)
+                           SequenceGroupMetadata, SequenceStage,
+                           SequenceStatus)
 
 logger = init_logger(__name__)
 
@@ -50,16 +51,25 @@ class SchedulingBudget:
     """
     token_budget: int
     max_num_seqs: int
+    logprob_budget: int = 256  # Added to prevent OOM in calculating logprobs
     _requeset_ids_num_batched_tokens: Set[str] = field(default_factory=set)
     _requeset_ids_num_curr_seqs: Set[str] = field(default_factory=set)
+    _requeset_ids_num_batched_logprobs: Set[str] = field(default_factory=set)
     _num_batched_tokens: int = 0
     _num_curr_seqs: int = 0
+    _num_batched_logprobs: int = 0
 
-    def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):
+    def can_schedule(self,
+                     *,
+                     num_new_tokens: int,
+                     num_new_seqs: int,
+                     num_new_logprobs: int = 0) -> bool:
         assert num_new_tokens != 0
-        assert num_new_seqs != 0
+        # We don't assert num_new_seqs != 0 here for decoding.
         return (self.num_batched_tokens + num_new_tokens <= self.token_budget
-                and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs)
+                and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs
+                and self.num_batched_logprobs + num_new_logprobs <=
+                self.logprob_budget)
 
     def remaining_token_budget(self):
         return self.token_budget - self.num_batched_tokens
@@ -89,6 +99,22 @@ def subtract_num_seqs(self, req_id: str, num_curr_seqs: int):
             self._requeset_ids_num_curr_seqs.remove(req_id)
             self._num_curr_seqs -= num_curr_seqs
 
+    def remaining_logprob_budget(self):
+        return self.logprob_budget - self._num_batched_logprobs
+
+    def add_num_batched_logprobs(self, req_id: str, num_batched_logprobs: int):
+        if req_id in self._requeset_ids_num_batched_logprobs:
+            return
+
+        self._requeset_ids_num_batched_logprobs.add(req_id)
+        self._num_batched_logprobs += num_batched_logprobs
+
+    def subtract_num_batched_logprobs(self, req_id: str,
+                                      num_batched_logprobs: int):
+        if req_id in self._requeset_ids_num_batched_logprobs:
+            self._requeset_ids_num_batched_logprobs.remove(req_id)
+            self._num_batched_logprobs -= num_batched_logprobs
+
     @property
     def num_batched_tokens(self):
         return self._num_batched_tokens
@@ -97,6 +123,10 @@ def num_batched_tokens(self):
     def num_curr_seqs(self):
         return self._num_curr_seqs
 
+    @property
+    def num_batched_logprobs(self):
+        return self._num_batched_logprobs
+
 
 @dataclass
 class ScheduledSequenceGroup:
@@ -117,6 +147,8 @@ class SchedulerOutputs:
     num_prefill_groups: int
     # Total number of batched tokens.
     num_batched_tokens: int
+    # Total number of batched logprobs.
+    num_batched_logprobs: int
     # Blocks to swap in. List of CPU -> GPU block number.
     blocks_to_swap_in: List[Tuple[int, int]]
     # Blocks to swap out. List of GPU -> CPU block number.
@@ -412,14 +444,27 @@ def _schedule_running(
             seq_group = running_queue[0]
             num_running_tokens = self._get_num_new_tokens(
                 seq_group, SequenceStatus.RUNNING, enable_chunking, budget)
+            num_running_logprobs = self._get_num_new_logprobs(
+                seq_group, SequenceStatus.RUNNING, enable_chunking, budget)
 
             if num_running_tokens == 0:
                 break
 
+            # Check if token and logprob budget is exceeded.
+            # This check is necessary because we do not guarantee that
+            # max_num_batched_tokens and max_num_batched_logprobs are
+            # less than max_seqs.
+            if not budget.can_schedule(num_new_tokens=num_running_tokens,
+                                       num_new_seqs=0,
+                                       num_new_logprobs=num_running_logprobs):
+                break
+
             running_queue.popleft()
             while not self._can_append_slots(seq_group):
                 budget.subtract_num_batched_tokens(seq_group.request_id,
                                                    num_running_tokens)
+                budget.subtract_num_batched_logprobs(seq_group.request_id,
+                                                     num_running_logprobs)
                 num_running_seqs = seq_group.get_max_num_running_seqs()
                 budget.subtract_num_seqs(seq_group.request_id,
                                          num_running_seqs)
@@ -459,6 +504,8 @@ def _schedule_running(
                                                token_chunk_size=1))
                 budget.add_num_batched_tokens(seq_group.request_id,
                                               num_running_tokens)
+                budget.add_num_batched_logprobs(seq_group.request_id,
+                                                num_running_logprobs)
                 # OPTIMIZATION:  Note that get_max_num_running_seqs is
                 # expensive. For the default scheduling chase where
                 # enable_chunking is False, num_seqs are updated before running
@@ -559,10 +606,13 @@ def _schedule_swapped(
             num_new_tokens = self._get_num_new_tokens(seq_group,
                                                       SequenceStatus.SWAPPED,
                                                       enable_chunking, budget)
+            num_new_logprobs = self._get_num_new_logprobs(
+                seq_group, SequenceStatus.SWAPPED, enable_chunking, budget)
 
-            if (num_new_tokens == 0
-                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
-                                               num_new_seqs=num_new_seqs)):
+            if (num_new_tokens == 0 or not budget.can_schedule(
+                    num_new_tokens=num_new_tokens,
+                    num_new_seqs=num_new_seqs,
+                    num_new_logprobs=num_new_logprobs)):
                 break
 
             if lora_int_id > 0 and curr_loras is not None:
@@ -579,6 +629,8 @@ def _schedule_swapped(
                 decode_seq_groups.append(
                     ScheduledSequenceGroup(seq_group, token_chunk_size=1))
             budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
+            budget.add_num_batched_logprobs(seq_group.request_id,
+                                            num_new_logprobs)
             budget.add_num_seqs(seq_group.request_id, num_new_seqs)
 
         swapped_queue.extendleft(leftover_swapped)
@@ -599,6 +651,11 @@ def _get_prompt_limit(self, seq_group: SequenceGroup) -> int:
         else:
             prompt_limit = min(self.scheduler_config.max_model_len,
                                self.scheduler_config.max_num_batched_tokens)
+            if (seq_group.sampling_params is not None
+                    and seq_group.sampling_params.prompt_logprobs):
+                prompt_limit = min(
+                    prompt_limit,
+                    self.scheduler_config.max_num_batched_logprobs)
 
         # Model is fine tuned with long context. Return the fine tuned max_len.
         if (seq_group.lora_request
@@ -662,6 +719,10 @@ def _schedule_prefills(
                 num_prompt_tokens = waiting_seqs[0].get_len()
                 assert num_new_tokens == num_prompt_tokens
 
+            num_new_logprobs = self._get_num_new_logprobs(
+                seq_group, SequenceStatus.WAITING, enable_chunking, budget)
+
+            # Check of num_tokens and num_logprobs are included
             prompt_limit = self._get_prompt_limit(seq_group)
             if num_new_tokens > prompt_limit:
                 logger.warning(
@@ -703,9 +764,10 @@ def _schedule_prefills(
                     continue
 
             num_new_seqs = seq_group.get_max_num_running_seqs()
-            if (num_new_tokens == 0
-                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
-                                               num_new_seqs=num_new_seqs)):
+            if (num_new_tokens == 0 or not budget.can_schedule(
+                    num_new_tokens=num_new_tokens,
+                    num_new_seqs=num_new_seqs,
+                    num_new_logprobs=num_new_logprobs)):
                 break
 
             # Can schedule this request.
@@ -717,6 +779,8 @@ def _schedule_prefills(
                 ScheduledSequenceGroup(seq_group=seq_group,
                                        token_chunk_size=num_new_tokens))
             budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
+            budget.add_num_batched_logprobs(seq_group.request_id,
+                                            num_new_logprobs)
             budget.add_num_seqs(seq_group.request_id, num_new_seqs)
 
         # Queue requests that couldn't be scheduled.
@@ -741,12 +805,14 @@ def _schedule_default(self) -> SchedulerOutputs:
         budget = SchedulingBudget(
             token_budget=self.scheduler_config.max_num_batched_tokens,
             max_num_seqs=self.scheduler_config.max_num_seqs,
+            logprob_budget=self.scheduler_config.max_num_batched_logprobs,
         )
         # Make sure we include num running seqs before scheduling prefill,
         # so that we don't schedule beyond max_num_seqs for prefill.
         for seq_group in self.running:
             budget.add_num_seqs(seq_group.request_id,
                                 seq_group.get_max_num_running_seqs())
+            # budget.add_num_batched_logprobs
         curr_loras = set(
             seq_group.lora_int_id for seq_group in self.running
             if seq_group.lora_int_id > 0) if self.lora_enabled else None
@@ -784,6 +850,8 @@ def _schedule_default(self) -> SchedulerOutputs:
 
         assert (budget.num_batched_tokens <=
                 self.scheduler_config.max_num_batched_tokens)
+        assert (budget.num_batched_logprobs <=
+                self.scheduler_config.max_num_batched_logprobs)
         assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs
 
         # Update waiting requests.
@@ -812,6 +880,7 @@ def _schedule_default(self) -> SchedulerOutputs:
                                   swapped_in.decode_seq_groups),
             num_prefill_groups=len(prefills.seq_groups),
             num_batched_tokens=budget.num_batched_tokens,
+            num_batched_logprobs=budget.num_batched_logprobs,
             blocks_to_swap_in=swapped_in.blocks_to_swap_in,
             blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
             blocks_to_copy=running_scheduled.blocks_to_copy +
@@ -840,7 +909,7 @@ def _schedule_chunked_prefill(self):
         budget = SchedulingBudget(
             token_budget=self.scheduler_config.max_num_batched_tokens,
             max_num_seqs=self.scheduler_config.max_num_seqs,
-        )
+            logprob_budget=self.scheduler_config.max_num_batched_logprobs)
         curr_loras: Set[int] = set()
 
         remaining_waiting, prefills = (self.waiting,
@@ -872,6 +941,8 @@ def _schedule_chunked_prefill(self):
 
         assert (budget.num_batched_tokens <=
                 self.scheduler_config.max_num_batched_tokens)
+        assert (budget.num_batched_logprobs <=
+                self.scheduler_config.max_num_batched_logprobs)
         assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs
 
         # Update waiting requests.
@@ -901,6 +972,7 @@ def _schedule_chunked_prefill(self):
                                 len(swapped_in.prefill_seq_groups) +
                                 len(running_scheduled.prefill_seq_groups)),
             num_batched_tokens=budget.num_batched_tokens,
+            num_batched_logprobs=budget.num_batched_logprobs,
             blocks_to_swap_in=swapped_in.blocks_to_swap_in,
             blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
             blocks_to_copy=running_scheduled.blocks_to_copy +
@@ -1198,3 +1270,27 @@ def _get_num_new_tokens(self, seq_group: SequenceGroup,
             num_new_tokens = min(num_new_tokens,
                                  budget.remaining_token_budget())
         return num_new_tokens
+
+    def _get_num_new_logprobs(self, seq_group: SequenceGroup,
+                              status: SequenceStatus, enable_chunking: bool,
+                              budget: SchedulingBudget) -> int:
+        if seq_group.sampling_params is None:
+            return 0
+
+        decode_logprobs = seq_group.sampling_params.logprobs is not None
+        prompt_logprobs = seq_group.sampling_params.prompt_logprobs is not None
+        num_new_logprobs = 0
+        seqs = seq_group.get_seqs(status=status)
+
+        for seq in seqs:
+            if (seq.data.stage == SequenceStage.DECODE
+                    and decode_logprobs) or (seq.data.stage
+                                             == SequenceStage.PREFILL
+                                             and prompt_logprobs):
+                num_new_logprobs += seq.get_num_new_tokens()
+
+        # The same logic as _get_num_new_tokens.
+        if enable_chunking and len(seqs) == 1:
+            num_new_logprobs = min(num_new_logprobs,
+                                   budget.remaining_logprob_budget())
+        return num_new_logprobs

vllm/engine/arg_utils.py

@@ -49,6 +49,7 @@ class EngineArgs:
     max_num_batched_tokens: Optional[int] = None
     max_num_seqs: int = 256
     max_logprobs: int = 5  # OpenAI default value
+    max_num_batched_logprobs: Optional[int] = None
     disable_log_stats: bool = False
     revision: Optional[str] = None
     code_revision: Optional[str] = None
@@ -380,6 +381,11 @@ def add_cli_args(
             default=EngineArgs.max_logprobs,
             help=('Max number of log probs to return logprobs is specified in'
                   ' SamplingParams.'))
+        parser.add_argument(
+            "--max-num-batched-logprobs",
+            type=int,
+            default=EngineArgs.max_num_batched_logprobs,
+            help='Maximum number of batched logprobs per iteration.')
         parser.add_argument('--disable-log-stats',
                             action='store_true',
                             help='Disable logging statistics.')
@@ -676,6 +682,7 @@ def create_engine_config(self, ) -> EngineConfig:
             enable_chunked_prefill=self.enable_chunked_prefill,
             embedding_mode=model_config.embedding_mode,
             preemption_mode=self.preemption_mode,
+            max_num_batched_logprobs=self.max_num_batched_logprobs,
         )
         lora_config = LoRAConfig(
             max_lora_rank=self.max_lora_rank,