[CB] Changes for long generation

src/transformers/generation/configuration_utils.py

@@ -1556,8 +1556,10 @@ class ContinuousBatchingConfig:
             Number of blocks in the KV cache. Auto-inferred from GPU memory when `None`.
         max_batch_tokens (`int`, *optional*):
             Maximum number of tokens in a batch. Auto-inferred from GPU memory when `None`.
-        max_memory_percent (`float`, *optional*, defaults to 0.8):
-            Maximum percentage of free GPU memory (after the model is loaded) to use for the KV cache.
+        max_memory_percent (`float`, *optional*):
+            Maximum percentage of free GPU memory (after the model is loaded) to use for the KV cache. When `None`,
+            resolved at runtime to 0.9 if there is no logit processing and 0.8 if there is, to leave headroom for
+            vocabulary-sized temporary tensors.
         max_blocks_per_request (`int`, *optional*, defaults to 0):
             Maximum blocks per request, used in the `flash_attn_with_kvcache` fast decode path to dimension
             the block table. Setting this to 0 disables the fast decode path.
@@ -1607,8 +1609,9 @@ class ContinuousBatchingConfig:
     num_blocks: int | None = None
     max_batch_tokens: int | None = None
 
-    # The max percentage of free GPU memory (after the model is loaded) to use for the KV cache.
-    max_memory_percent: float = 0.8
+    # The max percentage of free GPU memory (after the model is loaded) to use for the KV cache. If None, auto resolved
+    # to 0.9 (no logit processing) or 0.8 (logit processing) to leave headroom for temporary tensors.
+    max_memory_percent: float | None = None
 
     # This is only used in the flash_attn_with_kvcache fast decode path to dimension the block table. If it is set to 0,
     # the fast decode path will not be used. Currently turned off by default.
@@ -1773,6 +1776,13 @@ def decide_use_async_batching(self, is_attn_mask_needed: bool) -> bool:
             )
         return self.use_async_batching
 
+    def resolve_max_memory_percent(self, has_logit_processors: bool) -> None:
+        """Resolves `max_memory_percent` when unset: 0.9 without logit processors, 0.8 with them. Active processors
+        materialize `[N, V]` intermediates (e.g. top-p sort, softmax) that get captured into the CUDA graph pool, so
+        the cache has to cede some budget to that pool."""
+        if self.max_memory_percent is None:
+            self.max_memory_percent = 0.8 if has_logit_processors else 0.9
+
     def resolve_sentinel_values(self) -> None:
         """For some parameters (padding intervals and max cached graphs), the default is a sentinel value of 0: that
         way, if the user specifies a value for those parameters, we know they want it used, ie. we turn on cuda graphs.

src/transformers/generation/continuous_batching/cb_logits_processors.py

@@ -319,6 +319,8 @@ def __call__(self, scores: torch.FloatTensor, tensor_arg: torch.Tensor) -> torch
         return scores.masked_fill(indices_to_remove, self.filter_value)
 
 
+# TODO: add non-per-request CB variants so the memory-efficient warpers work when `per_request_processors=False`.
+# TODO: fuse temperature + top-k + top-p into a single pass to reuse the softmax/sort and cut activation peak.
 CLASSIC_TO_CB_PROCESSORS_MAP = {
     TemperatureLogitsWarper: ContinuousBatchingTemperatureLogitsWarper,
     TopKLogitsWarper: ContinuousBatchingTopKLogitsWarper,

src/transformers/generation/continuous_batching/continuous_api.py

@@ -623,26 +623,18 @@ def _sample(self, scores: torch.Tensor, logits_indices: torch.Tensor, output_ids
             output_ids[1, :tokens].copy_(logprobs.view(dtype=torch.int32))
 
     @torch.inference_mode()
-    def warmup(
-        self,
-        model: nn.Module,
-        logit_processor: LogitsProcessorList,
-        num_query_tokens: int = 0,
-        num_cache_tokens: int = 0,
-    ) -> None:
+    def warmup(self, model: nn.Module) -> None:
         """Pre-capture CUDA graphs (or trigger compile warmup) for varlen and decode paths. In async mode, both IO
-        pairs are warmed up since each has its own graph buffer and static tensors."""
+        pairs are warmed up since each has its own graph buffer and static tensors. The varlen path is warmed up at
+        the largest possible `(q, kv)` sizes so subsequent captures fit inside it without growing the pool."""
 
         if not self._pad_inputs:
             logger.info("CUDA graphs and compile are disabled, skipping warmup.")
             return None
 
