dotcache/integrations/vllm_adapter/adapter.py

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any

import numpy as np

from ...page_cache import PreparedPageCache
from ...tracing import ExecutionTrace
from .block_cache import VllmPagedKVCache
from .compat import VLLM_V1_MULTIPROCESSING_ENV, require_supported_vllm_version
from .config import VllmAdapterConfig, VllmAdapterMode

try:  # pragma: no cover - torch is optional for the base repo
    import torch
    import torch.nn as nn
except ImportError:  # pragma: no cover
    torch = None
    nn = object  # type: ignore[assignment]


def _require_torch() -> None:
    if torch is None:
        raise RuntimeError("torch is required for the vLLM adapter path")


def _torch_backend_matches_device(backend: str, device_type: str) -> bool:
    if device_type == "cuda":
        return backend in {"torch_cuda", "auto"}
    return False


def _looks_like_vllm_llama_attention(module: Any) -> bool:
    required = ("qkv_proj", "o_proj", "rotary_emb", "q_size", "kv_size", "num_heads", "num_kv_heads", "head_dim", "scaling")
    return all(hasattr(module, name) for name in required)


def _looks_like_vllm_llama_model(model: Any) -> bool:
    if not hasattr(model, "model") or not hasattr(model.model, "layers"):
        return False
    layers = getattr(model.model, "layers")
    if not layers:
        return False
    first_layer = layers[0]
    return hasattr(first_layer, "self_attn") and _looks_like_vllm_llama_attention(first_layer.self_attn)


def _extract_qkv(base_attention: Any, hidden_states) -> Any:
    projected = base_attention.qkv_proj(hidden_states)
    if isinstance(projected, tuple):
        projected = projected[0]
    return projected


def _split_qkv(base_attention: Any, qkv) -> tuple[Any, Any, Any]:
    return qkv.split([int(base_attention.q_size), int(base_attention.kv_size), int(base_attention.kv_size)], dim=-1)


def _apply_rope(base_attention: Any, positions, q, k) -> tuple[Any, Any]:
    rotated = base_attention.rotary_emb(positions, q, k)
    if isinstance(rotated, tuple) and len(rotated) == 2:
        return rotated
    raise ValueError("vLLM rotary_emb must return (query, key)")


def _project_dotcache_output(base_attention: Any, context_tensor) -> Any:
    projected = base_attention.o_proj(context_tensor)
    if isinstance(projected, tuple):
        return projected[0]
    return projected


def _resolve_token_positions(positions) -> np.ndarray:
    if torch is not None and torch.is_tensor(positions):
        return positions.reshape(-1).detach().cpu().numpy().astype(np.int64, copy=False)
    return np.asarray(positions, dtype=np.int64).reshape(-1)


