# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Worker-side wiring of HiSparse caches to attention layers."""

from collections.abc import Mapping
from copy import copy
from typing import TYPE_CHECKING, Any

import torch

from vllm.config import VllmConfig
from vllm.v1.hisparse.layout import (
    HISPARSE_HOT_SUFFIX,
    HISPARSE_RESIDENT_SUFFIX,
    get_hisparse_kv_cache_groups,
)
from vllm.v1.hisparse.runtime import (
    HiSparseCacheHandle,
    HiSparseHostPool,
    initialize_hisparse_runtime_buffers,
    release_pinned_state,
)
from vllm.v1.kv_cache_interface import (
    HiSparseHotSpec,
    HiSparseResidentSpec,
    KVCacheConfig,
    KVCacheLayout,
    KVCacheSpec,
    MLAAttentionSpec,
    UniformTypeKVCacheSpecs,
    create_kv_cache_views,
)
from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size

if TYPE_CHECKING:
    from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
    from vllm.v1.worker.gpu.block_table import BlockTables


def resolve_hisparse_specs(
    vllm_config: VllmConfig,
    kv_cache_spec: dict[str, KVCacheSpec],
    attn_layers: Mapping[str, "AttentionLayerBase"],
) -> dict[str, KVCacheSpec]:
    """Resolve a common kernel block size for HiSparse MLA specs, and add the
    resident/hot specs HiSparse derives from them."""
    mla_specs = {
        name: spec
        for name, spec in kv_cache_spec.items()
        if isinstance(spec, MLAAttentionSpec)
    }
    if not mla_specs:
        return kv_cache_spec
    block_sizes = {spec.block_size for spec in mla_specs.values()}
    if len(block_sizes) != 1:
        raise ValueError("HiSparse requires one scheduler block size.")
    backends = [attn_layers[name].get_attn_backend() for name in mla_specs]
    specs = list(mla_specs.values())
    try:
        block_size = select_common_block_size(block_sizes.pop(), backends, specs)
    except ValueError as error:
        raise ValueError(
            "HiSparse requires a GPU block size supported by every sparse "
            f"attention and indexer backend: {error}"
        ) from error
    kv_cache_spec = kv_cache_spec | {
        name: spec.copy_with_new_block_size(block_size)
        for name, spec in mla_specs.items()
    }
    groups = get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec)
    assert groups is not None
    return kv_cache_spec | {
        layer_name: group.kv_cache_spec
        for group in groups
        if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
        for layer_name in group.layer_names
    }


def allocate_hisparse_kv_caches(
    kv_cache_config: KVCacheConfig,
    device: torch.device,
    layout: KVCacheLayout,
    kernel_block_sizes: list[int],
    host_pool: HiSparseHostPool,
) -> dict[str, torch.Tensor]:
    """Allocate the host pool separately from the shared device backing."""
    device_config = copy(kv_cache_config)
    device_config.kv_cache_tensors = [
        tensor
        for tensor in kv_cache_config.kv_cache_tensors
        if not tensor.host_resident
    ]
    kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes)
    host_tensors = [
        tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident
    ]
    (host_size,) = {tensor.size for tensor in host_tensors}
    backing = host_pool.allocate(host_size)
    (host_group_id,) = kv_cache_config.host_group_ids
    host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec
    for tensor in host_tensors:
        spec = (
            host_spec.kv_cache_specs[tensor.layers[0]]
            if isinstance(host_spec, UniformTypeKVCacheSpecs)
            else host_spec
        )
        views = create_kv_cache_views(
            backing,
            spec,
            kv_cache_config.num_blocks_of(tensor),
            layout,
            tensor,
            kernel_block_size=kernel_block_sizes[host_group_id],
        )
        kv_caches.update(zip(tensor.layers, views))
    return kv_caches


def init_hisparse_kv_cache(
    kv_cache_config: KVCacheConfig,
    device: torch.device,
    kernel_block_sizes: list[int],
    vllm_config: VllmConfig,
    forward_context: dict[str, Any],
    block_tables: "BlockTables",
) -> dict[str, torch.Tensor]:
    """Allocate and bind HiSparse caches within the caller's allocation context."""
    host_pool = HiSparseHostPool(vllm_config, kv_cache_config)
    try:
        kv_caches = allocate_hisparse_kv_caches(
            kv_cache_config,
            device,
            vllm_config.cache_config.get_resolved_kv_cache_layout(),
            kernel_block_sizes,
            host_pool,
        )
        cache_handles = bind_hisparse_kv_caches(
            forward_context=forward_context,
            kv_cache_config=kv_cache_config,
            kv_caches=kv_caches,
            block_tables=block_tables,
            host_pool=host_pool,
        )
        initialize_hisparse_runtime_buffers(
            cache_handles,
            max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
        )
        return kv_caches
    except Exception:
        if host_pool.shared_region is not None:
            host_pool.shared_region.cleanup()
        raise


def _get_hisparse_cache(
    forward_context: dict[str, Any], layer_name: str
) -> HiSparseCacheHandle:
    attention_layer = forward_context[layer_name]
    hisparse_cache = attention_layer.hisparse_cache
    assert hisparse_cache is not None
    return hisparse_cache


