diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index 308c42564295..f601a97959c6 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/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. diff --git a/src/transformers/generation/continuous_batching/cache.py b/src/transformers/generation/continuous_batching/cache.py index 9fd0d3afba11..59de60bc957c 100644 --- a/src/transformers/generation/continuous_batching/cache.py +++ b/src/transformers/generation/continuous_batching/cache.py @@ -182,15 +182,30 @@ def __init__( else: num_attention_masks = 1 + # Peak activations coefficients (for number of blocks and number of batch tokens) + q_bytes_per_token = config.num_attention_heads * self.head_dim + lm_head_peak = ( + 0, # number of blocks does not affect the LM head peak activation + config.hidden_size + 2 * config.vocab_size, # hidden states + logits + ) + attention_peak = ( + 2 * page_size, # old K and V, read from cache (in the worst case scenario: whole cache is read) + config.hidden_size + q_bytes_per_token + 2 * page_size, # hidden state + Q + new K and V + ) + memory_handler = PagedAttentionMemoryHandler( - block_size=self.block_size, + continuous_batching_config=continuous_batching_config, page_size=page_size, num_groups=self.num_groups, group_size=group_size, - peak_activation_per_token=(config.hidden_size + config.vocab_size), + activation_peaks=[lm_head_peak, attention_peak], num_attention_masks=num_attention_masks, - continuous_batching_config=continuous_batching_config, ) + + # If somehow the max memory percent is not yet resolved, resolve it conservatively + if continuous_batching_config.max_memory_percent is None: + continuous_batching_config.resolve_max_memory_percent(has_logit_processors=True) + num_blocks, max_batch_tokens = memory_handler.infer_num_blocks_and_max_batch_tokens( num_blocks=continuous_batching_config.num_blocks, max_batch_tokens=continuous_batching_config.max_batch_tokens, @@ -316,17 +331,20 @@ def extend_read_and_write_indices( request_id: str, past_length: int, query_length: int, - read_index: list[list[int]], + read_index: list[list[int]] | None, write_index: list[list[int]], ) -> None: """Retrieve physical cache indices for reading KV states in the cache across all layer groups. This method coordinates with all cache managers to build the complete set of read indices needed for attention computation. + When read_index is None, the batch has no cache reads and we only compute the write indices. """ - for cm, read_indices, write_indices in zip(self.group_cache_managers, read_index, write_index): - indices = cm.get_read_indices(request_id, past_length, query_length) - read_indices.extend(indices) - indices = cm.get_write_indices(request_id, past_length, query_length) - write_indices.extend(indices) + # Write indices are always computed + for cm, write_indices in zip(self.group_cache_managers, write_index): + write_indices.extend(cm.get_write_indices(request_id, past_length, query_length)) + # Read indices are only computed if there are cache indices + if read_index is not None: + for cm, read_indices in zip(self.group_cache_managers, read_index): + read_indices.extend(cm.get_read_indices(request_id, past_length, query_length)) def fill_block_table( self, request_id: str, past_length: int, query_length: int, block_table: torch.Tensor @@ -355,26 +373,34 @@ def update( read_index: list[torch.Tensor], # shape [num_layer_groups, seqlen_kv + past_length] write_index: list[torch.Tensor], # shape [num_layer_groups, seqlen_q] ) -> tuple[torch.Tensor, torch.Tensor]: # shape [seqlen_kv + past_length, num_kv_heads, head_dim] - """Update the cache with new key-value states for a specific layer. This method writes new KV states to the - appropriate cache locations. The behavior differs based on the layer's attention type: + """Update the cache with new key-value states for a specific layer, and retrieves the relevant KV states from + the cache for attention computation. The behavior differs based on the layer's attention type: - Full attention: New KV states are written to cache, then complete sequence is read from cache - Sliding window: Old KV is read from cache along with extra spaces for the new KV, then new KV is written to cache. This is because new KV might overwrite the old KV, so we need to read the old KV first. + When the layer's read index is empty, the batch has no cache reads (all requests are non-chunked prefills): we + only write to the cache and return the input KV states directly, skipping the index_select read-back. + Returns the complete KV states (cached + new) for attention computation. """ - # Retrieve the layer read and write indices + # Retrieve the layer write index and the relevant cache tensors group_idx, layer_idx_in_group = self.layer_index_to_group_indices[layer_idx] layer_read_index = read_index[group_idx] layer_write_index = write_index[group_idx] - # Select the correct cache k_cache = self.key_cache[layer_idx_in_group] v_cache = self.value_cache[layer_idx_in_group] # Transpose the key and value states to match the cache shape, after which shape is [seqlen_kv, num_kv_heads, head_dim] key_states = key_states.transpose(1, 2).squeeze(0) value_states = value_states.transpose(1, 2).squeeze(0) + # Case: write-only, no cache read. The input KV states already contain everything the attention needs. + if layer_read_index.numel() == 0: + k_cache.index_copy_(0, layer_write_index, key_states) + v_cache.index_copy_(0, layer_write_index, value_states) + return key_states, value_states + # Case: full attention sliding_window = self.sliding_windows[layer_idx] if sliding_window == 1: @@ -509,25 +535,26 @@ class PagedAttentionMemoryHandler: _activation_dtype = torch.bfloat16 _input_dtype = torch.int32 - _upper_bound_max_batch_tokens = 256 + _upper_bound_max_batch_tokens = 1024 _upper_bound_num_blocks = 4096 def __init__( self, - block_size: int, + continuous_batching_config: ContinuousBatchingConfig, page_size: int, num_groups: int, group_size: int, - peak_activation_per_token: int, + activation_peaks: list[tuple[int, int]], num_attention_masks: int, - continuous_batching_config: ContinuousBatchingConfig, ) -> None: - """Initialize the memory handler.""" - self.block_size = block_size + """Initialize the memory handler. `activation_peaks` is a list of `(Δcn, Δcm)` pairs giving the activation memory + contributions proportional to N (pages) and M (batch tokens) for each peak. Memory must satisfy the constraint + at every peak, so we solve each polynomial independently and take the most restrictive result.""" + self.block_size = continuous_batching_config.block_size self.page_size = page_size self.num_groups = num_groups self.group_size = group_size - self.peak_activation_per_token = peak_activation_per_token + self.activation_peaks = activation_peaks self.num_attention_masks = num_attention_masks self.max_blocks_per_request = continuous_batching_config.max_blocks_per_request or 0 # This is the number of output rows for the output_ids tensor @@ -545,23 +572,29 @@ def get_available_memory(max_memory_percent: float = 1.0) -> int: # Formatting is disabled because of comment indentation, which improves readability. # fmt: off - def _equation_coefficients(self, cache_dtype: torch.dtype) -> tuple[int, int, int, int]: - """Returns (coeff_n, coeff_m, coeff_nm, coeff_mm) for the memory polynomial. Each addend is annotated with - the tensor it corresponds to in `ContinuousBatchingIOs._setup_static_tensors`. + def _equation_coefficients( + self, peak: tuple[int, int], cache_dtype: torch.dtype + ) -> tuple[int, int, int, int]: + """Returns `(coeff_n, coeff_m, coeff_nm, coeff_mm)` for the memory polynomial of a single activation peak. + `peak = (Δcn, Δcm)` is the peak-specific activation contribution; the rest of the coefficients are shared + across peaks. Each addend is annotated with the tensor it corresponds to in + `ContinuousBatchingIOs._setup_static_tensors` (or the forward pass, for activation terms). """ i = self._input_dtype.itemsize # int32 a = self._activation_dtype.itemsize # bfloat16 c = cache_dtype.itemsize k = self.io_multiplier # 1 sync, 2 async (IO tensors only) + delta_n, delta_m = peak # -- N terms: cost per cache page -------------------------------------------------- coeff_n = ( 2 * self.group_size * self.page_size * c # kv_cache: 2 * group_size * [N, page_size] * cache_dtype + k * self.num_groups * 8 # read_index: [num_groups, N + M] (N part only, int64) + + delta_n * a # activation peak: N-proportional part ) # -- M terms: cost per batch token ------------------------------------------------- coeff_m = ( - self.peak_activation_per_token * a # activation peak (largest hidden state per token) + delta_m * a # activation peak: M-proportional part + k * 7 * i # bulk_input: [7, M] int32, packed as 7 rows + k * self.num_output_rows * i # output_ids: [num_output_rows, M] int32 + k * self.num_groups # block_table: [bt_groups, M, max_blocks_per_req] int32 @@ -569,9 +602,9 @@ def _equation_coefficients(self, cache_dtype: torch.dtype) -> tuple[int, int, in + k * self.num_groups * 8 # write_index: [num_groups, M] int64 + k * self.num_groups * 8 # read_index: [num_groups, N + M] (M part only, int64) ) - # -- N·M terms: cost per (page × batch token) ------------------------------------- + # -- N·M terms: cost per (page × batch token) -------------------------------------- coeff_nm = k * self.num_attention_masks * a # attention_mask: [1, 1, M, N + M] (N·M part only) - # -- M² terms: cost per (batch token squared) ------------------------------------- + # -- M² terms: cost per (batch token squared) -------------------------------------- coeff_mm = k * self.num_attention_masks * a # attention_mask: [1, 1, M, N + M] (M² part only) return coeff_n, coeff_m, coeff_nm, coeff_mm @@ -590,55 +623,80 @@ def _solve_quadratic(a: float, b: float, c: float) -> float: raise ValueError(f"No positive solution (root = {root})") return root - def infer_num_blocks_and_max_batch_tokens( + def _solve_for_peak( self, - num_blocks: int | None = None, - max_batch_tokens: int | None = None, - max_memory_percent: float = 0.8, # FIXME: it seems we overcommit memory, was changed from 0.9 which caused OOMs in our benchmarking CI - cache_dtype: torch.dtype = torch.float16, + peak: tuple[int, int], + available: int, + num_blocks: int | None, + max_batch_tokens: int | None, + cache_dtype: torch.dtype, ) -> tuple[int, int]: - """Solve for the missing variable(s) in the memory polynomial (see ``_equation_coefficients``). When both - are unknown, assumes M = m·N (m = 0.01, i.e. one batch fills ~1 % of the cache) and solves the resulting - quadratic in N. - """ - available = self.get_available_memory(max_memory_percent) - coeff_n, coeff_m, coeff_nm, coeff_mm = self._equation_coefficients(cache_dtype) - logger.info(f"Cache memory: {available}") + """Solve for `(num_blocks, max_batch_tokens)` against one activation peak's memory polynomial. Clamps to upper + bounds. Either input may be None; whichever is None is solved for.""" + cn, cm, cnm, cmm = self._equation_coefficients(peak, cache_dtype) if num_blocks is None and max_batch_tokens is None: # Substitute M = m·N → (coeff_nm·m + coeff_mm·m²)·N² + (coeff_n + coeff_m·m)·N − avail = 0 m = 0.01 - num_pages = self._solve_quadratic( - coeff_nm * m + coeff_mm * m**2, - coeff_n + coeff_m * m, - -available, - ) - num_blocks = min(floor(num_pages) // self.block_size, self._upper_bound_num_blocks) - max_batch_tokens = min(int(num_pages * m), self._upper_bound_max_batch_tokens) - - elif num_blocks is None: + num_pages = self._solve_quadratic(cnm * m + cmm * m**2, cn + cm * m, -available) + max_batch_tokens = int(num_pages * m) + if max_batch_tokens > self._upper_bound_max_batch_tokens: + max_batch_tokens = self._upper_bound_max_batch_tokens + # If max_batch_tokens is clamped, we recompute num_blocks below to get a higher value + num_blocks = None + else: + num_blocks = min(floor(num_pages) // self.block_size, self._upper_bound_num_blocks) + + if num_blocks is None: # M given → linear in N: (coeff_n + coeff_nm·M)·N = avail − coeff_m·M − coeff_mm·M² M = max_batch_tokens - num_pages = floor((available - coeff_m * M - coeff_mm * M**2) / (coeff_n + coeff_nm * M)) + num_pages = floor((available - cm * M - cmm * M**2) / (cn + cnm * M)) num_blocks = min(num_pages // self.block_size, self._upper_bound_num_blocks) - elif max_batch_tokens is None: # N given → quadratic in M: coeff_mm·M² + (coeff_m + coeff_nm·N)·M + (coeff_n·N − avail) = 0 N = num_blocks * self.block_size - M = self._solve_quadratic(coeff_mm, coeff_m + coeff_nm * N, coeff_n * N - available) + M = self._solve_quadratic(cmm, cm + cnm * N, cn * N - available) max_batch_tokens = min(floor(M), self._upper_bound_max_batch_tokens) + return num_blocks, max_batch_tokens + + def infer_num_blocks_and_max_batch_tokens( + self, + num_blocks: int | None = None, + max_batch_tokens: int | None = None, + max_memory_percent: float = 0.9, + cache_dtype: torch.dtype = torch.float16, + ) -> tuple[int, int]: + """Solve for the missing variable(s) in the memory polynomial (see ``_equation_coefficients``). There is one + polynomial per activation peak; we solve each independently and take the most restrictive (smallest) result. + When both `N` and `M` are unknown, assumes `M = m·N` (m = 0.01, i.e. one batch fills ~1 % of the cache) and + solves the resulting quadratic in N. + """ + available = self.get_available_memory(max_memory_percent) + logger.info(f"Cache memory: {available}") + # Solve each peak independently, then take the element-wise min (tightest constraint wins) + acc_num_blocks = float("inf") + acc_max_batch_tokens = float("inf") + for peak in self.activation_peaks: + n_blocks, m_batch_tokens = self._solve_for_peak(peak, available, num_blocks, max_batch_tokens, cache_dtype) + acc_num_blocks = min(acc_num_blocks, n_blocks) + acc_max_batch_tokens = min(acc_max_batch_tokens, m_batch_tokens) + # Now update the value (cannot update in loop, it would overwrite the user-passed values) + num_blocks, max_batch_tokens = acc_num_blocks, acc_max_batch_tokens # Validate - memory_footprint = self.compute_memory_footprint( - max_batch_tokens=max_batch_tokens, num_blocks=num_blocks, cache_dtype=cache_dtype - ) + memory_footprint = self.compute_memory_footprint(num_blocks, max_batch_tokens, cache_dtype) if memory_footprint > available: raise MemoryError(f"Memory footprint {memory_footprint} is more than available memory {available}") return num_blocks, max_batch_tokens def compute_memory_footprint(self, num_blocks: int, max_batch_tokens: int, cache_dtype: torch.dtype) -> int: - """Evaluate the memory polynomial at concrete (N, M) values.""" + """Evaluate the memory polynomial at concrete (N, M) values, taking the max across activation peaks.""" N = num_blocks * self.block_size M = max_batch_tokens - cn, cm, cnm, cmm = self._equation_coefficients(cache_dtype) - return cn * N + cm * M + cnm * N * M + cmm * M * M + + max_memory_footprint = 0 + for peak in self.activation_peaks: + cn, cm, cnm, cmm = self._equation_coefficients(peak, cache_dtype) + memory_footprint = cn * N + cm * M + cnm * N * M + cmm * M * M + max_memory_footprint = max(max_memory_footprint, memory_footprint) + return max_memory_footprint diff --git a/src/transformers/generation/continuous_batching/cb_logits_processors.py b/src/transformers/generation/continuous_batching/cb_logits_processors.py index 3a5f7eb8df26..619d9fefea5e 100644 --- a/src/transformers/generation/continuous_batching/cb_logits_processors.py +++ b/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, diff --git a/src/transformers/generation/continuous_batching/continuous_api.py b/src/transformers/generation/continuous_batching/continuous_api.py index 47290b9d70b6..0521c6402ca9 100644 --- a/src/transformers/generation/continuous_batching/continuous_api.py +++ b/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]) diff --git a/src/transformers/generation/continuous_batching/input_outputs.py b/src/transformers/generation/continuous_batching/input_outputs.py index 134941c2526f..fbe7890a15b9 100644 --- a/src/transformers/generation/continuous_batching/input_outputs.py +++ b/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() diff --git a/src/transformers/generation/continuous_batching/requests.py b/src/transformers/generation/continuous_batching/requests.py index 05bf65725c5a..381c94bc2dc9 100644 --- a/src/transformers/generation/continuous_batching/requests.py +++ b/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() diff --git a/src/transformers/generation/continuous_batching/scheduler.py b/src/transformers/generation/continuous_batching/scheduler.py index f35d2e968342..284c202267c5 100644 --- a/src/transformers/generation/continuous_batching/scheduler.py +++ b/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 diff --git a/tests/generation/test_continuous_batching.py b/tests/generation/test_continuous_batching.py index ff3e54be374f..cd7c95f7bf4e 100644 --- a/tests/generation/test_continuous_batching.py +++ b/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