-        num_query_tokens = num_query_tokens if num_query_tokens > 0 else self.max_batch_tokens
-        num_query_tokens = min(num_query_tokens, self.max_batch_tokens)
-        num_cache_tokens = num_cache_tokens if num_cache_tokens > 0 else self.cache.block_size * num_query_tokens
-        num_cache_tokens = min(num_cache_tokens, self.cache.num_blocks * self.cache.block_size)
-
+        num_query_tokens = self.max_batch_tokens
         num_pages = self.cache.num_blocks * self.cache.block_size
+        num_cache_tokens = num_pages - num_query_tokens
         compute_stream = self.inputs_and_outputs.compute_stream
 
         # In async mode, each IO pair has its own graph buffer and static tensors, so we warm up both
@@ -677,7 +669,7 @@ def warmup(
                         forward_fn(*forward_fn_args)
                 logger.info(f"Varlen warmup completed in {perf_counter() - start:.2f}s")
             except Exception as e:
-                logger.warning(f"Failed to warm up varlen path: {e}")
+                logger.warning(f"Failed to warm up varlen path: {e}. Graph pool may fragment and OOM under load.")
             finally:
                 for fs in future_states:
                     self.cache.free_blocks(fs.state.request_id)
@@ -811,12 +803,12 @@ def is_running(self) -> bool:
         """Check if the background generation thread is running."""
         return self._generation_thread is not None and self._generation_thread.is_alive()
 
-    def warmup(self, num_query_tokens: int = 0, num_cache_tokens: int = 0) -> None:
+    def warmup(self) -> None:
         """Pre-capture CUDA graphs for varlen and decode paths by running dummy batches. Initializes the batch
         processor if not already done."""
         if self.batch_processor is None:
             self.batch_processor = self._create_batch_processor()
-        self.batch_processor.warmup(self.model, self.logit_processor, num_query_tokens, num_cache_tokens)
+        self.batch_processor.warmup(self.model)
         self.warmed_up = True
 
     # NOTE: don't forget to update `continuous_batching_context_manager` when changing this method's definition
@@ -1040,6 +1032,8 @@ def _generation_step(self) -> None:
         self.batch_processor._generation_step(self.model)
 
     def _create_batch_processor(self) -> ContinuousBatchProcessor:
+        # Resolve max_memory_percent now that we know whether any logit processors are active.
+        self.continuous_batching_config.resolve_max_memory_percent(self.logit_processor.do_processing)
         # Create the PagedAttentionCache
         paged_attention_cache = PagedAttentionCache(
             self.model.config,
@@ -1225,25 +1219,25 @@ def continuous_batching_context_manager(
         timeout: float | None = None,
         continuous_batching_config: ContinuousBatchingConfig | None = None,
         persistent_manager: bool = False,
-        warmup_requests: int | None = 0,
+        warmup: bool = True,
         **deprecated_kwargs,
     ) -> Generator[ContinuousBatchingManager]:
         """A context manager to safely use the continuous batching manager. Arguments are similar to the ones of
         `init_continuous_batching`, except for:
             - block: whether to block the thread when stopping the manager. Default is True.
             - timeout: maximum time to wait for the thread to stop. Default is None (no timeout).
-            - warmup_query_tokens: the number of expected requests for which to warmup. 0 is auto, None is no warmup.
+            - warmup: whether to pre-capture CUDA graphs at the largest sizes before running. Default is True.
         """
         manager = self.init_continuous_batching(
             generation_config=generation_config,
             continuous_batching_config=continuous_batching_config,
             **deprecated_kwargs,
         )
-        if not (warmup_requests is None or manager.warmed_up):
+        if warmup and not manager.warmed_up:
             # Warmup is long (~30 sec): best to signal the user it's happening than let them think the manager is stuck
-            logger.warning("Warming up for coninuous batching...")
+            logger.warning("Warming up for continuous batching...")
             start = perf_counter()
-            manager.warmup(num_query_tokens=warmup_requests, num_cache_tokens=0)
+            manager.warmup()
             logger.warning(f"Warming up completed in {perf_counter() - start:.2f}s.")
         manager.start()
         try:
@@ -1320,7 +1314,7 @@ def generate_batch(
             block=True,
             timeout=5,
             persistent_manager=persistent_manager,
-            warmup_requests=len(inputs) if warmup else None,
+            warmup=warmup,
             **deprecated_kwargs,
         )
         logging_cm = logging_redirect_tqdm([logger])

src/transformers/generation/continuous_batching/input_outputs.py

@@ -14,7 +14,6 @@
 from contextlib import nullcontext
 from dataclasses import dataclass
 from functools import partial
-from itertools import count
 from typing import Any
 
 import torch
@@ -250,10 +249,11 @@ def _transfer_inputs(
             # Only transfer block_table for decode-only batches (when it's actually used)
             if self.use_block_table:
                 other.block_table.copy_(self.block_table, non_blocking=non_blocking)
-            # Otherwise, we transfer the read and write indices
+            # Otherwise, we transfer the write indices (and read indices if the batch uses any cache reads)
             else:
                 other.write_index_storage.copy_(self.write_index_storage, non_blocking=non_blocking)
-                other.read_index_storage.copy_(self.read_index_storage, non_blocking=non_blocking)
+                if self.max_kv_read > 0:
+                    other.read_index_storage.copy_(self.read_index_storage, non_blocking=non_blocking)
             # Transfer the attention masks if needed
             if self.attention_mask is not None and other.attention_mask is not None:
                 for layer_type in self.attention_mask.keys():
@@ -373,14 +373,15 @@ def prepare_batch_tensors(
         self.requests_in_batch = []
         self.req_id_to_new_token_position = {}
 
-        # Prepare accumulators
+        # Prepare accumulators. For batches with no past cache to read, we leave read_index empty: the cache.update
+        # will detect the 0-size indices and skip the read.
         input_ids = []
         position_ids = []
         cumulative_seqlens_q = [0]
         logits_indices = []
         cumulative_seqlens_k = {layer_type: [0] for layer_type in self.cumulative_seqlens_k.keys()}
-        read_index = [[] for _ in range(self.cache.num_groups)]
         write_index = [[] for _ in range(self.cache.num_groups)]
+        read_index = None if self.max_kv_read == 0 else [[] for _ in range(self.cache.num_groups)]
 
         # Go through all the requests in the batch
         for i, future_state in enumerate(requests_in_batch):
@@ -448,14 +449,16 @@ def prepare_batch_tensors(
                     sliding_window=self.sliding_window if layer_type == "sliding_attention" else 1,
                 )
 
-        # If we are not using the block table, we populate the read and write indices
+        # If we are not using the block table, we populate the write indices (and maybe the read indices)
         if not self.use_block_table:
             to_index_tensor = partial(torch.tensor, dtype=torch.int64, device=self.device)
-            for i, group_read_indices, group_write_indices in zip(count(), read_index, write_index):
-                self.read_index_storage[i, : len(group_read_indices)] = to_index_tensor(group_read_indices)
+            for i, group_write_indices in enumerate(write_index):
                 self.write_index_storage[i, : len(group_write_indices)] = to_index_tensor(group_write_indices)
-                self.true_read_sizes[i] = len(group_read_indices)
                 self.true_write_sizes[i] = len(group_write_indices)
+            if read_index is not None:
+                for i, group_read_indices in enumerate(read_index):
+                    self.read_index_storage[i, : len(group_read_indices)] = to_index_tensor(group_read_indices)
+                    self.true_read_sizes[i] = len(group_read_indices)
 
     def get_model_kwargs(self, use_padding: bool = False) -> dict[str, Any]:
         """Get model keyword arguments for the current batch, eventually padding the query dimension and KV dimensions
@@ -500,10 +503,14 @@ def get_model_kwargs(self, use_padding: bool = False) -> dict[str, Any]:
 
         # For the attributes that are lists of tensors, we construct list of tensor references
         for i in range(self.cache.num_groups):
-            read_index_size = kv_size if use_padding else self.true_read_sizes[i]
             write_index_size = q_size if use_padding else self.true_write_sizes[i]
-            kwargs.read_index.append(self.read_index_storage[i, :read_index_size])
             kwargs.write_index.append(self.write_index_storage[i, :write_index_size])
+            # If there is no cache to read, pass a list of empty tensors so `cache.update` uses the write-only fast path
+            if self.max_kv_read == 0:
+                read_index_size = 0
+            else:
+                read_index_size = kv_size if use_padding else self.true_read_sizes[i]
+            kwargs.read_index.append(self.read_index_storage[i, :read_index_size])
 
         # For the attributes that are dict of tensors, we first fill the dict with the actual values
         for layer_type, seqlens_k in self.cumulative_seqlens_k.items():
@@ -531,11 +538,11 @@ def get_cb_kwargs(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return self.carry_over_ids, self.output_ids, self.output_ids
 
     def _get_graph_key(self) -> tuple[int, ...]:
-        # Keys for varlen path
-        if self.max_kv_read > 0:
-            return (self.num_q_tokens, self.max_kv_read, *self.max_seqlen_k.values())
         # Keys for decode fast path
-        return (self.num_q_tokens,)
+        if self.use_block_table:
+            return (self.num_q_tokens,)
+        # Keys for varlen path
+        return (self.num_q_tokens, self.max_kv_read, *self.max_seqlen_k.values())
 
     def get_graph(self) -> torch.cuda.CUDAGraph | None:
         key = self._get_graph_key()

src/transformers/generation/continuous_batching/requests.py

@@ -27,6 +27,7 @@
     import psutil
 
 # This is a temporary token ID used to represent a token that is not yet generated
+# TODO: update this to 0 and check it breaks nothing + simplify carry over and time new logic
 TMP_TOKEN_ID = -1
 
 
@@ -45,9 +46,11 @@ def get_device_and_memory_breakdown() -> tuple[torch.device, int, int, int]:
         device = torch.device("cuda")
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
-        total_memory = torch.cuda.get_device_properties(device).total_memory
+        # Use mem_get_info to get actual free memory: device_properties().total_memory returns the physical device
+        # total which ignores CUDA context and driver overhead (~0.5 GiB), leading to overcommit.
+        free_memory, total_memory = torch.cuda.mem_get_info(device)
         reserved_memory = torch.cuda.memory_reserved(device)
-        allocated_memory = torch.cuda.memory_allocated(device)
+        allocated_memory = total_memory - free_memory
     elif is_torch_xpu_available():
         device = torch.device("xpu")
         torch.xpu.empty_cache()

src/transformers/generation/continuous_batching/scheduler.py

@@ -205,7 +205,7 @@ def _process_candidates(
         """
         scheduled_requests = []
         one_allocation_failed = False
-        decode_fast_path = True
+        decode_fast_path = self.cache.max_blocks_per_request > 0  # best way to check if decode fast path availability
         safety_margins = safety_margin * self.cache.num_blocks
         original_token_budget, original_cache_budget = token_budget, cache_budget
 
@@ -219,17 +219,22 @@ def _process_candidates(
                 )
                 break
 
-            # Check cache budget
+            # 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)
+
+            # This block checks cache budget: decode batches have infinite budget, but varlen batches don't, because KV
+            # cache is read through a fixed-sized index tensor. We keep track of the current budget in case the batch
+            # goes from decode to varlen
+            is_decode_eligible = request_len == 1 and state.position_offset < self.max_decode_fast_path_length
             read_cache_needed = state.current_len()
             if self.read_cache_limit is not None:
                 read_cache_needed = min(read_cache_needed, self.read_cache_limit)
-            if cache_budget < read_cache_needed:
+            # A request that would change the batch from decode to varlen is rejected if the cache budget is too low
+            if not (decode_fast_path and is_decode_eligible) and cache_budget < read_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)
 
@@ -273,7 +278,7 @@ def _process_candidates(
                 request_ids_to_remove_from_waiting.add(req_id)
 
             # Early exit of the loop if we have no budget left
-            if token_budget == 0 or cache_budget == 0:
+            if token_budget == 0 or (cache_budget <= 0 and not decode_fast_path):
                 break
 
         num_q_tokens = original_token_budget - token_budget

tests/generation/test_continuous_batching.py

@@ -1274,16 +1274,16 @@ def test_memory_prediction(
             max_blocks_per_request=max_bpr,
             return_logprobs=logprobs,
             use_async_batching=use_async_batching,
+            block_size=block_size,
         )
 
         handler = PagedAttentionMemoryHandler(
-            block_size=block_size,
+            continuous_batching_config=cb_config,
             page_size=page_size,
             num_groups=num_groups,
             group_size=group_size,
-            peak_activation_per_token=peak_act,
+            activation_peaks=[(0, peak_act)],
             num_attention_masks=num_attn_masks,
-            continuous_batching_config=cb_config,
         )
 
         N = self.NUM_BLOCKS * block_size  # num_pages