class DotCacheVllmLlamaAttention(nn.Module):  # type: ignore[misc]
    def __init__(self, base_attention: Any, adapter: "VllmDotCacheModelAdapter", *, layer_id: int) -> None:
        _require_torch()
        super().__init__()
        self.base_attention = base_attention
        self.adapter = adapter
        self.layer_id = int(layer_id)

    def forward(self, positions, hidden_states):
        token_count = int(hidden_states.shape[0])
        active_trace = self.adapter.active_trace if self.adapter.active_trace is not None else self.adapter.runtime_trace
        if self.adapter.mode == "dense":
            return self.base_attention(positions, hidden_states)

        qkv = _extract_qkv(self.base_attention, hidden_states)
        query_states, key_states, value_states = _split_qkv(self.base_attention, qkv)
        query_states, key_states = _apply_rope(self.base_attention, positions, query_states, key_states)
        query_rows = query_states.view(token_count, int(self.base_attention.num_heads), int(self.base_attention.head_dim))
        key_rows = key_states.view(token_count, int(self.base_attention.num_kv_heads), int(self.base_attention.head_dim))
        value_rows = value_states.view(token_count, int(self.base_attention.num_kv_heads), int(self.base_attention.head_dim))
        token_positions = _resolve_token_positions(positions)

        if token_count != 1:
            dense_output = self.base_attention(positions, hidden_states)
            encode_start = time.perf_counter()
            if _torch_backend_matches_device(self.adapter.backend, hidden_states.device.type):
                self.adapter.block_cache.append_tokens_torch(
                    self.layer_id,
                    key_rows.detach().to(dtype=torch.float32),
                    value_rows.detach().to(dtype=torch.float32),
                    positions,
                    trace=active_trace,
                )
            else:
                self.adapter.block_cache.append_step(
                    self.layer_id,
                    key_rows.detach().permute(1, 0, 2).to(dtype=torch.float32).cpu().numpy(),
                    value_rows.detach().permute(1, 0, 2).to(dtype=torch.float32).cpu().numpy(),
                    int(token_positions[0]),
                    trace=active_trace,
                )
            self.adapter.prefill_block_encode_ms_total += (time.perf_counter() - encode_start) * 1000.0
            return dense_output

        dense_output = self.base_attention(positions, hidden_states) if self.adapter.mode == "dotcache_shadow" else None
        append_start = time.perf_counter()
        if _torch_backend_matches_device(self.adapter.backend, hidden_states.device.type):
            self.adapter.block_cache.append_step_torch(
                self.layer_id,
                key_rows.detach().permute(1, 0, 2).to(dtype=torch.float32),
                value_rows.detach().permute(1, 0, 2).to(dtype=torch.float32),
                int(token_positions[0]),
                trace=active_trace,
            )
        else:
            self.adapter.block_cache.append_step(
                self.layer_id,
                key_rows.detach().permute(1, 0, 2).to(dtype=torch.float32).cpu().numpy(),
                value_rows.detach().permute(1, 0, 2).to(dtype=torch.float32).cpu().numpy(),
                int(token_positions[0]),
                trace=active_trace,
            )
        self.adapter.append_runtime_ms_total += (time.perf_counter() - append_start) * 1000.0

        decode_start = time.perf_counter()
        if _torch_backend_matches_device(self.adapter.backend, hidden_states.device.type):
            context_states = self.adapter.block_cache.decode_layer_torch(
                self.layer_id,
                query_rows[0].detach().to(dtype=torch.float32),
                self.adapter.q_head_to_kv_head,
                query_scale=float(self.base_attention.scaling),
                trace=active_trace,
            )
        else:
            context_states = self.adapter.block_cache.decode_layer(
                self.layer_id,
                query_rows[0].detach().to(dtype=torch.float32).cpu().numpy(),
                self.adapter.q_head_to_kv_head,
                query_scale=float(self.base_attention.scaling),
                trace=active_trace,
            )
        self.adapter.decode_runtime_ms_total += (time.perf_counter() - decode_start) * 1000.0

        if not torch.is_tensor(context_states):
            context_states = torch.as_tensor(context_states, dtype=torch.float32, device=hidden_states.device)
        dotcache_output = _project_dotcache_output(
            self.base_attention,
            context_states.to(dtype=hidden_states.dtype, device=hidden_states.device).reshape(1, -1),
        )

        self.adapter.record_last_dotcache_output(self.layer_id, dotcache_output)
        if self.adapter.mode == "dotcache_shadow":
            if dense_output is None:
                raise RuntimeError("shadow mode requires the dense attention output for comparison")
            self.adapter.record_shadow_output(dense_output, dotcache_output)
            return dense_output
        return dotcache_output


@dataclass
class VllmDotCacheModelAdapter:
    model: Any
    adapter_config: VllmAdapterConfig
    backend: str = "torch_cuda"
    cache: PreparedPageCache | None = None

    def __post_init__(self) -> None:
        self.cache = self.cache if self.cache is not None else PreparedPageCache()
        self.block_cache = VllmPagedKVCache(
            config=self.adapter_config.dotcache_config,
            num_hidden_layers=self.model.config.num_hidden_layers,
            num_attention_heads=self.model.config.num_attention_heads,
            num_key_value_heads=getattr(self.model.config, "num_key_value_heads", self.model.config.num_attention_heads),
            block_size=self.adapter_config.block_size,
            backend=self.backend,
            cache=self.cache,
        )
        self.q_head_to_kv_head = self.block_cache.default_q_head_to_kv_head.copy()
        self.mode: VllmAdapterMode = self.adapter_config.mode
        self.active_trace: ExecutionTrace | None = None
        self.runtime_trace = ExecutionTrace()
        self.prefill_block_encode_ms_total = 0.0
        self.append_runtime_ms_total = 0.0
        self.decode_runtime_ms_total = 0.0
        self.shadow_output_max_abs_error = 0.0
        self.shadow_output_max_rel_error = 0.0
        self._last_dotcache_outputs: dict[int, Any] = {}
        self._wrappers: list[DotCacheVllmLlamaAttention] = []
        self._install_wrappers()

    @property
    def device(self):
        return next(self.model.parameters()).device

    @property
    def resident_bytes(self) -> int:
        return self.block_cache.resident_bytes

    def _install_wrappers(self) -> None:
        for layer_id, layer in enumerate(self.model.model.layers[: self.model.config.num_hidden_layers]):
            wrapper = DotCacheVllmLlamaAttention(layer.self_attn, self, layer_id=layer_id)
            layer.self_attn = wrapper
            self._wrappers.append(wrapper)

    def set_mode(self, mode: VllmAdapterMode) -> None:
        self.mode = mode

    def clear(self) -> None:
        self.block_cache.clear()
        self.active_trace = None
        self._last_dotcache_outputs.clear()
        self.reset_runtime_metrics()

    def reset_runtime_metrics(self) -> None:
        self.runtime_trace = ExecutionTrace()
        self.prefill_block_encode_ms_total = 0.0
        self.append_runtime_ms_total = 0.0
        self.decode_runtime_ms_total = 0.0
        self.shadow_output_max_abs_error = 0.0
        self.shadow_output_max_rel_error = 0.0

    def record_shadow_output(self, dense_output, dotcache_output) -> None:
        dense = dense_output.detach().to(dtype=torch.float32).cpu().numpy()
        dotcache = dotcache_output.detach().to(dtype=torch.float32).cpu().numpy()
        delta = np.abs(dotcache - dense)
        denom = np.maximum(np.abs(dense), 1e-8)
        self.shadow_output_max_abs_error = max(self.shadow_output_max_abs_error, float(np.max(delta)))
        self.shadow_output_max_rel_error = max(self.shadow_output_max_rel_error, float(np.max(delta / denom)))

    def record_last_dotcache_output(self, layer_id: int, output) -> None:
        self._last_dotcache_outputs[int(layer_id)] = output.detach().clone()

    def last_dotcache_output(self, layer_id: int):
        return self._last_dotcache_outputs[int(layer_id)]


