# SPDX-License-Identifier: Apache-2.0
# Standard
from dataclasses import dataclass
from enum import Enum, auto
from functools import cache, wraps
from typing import Any, List, Optional, Tuple, Union
import abc
import ctypes
import os
import threading

# Third Party
from sortedcontainers import SortedList
import torch

# First Party
from lmcache import torch_dev
from lmcache import torch_device_type as torch_device_type  # noqa: F401
from lmcache.integration.vllm.utils import get_size_bytes
from lmcache.logging import init_logger
from lmcache.observability import LMCStatsMonitor
from lmcache.utils import _lmcache_nvtx_annotate
from lmcache.v1.pin_monitor import PinMonitor
from lmcache.v1.platform import current_device_spec as current_device_spec  # noqa: F401
from lmcache.v1.system_detection import NUMAMapping
import lmcache.c_ops as lmc_ops

logger = init_logger(__name__)


# Cache for ctypes ubyte-array types keyed by length.
#
# ctypes does not cache `(c_ubyte * N)` array types -- each call to the `*`
# operator builds a fresh heap type via PyCArrayType_from_ctype. The heap
# type metadata stays alive forever (held by the type system), so calling
# `(ctypes.c_ubyte * N).from_address(...)` on every TensorMemoryObj.byte_array
# access leaks ~1-2 kB per call. Under long-running remote-backend put/get
# workloads this is the dominant source of monotonic anonymous-memory growth
# (see https://github.com/LMCache/LMCache/issues/3767).
#
# Caching by length is safe: the `from_address(addr)` instance never owns the
# underlying buffer, only the metadata (length, item-type), and that metadata
# depends solely on `N`. ``functools.cache`` provides a thread-safe unbounded
# memoization primitive, so concurrent first-time accesses for the same `N`
# cannot race to create distinct heap types.
@cache
def _get_cached_ubyte_array_type(num_bytes: int) -> type[ctypes.Array[ctypes.c_ubyte]]:
    """Return a cached ``ctypes.c_ubyte * num_bytes`` array type.

    Args:
        num_bytes: The length of the array type in bytes.

    Returns:
        The cached ``ctypes.Array`` subclass for the given length. Subsequent
        calls with the same ``num_bytes`` return the same type object.
    """
    return ctypes.c_ubyte * num_bytes


# Helper functions for thread safety
def synchronized(lock_attr_name):
    """
    Decorator to make a method thread-safe by acquiring the lock
    specified by lock_attr_name on the instance.
    """

    def decorator(method):
        @wraps(method)
        def wrapper(self, *args, **kwargs):
            lock = getattr(self, lock_attr_name)
            with lock:
                return method(self, *args, **kwargs)

        return wrapper

    return decorator


class MemoryFormat(Enum):
    UNDEFINED = 0
    """[2, num_layers, num_tokens, hidden_dim]
    """
    # KV_BLOB = 1
    KV_2LTD = auto()
    """[num_tokens, 2, hidden_dim]
    """
    # LAYER_KV_BLOB = 2
    KV_T2D = auto()
    """[2, num_tokens, hidden_dim]
    """

    KV_2TD = auto()
    """Compressed binary array format
    """
    BINARY = auto()

    BINARY_BUFFER = auto()

    KV_MLA_FMT = auto()
    """[1, num_layers, num_tokens, aligned_head_size]
    """

    # This is for the encoder cache (EC) tensor format
    EC_TD = auto()
    """[num_tokens, hidden_dim]
    """

    # Hidden-state store (HS) tensor format. Same logical shape as EC_TD
    # ([num_tokens, hidden_dim]) but tagged separately so the allocator and
    # any future mp/serialization paths can distinguish encoder-cache entries
    # from hidden-state entries.
    HS_TD = auto()
    """[num_tokens, hidden_dim]
    """

    def token_dim(self) -> int:
        if self == MemoryFormat.KV_2LTD:
            return 2
        elif self == MemoryFormat.KV_T2D:
            return 1
        elif self == MemoryFormat.KV_2TD:
            return 0
        elif self == MemoryFormat.BINARY:
            return 0
        elif self == MemoryFormat.BINARY_BUFFER:
            return 0
        elif self == MemoryFormat.KV_MLA_FMT:
            return 2
        elif self == MemoryFormat.EC_TD:
            return 0
        elif self == MemoryFormat.HS_TD:
            return 0
        return 0


@dataclass
class FreeBlock:
    """Metadata class used by the memory allocators"""

    start: int
    size: int

    def can_be_coalesced(self, succ: "FreeBlock") -> bool:
        return self.start + self.size == succ.start


@dataclass
class MemoryObjMetadata:
    # TODO(chunxiaozheng): use shapes and dtypes to replace shape and dtype
    # The 'logical' shape of the tensor
    shape: torch.Size

    # The 'logical' dtype of the tensor
    dtype: Optional[torch.dtype]

    # The 'physical address' of the tensor
    address: int

    # The 'physical size' in bytes of the allocated memory
    phy_size: int

    # Reference count
    ref_count: int

    # Whether the object is pinned and cannot be evicted
    # lookup pins are temporary
    # cache controller pins are persistent
    pin_count: int = 0

    # The 'logical' format of the tensor
    fmt: MemoryFormat = MemoryFormat.UNDEFINED

    # Positions when the cache is stored
    cached_positions: Optional[torch.Tensor] = None

    # shapes and dtypes should be used in the future
    shapes: Optional[list[torch.Size]] = None
    dtypes: Optional[list[torch.dtype]] = None

    def to_dict(self):
        # Note(Kuntai): this is used for serializing MemoryObjMetadata via
        # msgpack.
        return {
            "__type__": "MemoryObjMetadata",
            "shape": list(self.shape),  # torch.Size -> list
            "dtype": str(self.dtype) if self.dtype else None,
            "address": self.address,
            "phy_size": self.phy_size,
            "ref_count": self.ref_count,
            "fmt": self.fmt.value,
            "shapes": [list(shape) for shape in self.shapes] if self.shapes else None,
            "dtypes": [str(dtype) for dtype in self.dtypes] if self.dtypes else None,
        }

    @staticmethod
    def from_dict(d):
        dtype_str = d["dtype"]
        dtype = getattr(torch, dtype_str.replace("torch.", "")) if dtype_str else None
        shapes_list = d["shapes"]
        shapes = [torch.Size(s) for s in shapes_list] if shapes_list else None
        dtypes_list = d["dtypes"]
        dtypes = (
            [getattr(torch, d_str.replace("torch.", "")) for d_str in dtypes_list]
            if dtypes_list
            else None
        )
        return MemoryObjMetadata(
            shape=torch.Size(d["shape"]),
            dtype=dtype,
            address=d["address"],
            phy_size=d["phy_size"],
            ref_count=d["ref_count"],
            fmt=MemoryFormat(d["fmt"]),
            shapes=shapes,
            dtypes=dtypes,
        )

    def get_size(self) -> int:
        if self.shapes is not None and self.dtypes is not None:
            return get_size_bytes(self.shapes, self.dtypes)
        return self.shape.numel() * self.dtype.itemsize  # type: ignore


