transformers 5.0.0rc1__py3-none-any.whl → 5.0.0rc3__py3-none-any.whl

transformers/generation/continuous_batching/scheduler.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,7 +17,7 @@ from collections import deque
 
 from ...utils.metrics import attach_tracer, traced
 from .cache import PagedAttentionCache
-from .requests import RequestState, RequestStatus
+from .requests import RequestState, RequestStatus, logger
 
 
 class Scheduler(ABC):
@@ -36,6 +35,11 @@ class Scheduler(ABC):
         self.retain_cache_on_finish = retain_cache_on_finish
         self._cancellation_lock = threading.Lock()
         self._requests_to_cancel: set[str] = set()
+        self._requests_to_fork: list[RequestState] = []
+        # This state is used to avoid infinite loops when offloading requests
+        self.block_new_requests = False
+        # This is to compute the cache used by a new request being scheduled
+        self.cache_budget_module = None if cache.num_full_attention_groups else cache.config.sliding_window
 
     @traced
     def add_waiting_request(self, state: RequestState):
@@ -51,10 +55,11 @@ class Scheduler(ABC):
         self.waiting_requests_order.append(state.request_id)
 
     @abstractmethod
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
-        """Schedules requests for the next batch based on available token budget. This method selects which requests
-        should be processed in the current batch, considering the token budget and the scheduler's prioritization rules.
-        The token_budget is the maximum number of tokens that can be processed in this batch."""
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState]:
+        """Schedules requests for the next batch based on available token and cache budgets. This method selects which
+        requests should be processed in the current batch, considering the budgets and the scheduler's prioritization
+        rules. The token_budget is the maximum number of tokens that can be processed in a batch, and the cache_budget
+        is the maximum number of KV cache entries that can be read in a batch."""
 
     @traced
     def has_pending_requests(self) -> bool:
@@ -62,14 +67,13 @@ class Scheduler(ABC):
         return len(self.active_requests) or len(self.waiting_requests)
 
     @traced
-    def finish_request(self, request_id: str, evict_from_cache: bool = True):
+    def finish_request(self, request_id: str, evict_from_cache: bool = True) -> None:
         """Completes processing of a request and optionally frees its allocated cache blocks. This method is called
         when a request has finished generation or encountered an error.
         """
         if evict_from_cache:
             self.cache.free_blocks(request_id)
-            if request_id in self.active_requests:
-                del self.active_requests[request_id]
+            self.active_requests.pop(request_id, None)
 
     @traced
     def get_active_request_static_outputs(self, request_id: str) -> list[int]:
@@ -89,10 +93,8 @@ class Scheduler(ABC):
         """Remove all cancelled requests from active and waiting queues."""
         with self._cancellation_lock:
             for request_id in self._requests_to_cancel:
-                if request_id in self.active_requests:
-                    del self.active_requests[request_id]
-                if request_id in self.waiting_requests:
-                    del self.waiting_requests[request_id]
+                self.active_requests.pop(request_id, None)
+                self.waiting_requests.pop(request_id, None)
                 if request_id in self.waiting_requests_order:
                     self.waiting_requests_order.remove(request_id)
                 self.cache.free_blocks(request_id)
@@ -106,7 +108,7 @@ class Scheduler(ABC):
         )
 
     @traced
