# Copyright (c) DeepSpeed Team
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
"""CUDA-graph-compatible static KV cache for hybrid engine rollout.

Derived from HuggingFace transformers ``StaticCache`` / ``StaticLayer``, but
with a critical difference: the write position is supplied externally via a
shared tensor instead of an internal ``cumulative_length`` counter.

Why this matters
----------------
Transformers' ``StaticLayer.update()`` maintains its own ``cumulative_length``
tensor that advances on every call.  During CUDA graph capture the captured
forward "freezes" this counter at whatever value it had at capture time.
On replay the counter does *not* advance, so subsequent KV writes go to the
wrong positions and the model silently produces incorrect logits.

Our ``DeepSpeedStaticCache`` instead reads the write position from a shared
tensor (``write_position``) that the caller updates in-place before each graph
replay.  Because ``write_position`` is a real tensor at a fixed address, CUDA
graph replays read the current value each time.

The caller (HybridEngineRollout) must call ``cache.set_write_position(pos)``
before each replay, where ``pos`` is a scalar ``torch.long`` tensor on the
correct device.
"""

import torch


class DeepSpeedStaticLayer:
    """A single layer's static KV cache whose write position is externally set.

    Parameters
    ----------
    max_cache_len : int
        Maximum number of tokens the cache can hold (last dim size).
    """

    is_compileable = True
    is_sliding = False

    def __init__(self, max_cache_len: int):
        self.max_cache_len = max_cache_len
        self.keys: torch.Tensor | None = None
        self.values: torch.Tensor | None = None
        self.is_initialized = False
        self._write_position: torch.Tensor | None = None

    def set_write_position(self, pos: torch.Tensor):
        self._write_position = pos

    def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:
        self.dtype = key_states.dtype
        self.device = key_states.device
        max_batch_size, num_heads = key_states.shape[:2]
        self.max_batch_size = max_batch_size
        self.num_heads = num_heads
        self.k_head_dim = key_states.shape[-1]
        self.v_head_dim = value_states.shape[-1]

        self.keys = torch.zeros(
            (max_batch_size, num_heads, self.max_cache_len, self.k_head_dim),
            dtype=self.dtype,
            device=self.device,
        )
        self.values = torch.zeros(
            (max_batch_size, num_heads, self.max_cache_len, self.v_head_dim),
            dtype=self.dtype,
            device=self.device,
        )
        torch._dynamo.mark_static_address(self.keys)
        torch._dynamo.mark_static_address(self.values)
        self.is_initialized = True

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        *args,
        **kwargs,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if not self.is_initialized:
            self.lazy_initialization(key_states, value_states)

        kv_length = key_states.shape[-2]

        if self._write_position is not None:
            cache_position = torch.arange(kv_length, device=self.device) + self._write_position
        else:
            cache_position = torch.arange(kv_length, device=self.device)

        try:
            self.keys.index_copy_(2, cache_position, key_states)
            self.values.index_copy_(2, cache_position, value_states)
        except NotImplementedError:
            self.keys[:, :, cache_position] = key_states
            self.values[:, :, cache_position] = value_states

        return self.keys, self.values

    def get_mask_sizes(self, query_length: int) -> tuple[int, int]:
        return self.max_cache_len, 0

    def get_seq_length(self) -> int:
        if not self.is_initialized:
            return 0
        if self._write_position is not None:
            return self._write_position + 1
        return 0

    def get_max_cache_shape(self) -> int:
        return self.max_cache_len

    def reset(self) -> None:
        if self.is_initialized:
            self.keys.zero_()
            self.values.zero_()

    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
        if self.is_initialized:
            self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
            self.values = self.values.index_select(0, beam_idx.to(self.values.device))


class DeepSpeedStaticCache:
    """CUDA-graph-compatible static KV cache.

    Drop-in replacement for ``transformers.StaticCache`` in the graph-capture
    decode path of ``HybridEngineRollout``.  All layers share a single
    ``write_position`` tensor that the caller updates before each graph replay.

    Parameters
    ----------
    config : PreTrainedConfig
        HuggingFace model config (used to determine number of layers and head
        dimensions).
    batch_size : int
        Batch size for eager initialization.
    max_cache_len : int
        Maximum sequence length (prompt + generated tokens).
    device : torch.device | int | str | None
        Device for eager initialization.
    dtype : torch.dtype | None
        Dtype for eager initialization.
    """

    def __init__(
        self,
        config,
        batch_size: int = 1,
        max_cache_len: int = 4096,
        device=None,
        dtype=None,
    ):
        self.config = config
        text_config = getattr(config, "text_config", config)
        num_layers = getattr(text_config, "num_hidden_layers", 1)
        self._layers = [DeepSpeedStaticLayer(max_cache_len) for _ in range(num_layers)]
        self._max_cache_len = max_cache_len
        self._write_position: torch.Tensor | None = None

        if dtype is not None and device is not None and batch_size > 0:
            num_heads = getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))
            head_dim = getattr(text_config, "hidden_size", 1) // getattr(text_config, "num_attention_heads", 1)
            self.early_initialization(batch_size, num_heads, head_dim, dtype, device)

    @property
    def layers(self):
        return self._layers

    def set_write_position(self, pos: torch.Tensor):
        """Set the write position shared by all layers.

        Must be called before each graph replay with the decode step position
        as a scalar ``torch.long`` tensor on the correct device.  The tensor is
        stored by reference so subsequent in-place updates (e.g.
        ``pos.fill_(new_val)``) are immediately visible to all layers.
        """
        self._write_position = pos
        for layer in self._layers:
            layer.set_write_position(pos)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        *args,
        **kwargs,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if layer_idx >= len(self._layers):
            raise IndexError(f"layer_idx {layer_idx} out of range (cache has {len(self._layers)} layers)")
        return self._layers[layer_idx].update(key_states, value_states, *args, **kwargs)

    def early_initialization(
        self,
        batch_size: int,
        num_heads: int,
        head_dim: int,
        dtype: torch.dtype,
        device,
    ):
        for layer in self._layers:
            fake_k = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)
            fake_v = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)
            layer.lazy_initialization(fake_k, fake_v)

    def get_seq_length(self, layer_idx: int = 0) -> int:
        if layer_idx >= len(self._layers):
            return 0
        return self._layers[layer_idx].get_seq_length()

    def get_max_cache_shape(self, layer_idx: int = 0) -> int:
        if layer_idx >= len(self._layers):
            return self._max_cache_len
        return self._layers[layer_idx].get_max_cache_shape()

    def get_mask_sizes(self, query_length: int, layer_idx: int = 0) -> tuple[int, int]:
        if layer_idx >= len(self._layers):
            return self._max_cache_len, 0
        return self._layers[layer_idx].get_mask_sizes(query_length)

    def reset(self):
        for layer in self._layers:
            layer.reset()

    def __len__(self):
        return len(self._layers)