# Copyright 2025 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from math import floor, gcd, sqrt
from typing import Any

import torch

from ...configuration_utils import PreTrainedConfig
from ...generation.configuration_utils import ContinuousBatchingConfig
from ...utils.generic import is_flash_attention_requested
from .cache_manager import BlockManager, CacheAllocator, FullAttentionCacheAllocator, SlidingAttentionCacheAllocator
from .distributed import DistributedHelper
from .initialization import resolve_max_memory_percent
from .requests import RequestState, RequestStatus, get_device_and_memory_breakdown, logger


def find_num_kv_heads(config: PreTrainedConfig) -> int:
    """Finds the number of key-value heads for the given config."""
    # If the model supports GQA, we leverage it by using the num_key_value_heads attribute
    kv_heads = getattr(config, "num_key_value_heads", None)
    if kv_heads is not None:
        return kv_heads
    # Otherwise, the number of KV heads is the same as the number of attention heads
    kv_heads = getattr(config, "num_attention_heads", None)
    if kv_heads is not None:
        return kv_heads
    raise ValueError(f"num_key_value_heads or num_attention_heads could not be found in the config:\n{config}")


def find_head_dim(config: PreTrainedConfig) -> int:
    """Finds the head dimension for the given config."""
    # If the model has the head_dim attribute, there is nothing to do but return it
    head_dim = getattr(config, "head_dim", None)
    if head_dim is not None:
        return head_dim
    # If it is missing, we may reconstruct it from the hidden size and the number of attention heads
    hidden_size = getattr(config, "hidden_size", None)
    num_attention_heads = getattr(config, "num_attention_heads", None)
    if hidden_size is not None and num_attention_heads is not None:
        return hidden_size // num_attention_heads
    raise ValueError(f"head_dim or (hidden_size and num_attention_heads) could not be found in the config:\n{config}")


def group_layers_by_attn_type(config: PreTrainedConfig) -> tuple[list[list[int]], list[str]]:
    """
    Group layers depending on the attention mix, according to VLLM's hybrid allocator rules:
        - Layers in each group need to have the same type of attention
        - All groups have the same number of layers

    For a model with the following layer types: ["sliding", "full", "full", "sliding", "full", "full", "full", "full"]
    We would get four groups: [0, 3], [1, 2], [4,5] and [6,7].
    """
    # If the config has no layer_type attribute, it means all layers are the same attention type
    layer_types = getattr(config, "layer_types", None)
    if layer_types is None:
        attn_type = "sliding_attention" if getattr(config, "sliding_window", None) is not None else "full_attention"
        layer_types = [attn_type for _ in range(config.num_hidden_layers)]

    # We then count the number of layers of each type
    layer_counts = {}
    for i, layer_type in enumerate(layer_types):
        layer_counts[layer_type] = layer_counts.get(layer_type, []) + [i]

    # The size of all groups is the greatest common divisor of the number of layers of each type
    group_size = gcd(*[len(indices) for indices in layer_counts.values()])

    # We then group the layers by type
    layer_groups = []
    for layer_type, indices in layer_counts.items():
        for i in range(0, len(indices), group_size):
            layer_groups.append(indices[i : i + group_size])
    # And note the layer types
    group_types = [layer_types[lg[0]] for lg in layer_groups]
    return layer_groups, group_types