class MemoryObj(metaclass=abc.ABCMeta):
    """
    MemoryObj interface.
    """

    # subclasses should expose raw_data differently
    raw_data: Any

    def __init__(self, metadata: MemoryObjMetadata):
        self.meta = metadata

    @abc.abstractmethod
    def invalidate(self):
        """
        Invalidate the MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def is_valid(self):
        """
        Check if the MemoryObj is valid.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_size(self) -> int:
        """
        Get the size of the MemoryObj in bytes.
        Note that this number could be smaller than the physical size.
        The physical size is aligned to the allocator's alignment.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_shape(self) -> torch.Size:
        """
        Get the shape of the MemoryObj.
        """
        raise NotImplementedError

    def get_dtype(self) -> Optional[torch.dtype]:
        """
        Get the dtype of the MemoryObj.
        """
        return None

    @abc.abstractmethod
    def get_shapes(self) -> list[torch.Size]:
        """
        Get the shapes of the MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_dtypes(self) -> list[torch.dtype]:
        """
        Get the dtypes of the MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_memory_format(self) -> MemoryFormat:
        """
        Get the memory format of the MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_physical_size(self) -> int:
        """
        Get the physical size of the MemoryObj in bytes.
        """
        raise NotImplementedError

    def set_used_size(self, n: int) -> None:  # noqa: B027
        """Narrow this buffer's logical size to the first ``n`` bytes.

        Optional hook for callers that have just written ``n`` bytes
        into a buffer originally allocated with an upper-bound size
        (e.g. the async serde processor, where the destination is sized
        from ``estimate_serialized_size`` but ``serialize`` returns the
        actual ``n`` it wrote). After this call, ``get_size()`` /
        ``byte_array`` / any downstream L2 adapter that reads the
        logical size will see exactly ``n`` bytes.

        Default is a no-op so subclasses without a "used vs allocated"
        distinction (e.g. :class:`BytesBufferMemoryObj`, where the raw
        bytes already are the actual contents) keep working unchanged.

        Args:
            n: bytes actually used in this buffer. Subclasses that
                implement this must validate ``n`` and raise
                ``ValueError`` on out-of-range or unsupported layouts.
        """
        pass

    @abc.abstractmethod
    def pin(self) -> bool:
        """
        Pin the memory obj so that it will not be evicted.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def ref_count_up(self):
        """
        Increase ref count for the given MemoryObj by one.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def unpin(self) -> bool:
        """
        Unpin the memory obj so that it can be evicted.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def ref_count_down(self):
        """
        Decrease ref count for the given MemoryObj by one.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_ref_count(self) -> int:
        """
        Get ref count for the given MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_num_tokens(self) -> int:
        """
        Get token number for the given MemoryObj.
        """
        raise NotImplementedError

    @property
    def shm_offset(self) -> int:
        """Return the byte offset of this object inside the SHM pool."""
        return self.meta.address

    @property
    def shm_byte_length(self) -> int:
        """Return the byte length of this object inside the SHM pool."""
        return self.get_size()

    @property
    @abc.abstractmethod
    def metadata(self) -> MemoryObjMetadata:
        """
        Get the metada of the MemoryObj.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def tensor(self) -> Optional[torch.Tensor]:
        """
        Get the tensor from the MemoryObj.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def byte_array(self) -> bytes:
        """
        Get the byte array from the MemoryObj.
        The size is will be the physical size instead of the unaligned size.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def data_ptr(self) -> int:
        """
        Get the data pointer of the MemoryObj.
        This is used to access the raw data in the memory.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def is_pinned(self) -> bool:
        """
        Check whether the memory obj is pinned.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def can_evict(self) -> bool:
        """
        Check whether the memory obj can be evicted.
        """
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def raw_tensor(self) -> Optional[torch.Tensor]:
        """
        Get the raw tensor from the MemoryObj.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_tensor(self, index: int) -> Optional[torch.Tensor]:
        """
        Get the tensor from the MemoryObj at the given index(group).
        """
        raise NotImplementedError

    @abc.abstractmethod
    def parent(self) -> Optional["MemoryAllocatorInterface"]:
        """
        Get the allocator that allocates this memory object
        """
        raise NotImplementedError


@dataclass
class PinnedAllocFree:
    """Resolved alloc/free function pair for pinned CPU memory."""

    alloc_fn: Any
    alloc_args: tuple
    free_fn: Any
    free_args: tuple

    def alloc(self) -> int:
        """Allocate pinned memory and return the raw pointer."""
        return self.alloc_fn(*self.alloc_args)

    def free(self, ptr: int) -> None:
        """Free a previously allocated pinned-memory pointer."""
        self.free_fn(ptr, *self.free_args)


def _resolve_pinned_alloc_free(
    numa_mapping: Optional[NUMAMapping] = None,
    shm_name: Optional[str] = None,
    size: Optional[int] = None,
    use_hugepages: bool = False,
) -> PinnedAllocFree:
    """Resolve the alloc/free function pair based on memory type.

    Returns:
        A PinnedAllocFree with the resolved functions and their extra
        arguments.  Call ``ptr = resolved.alloc()`` and ``resolved.free(ptr)``.
    """
    if shm_name:
        if use_hugepages:
            raise ValueError("Hugepages are not supported with shared memory (shm)")
        return PinnedAllocFree(
            alloc_fn=lmc_ops.alloc_shm_pinned_ptr,
            alloc_args=(size, shm_name),
            free_fn=lmc_ops.free_shm_pinned_ptr,
            free_args=(size, shm_name),
        )
    elif numa_mapping:
        if torch_dev.is_available():
            current_device_id = torch_dev.current_device()
        else:
            current_device_id = 0
        gpu_to_numa_mapping = numa_mapping.gpu_to_numa_mapping
        assert current_device_id in gpu_to_numa_mapping, (
            f"Current device {current_device_id} is not in the GPU NUMA mapping."
        )
        numa_id = gpu_to_numa_mapping[current_device_id]
        if use_hugepages:
            return PinnedAllocFree(
                alloc_fn=lmc_ops.alloc_hugepage_pinned_numa_ptr,
                alloc_args=(size, numa_id),
                free_fn=lmc_ops.free_hugepage_pinned_numa_ptr,
                free_args=(size,),
            )
        else:
            return PinnedAllocFree(
                alloc_fn=lmc_ops.alloc_pinned_numa_ptr,
                alloc_args=(size, numa_id),
                free_fn=lmc_ops.free_pinned_numa_ptr,
                free_args=(size,),
            )
    else:
        flags = 0
        if use_hugepages:
            return PinnedAllocFree(
                alloc_fn=lmc_ops.alloc_hugepage_pinned_ptr,
                alloc_args=(size, flags),
                free_fn=lmc_ops.free_hugepage_pinned_ptr,
                free_args=(size,),
            )
        else:
            return PinnedAllocFree(
                alloc_fn=lmc_ops.alloc_pinned_ptr,
                alloc_args=(size, flags),
                free_fn=lmc_ops.free_pinned_ptr,
                free_args=(),
            )


def _read_hugepage_info() -> Optional[Tuple[int, int, int]]:
    """Read hugepage pool stats from sysfs.

    NOTE: We only use 2 MiB hugepages, so the pool stats are taken from
    the 2 MiB pool directly rather than the system default pool reported in
    ``/proc/meminfo`` (which can be 1 GiB on some hosts).

    Returns:
        ``(nr_hugepages, free_hugepages, page_size_mb)`` for the hugepage
        pool, or ``None`` if the sysfs entries are unavailable.
    """
    base = "/sys/kernel/mm/hugepages/hugepages-2048kB"
    try:
        with open(f"{base}/nr_hugepages") as f:
            total = int(f.read().strip())
        with open(f"{base}/free_hugepages") as f:
            free = int(f.read().strip())
        return total, free, 2
    except (OSError, ValueError):
        return None


def _allocate_cpu_memory(
    size: int,
    numa_mapping: Optional[NUMAMapping] = None,
    shm_name: Optional[str] = None,
    use_hugepages: bool = False,
) -> torch.Tensor:
    if size == 0:
        return torch.empty(0, dtype=torch.uint8)

    resolved = _resolve_pinned_alloc_free(
        numa_mapping,
        shm_name,
        size,
        use_hugepages,
    )

    try:
        ptr = resolved.alloc()
    except RuntimeError as e:
        if use_hugepages and "mmap failed" in str(e):
            diag = _read_hugepage_info()
            if diag is not None:
                total, free, page_mb = diag
                page_bytes = page_mb * 1024 * 1024
                needed = (size + page_bytes - 1) // page_bytes
                logger.error(
                    "Failed to allocate huge pages. "
                    "Pool has %d pages (%d free, each %d MiB). "
                    "Requested %d bytes (%d pages). "
                    "Please grow the %d MiB hugepage pool.",
                    total,
                    free,
                    page_mb,
                    size,
                    needed,
                    page_mb,
                )
            else:
                logger.error(
                    "Failed to allocate huge pages. "
                    "Please grow the 2 MiB hugepage pool."
                )
        raise

    array_type = ctypes.c_uint8 * size
    buf = array_type.from_address(ptr)
    buffer = torch.frombuffer(buf, dtype=torch.uint8)

    return buffer


def _free_cpu_memory(
    buffer: torch.Tensor,
    size: int | None = None,
    numa_mapping: Optional[NUMAMapping] = None,
    shm_name: Optional[str] = None,
    use_hugepages: bool = False,
) -> None:
    if torch_dev.is_available():
        torch_dev.synchronize()

    resolved = _resolve_pinned_alloc_free(
        numa_mapping,
        shm_name,
        size,
        use_hugepages,
    )
    resolved.free(buffer.data_ptr())


def _allocate_gpu_memory(
    size: int,
    device: str,
) -> Tuple[torch.Tensor, torch.Tensor]:
    page_size = os.sysconf("SC_PAGESIZE")

    # Over-allocate
    base_buffer = torch.empty(size + page_size, dtype=torch.uint8, device=device)
    offset = -base_buffer.data_ptr() % page_size

    # Make aligned view
    aligned_buffer = base_buffer[offset : offset + size]

    # Need to return the base buffer as well in order to prevent GC
    return base_buffer, aligned_buffer


class TensorMemoryObj(MemoryObj):
    """
    Wraps a raw flat tensor with some metadata
    """

    monitor = LMCStatsMonitor.GetOrCreate()

    def __init__(
        self,
        raw_data: torch.Tensor,
        metadata: MemoryObjMetadata,
        parent_allocator: Optional["MemoryAllocatorInterface"],
    ):
        assert metadata.dtype is not None, "dtype must be specified for TensorMemoryObj"
        super().__init__(metadata)
        self.raw_data = raw_data
        self.valid = True
        self.lock = threading.Lock()
        self.parent_allocator = parent_allocator
        # ``None`` means "use the layout-derived size from
        # group_prefix_sum"; a non-None value narrows the logical view to
        # exactly that many bytes (see set_used_size).  Allocator reuse
        # paths must reset this to None along with the rest of the
        # per-allocation metadata.
        self._used_size_override: Optional[int] = None
        # Calculate the prefix sum of the group sizes
        # If there are two groups, the prefix sum will be
        # [0, size_of_group_1, size_of_group_1 + size_of_group_2]
        self.group_prefix_sum = [0]
        if self.meta.shapes is not None and self.meta.dtypes is not None:
            size_in_bytes = 0
            for shape, dtype in zip(self.meta.shapes, self.meta.dtypes, strict=True):
                size_in_bytes += shape.numel() * dtype.itemsize
                self.group_prefix_sum.append(size_in_bytes)
        else:
            self.group_prefix_sum.append(self.meta.get_size())

    def __del__(self):
        """
        Destructor to ensure memory is released when the object is garbage collected.
        This acts as a safety net to prevent memory leaks if ref_count_down() is not
        called properly somewhere in the code path.
        """
        if self.parent_allocator is not None and self.is_valid():
            if self.meta.ref_count > 0 or self.meta.pin_count > 0:
                logger.warning(
                    "MemoryObj at %s is being garbage collected "
                    "with ref_count=%d, pin_count=%d. "
                    "This indicates ref_count_down()/unpin() was not called properly.",
                    self.meta.address,
                    self.meta.ref_count,
                    self.meta.pin_count,
                )
            self.parent_allocator.free(self)

    def invalidate(self):
        self.valid = False

    def is_valid(self):
        return self.valid

    def get_size(self) -> int:
        if self._used_size_override is not None:
            return self._used_size_override
        return self.group_prefix_sum[-1]

    def set_used_size(self, n: int) -> None:
        """Narrow the logical size to ``n`` bytes after a write.

        After this call, ``get_size()`` returns ``n`` and ``byte_array``
        exposes exactly ``n`` bytes from the start of ``raw_data``.  The
        physical allocation (``get_physical_size``) and ``raw_data``
        buffer are unchanged.  Allocator reuse resets this override to
        ``None`` so a recycled block returns to its layout-derived size.

        Note: the ``tensor`` property still derives its shape from
        ``meta.shape``, so accessing ``.tensor`` on a buffer narrowed
        below its layout size will fail to reshape.  Use ``byte_array``
        (or read ``raw_data[: get_size()]`` directly) for downstream
        I/O that must honor the narrowed size.

        Args:
            n: bytes actually written.  Must satisfy
                ``0 <= n <= get_physical_size()``.

        Raises:
            ValueError: if ``n`` is outside the allowed range.
        """
        if n < 0 or n > self.meta.phy_size:
            raise ValueError(
                f"set_used_size: n={n} out of range [0, {self.meta.phy_size}]"
            )
        with self.lock:
            self._used_size_override = n

    # TODO(chunxiaozheng): use get_shapes and get_dtypes to replace
    #  get_shape and get_dtype
    def get_shape(self) -> torch.Size:
        return self.meta.shape

    def get_dtype(self) -> torch.dtype:
        assert self.meta.dtype is not None
        return self.meta.dtype

    def get_shapes(self) -> list[torch.Size]:
        assert self.meta.shapes is not None
        return self.meta.shapes

    def get_dtypes(self) -> list[torch.dtype]:
        assert self.meta.dtypes is not None
        return self.meta.dtypes

    def get_memory_format(self) -> MemoryFormat:
        with self.lock:
            return self.meta.fmt

    def get_physical_size(self) -> int:
        return self.meta.phy_size

    def ref_count_up(self):
        with self.lock:
            self.meta.ref_count += 1

    def ref_count_down(self):
        with self.lock:
            self.meta.ref_count -= 1
            if self.meta.ref_count < 0:
                logger.warning(
                    f"Ref count of MemoryObj {self.meta.address}"
                    f"is negative: {self.meta.ref_count}."
                    "Double free occurred somewhere."
                    "Setting ref count back to 0 as a hack but please find the bug."
                )
                self.meta.ref_count = 0
            if (
                self.meta.ref_count == 0
                and self.parent_allocator is not None
                and self.meta.pin_count == 0
            ):
                self.parent_allocator.free(self)

    def get_ref_count(self) -> int:
        with self.lock:
            return self.meta.ref_count

    def get_num_tokens(self) -> int:
        with self.lock:
            token_dim = self.meta.fmt.token_dim()
            return self.meta.shape[token_dim]

    def pin(self) -> bool:
        with self.lock:
            # if pin_count is 0, indicates that the object is pinned for the first time
            if self.meta.pin_count == 0:
                TensorMemoryObj.monitor.update_pinned_memory_objs_count(1)

            self.meta.pin_count += 1

            # Register/update with PinMonitor for timeout tracking on every pin
            pin_monitor = PinMonitor.GetOrCreate()
            pin_monitor.on_pin(self)
            return True

    def unpin(self) -> bool:
        with self.lock:
            self.meta.pin_count -= 1

            # if pin_count is 0, indicates that the object is unpinned
            if self.meta.pin_count == 0:
                TensorMemoryObj.monitor.update_pinned_memory_objs_count(-1)
                # Unregister from PinMonitor when fully unpinned
                pin_monitor = PinMonitor.GetOrCreate()
                pin_monitor.on_unpin(self)

            if self.meta.pin_count <= 0 and self.meta.ref_count <= 0:
                if self.parent_allocator is None:
                    logger.error(
                        "Parent allocator is None when trying to free MemoryObj."
                        "This could cause memory leak"
                    )
                else:
                    self.parent_allocator.free(self)

            if self.meta.pin_count < 0:
                logger.warning(
                    f"Pin count of MemoryObj {self.meta.address}"
                    f"is negative: {self.meta.pin_count}."
                    "Double unpin occurred somewhere."
                    "Setting pin count back to 0 as a hack but please find the bug."
                )
                self.meta.pin_count = 0
            return True

    @property
    def metadata(self) -> MemoryObjMetadata:
        with self.lock:
            return self.meta

    @property
    def tensor(self) -> Optional[torch.Tensor]:
        if not self.valid:
            logger.warning("Trying to access an invalidated MemoryObj")
            return None
        assert self.meta.dtype is not None
        if self._used_size_override is not None:
            # Narrowed byte buffer (see set_used_size): expose exactly
            # the used bytes as a flat uint8 view.  Reshaping to the
            # original meta.shape would raise -- fewer than shape-many
            # bytes are logically present -- so keep the view consistent
            # with get_size()/byte_array/shm_byte_length, which all
            # report the narrowed length.  Consumers that build SHM
            # transport slots from ``tensor.shape`` and
            # ``shm_byte_length`` then stay self-consistent.
            return self.raw_data[: self._used_size_override].view(torch.uint8)
        # TODO(Jiayi): consider caching the `get_size()`
        return (
            self.raw_data[: self.get_size()].view(self.meta.dtype).view(self.meta.shape)
        )

    @property
    def byte_array(self) -> memoryview:
        # TODO: consider using one of the alternatives

        # Alternative 1:
        # # PyTorch tensors support buffer protocol directly for CPU tensors
        # return memoryview(self.raw_data)

        # Alternative 2:
        # assert self.raw_data.device.type == 'cpu',
        #   "byte_array only works with CPU tensors"
        # return memoryview(self.raw_data.contiguous().numpy())

        # Use logical size (get_size) rather than raw_data physical size.
        # The raw_data buffer may include alignment padding (e.g. from
        # batched_allocate) that must not be exposed to callers such as
        # remote-backend put/get which rely on byte_array length matching
        # the metadata length.
        num_bytes = self.get_size()
        ptr = self.raw_data.data_ptr()
        # ctypes does not cache (c_ubyte * N) array types -- each `*` builds a
        # fresh heap type. With this property accessed once per remote put/get,
        # uncached creation leaks ~1-2 kB per call (heap-type metadata is held
        # by the type system and never reclaimed). Cache the array type per
        # size so steady-state usage reuses a fixed set of types. See
        # https://github.com/LMCache/LMCache/issues/3767.
        arr_type = _get_cached_ubyte_array_type(num_bytes)
        ubyte_ptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte))
        byte_array = arr_type.from_address(ctypes.addressof(ubyte_ptr.contents))
        return memoryview(byte_array)

    @property
    def data_ptr(self) -> int:
        return self.raw_data.data_ptr()

    @property
    def is_pinned(self) -> bool:
        return self.metadata.pin_count > 0

    @property
    def can_evict(self) -> bool:
        """
        Check whether the memory obj can be evicted.
        A memory obj can be evicted if it is not pinned and ref_count=1.
        """
        return not self.is_pinned and self.get_ref_count() == 1

    @property
    def raw_tensor(self) -> Optional[torch.Tensor]:
        if not self.valid:
            logger.warning("Trying to access an invalidated MemoryObj")
            return None
        return self.raw_data

    def get_tensor(self, index: int) -> Optional[torch.Tensor]:
        if not self.valid:
            logger.warning("Trying to access an invalidated MemoryObj")
            return None
        assert self.meta.shapes is not None
        assert self.meta.dtypes is not None
        begin = self.group_prefix_sum[index]
        end = self.group_prefix_sum[index + 1]
        return (
            self.raw_data[begin:end]
            .view(self.meta.dtypes[index])
            .view(self.meta.shapes[index])
        )

    def parent(self) -> Optional["MemoryAllocatorInterface"]:
        return self.parent_allocator


class BytesBufferMemoryObj(MemoryObj):
    """
    Wraps a raw flat tensor with some metadata
    """

    def __init__(self, raw_bytes: bytes, metadata: Optional[MemoryObjMetadata] = None):
        self.raw_data = raw_bytes
        if metadata is None:
            bytes_shape = torch.Size([len(self.raw_data), 0, 0, 0])
            metadata = MemoryObjMetadata(
                shape=bytes_shape,
                dtype=None,
                address=0,
                phy_size=0,
                ref_count=1,
                pin_count=0,
                fmt=MemoryFormat.BINARY_BUFFER,
            )
        super().__init__(metadata)
        self.valid = True

    def invalidate(self):
        self.valid = False

    def is_valid(self):
        return self.valid

    def get_size(self) -> int:
        return len(self.raw_data)

    def get_shape(self) -> torch.Size:
        return torch.Size([len(self.raw_data), 0, 0, 0])

    def get_dtype(self) -> Optional[torch.dtype]:
        return None

    def get_shapes(self) -> list[torch.Size]:
        return [self.get_shape()]

    def get_dtypes(self) -> list[torch.dtype]:
        return []

    def get_memory_format(self) -> MemoryFormat:
        return self.metadata.fmt

    def get_physical_size(self) -> int:
        return self.metadata.phy_size

    def pin(self) -> bool:
        self.metadata.pin_count += 1
        return True

    def unpin(self) -> bool:
        self.metadata.pin_count -= 1
        if self.metadata.pin_count < 0:
            logger.warning(
                f"Pin count of MemoryObj {self.meta.address}"
                f"is negative: {self.meta.pin_count}."
                "Double unpin occurred somewhere."
                "Setting pin count back to 0 as a hack but please find the bug."
            )
            self.metadata.pin_count = 0
        return True

    def ref_count_up(self):
        pass

    def ref_count_down(self):
        pass

    def get_ref_count(self) -> int:
        return 1

    def get_num_tokens(self) -> int:
        # TODO(Jiayi): record the number of tokens somehow
        return 1

    @property
    def metadata(self) -> MemoryObjMetadata:
        return self.meta

    @property
    def tensor(self) -> Optional[torch.Tensor]:
        if not self.valid:
            logger.warning("Trying to access an invalidated MemoryObj")
            return None
        return None

    @property
    def byte_array(self) -> bytes:
        return self.raw_data

    @property
    def data_ptr(self) -> int:
        mv = memoryview(self.raw_data)
        addr = ctypes.addressof(ctypes.c_char.from_buffer(mv))
        return addr

    @property
    def is_pinned(self) -> bool:
        return self.metadata.pin_count > 0

    @property
    def can_evict(self) -> bool:
        """
        Check whether the memory obj can be evicted.
        A buffer memory obj can be evicted if it is not pinned.
        """
        return not self.is_pinned

    @property
    def raw_tensor(self) -> Optional[torch.Tensor]:
        if not self.valid:
            logger.warning("Trying to access an invalidated MemoryObj")
            return None
        return None

    def get_tensor(self, index: int) -> Optional[torch.Tensor]:
        return None

    def parent(self) -> Optional["MemoryAllocatorInterface"]:
        # NOTE: BytesBufferMemoryObj may not be allocated by any allocator,
        # so just return None here
        return None


class GDSMemoryObject(MemoryObj):
    """A slab-anchored ``MemoryObj`` for the GDS L1 tier.

    The bytes live in the GDS slab file, not in host or device memory, so
    this object carries only the slab ``(offset, size)`` (in ``meta.address``
    / ``meta.phy_size``) and is otherwise a placeholder: ``tensor`` is always
    ``None`` and ``byte_array`` / ``data_ptr`` raise.
    """

    def __init__(self, metadata: MemoryObjMetadata) -> None:
        super().__init__(metadata)
        self.valid = True

    @property
    def slab_offset(self) -> int:
        """Byte offset of this chunk within the slab file (== ``meta.address``)."""
        return self.meta.address

    def invalidate(self) -> None:
        self.valid = False

    def is_valid(self) -> bool:
        return self.valid

    def get_size(self) -> int:
        return self.meta.phy_size

    def get_shape(self) -> torch.Size:
        return self.meta.shape

    def get_dtype(self) -> Optional[torch.dtype]:
        return self.meta.dtype

    def get_shapes(self) -> list[torch.Size]:
        raise NotImplementedError(
            "GDSMemoryObject.get_shapes: per-group shapes are not tracked on "
            "the GDS path (only the singular meta.shape is); use get_shape()"
        )

    def get_dtypes(self) -> list[torch.dtype]:
        raise NotImplementedError(
            "GDSMemoryObject.get_dtypes: per-group dtypes are not tracked on "
            "the GDS path (only the singular meta.dtype is); use get_dtype()"
        )

    def get_memory_format(self) -> MemoryFormat:
        return self.meta.fmt

    def get_physical_size(self) -> int:
        return self.meta.phy_size

    def ref_count_up(self) -> None:
        raise NotImplementedError(
            "GDSMemoryObject.ref_count_up: not used on the GDS path"
        )

    def ref_count_down(self) -> None:
        raise NotImplementedError(
            "GDSMemoryObject.ref_count_down: not used on the GDS path"
        )

    def get_ref_count(self) -> int:
        raise NotImplementedError(
            "GDSMemoryObject.get_ref_count: not used on the GDS path"
        )

    def get_num_tokens(self) -> int:
        raise NotImplementedError(
            "GDSMemoryObject.get_num_tokens: not used on the GDS path"
        )

    def pin(self) -> bool:
        raise NotImplementedError("GDSMemoryObject.pin: not used on the GDS path")

    def unpin(self) -> bool:
        raise NotImplementedError("GDSMemoryObject.unpin: not used on the GDS path")

    @property
    def metadata(self) -> MemoryObjMetadata:
        return self.meta

    @property
    def tensor(self) -> Optional[torch.Tensor]:
        return None

    @property
    def byte_array(self) -> bytes:
        raise NotImplementedError(
            f"GDSMemoryObject(slab_offset={self.slab_offset}).byte_array is not "
            "supported; bytes live in the GDS slab file and the staging buffer "
            "is registered VRAM (no buffer protocol)."
        )

    @property
    def data_ptr(self) -> int:
        raise NotImplementedError(
            f"GDSMemoryObject(slab_offset={self.slab_offset}).data_ptr is not "
            "supported; GDS reads/writes use gpu_buffer.data_ptr() via the "
            "gpu_ops dispatch, never the MemoryObj's data_ptr."
        )

    @property
    def is_pinned(self) -> bool:
        raise NotImplementedError("GDSMemoryObject.is_pinned: not used on the GDS path")

    @property
    def can_evict(self) -> bool:
        raise NotImplementedError("GDSMemoryObject.can_evict: not used on the GDS path")

    @property
    def raw_tensor(self) -> Optional[torch.Tensor]:
        return None

    def get_tensor(self, index: int) -> Optional[torch.Tensor]:
        return None

    def parent(self) -> Optional["MemoryAllocatorInterface"]:
        # The GDS slab is not a MemoryAllocatorInterface; dispatch in gpu_ops
        # keys off the GDSMemoryObject type, not the parent allocator.
        return None


class MemoryAllocatorInterface(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def allocate(
        self,
        shapes: Union[torch.Size, list[torch.Size]],
        dtypes: Union[torch.dtype, list[torch.dtype]],
        fmt: MemoryFormat = MemoryFormat.UNDEFINED,
        allocator_type: Optional[str] = None,
    ) -> Optional[MemoryObj]:
        """
        Allocates the memory to hold a tensor of the given shape.

        :param torch.Size shapes: The shape of the tensor to allocate.
        :param torch.dtype dtypes: The dtype of the tensor to allocate.
        :param MemoryFormat fmt: The format of the memory to allocate.

        :return: A MemoryObj wrapping the allocated memory. Returns
            None if the allocation failed.

        :rtype: Optional[MemoryObj]
        """
        raise NotImplementedError

    @abc.abstractmethod
    def batched_allocate(
        self,
        shapes: Union[torch.Size, list[torch.Size]],
        dtypes: Union[torch.dtype, list[torch.dtype]],
        batch_size: int,
        fmt: MemoryFormat = MemoryFormat.UNDEFINED,
        allocator_type: Optional[str] = None,
    ) -> Optional[List[MemoryObj]]:
        """
        Batched allocate the memory to hold a tensor of the given shape.

        :param torch.Size shapes: The shape of the tensor to allocate.
        :param torch.dtype dtypes: The dtype of the tensor to allocate.
        :param int batch_size: The number of tensors to allocate.
        :param MemoryFormat fmt: The format of the memory to allocate.

        :return: A list of MemoryObjs wrapping the allocated memory.
            Returns None if the allocation failed.

        :rtype: Optional[List[MemoryObj]]
        """
        raise NotImplementedError

    @abc.abstractmethod
    def free(
        self,
        memory_obj: MemoryObj,
        allocator_type: Optional[str] = None,
    ):
        """
        Frees the memory allocated for the given MemoryObj.
        Note that this function shouldn't be explicitly called.
        Instead, use `ref_count_down` to decrease ref count.

        :param MemoryObj memory_obj: The MemoryObj to free.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def batched_free(
        self,
        memory_objs: List[MemoryObj],
        allocator_type: Optional[str] = None,
        update_stats: bool = True,
    ):
        """
        Frees the memory allocated for the given list of MemoryObjs.

        :param List[MemoryObj] memory_objs: The list of MemoryObjs
            to free.
        """
        raise NotImplementedError

    def close(self):
        """
        Closes the memory allocator.
        This is called when the LMCacheEngine is closed.
        """
        return

    def memcheck(self) -> bool:
        """
        Checks the memory allocator for consistency.

        Returns:
            True if everything is fine otherwise False
        """
        return True

    # TODO(chunxiaozheng): remove if after all params replaced by shapes/dtypes
    def _adapt_shapes_and_dtypes(
        self,
        shapes: Union[torch.Size, list[torch.Size]],
        dtypes: Union[torch.dtype, list[torch.dtype]],
    ) -> Tuple[list[torch.Size], list[torch.dtype]]:
        if isinstance(shapes, torch.Size):
            shapes = [shapes]

        if isinstance(dtypes, torch.dtype):
            dtypes = [dtypes]

        assert len(shapes) == len(dtypes), (
            f"shapes and dtypes must have the same length, "
            f"got {len(shapes)} and {len(dtypes)}, "
            f"shapes: {shapes}, dtypes: {dtypes}"
        )
        return shapes, dtypes


