diff --git a/benchmark_v2/benchmark_scripts/continuous_batching_overall.py b/benchmark_v2/benchmark_scripts/continuous_batching_overall.py index afe246dbe427..9b2bb875b758 100644 --- a/benchmark_v2/benchmark_scripts/continuous_batching_overall.py +++ b/benchmark_v2/benchmark_scripts/continuous_batching_overall.py @@ -2,6 +2,7 @@ import json import re import subprocess +from datetime import datetime from pathlib import Path from tabulate import tabulate @@ -10,7 +11,7 @@ SCRIPT_LOCATION = (Path(__file__).parent.parent.parent / "examples/pytorch/continuous_batching.py").as_posix() COMMON_ARGS = "--log-level WARNING --seed 0 --force-max-length".split() ERROR_OUTPUT = {"time_seconds": "X", "num_tokens": "X", "throughput_tok_per_sec": "ERROR"} -RESULTS_FILE = Path(__file__).parent.parent / "benchmark_results/cb_overall_main.json" +RESULTS_DIR = Path(__file__).parent.parent / "benchmark_results/cb_overall/" def run_and_parse_cb_example(args: str) -> dict: @@ -35,30 +36,38 @@ def run_and_parse_cb_example(args: str) -> dict: return {"args": args, **ERROR_OUTPUT} -def build_comparison_table(results: list[dict], main_results: list[dict]) -> list[dict]: - """Build a table comparing current results against saved main results.""" - main_by_args = {r["args"]: r for r in main_results} +def get_most_recent_file(prefix: str, exclude: Path | None = None) -> Path | None: + """Find the most recent results file in RESULTS_DIR matching the given prefix, optionally excluding one.""" + candidates = sorted(RESULTS_DIR.glob(f"{prefix}__*.json")) + if exclude: + candidates = [c for c in candidates if c != exclude] + return candidates[-1] if candidates else None + + +def build_comparison_table(results: list[dict], baseline_results: list[dict], baseline_label: str) -> list[dict]: + """Build a table comparing current results against baseline results.""" + baseline_by_args = {r["args"]: r for r in baseline_results} comparison = [ { "args": "Arguments", - "main_tok_per_sec": "Main (tok/s)", + "baseline_tok_per_sec": f"{baseline_label} (tok/s)", "current_tok_per_sec": "Current (tok/s)", "diff_percent": "Diff (%)", } ] for result in results: - main = main_by_args.get(result["args"]) - main_tp = main["throughput_tok_per_sec"] if main else None + baseline = baseline_by_args.get(result["args"]) + baseline_tp = baseline["throughput_tok_per_sec"] if baseline else None current_tp = result["throughput_tok_per_sec"] - if isinstance(main_tp, (int, float)) and isinstance(current_tp, (int, float)): - diff = (current_tp - main_tp) / main_tp * 100 + if isinstance(baseline_tp, (int, float)) and isinstance(current_tp, (int, float)): + diff = (current_tp - baseline_tp) / baseline_tp * 100 diff_str = f"{diff:+.1f}%" else: diff_str = "N/A" comparison.append( { "args": result["args"], - "main_tok_per_sec": main_tp if main_tp is not None else "N/A", + "baseline_tok_per_sec": baseline_tp if baseline_tp is not None else "N/A", "current_tok_per_sec": current_tp, "diff_percent": diff_str, } @@ -114,12 +123,27 @@ def build_comparison_table(results: list[dict], main_results: list[dict]) -> lis # The header row is results[0], data rows are results[1:] data_results = results[1:] + # Always save results to a new timestamped file + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + prefix = "main" if args.main else "run" + results_file = RESULTS_DIR / f"{prefix}__{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + results_file.write_text(json.dumps(data_results, indent=2)) + print(f"\nResults saved to {results_file}") + + # Compare against baseline if args.main: - RESULTS_FILE.write_text(json.dumps(data_results, indent=2)) - print(f"\nResults saved to {RESULTS_FILE}") + # Compare against the previous main baseline (the one that was most recent before this new file) + baseline_file = get_most_recent_file("main", exclude=results_file) + baseline_label = "Previous main" + else: + # Compare against the most recent main baseline + baseline_file = get_most_recent_file("main") + baseline_label = "Main" + + if baseline_file: + baseline_results = json.loads(baseline_file.read_text()) + comparison = build_comparison_table(data_results, baseline_results, baseline_label) + print(f"\nComparing against: {baseline_file.name}") + print(tabulate(comparison, tablefmt="github")) else: - if RESULTS_FILE.exists(): - main_results = json.loads(RESULTS_FILE.read_text()) - comparison = build_comparison_table(data_results, main_results) - print() - print(tabulate(comparison, tablefmt="github")) + print("\nNo baseline results found for comparison.") diff --git a/src/transformers/generation/continuous_batching/cache.py b/src/transformers/generation/continuous_batching/cache.py index 3600efc6ec12..9fd0d3afba11 100644 --- a/src/transformers/generation/continuous_batching/cache.py +++ b/src/transformers/generation/continuous_batching/cache.py @@ -189,6 +189,7 @@ def __init__( group_size=group_size, peak_activation_per_token=(config.hidden_size + config.vocab_size), num_attention_masks=num_attention_masks, + continuous_batching_config=continuous_batching_config, ) num_blocks, max_batch_tokens = memory_handler.infer_num_blocks_and_max_batch_tokens( num_blocks=continuous_batching_config.num_blocks, @@ -220,8 +221,12 @@ def __init__( # Initialize the cache self.key_cache: list[torch.Tensor] = [] self.value_cache: list[torch.Tensor] = [] - # We add two extra tokens to the cache to handle padding and generally discard unwanted tokens + # We add two extra blocks to the cache as a padding zone that no BlockManager ever allocates from: one for the + # sentinel index (marks the spot of a new token in the read indices) and one for the trash index (for padding, + # block is never used so writes are silently discarded) self.cache_shape = ((num_blocks + 2) * self.block_size, self.num_key_value_heads, self.head_dim) + self.sentinel_index = self.cache_shape[0] - 1 + self.trash_index = self.sentinel_index - 1 for _ in range(group_size): new_layer_key_cache = torch.empty(self.cache_shape, dtype=self.dtype, device=self.device) new_layer_value_cache = torch.empty(self.cache_shape, dtype=self.dtype, device=self.device) @@ -243,7 +248,9 @@ def __init__( cm = FullAttentionCacheAllocator(i, self.block_size, allow_block_sharing=self.allow_block_sharing) self.num_full_attention_groups += 1 elif group_type == "sliding_attention": - cm = SlidingAttentionCacheAllocator(i, self.block_size, config.sliding_window) + cm = SlidingAttentionCacheAllocator( + i, self.block_size, config.sliding_window, self.sentinel_index, self.trash_index + ) self.num_sliding_attention_groups += 1 self.max_sliding_window_blocks_per_request = cm._max_blocks_per_request else: @@ -357,7 +364,7 @@ def update( Returns the complete KV states (cached + new) for attention computation. """ - # Retrieve the layer read and write indices, and if there is a sliding window + # Retrieve the layer read and write indices 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] @@ -371,23 +378,24 @@ def update( # Case: full attention sliding_window = self.sliding_windows[layer_idx] if sliding_window == 1: - k_cache[layer_write_index, :, :] = key_states - v_cache[layer_write_index, :, :] = value_states - key_states_with_cache = k_cache[layer_read_index, :, :] - value_states_with_cache = v_cache[layer_read_index, :, :] + k_cache.index_copy_(0, layer_write_index, key_states) + v_cache.index_copy_(0, layer_write_index, value_states) + key_states_with_cache = torch.index_select(k_cache, 0, layer_read_index) + value_states_with_cache = torch.index_select(v_cache, 0, layer_read_index) # Case: sliding window -- we need to be careful of read/write order because of chunked prefill, because it's # the only case where you may write over cache you need to use else: - # Add the cache to the key and value states - mask = (layer_read_index == -1).unsqueeze(-1).unsqueeze(-1) # TODO: should this be precomputed? - key_states_with_cache = k_cache[layer_read_index, :, :] + # Sentinel positions in read_index mark new-token slots; index_select reads garbage there, + # then masked_scatter_ overwrites them with the actual new key/value states. + mask = (layer_read_index == self.sentinel_index).unsqueeze(-1).unsqueeze(-1) + key_states_with_cache = torch.index_select(k_cache, 0, layer_read_index) key_states_with_cache.masked_scatter_(mask, key_states) - value_states_with_cache = v_cache[layer_read_index, :, :] + value_states_with_cache = torch.index_select(v_cache, 0, layer_read_index) value_states_with_cache.masked_scatter_(mask, value_states) - # Write new KV values to the cache - k_cache[layer_write_index, :, :] = key_states - v_cache[layer_write_index, :, :] = value_states + # Write new KV values to the cache (padding slots in write_index point to the trash position) + k_cache.index_copy_(0, layer_write_index, key_states) + v_cache.index_copy_(0, layer_write_index, value_states) # Return the new KV values return key_states_with_cache, value_states_with_cache @@ -487,28 +495,16 @@ def free_all_requests(self) -> None: # TODO: rework computation with the groups and their sizes class PagedAttentionMemoryHandler: - """A helper class to determine the best number of pages and maximum number of tokens per batch for the paged - attention cache, providing automatic sizing based on available GPU memory. - The helper works using the number of pages, which is tied to the number of blocks by: - num_blocks = num_pages // block_size - - The memory footprint consists of three main components: - - Cache memory: the space needed to store the cache tensors: - 2 * layer_group_size * [num_pages, page_size] * cache_dtype - - Activation memory: the space temporarily taken by the largest activation during the model forward pass: - peak_activation_per_token * max_tokens_per_batch * activation_dtype_size - - Static tensors: the space taken by the input/output buffers and metadata tensors for batch processing, sum of: - - inputs_ids + outputs_ids + position_ids + logits_indices: 4 * max_tokens_per_batch * int32_size - - attention_mask: num_attention_masks * num_pages * max_tokens_per_batch * activation_dtype_size - - cumulative_seqlens_q + cumulative_seqlens_k: (1 + 2) * max_tokens_per_batch * int32_size - - write_index_tensor: num_groups * max_tokens_per_batch * int32_size - - read_index_tensor: num_groups * (num_pages + max_tokens_per_batch) * int32_size - - The handler can operate in three modes: - 1. Auto-sizing: Determines both number of pages and maximum number of tokens per batch using quadratic optimization - 2. Fixed cache: Calculates max batch tokens given a fixed number of pages - 3. Fixed batch: Calculates number of pages given a fixed maximum batch size + """Determines the optimal number of pages (N) and max batch tokens (M) for the paged attention cache, given + available GPU memory. The relation between N and number of blocks is: num_blocks = N // block_size. + + The memory footprint is a polynomial in N and M, where each term maps to a tensor allocated in + ``ContinuousBatchingIOs._setup_static_tensors`` or ``PagedAttentionCache.__init__``: + memory(N, M) = coeff_n · N + coeff_m · M + coeff_nm · N·M + coeff_mm · M² + + See ``_equation_coefficients`` for the breakdown. All three solving modes (auto, fixed-N, fixed-M) reduce to + solving this equation, which is at most quadratic in one variable. """ _activation_dtype = torch.bfloat16 @@ -524,41 +520,76 @@ def __init__( group_size: int, peak_activation_per_token: int, num_attention_masks: int, + continuous_batching_config: ContinuousBatchingConfig, ) -> None: - """Initialize the memory handler with the parameters that cannot be automatically inferred. - - Args: - block_size: Size of the cache blocks - page_size: Size of the cache pages - num_groups: Number of layer groups - group_size: Number of layers per layer group - peak_activation_per_token: Maximum size of activation tensor per token, = hidden_size + vocab_size - num_attention_masks: Number of attention masks, 0 if no attention mask is used, 2 if hybrid model, else 1 - """ + """Initialize the memory handler.""" self.block_size = 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.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 + self.num_output_rows = 2 if continuous_batching_config.return_logprobs else 1 + # This account for the set of 2 IOs if async batching is used + self.io_multiplier = 2 if continuous_batching_config.use_async_batching else 1 @staticmethod def get_available_memory(max_memory_percent: float = 1.0) -> int: - """Calculate available GPU memory for cache allocation, accounting for already allocated tensors. - This method queries the current memory state and applies the specified percentage limit to determine - how much memory can be safely used for the paged attention cache. - - Args: - max_memory_percent: Fraction of available memory to use (0.0-1.0). 1.0 means use all available memory. - - Returns: - int: Available memory in bytes for cache allocation - """ + """Calculate available GPU memory for cache allocation, accounting for already allocated tensors.""" _, total, reserved, allocated = get_device_and_memory_breakdown() available_memory = total - max(allocated, reserved) available_memory = int(available_memory * max_memory_percent) return available_memory + # 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`. + """ + 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) + + # -- 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) + ) + # -- M terms: cost per batch token ------------------------------------------------- + coeff_m = ( + self.peak_activation_per_token * a # activation peak (largest hidden state per token) + + 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 + * self.max_blocks_per_request * i # (zero when fast-decode is off) + + 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) ------------------------------------- + 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) ------------------------------------- + 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 + # fmt: on + + @staticmethod + def _solve_quadratic(a: float, b: float, c: float) -> float: + """Largest positive root of a·x² + b·x + c = 0. Falls back to linear when a == 0.""" + if a == 0: + return -c / b + discriminant = b**2 - 4 * a * c + if discriminant < 0: + raise ValueError(f"No real solution (discriminant = {discriminant})") + root = (-b + sqrt(discriminant)) / (2 * a) + if root < 0: + raise ValueError(f"No positive solution (root = {root})") + return root + def infer_num_blocks_and_max_batch_tokens( self, num_blocks: int | None = None, @@ -566,196 +597,48 @@ def infer_num_blocks_and_max_batch_tokens( 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, ) -> tuple[int, int]: - """Determine optimal number of blocks and maximum number of tokens per batch based on available memory and - constraints. Check the class docstring for more details. Naming the number of pages as N and the maximum number - of tokens per batch as M, the equation solved is: - - available_memory = sum([ - MN * num_attention_masks * activation_dtype_size, - 2N * (layer_group_size * page_size * cache_dtype + 2 * num_group), - M * (peak_activation_per_token * activation_dtype + 28 + 4 * num_group), - ]) - - where we already simplified int32_size = 4. + """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. """ - if num_blocks is None: - if max_batch_tokens is None: - # If neither num_blocks nor max_batch_tokens are provided, we use a second-order polynomial - num_blocks, max_batch_tokens = self.compute_num_blocks_and_max_batch_tokens( - max_memory_percent, cache_dtype - ) - else: - # If only max_batch_tokens is provided, we infer the num_blocks - num_blocks = self.compute_num_blocks(max_batch_tokens, max_memory_percent, cache_dtype) + 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}") + + 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: + # 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_blocks = min(num_pages // self.block_size, self._upper_bound_num_blocks) + elif max_batch_tokens is None: - # If only num_blocks is provided, we infer the max_batch_tokens - max_batch_tokens = self.compute_max_batch_tokens(num_blocks, max_memory_percent, cache_dtype) - else: - # If both num_blocks and max_batch_tokens are provided, we use them (useless, but helps with typing) - max_batch_tokens = max_batch_tokens + # 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) + max_batch_tokens = min(floor(M), self._upper_bound_max_batch_tokens) - # We check if the memory footprint is too large in all cases - available_memory = self.get_available_memory(max_memory_percent) + # Validate memory_footprint = self.compute_memory_footprint( max_batch_tokens=max_batch_tokens, num_blocks=num_blocks, cache_dtype=cache_dtype ) - if memory_footprint > available_memory: - raise MemoryError(f"Memory footprint {memory_footprint} is more than available memory {available_memory}") + 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_num_blocks_and_max_batch_tokens( - self, - max_memory_percent: float, - cache_dtype: torch.dtype = torch.float16, - m: float = 0.01, - ) -> tuple[int, int]: - """Calculate optimal number of blocks and maximum number of tokens per batch using quadratic optimization when - neither is fixed. This method assumes a relationship M = m * N where m is a small ratio below 1 and solves the - resulting quadratic equation to find the optimal N that maximizes utilization within memory constraints. m is - the amount of cache we can fill with one batch: m=0.01 means a batch fills at most 1% of the cache. The equation - to solve is: - - available_memory = sum([ - m * N^2 * num_attention_masks * activation_dtype_size, - 2N * (layer_group_size * page_size * cache_dtype + 2 * num_group), - m * N * (peak_activation_per_token * activation_dtype + 28 + 4 * num_group), - ]) - - If num_attention_masks is 0, the equation simplifies to a 1st degree polynomial. - """ - cache_memory = self.get_available_memory(max_memory_percent) - logger.info(f"Cache memory: {cache_memory}") - - # Compute second-degree polynomial coefficients - a = m * self.num_attention_masks * self._activation_dtype.itemsize - b = 2 * (self.group_size * self.page_size * cache_dtype.itemsize + 2 * self.num_groups) - b += m * (self.peak_activation_per_token * self._activation_dtype.itemsize + 28 + 4 * self.num_groups) - c = -cache_memory - logger.debug(f"Coefficients of 2nd degree polynomial: {a = }, {b = }, {c = }") - - # If num_attention_masks is 0, the equation simplifies to a 1st degree polynomial - if self.num_attention_masks == 0: - greatest_solution = -c / b - # Otherwise, we solve the quadratic equation - else: - discriminant = b**2 - 4 * a * c - if discriminant < 0: - raise ValueError(f"Discriminant is negative: {discriminant = }") - greatest_solution = (-b + sqrt(discriminant)) / (2 * a) - - if greatest_solution < 0: - raise ValueError(f"Greatest solution is negative: {greatest_solution = }") - - # Infer number of blocks and max batch tokens - num_pages = floor(greatest_solution) - num_blocks = num_pages // self.block_size - if num_blocks > self._upper_bound_num_blocks: - logger.info(f"{num_blocks = } is too large, setting to {self._upper_bound_num_blocks = }") - num_blocks = self._upper_bound_num_blocks - max_batch_tokens = int(greatest_solution * m) - if max_batch_tokens > self._upper_bound_max_batch_tokens: - logger.info(f"{max_batch_tokens = } is too large, setting to {self._upper_bound_max_batch_tokens = }") - max_batch_tokens = self._upper_bound_max_batch_tokens - return num_blocks, max_batch_tokens - - def compute_max_batch_tokens( - self, - num_blocks: int, - max_memory_percent: float, - cache_dtype: torch.dtype = torch.float16, - ) -> int: - """Calculate maximum batch tokens M given a fixed number of cache blocks. The formula for M is given by: - - M = (available_memory - 2N * (layer_group_size * page_size * cache_dtype + 2 * num_group)) - / (activation_dtype_size * (N * num_attention_masks + peak_activation_per_token) + 28 + 4 * num_group) - """ - cache_memory = self.get_available_memory(max_memory_percent) - num_pages = num_blocks * self.block_size - # Compute numerator - num = cache_memory - num -= 2 * num_pages * (self.group_size * self.page_size * cache_dtype.itemsize + 2 * self.num_groups) - # Compute denominator - denum = self._activation_dtype.itemsize * ( - num_pages * self.num_attention_masks + self.peak_activation_per_token - ) - denum += 28 + 4 * self.num_groups - # Compute max batch tokens and return - max_batch_tokens = floor(num / denum) - if max_batch_tokens > self._upper_bound_max_batch_tokens: - logger.info(f"{max_batch_tokens = } is too large, setting to {self._upper_bound_max_batch_tokens = }") - max_batch_tokens = self._upper_bound_max_batch_tokens - return max_batch_tokens - - def compute_num_blocks( - self, - max_batch_tokens: int, - max_memory_percent: float, - cache_dtype: torch.dtype = torch.float16, - ) -> int: - """Calculate number of cache blocks N given a fixed maximum token per token M. The formula for N is given by: - - N = (available_memory - M * (peak_activation_per_token * activation_dtype + 28 + 4 * num_group)) - / (2 * (layer_group_size * page_size * cache_dtype + 2 * num_group) + M * (num_attention_masks * activation_dtype_size)) - """ - cache_memory = self.get_available_memory(max_memory_percent) - # Compute numerator - num = cache_memory - num -= max_batch_tokens * self.peak_activation_per_token * self._activation_dtype.itemsize - num -= max_batch_tokens * (28 + 4 * self.num_groups) - # Compute denominator - denum = 2 * (self.group_size * self.page_size * cache_dtype.itemsize + 2 * self.num_groups) - denum += max_batch_tokens * (self.num_attention_masks * self._activation_dtype.itemsize) - denum += max_batch_tokens * self._activation_dtype.itemsize - # Compute cache size and return number of blocks - num_pages = floor(num / denum) - num_blocks = num_pages // self.block_size - if num_blocks > self._upper_bound_num_blocks: - logger.info(f"{num_blocks = } is too large, setting to {self._upper_bound_num_blocks = }") - num_blocks = self._upper_bound_num_blocks - return num_blocks - - def compute_memory_footprint( - self, - num_blocks: int, - max_batch_tokens: int, - cache_dtype: torch.dtype, - ) -> int: - """Calculate the memory footprint breakdown for a given number of blocks and maximum batch tokens. The memory - footprint is given by: - - available_memory = sum([ - MN * num_attention_masks * activation_dtype_size, - 2N * (layer_group_size * page_size * cache_dtype + 2 * num_group), - M * (peak_activation_per_token * activation_dtype + 28 + 4 * num_group), - ]) - but is broken down below. - """ - num_pages = num_blocks * self.block_size - - cache_memory_footprint = 2 * self.group_size * num_pages * self.page_size * cache_dtype.itemsize - - activation_memory_footprint = self.peak_activation_per_token * self._activation_dtype.itemsize - activation_memory_footprint *= max_batch_tokens - - inputs_outputs_positions_and_logits_memory_footprint = 4 * max_batch_tokens * 4 # second 4 is for int32 size - - attention_memory_footprint = self.num_attention_masks * self._activation_dtype.itemsize - attention_memory_footprint *= num_pages * max_batch_tokens - - cumulative_seqlens_memory_footprint = 3 * max_batch_tokens * 4 # 4 is for int32 size - - write_index_memory_footprint = self.num_groups * max_batch_tokens * 4 # 4 is for int32 size - read_index_memory_footprint = self.num_groups * (num_pages + max_batch_tokens) * 4 # 4 is for int32 size - - total_memory_footprint = sum( - [ - cache_memory_footprint, - activation_memory_footprint, - inputs_outputs_positions_and_logits_memory_footprint, - attention_memory_footprint, - cumulative_seqlens_memory_footprint, - write_index_memory_footprint, - read_index_memory_footprint, - ] - ) - return total_memory_footprint + 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.""" + 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 diff --git a/src/transformers/generation/continuous_batching/cache_manager.py b/src/transformers/generation/continuous_batching/cache_manager.py index a255cc2aa0da..71d0d3c23f71 100644 --- a/src/transformers/generation/continuous_batching/cache_manager.py +++ b/src/transformers/generation/continuous_batching/cache_manager.py @@ -441,17 +441,19 @@ def fill_block_table( class SlidingAttentionCacheAllocator(CacheAllocator): """Cache manager for sliding window attention layers.""" - def __init__(self, index: int, block_size: int, sliding_window: int) -> None: - """Initializes the cache manager for a group of sliding window attention layers. - Args: - - index: the index of the associated layer group - - block_size: the size of the blocks in the cache - - sliding_window: the size of the sliding window + def __init__( + self, index: int, block_size: int, sliding_window: int, sentinel_index: int, trash_index: int + ) -> None: + """Initializes the cache manager for a group of sliding window attention layers. ``sentinel_index`` and + ``trash_index`` are valid cache positions in the padding zone, used instead of -1 in read and write indices + respectively so that index_select/index_copy_ never receive negative values. """ self._index = index self.uses_block_sharing = False self.block_size = block_size self.sliding_window = sliding_window + self.sentinel_index = sentinel_index + self.trash_index = trash_index self._max_blocks_per_request = ceil(self.sliding_window / self.block_size) self.block_table = {} @@ -481,8 +483,8 @@ def get_read_indices(self, request_id: str, past_length: int, query_length: int) """Returns the physical indices of where to read request_id's cache in the cache tensor. For a group of sliding window attention layers, we read from the cache tensor before writing on it, because the new cache can overwrite the old one. To form the cache + new key / values states, we read the at most - sliding_window - 1 cache page and then manually add the new key / values states after. Hence the -1 indices - which indicate where to store the new key or values indices.""" + sliding_window - 1 cache page and then manually add the new key / values states after. Hence the sentinel + indices which indicate where to store the new key or values indices.""" # Retrieve the block table for the request and raise an error if it doesn't exist block_table = self.block_table.get(request_id) if block_table is None: @@ -498,7 +500,7 @@ def get_read_indices(self, request_id: str, past_length: int, query_length: int) block_offset = i % self.block_size physical_index = block_table[block_idx] * self.block_size + block_offset physical_indices.append(physical_index) - return physical_indices + [-1] * query_length + return physical_indices + [self.sentinel_index] * query_length def get_write_indices(self, request_id: str, past_length: int, query_length: int) -> list[int]: """Returns the physical indices of where to write request_id's cache in the cache tensor. For a group of @@ -521,7 +523,7 @@ def get_write_indices(self, request_id: str, past_length: int, query_length: int physical_index = block_table[block_idx] * self.block_size + block_offset physical_indices.append(physical_index) if padding_length > 0: - physical_indices = [-1] * padding_length + physical_indices + physical_indices = [self.trash_index] * padding_length + physical_indices return physical_indices # TODO: implement this diff --git a/src/transformers/generation/continuous_batching/input_outputs.py b/src/transformers/generation/continuous_batching/input_outputs.py index abd236251e69..e53020e43842 100644 --- a/src/transformers/generation/continuous_batching/input_outputs.py +++ b/src/transformers/generation/continuous_batching/input_outputs.py @@ -124,6 +124,7 @@ def __init__( self.requests_in_batch: list[FutureRequestState] = [] self.req_id_to_new_token_position: dict[str, int] = {} # only used for async API self.graphs: CudaGraphBuffer = CudaGraphBuffer(max_graphs) + self._trash_index = cache.trash_index # Setup static tensors and compute stream self._setup_static_tensors() self._reset_static_tensors(full_reset=True) @@ -204,12 +205,12 @@ def _setup_static_tensors(self) -> None: # For other kwargs, we need a list of tensors with as many tensors as there are groups self.write_index_storage = torch.empty( - (num_groups, max_batch_tokens), dtype=torch.int32, device=self.device, pin_memory=pin_memory + (num_groups, max_batch_tokens), dtype=torch.int64, device=self.device, pin_memory=pin_memory ) self.read_index_storage = torch.empty( - (num_groups, num_pages + max_batch_tokens), dtype=torch.int32, device=self.device, pin_memory=pin_memory + (num_groups, num_pages + max_batch_tokens), dtype=torch.int64, device=self.device, pin_memory=pin_memory ) - # For read index, the +T is because there are -1 for seqlen_q when model uses a sliding window + # For read index, the +T is because there are sentinel indices for seqlen_q when model uses a sliding window def _transfer_inputs( self, other: "ContinuousBatchingIOs", stream: torch.cuda.Stream, non_blocking: bool = False @@ -271,15 +272,15 @@ def _reset_static_tensors(self, full_reset: bool = False) -> None: # If this is a full reset, we reset every tensors if full_reset: self.block_table[:, :q_len].fill_(-1) - self.write_index_storage[:, :q_len].fill_(-2) # -1 is used to let the cache where new states go - self.read_index_storage[:, : q_len + kv_len].fill_(-2) # same + self.write_index_storage[:, :q_len].fill_(self._trash_index) + self.read_index_storage[:, : q_len + kv_len].fill_(self._trash_index) # If this is not a full reset, and we are going to use the block table, we only reset it elif self.use_block_table: self.block_table[:, :q_len].fill_(-1) # Otherwise, the read and write indices are the ones used, so we reset them else: - self.write_index_storage[:, :q_len].fill_(-2) # -1 is used to let the cache where new states go - self.read_index_storage[:, : q_len + kv_len].fill_(-2) # same + self.write_index_storage[:, :q_len].fill_(self._trash_index) + self.read_index_storage[:, : q_len + kv_len].fill_(self._trash_index) def reset(self) -> None: """Reset all relevant states for a new generation loop.""" @@ -423,9 +424,10 @@ def prepare_batch_tensors( # If we are not using the block table, we populate the read and write 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_tensor(group_read_indices) - self.write_index_storage[i, : len(group_write_indices)] = to_tensor(group_write_indices) + self.read_index_storage[i, : len(group_read_indices)] = to_index_tensor(group_read_indices) + 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) diff --git a/tests/generation/test_continuous_batching.py b/tests/generation/test_continuous_batching.py index 89857a220111..464b9564079a 100644 --- a/tests/generation/test_continuous_batching.py +++ b/tests/generation/test_continuous_batching.py @@ -35,6 +35,7 @@ ) from transformers.generation.continuous_batching.cache import ( PagedAttentionCache, + PagedAttentionMemoryHandler, SlidingAttentionCacheAllocator, group_layers_by_attn_type, ) @@ -48,6 +49,7 @@ require_flash_attn, require_kernels, require_torch_accelerator, + require_torch_gpu, slow, torch_device, ) @@ -1095,3 +1097,116 @@ def test_flash_attn_with_kvcache_parity(self, use_cuda_graph: bool, use_async: b text_fa2 = tokenizer.decode(out_fa2.generated_tokens, skip_special_tokens=True) text_fa3 = tokenizer.decode(out_fa3.generated_tokens, skip_special_tokens=True) self.assertEqual(text_fa2, text_fa3, f"Mismatch:\nFA2: {text_fa2}\nFA3: {text_fa3}") + + +@require_torch_gpu +class TestMemoryHandlerPrediction(unittest.TestCase): + """Verifies that ``PagedAttentionMemoryHandler.compute_memory_footprint`` matches real GPU memory usage. + + For each configuration we allocate tensors at the *idealized* sizes modeled by the handler (same shapes, same + dtypes, no alignment padding or extra blocks) and compare the CUDA memory delta to the handler's prediction. + """ + + # (block_size, page_size, num_groups, group_size, peak_act, num_attn_masks, max_bpr, logprobs, cache_dtype, use_async_batching) + CONFIGS = [ + (32, 256, 1, 22, 34048, 1, 0, False, torch.float16, False), # sdpa-like, 1 attn mask + (256, 256, 1, 22, 34048, 0, 0, False, torch.float16, False), # flash-like, no attn mask + (32, 256, 2, 14, 34048, 2, 0, False, torch.bfloat16, False), # hybrid model, 2 groups + 2 masks + (32, 128, 1, 16, 8192, 1, 8, True, torch.float16, False), # with block_table + logprobs + (32, 128, 1, 16, 8192, 1, 8, True, torch.float16, True), # with block_table + logprobs + async batching + ] + + NUM_BLOCKS = 4 + MAX_BATCH_TOKENS = 64 + + @parameterized.expand(CONFIGS) + def test_memory_prediction( + self, + block_size: int, + page_size: int, + num_groups: int, + group_size: int, + peak_act: int, + num_attn_masks: int, + max_bpr: int, + logprobs: bool, + cache_dtype: torch.dtype, + use_async_batching: bool, + ) -> None: + cb_config = ContinuousBatchingConfig( + max_blocks_per_request=max_bpr, + return_logprobs=logprobs, + use_async_batching=use_async_batching, + ) + + handler = PagedAttentionMemoryHandler( + block_size=block_size, + page_size=page_size, + num_groups=num_groups, + group_size=group_size, + peak_activation_per_token=peak_act, + num_attention_masks=num_attn_masks, + continuous_batching_config=cb_config, + ) + + N = self.NUM_BLOCKS * block_size # num_pages + M = self.MAX_BATCH_TOKENS + predicted = handler.compute_memory_footprint(self.NUM_BLOCKS, M, cache_dtype) + num_output_rows = 2 if logprobs else 1 + act_dtype = handler._activation_dtype + i32 = handler._input_dtype + + # -- Allocate tensors at the exact idealized sizes the handler models -- + device = "cuda" + torch.cuda.empty_cache() + baseline = torch.cuda.memory_allocated(device) + + k = handler.io_multiplier # 1 sync, 2 async -- scales IO tensors only + tensors = [] + # kv_cache: 2 * group_size tensors of [N, page_size] (not scaled by k) + for _ in range(group_size): + tensors.append(torch.empty((N, page_size), dtype=cache_dtype, device=device)) + tensors.append(torch.empty((N, page_size), dtype=cache_dtype, device=device)) + # activation peak: flat tensor of peak_act * M elements (not scaled by k) + tensors.append(torch.empty(peak_act * M, dtype=act_dtype, device=device)) + # IO tensors below are allocated k times (once per IO instance) + for _ in range(k): + # bulk_input: [7, M] + tensors.append(torch.empty((7, M), dtype=i32, device=device)) + # output_ids: [num_output_rows, M] + tensors.append(torch.empty((num_output_rows, M), dtype=i32, device=device)) + # attention_mask: [1, 1, M, N + M] per mask type + for _ in range(num_attn_masks): + tensors.append(torch.empty((1, 1, M, N + M), dtype=act_dtype, device=device)) + # block_table: [num_groups, M, max_bpr] (empty when max_bpr == 0) + if max_bpr > 0: + tensors.append(torch.empty((num_groups, M, max_bpr), dtype=i32, device=device)) + # write_index: [num_groups, M] + tensors.append(torch.empty((num_groups, M), dtype=torch.int64, device=device)) + # read_index: [num_groups, N + M] + tensors.append(torch.empty((num_groups, N + M), dtype=torch.int64, device=device)) + + actual_cuda = torch.cuda.memory_allocated(device) - baseline + expected_nbytes = sum(t.nbytes for t in tensors) + num_allocations = len(tensors) + + del tensors + torch.cuda.empty_cache() + + # 1) Exact check: prediction must equal the sum of tensor nbytes. This validates the polynomial + # coefficients against the tensor shapes, with zero tolerance. + self.assertEqual( + predicted, + expected_nbytes, + f"Prediction ({predicted}) != sum of tensor nbytes ({expected_nbytes})", + ) + + # 2) GPU memory check: CUDA's caching allocator rounds each allocation up (typically to 512 bytes). + # We allow up to 512 bytes of overhead per allocation. + max_cuda_overhead = num_allocations * 512 + self.assertLessEqual( + abs(actual_cuda - predicted), + max_cuda_overhead, + f"CUDA delta ({actual_cuda}) too far from prediction ({predicted}), " + f"allowed overhead = {max_cuda_overhead} ({num_allocations} allocs × 512B)", + ) diff --git a/tests/generation/test_paged_attention.py b/tests/generation/test_paged_attention.py index 015dab8febff..c39325332a2b 100644 --- a/tests/generation/test_paged_attention.py +++ b/tests/generation/test_paged_attention.py @@ -4,6 +4,7 @@ from parameterized import parameterized from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig +from transformers.generation.configuration_utils import ContinuousBatchingConfig from transformers.testing_utils import Expectations, slow @@ -30,7 +31,12 @@ "orange.\n\n## Step 1: Identify the key characteristics of the fruit\nThe fruit is described as being orange in color and round in shape.\n\n##", "This riddle is a classic example of a lateral thinking puzzle, which requires the test-taker to think creatively and consider multiple possibilities. The answer", "get in touch with us. We will respond to your message as soon as possible.\n\n[Your Name]\n[Your Email]\n[Your Phone Number]", - "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train", + # The last prompt sits on a numerical boundary: eager/flex produce "does", sdpa/fa2 produce "will". + # We use a tuple to accept either variant. + ( + "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train", + "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers will the train", + ), ], } ) @@ -48,7 +54,8 @@ def setUpClass(cls): if cls.tokenizer.pad_token is None: cls.tokenizer.pad_token = cls.tokenizer.eos_token - cls.model.config.pad_token_id = cls.model.config.eos_token_id + eos_id = cls.model.config.eos_token_id + cls.model.config.pad_token_id = eos_id[0] if isinstance(eos_id, list) else eos_id cls.model.use_cache = False @@ -63,15 +70,17 @@ def setUpClass(cls): def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max_batch_tokens): self.model.config.attn_implementation = attn_impl + cb_config = ContinuousBatchingConfig( + num_blocks=num_blocks, + block_size=block_size, + max_batch_tokens=max_batch_tokens, + ) generation_config = GenerationConfig( max_new_tokens=30, top_k=0, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, use_cache=False, - num_blocks=num_blocks, - block_size=block_size, - max_batch_tokens=max_batch_tokens, ) tokenized = self.tokenizer(_TEST_PROMPTS, truncation=True, max_length=512) @@ -80,17 +89,18 @@ def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max batch_outputs = self.model.generate_batch( inputs=batch_inputs, generation_config=generation_config, + continuous_batching_config=cb_config, ) expected_outputs = _EXPECTED_OUTPUTS.get_expectation() for i, (output, expected_output) in enumerate(zip(batch_outputs.values(), expected_outputs)): generated = self.tokenizer.decode(output.generated_tokens, skip_special_tokens=False).strip() - expected = expected_output.strip() - self.assertEqual( + expected_output = (expected_output.strip(),) if isinstance(expected_output, str) else expected_output + self.assertIn( generated, - expected, - msg=f"[{attn_impl}] Mismatch in request {i}:\nExpected: {expected}\nGot: {generated}", + [e.strip() for e in expected_output], + msg=f"[{attn_impl}] Mismatch in request {i}:\nExpected one of: {expected_output}\nGot: {generated}", ) @parameterized.expand( @@ -105,6 +115,11 @@ def test_generate_batch_with_sampling(self, attn_impl, num_blocks, block_size, m """Test batch generation with do_sampling=True to verify sampling works correctly.""" self.model.config.attn_implementation = attn_impl + cb_config = ContinuousBatchingConfig( + num_blocks=num_blocks, + block_size=block_size, + max_batch_tokens=max_batch_tokens, + ) generation_config = GenerationConfig( max_new_tokens=30, do_sample=True, @@ -114,9 +129,6 @@ def test_generate_batch_with_sampling(self, attn_impl, num_blocks, block_size, m eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, use_cache=False, - num_blocks=num_blocks, - block_size=block_size, - max_batch_tokens=max_batch_tokens, ) tokenized = self.tokenizer(_TEST_PROMPTS, truncation=True, max_length=512) # Use fewer prompts for faster test @@ -126,6 +138,7 @@ def test_generate_batch_with_sampling(self, attn_impl, num_blocks, block_size, m batch_outputs = self.model.generate_batch( inputs=batch_inputs, generation_config=generation_config, + continuous_batching_config=cb_config, ) end = time.time() print(