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

transformers/generation/configuration_utils.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2022 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -101,6 +100,10 @@ class GenerationConfig(PushToHubMixin):
 
     </Tip>
 
+    Note: the configuration field that are still `None` will be overriden by `GenerationConfig._get_default_generation_params()`
+    during the generation loop. If you want to use different values for these fields, make sure to explicitly set them in the
+    generation config.
+
     Arg:
         > Parameters that control the length of the output
 
@@ -129,14 +132,14 @@ class GenerationConfig(PushToHubMixin):
 
         > Parameters that control the generation strategy used
 
-        do_sample (`bool`, defaults to `False`):
+        do_sample (`bool`):
             Whether or not to use sampling ; use greedy decoding otherwise.
         num_beams (`int`, *optional*):
             Number of beams for beam search. 1 means no beam search.
 
         > Parameters that control the cache
 
-        use_cache (`bool`, defaults to `True`):
+        use_cache (`bool`):
             Whether or not the model should use the past last key/values attentions (if applicable to the model) to
             speed up decoding.
         cache_implementation (`str`, *optional*):
@@ -206,7 +209,7 @@ class GenerationConfig(PushToHubMixin):
         bad_words_ids (`list[list[int]]`, *optional*):
             List of list of token ids that are not allowed to be generated. Check
             [`~generation.NoBadWordsLogitsProcessor`] for further documentation and examples.
-        renormalize_logits (`bool`, defaults to `False`):
+        renormalize_logits (`bool`):
             Whether to renormalize the logits after applying all the logits processors (including the custom
             ones). It's highly recommended to set this flag to `True` as the search algorithms suppose the score logits
             are normalized but some logit processors break the normalization.