class AddressManager:
    """
    Manages a virtual address space starting from 0 for memory allocation.

    Key interfaces:
    - allocate(size): Allocate a block of memory of the given size. The starting
      address and the actual allocated size will be aligned.

    - free(address, size): Free a previously allocated region. Note that if the
      region is not "allocated" before, it may have internal errors.

    - sbrk(size): Expand the virtual address space by the given size. The size
      will be aligned internally.

    Core assumptions:
    - The allocated size should be aligned with ALIGN_BYTES.
    """

    ALIGN_BYTES = 4096

    def __init__(self, size: int, align_bytes: int = ALIGN_BYTES):
        """
        Initializes the AddressManager with a given size.

        Args:
            size: The initial size of the virtual address space.
            align_bytes: The alignment requirement for allocations.
        """
        self._size = size
        self._align = align_bytes

        # Current implementation: explicit list
        self._explicit_list: SortedList[FreeBlock] = SortedList(key=lambda x: x.start)
        self._explicit_list.add(FreeBlock(start=0, size=size))

        # thread safe lock
        self._lock = threading.Lock()

        # For debugging purposes
        self.total_allocated_size = 0

    def compute_aligned_size(self, raw_size: int) -> int:
        """
        Helper function to compute the aligned size for a given raw size.

        Args:
            raw_size: The raw size to be aligned.

        Returns:
            The aligned size.
        """
        return (raw_size + self._align - 1) & ~(self._align - 1)

    def _can_merge_with_prev(
        self, curr_block: FreeBlock, prev_block: FreeBlock
    ) -> bool:
        """Hook: Check if curr_block can merge with prev_block."""
        return prev_block.can_be_coalesced(curr_block)

    def _can_merge_with_succ(
        self, curr_block: FreeBlock, succ_block: FreeBlock
    ) -> bool:
        """Hook: Check if curr_block can merge with succ_block."""
        return curr_block.can_be_coalesced(succ_block)

    @_lmcache_nvtx_annotate
    def _coalesce(
        self,
        curr_block: FreeBlock,
        prev_block: Optional[FreeBlock],
        succ_block: Optional[FreeBlock],
    ):
        """
        Coalesces the current block with the previous and/or successor block.
        This assumes the curr_block is NOT in self._explicit_list

        Returns True if the current block was coalesced, otherwise False.
        """
        merge_prev = prev_block is not None and self._can_merge_with_prev(
            curr_block, prev_block
        )
        merge_succ = succ_block is not None and self._can_merge_with_succ(
            curr_block, succ_block
        )

        if merge_prev and merge_succ:
            prev_block.size += curr_block.size + succ_block.size  # type: ignore
            self._explicit_list.remove(succ_block)
        elif merge_prev:
            prev_block.size += curr_block.size  # type: ignore
        elif merge_succ:
            # NOTE: logically, this won't change the order of the succ_block,
            #       so we don't need to do a "remove" and "reinsert" here
            self._explicit_list.remove(succ_block)
            succ_block.start -= curr_block.size  # type: ignore
            succ_block.size += curr_block.size  # type: ignore
            self._explicit_list.add(succ_block)

        return merge_prev or merge_succ

    @_lmcache_nvtx_annotate
    @synchronized("_lock")
    def allocate(self, size: int) -> tuple[int, int]:
        """
        Allocate a block of memory from the virtual address space of a given
        size. The actual allocated size could be larger than the requested size
        in order to satisfy alignment requirements.

        Args:
            size: The requested size of the memory block. Should be greater
                than 0.

        Returns:
            A tuple (address, allocated_size) where address is the starting
            address of the allocated block and allocated_size is the actual
            size of the allocated block.

        Raises:
            RuntimeError: If no memory is available to allocate.
        """
        aligned_size = self.compute_aligned_size(size)
        for block in self._explicit_list:
            if block.size >= aligned_size:
                break
        else:
            logger.warning(
                "Failed to allocate memory block of size %d "
                "because no memory is available",
                size,
            )
            raise RuntimeError(
                f"Failed to allocate memory block of size {size} "
                "because no memory is available"
            )

        self._explicit_list.remove(block)
        if block.size > aligned_size:
            self._explicit_list.add(
                FreeBlock(
                    start=block.start + aligned_size,
                    size=block.size - aligned_size,
                )
            )

        # For debug
        self.total_allocated_size += aligned_size

        return block.start, aligned_size

    @_lmcache_nvtx_annotate
    @synchronized("_lock")
    def batched_allocate(self, size: int, batch_size: int) -> list[tuple[int, int]]:
        """
        Allocate blocks of memory from the virtual address space of a given
        size and batch size. The actual allocated size could be larger than
        the requested size in order to satisfy alignment requirements.

        Args:
            size: The requested size of the memory block. Should be greater
                than 0.
            batch_size: The number of memory blocks to allocate.

        Returns:
            A list of tuple (address, allocated_size) where address is the starting
            address of the allocated block and allocated_size is the actual size of
            the allocated block.
            Note: the length of the return list is the same as the batch_size.

        Raises:
            RuntimeError: If no memory is available to allocate.
        """
        aligned_size = self.compute_aligned_size(size)
        remaining = batch_size
        allocate_result: list[tuple[int, int]] = []

        blocks_to_remove: list[FreeBlock] = []
        blocks_to_add: list[FreeBlock] = []

        for block in self._explicit_list:
            if remaining <= 0:
                break
            if block.size < aligned_size:
                continue

            # Greedily carve out as many aligned_size chunks as possible
            num_from_block = min(remaining, block.size // aligned_size)
            start = block.start
            for i in range(num_from_block):
                allocate_result.append((start + i * aligned_size, aligned_size))
            remaining -= num_from_block

            # Mark the original block for removal
            blocks_to_remove.append(block)

            # Keep the remaining tail as a new free block if any space is left
            used = num_from_block * aligned_size
            if block.size > used:
                blocks_to_add.append(
                    FreeBlock(start=block.start + used, size=block.size - used)
                )

        if remaining > 0:
            # Not enough memory; free list is untouched, no rollback needed
            logger.warning(
                "Failed to batched allocate %d memory blocks of size %d "
                "because no enough memory is available (short by %d blocks)",
                batch_size,
                size,
                remaining,
            )
            raise RuntimeError(
                f"Failed to batched allocate {batch_size} memory blocks "
                f"of size {size} because no enough memory is available"
            )
        if len(allocate_result) != batch_size:
            # The length of allocate_result is not equal to batch_size;
            # free list is untouched, no rollback needed
            logger.warning(
                "Failed to batched allocate %d memory blocks of size %d "
                "because the length of allocate_result %d is not equal to batch_size",
                batch_size,
                size,
                len(allocate_result),
            )
            raise RuntimeError(
                f"Failed to batched allocate {batch_size} memory blocks "
                f"of size {size} because the length of allocate_result "
                f"{len(allocate_result)} is not equal to batch_size"
            )

        # Allocation succeeded; batch-update the free list
        for block in blocks_to_remove:
            self._explicit_list.remove(block)
        for block in blocks_to_add:
            self._explicit_list.add(block)

        # Update debug statistics
        total_allocated = aligned_size * batch_size
        self.total_allocated_size += total_allocated

        return allocate_result

    @_lmcache_nvtx_annotate
    @synchronized("_lock")
    def free(self, address: int, size: int):
        """
        Free a previously allocated block of memory.

        Args:
            address: The starting address of the block to free.
            size: The size of the block to free. Should be greater than 0.
        """
        new_free_block = FreeBlock(start=address, size=size)
        index = self._explicit_list.bisect_left(new_free_block)
        prev_block = self._explicit_list[index - 1] if index > 0 else None
        succ_block = (
            self._explicit_list[index] if index < len(self._explicit_list) else None
        )

        coalesced = self._coalesce(new_free_block, prev_block, succ_block)
        if not coalesced:
            self._explicit_list.add(new_free_block)

        # For debug
        self.total_allocated_size -= size

    @synchronized("_lock")
    def sbrk(self, size: int):
        """
        Expand the virtual address space by a given size.

        Args:
            size: The size to expand the address space. Will be aligned internally
                with the ALIGN_BYTES
        """
        size = self.compute_aligned_size(size)
        new_block = FreeBlock(start=self._size, size=size)
        prev_block = self._explicit_list[-1] if len(self._explicit_list) > 0 else None
        succ_block = None
        coalesced = self._coalesce(new_block, prev_block, succ_block)
        if not coalesced:
            self._explicit_list.add(new_block)

        self._size += size

    def get_heap_size(self) -> int:
        """
        Get the total size of the address space.

        Returns:
            The total size in bytes.
        """
        return self._size

    def get_free_size(self) -> int:
        """
        Get the total free size in the address space.

        Returns:
            The total free size in bytes.
        """
        return self._size - self.total_allocated_size

    def check_consistency(self) -> bool:
        """
        Check if the address manager is consistent.

        Returns:
            True if consistent, False otherwise.
        """
        # Check if free blocks are properly coalesced
        for prev, succ in zip(
            self._explicit_list[:-1], self._explicit_list[1:], strict=False
        ):
            if prev.can_be_coalesced(succ):
                return False

        # Check if total size matches
        total_free_size = sum(block.size for block in self._explicit_list)
        if total_free_size + self.total_allocated_size != self._size:
            return False

        return True


_CORE_EXPORTS = [
    "AddressManager",
    "BytesBufferMemoryObj",
    "FreeBlock",
    "GDSMemoryObject",
    "MemoryAllocatorInterface",
    "MemoryFormat",
    "MemoryObj",
    "MemoryObjMetadata",
    "TensorMemoryObj",
    "torch_device_type",
]

__all__ = list(_CORE_EXPORTS)