-    def _allocate_blocks_if_needed(self, state: RequestState) -> bool:
+    def _allocate_blocks_if_needed(self, state: RequestState, len_next_tokens: int) -> bool:
         """Allocate additional cache blocks for a request if the currently allocated blocks are insufficient to
         accommodate the next tokens. It calculates how many blocks are needed based on the request's current
         cache occupancy and the number of tokens to be processed. The allocation itself is done by the CacheAllocator
@@ -115,20 +117,16 @@ class Scheduler(ABC):
         # 1. we check that the occupancy is less than the requested length
         # 2. we allocate enough blocks to cover the requested length
         current_len = state.current_len()
-        len_next_tokens = len(state.tokens_to_process)
         occupancy = state.allocated_blocks * self.cache.block_size - current_len
         if occupancy < len_next_tokens or state.allocated_blocks == 0:
             blocks_needed = ((len_next_tokens - occupancy + 1) // self.cache.block_size) + 1
-            allocated = self.cache.allocate_blocks(blocks_needed, state)
+            allocated = self.cache.allocate_blocks(blocks_needed, state.request_id, state.allocated_blocks)
             if allocated is None:
                 return False
             state.allocated_blocks += allocated
         return True
 
-    @traced(span_name="prepare_request")
-    def _prepare_request_for_processing(
-        self, state: RequestState, token_budget: int, request_ids_to_remove_from_waiting: set[str]
-    ) -> None:
+    def _infer_request_tokens(self, state: RequestState, request_ids_to_remove_from_waiting: set[str]) -> list[int]:
         """Prepares a request for processing in the current batch. If prefix sharing is enabled, and the request was
         pending, this is where we look for a prefix match and split the request if found."""
         # If prefix sharing is enabled, we look for a prefix match and split the request if found
@@ -138,6 +136,8 @@ class Scheduler(ABC):
                 self.active_requests[state.request_id] = state
                 request_ids_to_remove_from_waiting.add(state.request_id)
                 state.status = RequestStatus.SPLIT_PENDING_REMAINDER
+                # We keep track of the number of allocated blocks to avoid double allocation
+                state.allocated_blocks += prefill_length // self.cache.block_size
                 # Even if we match the whole request, we keep at least 1 token to start decoding
                 prefill_length = min(prefill_length, len(state.tokens_to_process) - 1)
                 state.remaining_prefill_tokens = state.tokens_to_process[prefill_length:]
@@ -150,9 +150,25 @@ class Scheduler(ABC):
         # Otherwise, the tokens to process are the prompt ids, which are the full prompt or the last predicted tokens
         else:
             request_tokens = state.tokens_to_process
-
+        return request_tokens
+
+    def _schedule_request(
+        self,
+        state: RequestState,
+        request_tokens: list[int],
+        token_budget: int,
+        request_ids_to_remove_from_waiting: set[str],
+    ) -> None:
+        """Schedules a request for the current batch, updating the request's status according to the token budget left.
+        If the request has children (for parallel decoding), it ensures at least one token remains before the request is
+        forked."""
+        # If the request has one or more children we make sure not to prefill it entirely
+        if state.num_children > 0 and token_budget >= len(request_tokens) - 1:
+            token_budget = len(request_tokens) - 1
+            self._requests_to_fork.append(state)
+
+        # Case: we can process the entire prompt/remainder
         if len(request_tokens) < token_budget:
-            # Can process the entire prompt/remainder
             if state.status == RequestStatus.PENDING:
                 self.active_requests[state.request_id] = state
                 state.status = RequestStatus.PREFILLING
@@ -161,8 +177,9 @@ class Scheduler(ABC):
                 state.status = RequestStatus.PREFILLING
                 state.tokens_to_process = state.remaining_prefill_tokens
                 state.remaining_prefill_tokens = []
+
+        # Otherwise: we need to split the request
         else:
-            # Need to split the request
             if state.status == RequestStatus.PENDING:
                 self.active_requests[state.request_id] = state
                 state.status = RequestStatus.PREFILLING_SPLIT
@@ -189,7 +206,7 @@ class FIFOScheduler(Scheduler):
         self.safety_margin = safety_margin
 
     @traced
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
         priority_states: list[RequestState] = []
         second_priority_states: list[RequestState] = []
         scheduled_requests = []
@@ -201,35 +218,62 @@ class FIFOScheduler(Scheduler):
                 second_priority_states.append(state)
 
         # Add waiting requests to second priority
-        for req_id in self.waiting_requests_order:
-            second_priority_states.append(self.waiting_requests[req_id])
+        if not self.block_new_requests:
+            for req_id in self.waiting_requests_order:
+                second_priority_states.append(self.waiting_requests[req_id])
 
         candidates = priority_states + second_priority_states
         request_ids_to_remove_from_waiting = set()
         safety_margins = self.safety_margin * self.cache.num_blocks
 
+        one_allocation_failed = False
+
         for state in candidates:
             # If we are out the safety margin, we only accept decoding requests or the first prefill request
             num_free_blocks = self.cache.get_num_free_blocks()
             outside_safety_margin = num_free_blocks < safety_margins
             if outside_safety_margin and scheduled_requests and state.status != RequestStatus.DECODING:
+                logger.info(
+                    f"Outside safety margin, breaking out of scheduling loop. {num_free_blocks = } {safety_margins = }"
+                )
                 break
 
-            self._prepare_request_for_processing(state, token_budget, request_ids_to_remove_from_waiting)
-            request_len = len(state.tokens_to_process)
-            # If we can't allocate blocks, do not schedule the request and break if the cache is full
-            if not self._allocate_blocks_if_needed(state):
-                if self.cache.get_num_free_blocks() == 0:
+            # Check cache budget
+            cache_needed = state.current_len()
+            cache_needed = (
+                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
+            )
+            if cache_budget < cache_needed:
+                continue
+
+            # Infer the tokens that will be present in the batch if token budget is enough
+            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
+            # Account for token budget
+            request_len = min(len(request_tokens), token_budget)
+            # Check there will be enough cache for the new tokens
+            allocation_successful = self._allocate_blocks_if_needed(state, request_len)
+
+            # If the allocation would not be successful, we move on to the next request
+            if not allocation_successful:
+                one_allocation_failed = True
+                # If we have reached a request that was waiting, all subsequent requests are also waiting, and will need
+                # allocation as well. So if there is no more free blocks, we can safely break out of the loop.
+                if num_free_blocks == 0 and state.request_id in self.waiting_requests:
+                    logger.info(f"Breaking mid-loop for request {state.request_id} because the cache is full")
                     break
                 continue
 
-            # Add the request to the scheduled requests
+            # If this point is reached, it means we can safely schedule the request
+            self._schedule_request(state, request_tokens, token_budget, request_ids_to_remove_from_waiting)
+            request_len = len(state.tokens_to_process)  # it may change after scheduling
             scheduled_requests.append(state)
 
-            # Update the token budget
+            # Update the token and cache budgets
             token_budget -= request_len
+            cache_budget -= cache_needed
+
             # If using prefix sharing, we make note of the blocks that will be computed in the forward pass
-            if self.cache.use_prefix_sharing:
+            if self.cache.allow_block_sharing:
                 tokens_in_current_block = state.current_len() % self.cache.block_size
                 tokens_after_forward = tokens_in_current_block + request_len
                 complete_blocks = tokens_after_forward // self.cache.block_size
@@ -241,18 +285,24 @@ class FIFOScheduler(Scheduler):
             if was_waiting:
                 request_ids_to_remove_from_waiting.add(req_id)
 
-            # Early exit of the loop if we have no token budget left
-            if token_budget == 0:
+            # Early exit of the loop if we have no budget left
+            if token_budget == 0 or cache_budget == 0:
                 break
 
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
         self.waiting_requests_order = deque(
             [req_id for req_id in self.waiting_requests_order if req_id not in request_ids_to_remove_from_waiting]
         )
 
+        # If no requests were scheduled and the cache is full, we signal it by returning None
+        if not scheduled_requests and one_allocation_failed:
+            return None
+
         return scheduled_requests
 
 
 # FIXME: prioritize adding from waiting reqs before scheduling `RequestStatus.DECODING` when cache space allows it
+# TODO: further consolidate the code by making more of it common. The reference Scheduler is FIFO, not this one.
 @attach_tracer()
 class PrefillFirstScheduler(Scheduler):
     """Scheduler that prioritizes split prefill requests over decoding requests. This scheduler ensures that split