class PagedAttentionCache:
    """
    Manages the cache for a paged attention mechanism, inspired by VLLM's hybrid allocator. The cache relies on making
    groups of layers to reduce the complexity of cache management and fragmentation.

    The cache uses a three-level hierarchy:
    - Pages: The smallest unit of cache, a page has a size of [num_heads, head_size], which is the space needed to
        store the key or value states for one token and one layer. For a model with only full-attention layers, to store
        the KV cache of one token, we need `2 * num_layers` pages: key and values each take `num_layers` pages.
        Pages are grouped into blocks:
    - Blocks: A block is a collection of `block_size` pages, serving as the allocation unit to reduce management
        complexity and fragmentation. Cache is allocated and freed block by block, not page by page. One block is
        allocated to one layer group, which only has one attention type, like full-attention or sliding-attention.
        If all layers in the model have the same attention type, then all layers will be in the same group. There is
        more than one group if and only if the model has a mixed attention types, like layers with full-attention and
        layers with sliding-attention.
    - Cache tensors: The physical supports for the cache. There are as many cache tensors as there are layer in a
        layer group, and the shape of the cache tensor is `[num_blocks * block_size, num_heads, head_size]`.

    Grouping layers into groups is useful because when we allocate one block to a group N, the block allocated is the
        same for all layers in group N, equivalently it is allocated across all cache tensors. This allows us to
        efficiently allocate and free blocks, and to efficiently read and write key and value states.

    For instance, imagine we have 8 blocks of cache and a model with two layer groups: a full-attention group with 3
    layers and a sliding-attention group with 3 layers. At creation time, the physical cache tensors look like this:

    cache_tensor_0: □ □ □ □ □ □ □ □
    cache_tensor_1: □ □ □ □ □ □ □ □
    cache_tensor_2: □ □ □ □ □ □ □ □

    where □ means the blocks is not allocated to any layer group yet. We have 3 cache tensors because there are
    3 layers per group.
    We allocate 1 block to each group, after allocation, the cache tensors look like this:

    cache_tensor_0: ✖ ◉ □ □ □ □ □ □
    cache_tensor_1: ✖ ◉ □ □ □ □ □ □
    cache_tensor_2: ✖ ◉ □ □ □ □ □ □

    where ✖ means the block is allocated to the full-attention group, and ◉ means the block is allocated to the
    sliding-attention group.
    Now, if we continue to generate, and the sliding window has been reached, we only need to allocate a new block
    for the full-attention group, and the cache tensors look like this:

    cache_tensor_0: ✖ ◉ ✖ □ □ □ □ □
    cache_tensor_1: ✖ ◉ ✖ □ □ □ □ □
    cache_tensor_2: ✖ ◉ ✖ □ □ □ □ □

    And after further generation, when we need a new block allocated:

    cache_tensor_0: ✖ ◉ ✖ ✖ □ □ □ □
    cache_tensor_1: ✖ ◉ ✖ ✖ □ □ □ □
    cache_tensor_2: ✖ ◉ ✖ ✖ □ □ □ □

    This would not have been possible if all layers were in the same group: we would have had to allocate a new block
    for the sliding-attention group, although it is not needed.
    """

    _min_block_size = 4

    def __init__(
        self,
        config: PreTrainedConfig,
        continuous_batching_config: ContinuousBatchingConfig,
        device: torch.device | str,
        distributed_helper: DistributedHelper,
        tp_plan: dict[str, Any],
        dtype: torch.dtype = torch.float16,
    ) -> None:
        """Initialize a paged attention cache for efficient memory usage. Also turns in prefix sharing if the model has
        only full attention layers.

        Args:
            config: Model configuration
            continuous_batching_config: Continuous batching configuration containing cache parameters
            device: Device for the cache tensors
            distributed_helper: TP-aware helper. Used to dispatch attention heads and ensure coherent cache size
            tp_plan: Tensor parallelism plan
            dtype: Data type of the activation and the cache (for now, these are the same)
        """
        self.config = config
        self.dtype = dtype
        self.device = device

        # Extract model dimensions
        self.num_key_value_heads: int = find_num_kv_heads(config)
        self.head_dim: int = find_head_dim(config)

        # Extract cache dimensions. Default used to be 32, now it's 256 to be compatible with flash_with_kvcache.
        self.block_size = continuous_batching_config.block_size
        if self.block_size < self._min_block_size:
            raise ValueError(f"Block size must be at least {self._min_block_size}, but got {self.block_size}")

        # Group layers depending on the attention mix
        layer_groups, group_types = group_layers_by_attn_type(config)
        group_size = len(layer_groups[0])
        self.num_groups = len(layer_groups)

        self.sliding_windows = {}
        self.layer_index_to_group_indices = {}
        for i, group in enumerate(layer_groups):
            sliding_window = config.sliding_window if group_types[i] == "sliding_attention" else 1
            for j, layer in enumerate(group):
                self.layer_index_to_group_indices[layer] = (i, j)
                self.sliding_windows[layer] = sliding_window

        # Check if the KV heads are part of the TP plan. If they are not, the cache does not need plan for TP.
        # TODO: this is fragile. If your model fails to TP properly because of this, please open an issue.
        kv_is_tp = True
        for key in ["layers.*.self_attn.k_proj", "layers.*.self_attn.v_proj"]:
            if not (key in tp_plan or "model." + key in tp_plan):
                kv_is_tp = False
                break

        # If the KV heads are TP'ed, each KV head is dispatched to a different GPU, so the effective number of KV heads
        # per GPU is simply divided by the TP size
        tp_size = distributed_helper.tp_size
        if tp_size > 1 and kv_is_tp:
            if self.num_key_value_heads % tp_size != 0:
                raise ValueError(
                    f"Number of key value heads {self.num_key_value_heads} must be divisible by tensor parallel size {tp_size}."
                )
            self.num_key_value_heads //= tp_size

        # If somehow the max memory percent is not yet resolved, resolve it conservatively
        if continuous_batching_config.max_memory_percent is None:
            resolve_max_memory_percent(cb_config=continuous_batching_config, has_logit_processors=True)

        max_batch_tokens, num_blocks = PagedAttentionMemoryHandler(
            config=config,
            continuous_batching_config=continuous_batching_config,
            dtype=self.dtype,
            group_types=group_types,
            group_size=group_size,
        ).infer_max_batch_tokens_and_num_blocks()

        # For TP, align max_batch_tokens and num_blocks to the minimal value across the TP group
        if tp_size > 1:
            sync = torch.tensor([max_batch_tokens, num_blocks], device=self.device, dtype=torch.int64)
            distributed_helper.tp_all_reduce_min(sync)
            max_batch_tokens, num_blocks = int(sync[0].item()), int(sync[1].item())

        # Add the inferred attributes to the class
        self.max_batch_tokens = max_batch_tokens
        self.num_blocks = num_blocks
        self.num_pages = self.num_blocks * self.block_size
        logger.info(f"Paged cache initialized: {self.max_batch_tokens = }, {self.num_blocks = }, {self.block_size = }")

        # If max_blocks_per_request is not set, initialize it to the non-zero fallback value
        max_blocks_per_request = continuous_batching_config.max_blocks_per_request
        if max_blocks_per_request is None:
            max_blocks_per_request = continuous_batching_config.fallback_max_blocks_per_request
        self.max_blocks_per_request = max_blocks_per_request

        # Initialize the cache
        self.key_cache: list[torch.Tensor] = []
        self.value_cache: list[torch.Tensor] = []
        # We add two extra blocks to the cache as a padding zone that no BlockManager ever allocates from.
        # The first one is zeroed and then never written to. Its first index is the read trash, from which padding
        # tokens read their KV cache, and its second index is the sentinel index, to indicate where to store the new key
        # or values indices for sliding window attention groups.
        # The second is the write trash, where padding tokens can safely write their KV cache (it's never read from).
        block_based_shape = (num_blocks + 2, self.block_size, self.num_key_value_heads, self.head_dim)

        self.cache_shape = ((num_blocks + 2) * self.block_size, self.num_key_value_heads, self.head_dim)
        self.read_trash_index = num_blocks * self.block_size
        self.sentinel_index = num_blocks * self.block_size + 1  # since block size >= 4 >= 2, this is safe
        self.write_trash_index = (num_blocks + 1) * self.block_size
        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)
            torch._dynamo.mark_static_address(new_layer_key_cache)
            torch._dynamo.mark_static_address(new_layer_value_cache)
            self.key_cache.append(new_layer_key_cache)
            self.value_cache.append(new_layer_value_cache)
            # Write 0s in the read trash block so that the padding tokens read always 0-valued KV cache
            new_layer_key_cache.view(block_based_shape)[num_blocks].fill_(0)
            new_layer_value_cache.view(block_based_shape)[num_blocks].fill_(0)
        logger.info(f"{self.cache_shape = } {self.key_cache[0].shape = } {self.key_cache[0].numel() = }")

        # Block management data structures
        self.allow_block_sharing = continuous_batching_config.allow_block_sharing
        self.group_cache_managers: list[CacheAllocator] = []
        self.num_full_attention_groups = 0
        self.num_sliding_attention_groups = 0
        self.max_sliding_window_blocks_per_request = 0

        for i, group_type in enumerate(group_types):
            if group_type == "full_attention":
                cm = FullAttentionCacheAllocator(i, self.block_size, allow_block_sharing=self.allow_block_sharing)
                self.num_full_attention_groups += 1
            elif group_type == "sliding_attention":
                cm = SlidingAttentionCacheAllocator(
                    i, self.block_size, config.sliding_window, self.sentinel_index, self.write_trash_index
                )
                self.num_sliding_attention_groups += 1
                self.max_sliding_window_blocks_per_request = cm._max_blocks_per_request
            else:
                raise ValueError(f"Invalid group type: {group_type}")
            self.group_cache_managers.append(cm)

        # We only use prefix sharing if the whole model has only full attention layers and block sharing is allowed
        self.use_prefix_sharing = self.allow_block_sharing and group_types == ["full_attention"]
        self._block_manager = BlockManager(num_blocks, self.block_size, tp_on=tp_size > 1)
        self._total_prefix_length: int = 0  # a counter to measure the impact of prefix sharing, also used in tests

        # For block table support, we lazy init the name of the block table key
        self._block_table_key = None

    def blocks_needed(self, num_requested_blocks: int, allocated_blocks: int) -> int:
        """Returns the number of physical blocks needed to allocate (num_requested_blocks) blocks to a request that
        already has (allocated_blocks) blocks. The number of newly allocated blocks needed is predicted by the
        following rules:
        - for full attention groups: since there is no sliding window for full attention layers, one requested block is
            always equivalent to one newly allocated block for EACH full attention group
        - for sliding window groups: because of the sliding window, the number of blocks allocated to a request is
            capped. Using the number of already (allocated_blocks) we can compute the number of new blocks to actually
            allocate to the request, which can be lower than the number of requested blocks. That number is the same for
            all sliding window groups, as only one sliding window size is supported.
        """
        # This is not in a branch, because it is very rare to have zero full attention layer
        needed_blocks = num_requested_blocks * self.num_full_attention_groups
        # Only take this branch if the model has sliding window attention layers
        if self.num_sliding_attention_groups:
            blocks_left = max(self.max_sliding_window_blocks_per_request - allocated_blocks, 0)
            needed_blocks += min(blocks_left, num_requested_blocks) * self.num_sliding_attention_groups
        return needed_blocks

    def will_allocation_be_successful(self, num_requested_blocks: int, allocated_blocks: int) -> bool:
        """Returns a boolean indicating if the allocation of (num_requested_blocks) blocks will be successful."""
        return self.blocks_needed(num_requested_blocks, allocated_blocks) <= self.get_num_free_blocks()

    def blocks_in_use(self, request_id: str) -> int:
        """Returns the total number of physical blocks currently referenced by a request across all layer groups."""
        return sum(len(cm.block_table.get(request_id, ())) for cm in self.group_cache_managers)

    def allocate_blocks(self, n_blocks: int, request_id: str, allocated_blocks: int) -> int | None:
        """Allocate cache blocks across all layer groups for a given request. Actual allocation is done by the cache
        managers, and this method only returns the maximum number of blocks actually allocated across all managers."""
        # First check allocation will be successful before starting, to avoid partial allocations
        if not self.will_allocation_be_successful(n_blocks, allocated_blocks):
            return None
        # Allocate blocks across all cache managers
        max_allocated = 0
        for cm in self.group_cache_managers:
            num_allocated_blocks = cm.allocate_blocks(n_blocks, request_id, self._block_manager)
            if num_allocated_blocks is None:
                raise ValueError(f"Failed to allocate {n_blocks} blocks for request {request_id}")
            max_allocated = max(max_allocated, num_allocated_blocks)
        return max_allocated

    def free_blocks(self, request_id: str) -> None:
        """Free all allocated cache blocks for a given request across all layer groups. Actual deallocation is done
        by the cache managers."""
        for cm in self.group_cache_managers:
            cm.free_blocks(request_id, self._block_manager)

    def get_num_free_blocks(self) -> int:
        """Get the current number of unallocated blocks available for new requests."""
        return self._block_manager.num_free_blocks

    def extend_read_and_write_indices(
        self,
        request_id: str,
        past_length: int,
        query_length: 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.
        """
        # 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
    ) -> None:
        for i, cm in enumerate(self.group_cache_managers):
            cm.fill_block_table(request_id, past_length, query_length, block_table[i])

    def get_seqlens_k(self, past_length: int, query_length: int) -> dict[str, int]:
        """Retrieve the key sequence length for the given request_id across all layer types. Returns a dictionary of
        layer types to their corresponding key sequence lengths."""
        seqlens_k = {}
        if self.num_full_attention_groups > 0:
            seqlens_k["full_attention"] = past_length + query_length
        if self.num_sliding_attention_groups > 0:
            seqlens_k["sliding_attention"] = query_length + min(past_length, self.config.sliding_window - 1)
        # NOTE: when we add more attention types / different sliding windows, we can go back to looping over CMs
        return seqlens_k

    def update(
        self,
        key_states: torch.Tensor,  # shape [1, num_kv_heads, seqlen_kv, head_dim]
        value_states: torch.Tensor,  # shape [1, num_kv_heads, seqlen_kv, head_dim]
        layer_idx: int,
        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, 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 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]
        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:
            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:
            # 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 = 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 (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

    def get_block_table_key(self, flash_attn_with_kvcache_fn: Any) -> str:
        """A function to get the name of the block table key for the given flash_attn_with_kvcache_fn. The function's
        signature is only inspected once. This is necessary because different version of flash have different names for
        the block table key."""
        if self._block_table_key is None:
            kwarg_names = inspect.signature(flash_attn_with_kvcache_fn).parameters.keys()
            if "block_table" in kwarg_names:
                self._block_table_key = "block_table"
            elif "page_table" in kwarg_names:
                self._block_table_key = "page_table"
            else:
                raise ValueError(
                    f"flash_attn_with_kvcache_fn does not have a block_table or page_table argument: {inspect.signature(flash_attn_with_kvcache_fn)}"
                )
        return self._block_table_key

    def search_prefix_match(self, request_id: str, prompt_ids: list[int]) -> int:
        """Searches for a prefix match in the cache for the given (prompts_ids). If one is found, we reference the
        matching blocks in the (request_id), increase the reference count of the blocks and return the number of blocks
        that match. If no prefix match is found, we return 0."""
        current_hash = None
        allocated_blocks = []
        for b in range(len(prompt_ids) // self.block_size):
            tokens = prompt_ids[b * self.block_size : (b + 1) * self.block_size]
            # Prefix sharing is only supported when there is only one full attention layer group, so group_id=0.
            current_hash = self._block_manager.compute_hash(current_hash, tokens, group_id=0)
            block_id = self._block_manager._hash_to_id.get(current_hash)
            if block_id is not None:
                allocated_blocks.append(block_id)
                self._block_manager.increase_ref_count(block_id)
            else:
                break
        # If we found a matching prefix, we reference the blocks in the request
        if allocated_blocks:
            logger.debug(f"Found prefix match for request {request_id} with {len(allocated_blocks)} blocks")
            cm = self.group_cache_managers[0]
            cm.block_table[request_id] = allocated_blocks

        prefix_length = len(allocated_blocks) * self.block_size
        self._total_prefix_length += prefix_length
        return prefix_length

    def mark_shareable_blocks_as_complete(self, state: RequestState, num_complete_blocks: int) -> None:
        """Marks the blocks allocated to a request (state) as complete if they are shareable and they have been computed
        in the forward pass. A complete block is a block where the KV cache has been fully computed: if the block has
        enough space to hold the cache for N tokens, the block is marked as complete when the cache data is present for
        the N tokens. If block sharing is off, this is a no-op."""
        # The status can be FINISHED in async mode, because batch N+1 offloaded the request before batch N was over. So
        # we need to check for this case to avoid looking in the block table for blocks that no longer exist.
        if num_complete_blocks == 0 or state.status == RequestStatus.FINISHED:
            return None
        for cm in self.group_cache_managers:
            if cm.uses_block_sharing:
                self._block_manager.mark_shareable_blocks_as_complete(
                    num_complete_blocks=num_complete_blocks,
                    allocated_blocks=cm.block_table[state.request_id],
                    prompt_ids=(state.initial_tokens + state.generated_tokens),
                )

    def copy_cache(self, list_source_blocks: list[int], list_forked_blocks: list[int]) -> None:
        """Copy the cache from the source blocks to the forked blocks."""
        source_blocks = torch.tensor(list_source_blocks, device=self.device, dtype=torch.int32)
        forked_blocks = torch.tensor(list_forked_blocks, device=self.device, dtype=torch.int32)
        for key_cache, value_cache in zip(self.key_cache, self.value_cache):
            key_cache = key_cache.view(-1, self.block_size, self.num_key_value_heads, self.head_dim)
            value_cache = value_cache.view(-1, self.block_size, self.num_key_value_heads, self.head_dim)
            key_cache[forked_blocks] = key_cache[source_blocks]
            value_cache[forked_blocks] = value_cache[source_blocks]
        # FIXME: consolidate the cache into a single tensor of shape (group_size, 2, *self.k_or_v_cache_shape)
        # This will allow for  better .update and a single copy instead of one per cache tensor

    def compute_max_num_forks(self, source_request_id: str) -> int:
        """Computes the maximum number of children requests that can be forked from the source request."""
        # Count, across all groups, the new blocks each fork would have to allocate (i.e. non-shareable blocks)
        blocks_needed_per_fork = 0
        for cm in self.group_cache_managers:
            block_ids = cm.block_table[source_request_id]
            shareable_blocks = 0
            if cm.uses_block_sharing:
                for block_id in block_ids:
                    if not self._block_manager._id_to_block[block_id].is_complete:
                        break
                    shareable_blocks += 1
            blocks_needed_per_fork += len(block_ids) - shareable_blocks
        # If every block can be shared, no new allocations are needed and any number of forks is possible
        if blocks_needed_per_fork == 0:
            return 2**31  # absurdly large number, virtually infinite number of forks
        return self.get_num_free_blocks() // blocks_needed_per_fork

    def fork_request(self, source_request_id: str, destination_request_ids: list[str]) -> tuple[list[int], list[int]]:
        """Fork the cache of a request (state) into the one of a list of requests with the given (dst_request_ids)."""
        # These lists will be the accumulators for the source and destination blocks for the cache copy
        source_blocks, destination_blocks = [], []
        # Main fork loop
        for cm in self.group_cache_managers:
            src_blocks, dst_blocks = cm.fork_blocks(source_request_id, destination_request_ids, self._block_manager)
            source_blocks.extend(src_blocks)
            destination_blocks.extend(dst_blocks)
        return source_blocks, destination_blocks

    def free_all_requests(self) -> None:
        """Free all blocks allocated to requests across all cache managers. This preserves prefix hashes in the block
        manager (blocks become initialized rather than uninitialized if they were complete), allowing prefix sharing
        to work across generation sessions."""
        all_request_ids = set()
        for cm in self.group_cache_managers:
            all_request_ids.update(cm.block_table.keys())
        for request_id in all_request_ids:
            self.free_blocks(request_id)


class PagedAttentionMemoryHandler:
    """Determines the optimal max batch tokens (M) and number of blocks (N) 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 M and N, where each term maps to a tensor allocated in
    ``ContinuousBatchingIOs._setup_static_tensors`` or ``PagedAttentionCache.__init__``:

        memory(M, N)  =  coeff_m · M  +  coeff_n · N  +  coeff_mn · M·N  +  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.
    """

    _min_max_batch_tokens = 256
    _default_max_batch_tokens = 8192

    def __init__(
        self,
        config: PreTrainedConfig,
        continuous_batching_config: ContinuousBatchingConfig,
        dtype: torch.dtype,
        group_types: list[str],
        group_size: int,
    ) -> None:
        """Initialize the memory handler. Args:
        - config: the model configuration
        - continuous_batching_config: the continuous batching configuration
        - dtype: the data type of the activation and the cache
        - group_types: the list of all attention group types, formatted as strings
        - group_size: the size (in layers) of an attention group
        """
        self.config = config
        self.cb_config = continuous_batching_config
        self.cache_dtype = dtype
        self.activation_dtype = dtype
        self.block_size = continuous_batching_config.block_size
        self.page_size = find_head_dim(config) * find_num_kv_heads(config)
        self.num_groups = len(group_types)
        self.group_size = group_size

        # TODO: when we generalize to allow for block-attn, we can use `num_attention_masks=sum(set(group_types))`
        if is_flash_attention_requested(self.config):
            self.num_attention_masks = 0
        else:
            self.num_attention_masks = 2 if "sliding_attention" in group_types else 1

        self.max_blocks_per_request = continuous_batching_config.max_blocks_per_request
        if self.max_blocks_per_request is None:
            self.max_blocks_per_request = continuous_batching_config.fallback_max_blocks_per_request
        # 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
        self.available_memory = self.get_available_memory()

    @property
    def activation_peak(self) -> dict[str, tuple[int, ...]]:
        mem_per_q_token = self.config.num_attention_heads * find_head_dim(self.config)
        mem_per_k_or_v_token = self.page_size
        peaks = {}

        # LM head peak: this is when we turn the hidden states into logits
        delta_m = self.config.hidden_size * self.activation_dtype.itemsize  # hidden_shape, shape [M, hidden_size]
        delta_m += self.config.vocab_size * torch.float32.itemsize  # logits, shape [M, V], always in fp32
        peaks["lm_head"] = (delta_m, 0, 0, 0)

        # Attention peak: this is when we read the key and value states from the cache
        delta_m = self.activation_dtype.itemsize * (
            self.config.hidden_size  # hidden state, shape [M, hidden_size]
            + mem_per_q_token  # q_projection, shape [M, mem_per_q_token]
            + 2 * mem_per_k_or_v_token  # new K and V, shape [M, page_size]
        )
        # old K and V, read from cache (worst case scenario: whole cache is read)
        delta_n = 2 * mem_per_k_or_v_token * self.activation_dtype.itemsize
        peaks["attention"] = (delta_m, delta_n, 0, 0)

        return peaks

    def get_available_memory(self) -> int:
        """Calculate available GPU memory for cache allocation in bytes, accouting for the maximum memory percent limit
        fixed by the continuous batching config."""
        _, total, reserved, allocated = get_device_and_memory_breakdown()
        available_memory = total - max(allocated, reserved)
        available_memory = int(available_memory * self.cb_config.max_memory_percent)
        logger.info(f"Memory available for cache allocation: {available_memory // 1024**2} MB")
        return available_memory

    def infer_max_batch_tokens_and_num_blocks(self) -> tuple[int, int]:
        """Infers max_batch_tokens and num_blocks based on the available memory and the size of the activation peaks.
        If neither value is provided, we use a default value of 8192 for max_batch_tokens, apply bounds depending on the
        available VRAM, and solve for num_blocks. If one value is provided, the other is found using a linear solve."""
        max_batch_tokens = self.cb_config.max_batch_tokens
        num_blocks = self.cb_config.num_blocks

        # If both values are provided, just make sure they make sense
        if max_batch_tokens is not None and num_blocks is not None:
            return self._check_footprint(max_batch_tokens, num_blocks)

        # If one or more value is provided, solve for the other
        if max_batch_tokens is not None or num_blocks is not None:
            max_batch_tokens, num_blocks = self._solve_for_peaks(
                max_batch_tokens, num_blocks, cache_fill_per_batch=None
            )
            return self._check_footprint(max_batch_tokens, num_blocks)

        # If no value is provided, use the default value for max_batch_tokens w/ VRAM-based upper bound
        upper_bound_vram, _ = self._solve_for_peaks(
            max_batch_tokens=None,
            num_blocks=None,
            cache_fill_per_batch=0.1,  # each cache must fill 10% of the cache at most
        )
        max_batch_tokens = min(self._default_max_batch_tokens, upper_bound_vram)
        max_batch_tokens = max(max_batch_tokens, self._min_max_batch_tokens)
        # Then solve with that value
        max_batch_tokens, num_blocks = self._solve_for_peaks(max_batch_tokens, num_blocks, cache_fill_per_batch=None)
        return self._check_footprint(max_batch_tokens, num_blocks)

    def _solve_for_peaks(
        self,
        max_batch_tokens: int | None,
        num_blocks: int | None,
        cache_fill_per_batch: float | None,
    ) -> tuple[int, int]:
        """Returns max_batch_tokens and num_blocks so that their memory footprint is within the available memory for all
        activation peaks. If neither value is given, a value must be provided for cache_fill_per_batch: this means we
        solve for both varibles by saying each batch fill a certain percentage of the cache (eg, if cache_fill_per_batch
        is 0.01, each batch will fill 1% of the cache)."""
        solutions = []

        for peak_deltas in self.activation_peak.values():
            m, n = self._solve_for_peak(peak_deltas, max_batch_tokens, num_blocks, cache_fill_per_batch)
            solutions.append((m, n))

        final_m = min([solution[0] for solution in solutions])
        final_n = min([solution[1] for solution in solutions])
        return final_m, final_n

    def _solve_for_peak(
        self,
        peak: tuple[int, ...],
        max_batch_tokens: int | None,
        num_blocks: int | None,
        cache_fill_per_batch: float | None,
    ) -> tuple[int, int]:
        """Returns a couple of `(max_batch_tokens, num_blocks)` that satisfy the memory constraint for the given
        activation peak."""
        cm, cn, cmn, cmm = self._equation_coefficients(peak)

        # If neither variable is defined, use a quadratic solver
        if max_batch_tokens is None and num_blocks is None:
            # Substitute M = m·N → (coeff_nm·m + coeff_mm·m²)·N² + (coeff_n + coeff_m·m)·N − avail = 0
            if cache_fill_per_batch is None:
                raise ValueError("m must be provided if max_batch_tokens and num_blocks are None")
            m = cache_fill_per_batch  # as in, m is a substitute for big M, which is max_batch_tokens
            num_pages = self._solve_quadratic(cmn * m + cmm * m**2, cn + cm * m, -self.available_memory)
            max_batch_tokens = int(num_pages * m)
            num_blocks = int(num_pages) // self.block_size

        # Otherwise, use a linear solver
        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((self.available_memory - cm * M - cmm * M**2) / (cn + cmn * M))
            num_blocks = num_pages // self.block_size

        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
            max_batch_tokens = int(self._solve_quadratic(cmm, cm + cmn * N, cn * N - self.available_memory))

        return max_batch_tokens, num_blocks

    def _check_footprint(self, max_batch_tokens: int, num_blocks: int) -> tuple[int, int]:
        """Checks if the footprint of the cache is within the available memory."""
        memory_footprint = self.compute_memory_footprint(max_batch_tokens, num_blocks)
        if memory_footprint > self.available_memory:
            raise MemoryError(
                f"Memory footprint {memory_footprint} is more than available memory {self.available_memory}"
            )
        if max_batch_tokens <= 0 or num_blocks <= 0:
            raise ValueError(f"Invalid values: max_batch_tokens = {max_batch_tokens}, num_blocks = {num_blocks}")
        return max_batch_tokens, num_blocks

    def _solve_quadratic(self, a: float, b: float, c: float) -> int:
        """Largest positive root of a·x² + b·x + c = 0. Falls back to linear when a == 0. Rounded down."""
        if a == 0:
            return int(-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 int(floor(root))

    # Formatting is disabled because of comment indentation, which improves readability.
    # fmt: off
    def _equation_coefficients(self, peak_deltas: tuple[int, ...]) -> tuple[int, ...]:
        """Given some deltas corresponding to an activation peak, returns the coefficients for the memory polynomial of
        that peak. The memory polynomial is described in that class docstring."""
        delta_m, delta_n, delta_mm, delta_mn = peak_deltas

        i = torch.int32.itemsize             # size of int32 in bytes, used for index, input_ids, ...
        a = self.activation_dtype.itemsize             # for now, the cache and the activation have the same dtype
        c = self.cache_dtype.itemsize
        k = self.io_multiplier               # 1 sync, 2 async (IO tensors only)

        # -- N terms: cost per cache page --------------------------------------------------
        coeff_n = (
            delta_n                                      # activation peak: N-proportional part
            + 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 = (
            delta_m                                    # 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
            * 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)
        )
        # TODO: the above could be refined by introducing the max_requests_per_batch, but then there is a min() and this
        # is no longer a simple polynomial. Could be worth checking into.
        # -- M·N terms: cost per (page × batch token) --------------------------------------
        coeff_mn = (
            delta_mn                             # activation peak: M·N-proportional part
            + 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 = (
            delta_mm                            # activation peak: M²-proportional part
            + k * self.num_attention_masks * a  # attention_mask: [1, 1, M, N + M] (M² part only)
        )

        return coeff_m, coeff_n, coeff_mn, coeff_mm
    # fmt: on

    def compute_memory_footprint(self, max_batch_tokens: int, num_blocks: int) -> int:
        """Evaluate the memory polynomial at concrete (N, M) values, taking the max across activation peaks."""
        M = max_batch_tokens
        N = num_blocks * self.block_size

        max_memory_footprint = 0
        for peak in self.activation_peak.values():
            cm, cn, cmn, cmm = self._equation_coefficients(peak)
            memory_footprint = cn * N + cm * M + cmn * N * M + cmm * M * M
            max_memory_footprint = max(max_memory_footprint, memory_footprint)
        return max_memory_footprint