def install_dotcache_on_vllm_model(
    model: Any,
    dotcache_config,
    *,
    block_size: int,
    backend: str = "torch_cuda",
    mode: VllmAdapterMode = "dense",
    cache: PreparedPageCache | None = None,
) -> VllmDotCacheModelAdapter:
    if not _looks_like_vllm_llama_model(model):
        raise ValueError("target model is not a supported vLLM Llama-family executor model")
    adapter_config = VllmAdapterConfig(
        dotcache_config=dotcache_config,
        block_size=block_size,
        mode=mode,
        model_family="llama",
    )
    return VllmDotCacheModelAdapter(model=model, adapter_config=adapter_config, backend=backend, cache=cache)


def _infer_block_size_from_target(target: Any) -> int | None:
    for attr_name in ("cache_config", "vllm_config"):
        attr = getattr(target, attr_name, None)
        if attr is not None and hasattr(attr, "block_size"):
            return int(attr.block_size)
        cache_config = getattr(attr, "cache_config", None)
        if cache_config is not None and hasattr(cache_config, "block_size"):
            return int(cache_config.block_size)
    return None


def _search_for_model(target: Any, *, max_depth: int = 6, visited: set[int] | None = None) -> Any | None:
    if visited is None:
        visited = set()
    if target is None:
        return None
    target_id = id(target)
    if target_id in visited or max_depth < 0:
        return None
    visited.add(target_id)

    if _looks_like_vllm_llama_model(target):
        return target

    for attr_name in (
        "model",
        "runner",
        "model_runner",
        "driver_worker",
        "worker",
        "model_executor",
        "llm_engine",
        "engine",
        "engine_core",
        "executor",
    ):
        child = getattr(target, attr_name, None)
        found = _search_for_model(child, max_depth=max_depth - 1, visited=visited)
        if found is not None:
            return found
    return None


def install_dotcache_on_vllm_runtime(
    target: Any,
    dotcache_config,
    *,
    block_size: int | None = None,
    backend: str = "torch_cuda",
    mode: VllmAdapterMode = "dense",
    cache: PreparedPageCache | None = None,
) -> VllmDotCacheModelAdapter:
    require_supported_vllm_version()
    model = _search_for_model(target)
    if model is None:
        llm_engine = getattr(target, "llm_engine", None)
        engine_core = getattr(llm_engine, "engine_core", None)
        if engine_core is not None and engine_core.__class__.__name__ != "InprocClient":
            raise RuntimeError(
                "could not locate a supported vLLM Llama-family executor model inside the target runtime; "
                f"for vLLM 0.18.x use the in-process engine path by setting "
                f"{VLLM_V1_MULTIPROCESSING_ENV}=0 or calling "
                "configure_vllm_inprocess_runtime() before constructing vllm.LLM"
            )
        raise RuntimeError("could not locate a supported vLLM Llama-family executor model inside the target runtime")
    resolved_block_size = int(block_size) if block_size is not None else _infer_block_size_from_target(target)
    if resolved_block_size is None:
        raise RuntimeError("could not infer vLLM block_size from the target runtime; pass block_size explicitly")
    return install_dotcache_on_vllm_model(
        model,
        dotcache_config,
        block_size=resolved_block_size,
        backend=backend,
        mode=mode,
        cache=cache,
    )