@@ -260,7 +310,7 @@ class PrefillFirstScheduler(Scheduler):
     decoding requests."""
 
     @traced
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
         priority_states: list[RequestState] = []
         second_priority_states: list[RequestState] = []
         scheduled_requests = []
@@ -273,29 +323,49 @@ class PrefillFirstScheduler(Scheduler):
                 second_priority_states.append(state)
 
         # Add waiting requests to second priority
-        for req_id in self.waiting_requests_order:
-            second_priority_states.append(self.waiting_requests[req_id])
+        if not self.block_new_requests:
+            for req_id in self.waiting_requests_order:
+                second_priority_states.append(self.waiting_requests[req_id])
 
         candidates = priority_states + second_priority_states
-
         request_ids_to_remove_from_waiting = set()
+        one_allocation_failed = False
 
         for state in candidates:
-            self._prepare_request_for_processing(state, token_budget, request_ids_to_remove_from_waiting)
-            request_len = len(state.tokens_to_process)
-            # If we can't allocate blocks, do not schedule the request and break if the cache is full
-            if not self._allocate_blocks_if_needed(state):
-                if self.cache.get_num_free_blocks() == 0:
+            # Check cache budget
+            cache_needed = state.current_len()
+            cache_needed = (
+                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
+            )
+            if cache_budget < cache_needed:
+                continue
+
+            # Infer the tokens that will be present in the batch if token budget is enough
+            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
+            # Account for token budget
+            request_len = min(len(request_tokens), token_budget)
+            # Check there will be enough cache for the new tokens
+            allocation_successful = self._allocate_blocks_if_needed(state, request_len)
+
+            # If the allocation would not be successful, we move on to the next request
+            if not allocation_successful:
+                one_allocation_failed = True
+                # If the request was waiting, all requests afterwards will need allocation, so we break if the cache is full
+                if state.request_id in self.waiting_requests and self.cache.get_num_free_blocks() == 0:
                     break
                 continue
 
-            # Add the request to the scheduled requests
+            # If this point is reached, it means we can safely schedule the request
+            self._schedule_request(state, request_tokens, token_budget, request_ids_to_remove_from_waiting)
+            request_len = len(state.tokens_to_process)  # it may change after scheduling
             scheduled_requests.append(state)
 
-            # Update the token budget
+            # Update the token and cache budgets
             token_budget -= request_len
+            cache_budget -= cache_needed
+
             # If using prefix sharing, we make note of the blocks that will be computed in the forward pass
-            if self.cache.use_prefix_sharing:
+            if self.cache.allow_block_sharing:
                 tokens_in_current_block = state.current_len() % self.cache.block_size
                 tokens_after_forward = tokens_in_current_block + request_len
                 complete_blocks = tokens_after_forward // self.cache.block_size
@@ -303,18 +373,23 @@ class PrefillFirstScheduler(Scheduler):
 
             # Remove the request from the waiting queue and mark it as removed
             req_id = state.request_id
-            if req_id in self.waiting_requests:
-                del self.waiting_requests[req_id]
+            was_waiting = self.waiting_requests.pop(req_id, None) is not None
+            if was_waiting:
                 request_ids_to_remove_from_waiting.add(req_id)
 
-            # Early exit of the loop if we have no token budget left
-            if token_budget == 0:
+            # Early exit of the loop if we have no budget left
+            if token_budget == 0 or cache_budget == 0:
                 break
 
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
         self.waiting_requests_order = deque(
             [req_id for req_id in self.waiting_requests_order if req_id not in request_ids_to_remove_from_waiting]
         )
 
+        # If no requests were scheduled and the cache is full, we signal it by returning None
+        if not scheduled_requests and one_allocation_failed:
+            return None
+
         return scheduled_requests
 
 

transformers/generation/logits_process.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2024 The HuggingFace Inc. team and Google DeepMind.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1543,133 +1542,6 @@ class PrefixConstrainedLogitsProcessor(LogitsProcessor):
         return scores_processed
 
 
-class HammingDiversityLogitsProcessor(LogitsProcessor):
-    r"""
-    [`LogitsProcessor`] that enforces diverse beam search.
-    Note that this logits processor is only effective for [`PreTrainedModel.group_beam_search`]. See [Diverse Beam
-    Search: Decoding Diverse Solutions from Neural Sequence Models](https://huggingface.co/papers/1610.02424) for more
-    details.
-    Traditional beam search often generates very similar sequences across different beams.
-    `HammingDiversityLogitsProcessor` addresses this by penalizing beams that generate tokens already chosen by other
-    beams in the same time step.
-    Args:
-        diversity_penalty (`float`):
-            This value is subtracted from a beam's score if it generates a token same as any beam from other group at a
-            particular time. A higher `diversity_penalty` will enforce greater diversity among the beams. Adjusting
-            this value can help strike a balance between diversity and natural likelihood.
-        num_beams (`int`):
-            Number of beams for beam search. 1 means no beam search.
-        num_beam_groups (`int`):
-            Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams.
-            [this paper](https://huggingface.co/papers/1610.02424) for more details.
-    Examples:
-    ```python
-    >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-    >>> import torch
-    >>> # Initialize the model and tokenizer
-    >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
-    >>> model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
-    >>> # A long text about the solar system
-    >>> text = (
-    ...     "The Solar System is a gravitationally bound system comprising the Sun and the objects that orbit it, "
-    ...     "either directly or indirectly. Of the objects that orbit the Sun directly, the largest are the eight "
-    ...     "planets, with the remainder being smaller objects, such as the five dwarf planets and small Solar System "
-    ...     "bodies. The Solar System formed 4.6 billion years ago from the gravitational collapse of a giant "
-    ...     "interstellar molecular cloud."
-    ... )
-    >>> inputs = tokenizer("summarize: " + text, return_tensors="pt")
-    >>> # Generate diverse summary
-    >>> outputs_diverse = model.generate(
-    ...     **inputs,
-    ...     num_beam_groups=2,
-    ...     diversity_penalty=10.0,
-    ...     max_length=100,
-    ...     num_beams=4,
-    ...     num_return_sequences=2,
-    ... )
-    >>> summaries_diverse = tokenizer.batch_decode(outputs_diverse, skip_special_tokens=True)
-    >>> # Generate non-diverse summary
-    >>> outputs_non_diverse = model.generate(
-    ...     **inputs,
-    ...     max_length=100,
-    ...     num_beams=4,
-    ...     num_return_sequences=2,
-    ... )
-    >>> summary_non_diverse = tokenizer.batch_decode(outputs_non_diverse, skip_special_tokens=True)
-    >>> # With `diversity_penalty`, the resulting beams are much more diverse
-    >>> print(summary_non_diverse)
-    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
-    'the Solar System formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.']
-    >>> print(summaries_diverse)
-    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
-    'the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets. the rest of the objects are smaller objects, such as the five dwarf planets and small solar system bodies.']
-    ```
-    """
-
-    def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int):
-        logger.warning_once(
-            "`HammingDiversityLogitsProcessor` is deprecated and will be removed in v4.62.0, as constrained beam search has been moved to the Hub: https://hf.co/transformers-community/constrained-beam-search."
-        )
-        if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0):
-            raise ValueError("`diversity_penalty` should be a float strictly larger than 0.")
-        self._diversity_penalty = diversity_penalty
-        if not isinstance(num_beams, int) or num_beams < 2:
-            raise ValueError("`num_beams` should be an integer strictly larger than 1.")
-        self._num_beams = num_beams
-        if not isinstance(num_beam_groups, int) or num_beam_groups < 2:
-            raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.")
-        if num_beam_groups > num_beams:
-            raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.")
-        self._num_sub_beams = num_beams // num_beam_groups
-
-    def __call__(
-        self,
-        input_ids: torch.LongTensor,
-        scores: torch.FloatTensor,
-        current_tokens: torch.LongTensor,
-        beam_group_idx: int,
-    ) -> torch.FloatTensor:
-        r"""
-        Args:
-            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
-            scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
-                Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
-                beam search or log softmax for each vocabulary token when using beam search
-            current_tokens (`torch.LongTensor` of shape `(batch_size)`):
-                Indices of input sequence tokens in the vocabulary, corresponding to the tokens selected by the other
-                beam groups in the current generation step.
-            beam_group_idx (`int`):
-                The index of the beam group currently being processed.
-        Return:
-            `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
-                The processed prediction scores.
-        """
-        # hamming diversity: penalise using same token in current group which was used in previous groups at
-        # the same time step
-        batch_size = current_tokens.shape[0] // self._num_beams
-        group_start_idx = beam_group_idx * self._num_sub_beams
-        group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams)
-        group_size = group_end_idx - group_start_idx
-        vocab_size = scores.shape[-1]
-
-        if group_start_idx == 0:
-            return scores
-
-        scores_processed = scores.clone()
-        for batch_idx in range(batch_size):
-            # predicted tokens of last time step of previous groups
-            previous_group_tokens = current_tokens[
-                batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx
-            ]
-            token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device)
-            scores_processed[batch_idx * group_size : (batch_idx + 1) * group_size] -= (
-                self._diversity_penalty * token_frequency
-            )
-
-        return scores_processed
-
-
 class ForcedBOSTokenLogitsProcessor(LogitsProcessor):
     r"""
     [`LogitsProcessor`] that enforces the specified token as the first generated token. Used with encoder-decoder

transformers/generation/stopping_criteria.py

@@ -430,7 +430,7 @@ class StopStringCriteria(StoppingCriteria):
         initial_match = end_lengths > 0
 
         # Tokens continue the string if the cumsum() so far is one of the valid positions for that token
-        # Note that we're actually tracking one cumsum() for for each possible end_length
+        # Note that we're actually tracking one cumsum() for each possible end_length
         later_match = torch.any(cumsum[:, :-1, :, None] == valid_positions[:, :, :, :, None], axis=-2)
 
         # The match vector is a boolean vector that indicates which positions have valid tokens

transformers/generation/streamers.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2023 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");