def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None:
    cache_handles = [
        cache
        for layer in forward_context.values()
        if (cache := getattr(layer, "hisparse_cache", None)) is not None
    ]
    runtimes = {
        id(cache.runtime): cache.runtime
        for cache in cache_handles
        if hasattr(cache.runtime, "_host_cache")
    }
    if not runtimes:
        return

    shared_regions = {
        id(runtime.shared_host_region): runtime.shared_host_region
        for runtime in runtimes.values()
        if runtime.shared_host_region is not None
    }
    assert len(shared_regions) <= 1
    shared_region = next(iter(shared_regions.values()), None)
    registered_pools = (
        []
        if shared_region is not None
        else list(
            {
                runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool
                for runtime in runtimes.values()
            }.values()
        )
    )
    release_pinned_state(list(runtimes.values()), registered_pools, shared_region)


def bind_hisparse_kv_caches(
    *,
    forward_context: dict[str, Any],
    kv_cache_config: KVCacheConfig,
    kv_caches: dict[str, torch.Tensor],
    block_tables: "BlockTables",
    host_pool: HiSparseHostPool,
) -> list[HiSparseCacheHandle]:
    """Bind existing cache storage and block tables; return the bound handles."""
    assert host_pool.registered is not None
    tensor_configs = {
        name: (
            tensor_config,
            tensor_config.offset + layer_index * tensor_config.layer_stride,
        )
        for tensor_config in kv_cache_config.kv_cache_tensors
        for layer_index, name in enumerate(tensor_config.layers)
    }
    resident_groups = [
        (group_id, group)
        for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
        if isinstance(group.kv_cache_spec, HiSparseResidentSpec)
    ]
    residency: torch.Tensor | None = None
    if resident_groups:
        table_shapes = {
            block_tables.input_block_tables[group_id].shape
            for group_id, _ in resident_groups
        }
        assert len(table_shapes) == 1
        max_num_reqs, max_num_pages = table_shapes.pop()
        residency = torch.zeros(
            (max_num_reqs, len(resident_groups), max_num_pages),
            dtype=torch.int32,
            device=block_tables.input_block_tables[resident_groups[0][0]].device,
        )
    for resident_source_index, (group_id, group) in enumerate(resident_groups):
        assert residency is not None
        for cache_name in group.layer_names:
            assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
            layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
            tensor_config, byte_offset = tensor_configs[cache_name]
            assert not tensor_config.host_resident
            cache_handle = _get_hisparse_cache(forward_context, layer_name)
            cache_handle.bind_cache(
                kv_caches[cache_name],
                byte_offset=byte_offset,
                block_stride=tensor_config.block_stride,
                num_blocks=kv_cache_config.num_blocks,
                block_size=group.kv_cache_spec.block_size,
                block_table=residency[:, resident_source_index],
                slot_mapping=block_tables.slot_mappings[group_id],
            )
            cache_handle.residency = residency
            assert cache_handle.view is not None
            kv_caches[cache_name] = cache_handle.view.cache
            cache_handle.runtime.resident_source_index = resident_source_index

    hot_backing: torch.Tensor | None = None
    cache_handles: list[HiSparseCacheHandle] = []
    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
        if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
            continue
        for cache_name in group.layer_names:
            assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
            raw_tensor = kv_caches[cache_name]
            if hot_backing is None:
                hot_backing = raw_tensor
            elif hot_backing.untyped_storage().data_ptr() != (
                raw_tensor.untyped_storage().data_ptr()
            ):
                raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
            layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
            cache_handle = _get_hisparse_cache(forward_context, layer_name)
            tensor_config, byte_offset = tensor_configs[cache_name]
            assert not tensor_config.host_resident
            cache_handle.runtime.bind_hot_cache(
                raw_tensor,
                byte_offset=byte_offset,
                block_stride=tensor_config.block_stride,
                num_blocks=kv_cache_config.num_blocks,
                block_size=group.kv_cache_spec.block_size,
                block_table=block_tables.input_block_tables[group_id],
            )
            resident = cache_handle.view
            hot = cache_handle.runtime.hot
            assert resident is not None
            if (
                resident.cache.untyped_storage().data_ptr()
                != hot.cache.untyped_storage().data_ptr()
                or resident.cache.data_ptr() != hot.cache.data_ptr()
                or resident.cache.stride() != hot.cache.stride()
            ):
                raise RuntimeError("HiSparse resident and hot layouts must match.")
            source_cache = kv_caches[layer_name]
            assert source_cache.untyped_storage().data_ptr() == (
                host_pool.registered.untyped_storage().data_ptr()
            )
            cache_handle.draft_layer = forward_context[layer_name].is_draft_layer
            cache_handle.runtime.shared_host_region = host_pool.shared_region
            cache_handle.runtime.bind_source_cache(
                source_cache,
                registered_host_pool=host_pool.registered,
            )
            cache_handles.append(cache_handle)
            # The hot slab is raw storage owned by the runtime, not a layer
            # cache: nothing downstream binds or registers it.
            del kv_caches[cache_name]

    if hot_backing is None or not cache_handles:
        raise RuntimeError("HiSparse found no hot-cache handles.")
    (source_group_id,) = kv_cache_config.host_group_ids
    for cache_handle in cache_handles:
        cache_handle.source_block_table = block_tables.input_block_tables[
            source_group_id
        ]
        cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id]
    return cache_handles