@@ -217,7 +220,7 @@ class GenerationConfig(PushToHubMixin):
         forced_eos_token_id (`int` or list[int]`, *optional*, defaults to `model.config.forced_eos_token_id`):
             The id of the token to force as the last generated token when `max_length` is reached. Optionally, use a
             list to set multiple *end-of-sequence* tokens.
-        remove_invalid_values (`bool`, defaults to `model.config.remove_invalid_values`):
+        remove_invalid_values (`bool`):
             Whether to remove possible *nan* and *inf* outputs of the model to prevent the generation method to crash.
             Note that using `remove_invalid_values` can slow down generation.
         exponential_decay_length_penalty (`tuple(int, float)`, *optional*):
@@ -234,7 +237,7 @@ class GenerationConfig(PushToHubMixin):
             Dictionary that maps a sequence of tokens to its bias term. Positive biases increase the odds of the
             sequence being selected, while negative biases do the opposite. Check
             [`~generation.SequenceBiasLogitsProcessor`] for further documentation and examples.
-        token_healing (`bool`, defaults to `False`):
+        token_healing (`bool`):
             Heal tail tokens of prompts by replacing them with their appropriate extensions.
             This enhances the quality of completions for prompts affected by greedy tokenization bias.
         guidance_scale (`float`, *optional*):
@@ -248,20 +251,20 @@ class GenerationConfig(PushToHubMixin):
 
         > Parameters that define the output variables of generate
 
-        num_return_sequences (`int`, *optional*, defaults to 1):
+        num_return_sequences (`int`, *optional*):
             The number of independently computed returned sequences for each element in the batch.
-        output_attentions (`bool`, defaults to `False`):
+        output_attentions (`bool`):
             Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
             tensors for more details.
-        output_hidden_states (`bool`, defaults to `False`):
+        output_hidden_states (`bool`):
             Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
             more details.
-        output_scores (`bool`, defaults to `False`):
+        output_scores (`bool`):
             Whether or not to return the prediction scores. See `scores` under returned tensors for more details.
-        output_logits (`bool`, defaults to `False`):
+        output_logits (`bool`):
             Whether or not to return the unprocessed prediction logit scores. See `logits` under returned tensors for
             more details.
-        return_dict_in_generate (`bool`, defaults to `False`):
+        return_dict_in_generate (`bool`):
             Whether or not to return a [`~utils.ModelOutput`], as opposed to returning exclusively the generated
             sequence. This flag must be set to `True` to return the generation cache (when `use_cache` is `True`)
             or optional outputs (see flags starting with `output_`)
@@ -286,7 +289,7 @@ class GenerationConfig(PushToHubMixin):
             (e.g. multilingual models with different target languages in one batch)
 
         > Generation parameters exclusive to assistant generation
-        is_assistant (`bool`, defaults to `False`):
+        is_assistant (`bool`):
             Whether the model is an assistant (draft) model.
         num_assistant_tokens (`int`, *optional*):
             Defines the number of _speculative tokens_ that shall be generated by the assistant model before being
@@ -327,7 +330,7 @@ class GenerationConfig(PushToHubMixin):
         compile_config (CompileConfig, *optional*):
             If using a compilable cache, this controls how `generate` will `compile` the forward pass for faster
             inference.
-        disable_compile (`bool`, defaults to `False`):
+        disable_compile (`bool`):
             Whether to disable the automatic compilation of the forward pass. Automatic compilation happens when
             specific criteria are met, including using a compilable cache. Please open an issue if you find the
             need to use this flag.
@@ -346,11 +349,11 @@ class GenerationConfig(PushToHubMixin):
         self.stop_strings = kwargs.pop("stop_strings", None)
 
         # Parameters that control the generation strategy used
-        self.do_sample = kwargs.pop("do_sample", False)
+        self.do_sample = kwargs.pop("do_sample", None)
         self.num_beams = kwargs.pop("num_beams", None)
 
         # Parameters that control the cache
-        self.use_cache = kwargs.pop("use_cache", True)
+        self.use_cache = kwargs.pop("use_cache", None)
         self.cache_implementation = kwargs.pop("cache_implementation", None)
         self.cache_config = kwargs.pop("cache_config", None)
 
@@ -368,15 +371,15 @@ class GenerationConfig(PushToHubMixin):
         self.length_penalty = kwargs.pop("length_penalty", None)
         self.no_repeat_ngram_size = kwargs.pop("no_repeat_ngram_size", None)
         self.bad_words_ids = kwargs.pop("bad_words_ids", None)
-        self.renormalize_logits = kwargs.pop("renormalize_logits", False)
+        self.renormalize_logits = kwargs.pop("renormalize_logits", None)
         self.forced_bos_token_id = kwargs.pop("forced_bos_token_id", None)
         self.forced_eos_token_id = kwargs.pop("forced_eos_token_id", None)
-        self.remove_invalid_values = kwargs.pop("remove_invalid_values", False)
+        self.remove_invalid_values = kwargs.pop("remove_invalid_values", None)
         self.exponential_decay_length_penalty = kwargs.pop("exponential_decay_length_penalty", None)
         self.suppress_tokens = kwargs.pop("suppress_tokens", None)
         self.begin_suppress_tokens = kwargs.pop("begin_suppress_tokens", None)
         self.sequence_bias = kwargs.pop("sequence_bias", None)
-        self.token_healing = kwargs.pop("token_healing", False)
+        self.token_healing = kwargs.pop("token_healing", None)
         self.guidance_scale = kwargs.pop("guidance_scale", None)
 
         self.watermarking_config = kwargs.pop("watermarking_config", None)
@@ -384,12 +387,12 @@ class GenerationConfig(PushToHubMixin):
             self.watermarking_config = WatermarkingConfig.from_dict(self.watermarking_config)
 
         # Parameters that define the output variables of `generate`
-        self.num_return_sequences = kwargs.pop("num_return_sequences", 1)
-        self.output_attentions = kwargs.pop("output_attentions", False)
-        self.output_hidden_states = kwargs.pop("output_hidden_states", False)
-        self.output_scores = kwargs.pop("output_scores", False)
-        self.output_logits = kwargs.pop("output_logits", False)
-        self.return_dict_in_generate = kwargs.pop("return_dict_in_generate", False)
+        self.num_return_sequences = kwargs.pop("num_return_sequences", None)
+        self.output_attentions = kwargs.pop("output_attentions", None)
+        self.output_hidden_states = kwargs.pop("output_hidden_states", None)
+        self.output_scores = kwargs.pop("output_scores", None)
+        self.output_logits = kwargs.pop("output_logits", None)
+        self.return_dict_in_generate = kwargs.pop("return_dict_in_generate", None)
 
         # Special tokens that can be used at generation time
         self.pad_token_id = kwargs.pop("pad_token_id", None)
@@ -401,7 +404,7 @@ class GenerationConfig(PushToHubMixin):
         self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None)
 
         # Assistant generation
-        self.is_assistant = kwargs.pop("is_assistant", False)
+        self.is_assistant = kwargs.pop("is_assistant", None)
         self.num_assistant_tokens = kwargs.pop("num_assistant_tokens", None)
         self.num_assistant_tokens_schedule = kwargs.pop("num_assistant_tokens_schedule", None)
         self.assistant_confidence_threshold = kwargs.pop("assistant_confidence_threshold", None)
@@ -413,7 +416,7 @@ class GenerationConfig(PushToHubMixin):
 
         # Performance
         self.compile_config = kwargs.pop("compile_config", None)
-        self.disable_compile = kwargs.pop("disable_compile", False)
+        self.disable_compile = kwargs.pop("disable_compile", None)
 
         # Deprecated (moved to the Hub). TODO remove for v5
         self.low_memory = kwargs.pop("low_memory", None)
@@ -483,7 +486,7 @@ class GenerationConfig(PushToHubMixin):
         if self.constraints is not None or self.force_words_ids is not None:
             generation_mode = GenerationMode.CONSTRAINED_BEAM_SEARCH
         elif self.num_beams is None or self.num_beams == 1:
-            if not self.do_sample:
+            if self.do_sample is not True:
                 if (
                     self.top_k is not None
                     and self.top_k > 1
@@ -498,7 +501,7 @@ class GenerationConfig(PushToHubMixin):
         else:
             if self.num_beam_groups is not None and self.num_beam_groups > 1:
                 generation_mode = GenerationMode.GROUP_BEAM_SEARCH
-            elif self.do_sample:
+            elif self.do_sample is True:
                 generation_mode = GenerationMode.BEAM_SAMPLE
             else:
                 generation_mode = GenerationMode.BEAM_SEARCH
@@ -537,6 +540,7 @@ class GenerationConfig(PushToHubMixin):
             "max_length": 20,
             "min_length": 0,
             "do_sample": False,
+            "use_cache": True,
             "early_stopping": False,
             "num_beams": 1,
             "temperature": 1.0,
@@ -616,7 +620,10 @@ class GenerationConfig(PushToHubMixin):
 
         # 2. Validation of attribute combinations
         # 2.1. detect sampling-only parameterization when not in sampling mode
-        if not self.do_sample:
+
+        # Note that we check `is not True` in purpose. Boolean fields can also be `None` so we
+        # have to be explicit. Value of `None` is same as having `False`, i.e. the default value
+        if self.do_sample is not True:
             greedy_wrong_parameter_msg = (
                 "`do_sample` is set not to set `True`. However, `{flag_name}` is set to `{flag_value}` -- this flag is only "
                 "used in sample-based generation modes. You should set `do_sample=True` or unset `{flag_name}`."
@@ -662,21 +669,25 @@ class GenerationConfig(PushToHubMixin):
                 )
 
         # 2.4. check `num_return_sequences`
-        if self.num_return_sequences > 1:
+        if self.num_return_sequences is not None and self.num_return_sequences > 1:
             if self.num_beams is None or self.num_beams == 1:
                 if not self.do_sample:
                     raise ValueError(
                         "Greedy methods (do_sample != True) without beam search do not support "
                         f"`num_return_sequences` different than 1 (got {self.num_return_sequences})."
                     )
-            elif self.num_beams is not None and self.num_return_sequences > self.num_beams:
+            elif (
+                self.num_beams is not None
+                and self.num_return_sequences is not None
+                and self.num_return_sequences > self.num_beams
+            ):
                 raise ValueError(
                     f"`num_return_sequences` ({self.num_return_sequences}) has to be smaller or equal to `num_beams` "
                     f"({self.num_beams})."
                 )
 
         # 2.5. check cache-related arguments
-        if not self.use_cache:
+        if self.use_cache is False:
             # In this case, all cache-related arguments should be unset. However, since `use_cache=False` is often used
             # passed to `generate` directly to hot-fix cache issues, let's raise a warning instead of an error
             # (otherwise a user might need to overwrite several parameters).
@@ -691,9 +702,9 @@ class GenerationConfig(PushToHubMixin):
                     )
 
         # 2.6. other incorrect combinations
-        if not self.return_dict_in_generate:
+        if self.return_dict_in_generate is not True:
             for extra_output_flag in self.extra_output_flags:
-                if getattr(self, extra_output_flag):
+                if getattr(self, extra_output_flag) is True:
                     minor_issues[extra_output_flag] = (
                         f"`return_dict_in_generate` is NOT set to `True`, but `{extra_output_flag}` is. When "
                         f"`return_dict_in_generate` is not `True`, `{extra_output_flag}` is ignored."
@@ -1182,7 +1193,7 @@ class GenerationConfig(PushToHubMixin):
         generation_config._original_object_hash = hash(generation_config)
         return generation_config
 
-    def update(self, defaults_only=False, **kwargs):
+    def update(self, defaults_only=False, allow_custom_entries=False, **kwargs):
         """
         Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes,
         returning all the unused kwargs.
@@ -1190,6 +1201,8 @@ class GenerationConfig(PushToHubMixin):
         Args:
             defaults_only (`bool`, *optional*, defaults to `False`):
                 Whether to update all keys in config with `kwargs` or only those that are set to `None` (i.e. default value).
+            allow_custom_entries (`bool`, *optional*, defaults to `False`):
+                Whether to allow updating custom entries into the config with `kwargs` if not present in the current config.
             kwargs (`dict[str, Any]`):
                 Dictionary of attributes to tentatively update this class.
 
@@ -1198,7 +1211,10 @@ class GenerationConfig(PushToHubMixin):
         """
         to_remove = []
         for key, value in kwargs.items():
-            if hasattr(self, key):
+            if allow_custom_entries and not hasattr(self, key):
+                setattr(self, key, value)
+                to_remove.append(key)
+            elif hasattr(self, key):
                 if not defaults_only or getattr(self, key) is None:
                     setattr(self, key, value)
                     to_remove.append(key)
@@ -1261,8 +1277,7 @@ class BaseWatermarkingConfig(ABC):
         return output
 
     def __iter__(self):
-        for attr, value in copy.deepcopy(self.__dict__).items():
-            yield attr, value
+        yield from copy.deepcopy(self.__dict__).items()
 
     def __repr__(self):
         return f"{self.__class__.__name__} {self.to_json_string()}"

transformers/generation/continuous_batching/__init__.py

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

transformers/generation/continuous_batching/cache.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -223,11 +222,18 @@ class PagedAttentionCache:
         # Block management data structures
         self.allow_block_sharing = allow_block_sharing
         self.group_cache_managers: list[CacheAllocator] = []
+        self.num_full_attention_groups = 0
+        self.num_sliding_attention_groups = 0
+        self.max_sliding_window_blocks_per_request = 0
+
         for i, group_type in enumerate(group_types):
             if group_type == "full_attention":
                 cm = FullAttentionCacheAllocator(i, self.block_size, allow_block_sharing=allow_block_sharing)
+                self.num_full_attention_groups += 1
             elif group_type == "sliding_attention":
                 cm = SlidingAttentionCacheAllocator(i, self.block_size, config.sliding_window)
+                self.num_sliding_attention_groups += 1
+                self.max_sliding_window_blocks_per_request = cm._max_blocks_per_request
             else:
                 raise ValueError(f"Invalid group type: {group_type}")
             self.group_cache_managers.append(cm)
@@ -238,16 +244,38 @@ class PagedAttentionCache:
         self.blocks_to_complete: dict[str, int] = {}
         self._total_prefix_length: int = 0  # a counter to measure the impact of prefix sharing, also used in tests
 
+    def will_allocation_be_successful(self, num_requested_blocks: int, allocated_blocks: int) -> bool:
+        """Returns a boolean indicating if the allocation of (num_requested_blocks) blocks will be successful. The
+        number of newly allocated blocks needed is predicted by the following rules:
+        - for full attention groups: since there is no sliding window for full attention layers, one requested block is
+            always equivalent to one newly allocated block for EACH full attention group
+        - for sliding window groups: because of the sliding window, the number of blocks allocated to a request is
+            capped. Using the number of already (allocated_blocks) we can compute the number of new blocks to actually
+            allocate to the request, which can be lower than the number of requested blocks. That number is the same for
+            all sliding window groups, as only one sliding window size is supported.
+        """
+        # This is not in a branch, because it is very rare to have zero full attention layer
+        needed_blocks = num_requested_blocks * self.num_full_attention_groups
+        # Only take this branch if the model has sliding window attention layers
+        if self.num_sliding_attention_groups:
+            blocks_left = max(self.max_sliding_window_blocks_per_request - allocated_blocks, 0)
+            needed_blocks += min(blocks_left, num_requested_blocks) * self.num_sliding_attention_groups
+        return needed_blocks <= self.get_num_free_blocks()
+
     @traced
-    def allocate_blocks(self, n_blocks: int, state: RequestState) -> int:
+    def allocate_blocks(self, n_blocks: int, request_id: str, allocated_blocks: int) -> int | None:
         """Allocate cache blocks across all layer groups for a given request. Actual allocation is done by the cache
         managers, and this method only returns the maximum number of blocks actually allocated across all managers."""
+        # First check allocation will be successful before starting, to avoid partial allocations
+        if not self.will_allocation_be_successful(n_blocks, allocated_blocks):
+            return None
+        # Allocate blocks across all cache managers
         max_allocated = 0
         for cm in self.group_cache_managers:
-            allocated = cm.allocate_blocks(n_blocks, state.request_id, self._block_manager)
-            if allocated is None:
-                return None
-            max_allocated = max(max_allocated, allocated)
+            num_allocated_blocks = cm.allocate_blocks(n_blocks, request_id, self._block_manager)
+            if num_allocated_blocks is None:
+                raise ValueError(f"Failed to allocate {n_blocks} blocks for request {request_id}")
+            max_allocated = max(max_allocated, num_allocated_blocks)
         return max_allocated
 
     @traced

transformers/generation/continuous_batching/cache_manager.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -207,6 +206,15 @@ class BlockManager:
         else:
             self._uninit_block_ids.extend(blocks)
 
+    def uninitialize_unshared_block(self, block_id: int) -> None:
+        """Marks a block as uninitialized. Raises an error if the block has more than one reference."""
+        # Make sure the block has only one reference and remove it from the block table
+        block = self._id_to_block.pop(block_id)
+        if block.ref_count > 1:
+            raise RuntimeError(f"Block {block_id} has more than one reference: {block.ref_count = }")
+        # Add the block to the uninitialized blocks queue
+        self._uninit_block_ids.append(block_id)
+
     def mark_shareable_blocks_as_complete(
         self, num_complete_blocks: int, allocated_blocks: list[int], prompt_ids: list[int]
     ) -> None:
@@ -242,13 +250,17 @@ class BlockManager:
             block.hash = self.compute_hash(parent_hash, tokens, block.group_id)
 
             existing_block_id = self._hash_to_id.get(block.hash)
-            # If the block hash is already in the hash to id mapping, we reference the existing block instead
+            # If their was a different block with the same hash, we reference the existing block instead
             if existing_block_id is not None:
-                logger.debug(f"Found existing block {existing_block_id} for block {block.id}")
-                allocated_blocks[i] = existing_block_id
-                self._id_to_block[existing_block_id].ref_count += 1
-                new_parent_id = existing_block_id
-                self.free_blocks([block.id], shareable=True)
+                if existing_block_id == block.id:
+                    # This should not happen, but is not a problem in itself, so we just log a warning
+                    logger.warning(f"Block {block.id} was marked as complete more than once")
+                else:
+                    logger.debug(f"Found existing block {existing_block_id} for block {block.id}")
+                    allocated_blocks[i] = existing_block_id
+                    new_parent_id = existing_block_id
+                    self.increase_ref_count(existing_block_id)
+                    self.uninitialize_unshared_block(block.id)
 
             # Otherwise, we add the completed block to the hash table
             else:
@@ -349,16 +361,17 @@ class FullAttentionCacheAllocator(CacheAllocator):
         allocated if successful and None otherwise. For group of full attention layers, we always allocate the number of
         requested blocks."""
         # Make sure the request_id is in the block table and get the first block id
-        if request_id not in self.block_table:
-            self.block_table[request_id] = []  # TODO: check the impact of making this a deque
-            last_block_id = None
+        block_table = self.block_table.get(request_id, [])
+        if block_table:
+            last_block_id = block_table[-1]
         else:
-            last_block_id = self.block_table[request_id][-1]
+            self.block_table[request_id] = block_table  # TODO: check the impact of making this a deque
+            last_block_id = None
         # Actual allocation, return early if failed
         allocated_blocks = block_manager.get_free_blocks(n_blocks, last_block_id, self.uses_block_sharing, self._index)
         if allocated_blocks is None:
             return None
-        self.block_table[request_id].extend(allocated_blocks)
+        block_table.extend(allocated_blocks)
         return n_blocks
 
     def get_read_indices(self, request_id: str, past_length: int, query_length: int) -> list[int]:

transformers/generation/continuous_batching/continuous_api.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2024 The HuggingFace Inc. team.
 # Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
 #
@@ -34,7 +33,7 @@ from ...generation.logits_process import LogitsProcessor
 from ...utils.logging import logging
 from ...utils.metrics import ContinuousBatchProcessorMetrics, attach_tracer, traced
 from .cache import PagedAttentionCache
-from .requests import GenerationOutput, RequestState, RequestStatus, get_device_and_memory_breakdown, logger
+from .requests import GenerationOutput, RequestState, RequestStatus, logger
 from .scheduler import SCHEDULER_MAPPING, FIFOScheduler, Scheduler
 
 
@@ -250,7 +249,7 @@ class ContinuousBatchProcessor:
     def setup_static_tensors(self, num_groups: int) -> None:
         """Setup the static tensors that are used for storage during the generation step. No other tensor will be
         allowed for the inputs or the outputs of the generation step."""
-        num_pages = self.cache.num_blocks * self.cache.block_size
+        self.num_pages = self.cache.num_blocks * self.cache.block_size
         self.tensor_metadata = {"dtype": torch.int32, "device": self.model_device}
 
         # Some tensors always have the same shape regardless of the model
@@ -275,7 +274,7 @@ class ContinuousBatchProcessor:
 
         if attn_mask_is_needed(self.config):
             attn_mask_kwargs = {
-                "size": (1, 1, self.max_batch_tokens, num_pages + self.max_batch_tokens),
+                "size": (1, 1, self.max_batch_tokens, self.num_pages + self.max_batch_tokens),
                 "dtype": self.model_dtype,
                 "device": self.model_device,
             }
@@ -288,7 +287,7 @@ class ContinuousBatchProcessor:
             torch.empty((self.max_batch_tokens,), **self.tensor_metadata) for _ in range(num_groups)
         ]
         self.read_index_storage = [
-            torch.empty((num_pages + self.max_batch_tokens), **self.tensor_metadata) for _ in range(num_groups)
+            torch.empty((self.num_pages + self.max_batch_tokens), **self.tensor_metadata) for _ in range(num_groups)
         ]
         # For read index, the +T is because there are -1 for seqlen_q when model uses a sliding window
 
@@ -428,6 +427,33 @@ class ContinuousBatchProcessor:
         self.metrics.record_request_completion(state.created_time, state.request_id)
         self.output_queue.put(state.to_generation_output())
 
+    # TODO: there should be a way to choose the offloading policy: biggest request, oldest request, etc.
+    # Including a policy to not allow offloading and crashing the generation
+    def soft_reset_one_request(self) -> None:
+        """Soft resets one active request by removing it from active requests and re-adding it to the waiting queue.
+
+        The generated tokens are kept as part of the new request's initial prompt. When `block_new_requests` is False,
+        the oldest request is offloaded; when True, the newest request is offloaded. This method also sets
+        `block_new_requests` to True to prevent infinite loops of offloading and re-scheduling requests.
+        """
+        # The offloaded request is the newest (resp. oldest) if block_new_requests is True (resp. False)
+        if self.scheduler.block_new_requests:
+            request_id, state = self.scheduler.active_requests.popitem()
+        else:
+            request_id, state = next(iter(self.scheduler.active_requests.items()))
+        logger.info(
+            f"Soft resetting request {request_id} with {len(state.initial_tokens)} initial tokens and "
+            f"{len(state.generated_tokens)} generated tokens"
+        )
+        # Create a copy of the offloaded request keeping the generated tokens as addition to the initial prompt
+        new_state = state.create_equivalent_initial_request()
+        # Actual offloading of the request
+        self.scheduler.finish_request(request_id, evict_from_cache=True)
+        self.scheduler.add_waiting_request(new_state)
+        # This flag blocks any new requests from being scheduled until one request is finished. This ensures that we
+        # don't enter an offload / schedule loop
+        self.scheduler.block_new_requests = True
+
     @traced
     def prepare_next_batch(self) -> bool:
         """Prepare tensors and metadata for the next model forward pass. Returns True if there are requests to process,
@@ -441,9 +467,18 @@ class ContinuousBatchProcessor:
         self.metrics.record_queue_metrics(len(self.scheduler.active_requests), len(self.scheduler.waiting_requests))
 
         # Schedule the next batch of requests, stop if there are no requests in the batch
-        self.requests_in_batch = self.scheduler.schedule_batch(self.max_batch_tokens)
+        self.requests_in_batch = self.scheduler.schedule_batch(self.max_batch_tokens, self.num_pages)
+
+        # If requests_in_batch is None, it means we need to offload some requests if possible
+        if self.requests_in_batch is None:
+            if len(self.scheduler.active_requests) > 1:
+                self.soft_reset_one_request()
+            else:
+                raise RuntimeError("No requests can be scheduled and no request can be offloaded.")
+        # If it's an empty list, it means we have no requests to process
         if not self.requests_in_batch:
             return False
+        # Otherwise, we can continue with the non-empty batch
         self.metrics.record_batch_metrics(self.requests_in_batch)
 
         # Reset the static tensors used for storage
@@ -592,6 +627,7 @@ class ContinuousBatchProcessor:
                 if is_finished:
                     self.metrics.record_request_completion(state.created_time, state.request_id)
                     self.scheduler.finish_request(state.request_id, evict_from_cache=(not self.manual_eviction))
+                    self.scheduler.block_new_requests = False
                 self._maybe_send_output(state)
             #  Otherwise, the request is still prefilling, but the prefill has been split
             elif state.status == RequestStatus.PREFILLING_SPLIT:
@@ -621,9 +657,6 @@ class ContinuousBatchProcessor:
         if copy_source:
             self.cache.copy_cache(copy_source, copy_destination)
 
-        if self.cache.get_num_free_blocks() == 0:
-            raise ValueError("No more free blocks")
-
     @traced
     def has_pending_requests(self) -> bool:
         """Check if there are any active or waiting requests."""
@@ -679,11 +712,8 @@ class ContinuousBatchProcessor:
         if self._pad_inputs:
             padded_q = pad_by_intervals(self.actual_query_length, self.max_batch_tokens, self.q_padding_intervals)
             max_read_index_size = max(self.actual_index_sizes[i][0] for i in range(self.cache.num_groups))
-            padded_read_index_size = pad_by_intervals(
-                max_read_index_size - self.max_batch_tokens,
-                self.cache.num_blocks * self.cache.block_size,
-                self.kv_padding_intervals,
-            )
+            # The space planned for query tokens will be added later, so we remove it from the space planned for KV
+            padded_read_index_size = pad_by_intervals(max_read_index_size, self.num_pages, self.kv_padding_intervals)
         else:
             padded_q, padded_read_index_size = 0, 0
         # Retrieve the model kwargs with or without padding
@@ -811,7 +841,8 @@ class ContinuousBatchingManager:
         self.log_prob_generation = getattr(generation_config, "log_prob_generation", False)
         self.do_sample = getattr(generation_config, "do_sample", True)
         self.logit_processor = self.model._get_logits_processor(generation_config)
-        self.num_return_sequences = getattr(generation_config, "num_return_sequences", 1)
+        num_return_sequences = getattr(generation_config, "num_return_sequences", None)
+        self.num_return_sequences = num_return_sequences if num_return_sequences is not None else 1
 
         # self.model.generation_config.top_p = None NOTE: figure out why this was here
 
@@ -1108,13 +1139,13 @@ class ContinuousBatchingManager:
         # Loop body ends if there is no requests in the batch
         if not batch_processor.prepare_next_batch():
             return
-        # Debug logging of the current memory usage
-        if logger.level <= logging.DEBUG:
-            device, total, reserved, allocated = get_device_and_memory_breakdown()
-            available_memory = total - max(allocated, reserved)
-            logger.debug(
-                f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}, Available: {available_memory}"
-            )
+        # Debug logging of the current memory usage -- commented out because it's often not used even in debug
+        # if logger.level < logging.DEBUG:
+        #     device, total, reserved, allocated = get_device_and_memory_breakdown()
+        #     available_memory = total - max(allocated, reserved)
+        #     logger.debug(
+        # f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}, Available: {available_memory}"
+        #     )
 
         self._generation_step()
         batch_processor.update_batch()
@@ -1262,7 +1293,7 @@ class ContinuousMixin:
         # Initialize manager with the batch inputs
         results = {}
         gen_cfg = self.generation_config if generation_config is None else generation_config
-        num_requests = len(inputs) * gen_cfg.num_return_sequences
+        num_requests = len(inputs) * (gen_cfg.num_return_sequences if gen_cfg.num_return_sequences is not None else 1)
         # Prepare context managers for the main loop
         manager_cm = self.continuous_batching_context_manager(
             generation_config=generation_config,