from __future__ import annotations import os import time from dataclasses import dataclass import json from pathlib import Path from typing import Any, Literal, Sequence import numpy as np from ..config import DotCacheConfig from ..model_kv_cache import ModelPagedKVCache from ..page_cache import PreparedPageCache from ..page_oracle import PageTraceRecord, save_page_trace from ..tracing import ExecutionTrace def transformers_available() -> bool: try: import torch # noqa: F401 import transformers # noqa: F401 except ImportError: return False return True if transformers_available(): import torch import torch.nn as nn import transformers.models.llama.modeling_llama as llama_mod from transformers import AutoModelForCausalLM, AutoTokenizer else: # pragma: no cover - exercised in environments without transformers torch = None nn = object # type: ignore[assignment] llama_mod = None AutoModelForCausalLM = None AutoTokenizer = None AttentionMode = Literal["dense", "dotcache"] @dataclass(slots=True) class LlamaReplayRecord: step_index: int layer_id: int token_index: int query_states: np.ndarray key_states: np.ndarray value_states: np.ndarray context_states: np.ndarray output_states: np.ndarray gate_states: np.ndarray | None = None @dataclass(slots=True) class LlamaLayerRuntimeProfile: layer_id: int call_count: int = 0 qkv_projection_ms_total: float = 0.0 append_ms_total: float = 0.0 decode_ms_total: float = 0.0 output_projection_ms_total: float = 0.0 def reset(self) -> None: self.call_count = 0 self.qkv_projection_ms_total = 0.0 self.append_ms_total = 0.0 self.decode_ms_total = 0.0 self.output_projection_ms_total = 0.0 def to_dict(self) -> dict[str, float | int]: return { "layer_id": self.layer_id, "call_count": self.call_count, "qkv_projection_ms_total": float(self.qkv_projection_ms_total), "append_ms_total": float(self.append_ms_total), "decode_ms_total": float(self.decode_ms_total), "output_projection_ms_total": float(self.output_projection_ms_total), "qkv_projection_ms_per_call": float(self.qkv_projection_ms_total / max(self.call_count, 1)), "append_ms_per_call": float(self.append_ms_total / max(self.call_count, 1)), "decode_ms_per_call": float(self.decode_ms_total / max(self.call_count, 1)), "output_projection_ms_per_call": float(self.output_projection_ms_total / max(self.call_count, 1)), } def _require_transformers() -> None: if not transformers_available(): raise RuntimeError("transformers and torch are required for the Llama integration path") def resolve_hf_auth_kwargs() -> dict[str, str]: for env_name in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN"): token = os.environ.get(env_name) if token is not None: token = token.strip() if token: return {"token": token} return {} def _torch_backend_matches_device(backend: str, device_type: str) -> bool: if device_type == "mps": return backend in {"torch_mps", "auto"} if device_type == "cuda": return backend in {"torch_cuda", "auto"} return False def _default_model_device() -> str: if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available(): return "mps" return "cpu" def _device_type(device: Any) -> str: if hasattr(device, "type"): return str(device.type) return str(device) def _synchronize_device(device: Any) -> None: device_type = _device_type(device) if device_type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(device=device) elif device_type == "mps" and torch.backends.mps.is_available(): torch.mps.synchronize() def _timed_call(fn, *, device: Any) -> tuple[Any, float]: _synchronize_device(device) start = time.perf_counter() result = fn() _synchronize_device(device) return result, (time.perf_counter() - start) * 1000.0 def _run_inference(fn): with torch.inference_mode(): return fn() def _prewarm_torch_decode_layers(adapter: "LlamaDotCacheModelAdapter", *, device: Any) -> None: device_type = _device_type(device) if device_type != "cuda" or not _torch_backend_matches_device(adapter.backend, device_type): return if adapter.model_kv_cache._torch_device_type is None: return zero_query = torch.zeros( (adapter.model.config.num_attention_heads, adapter.dotcache_config.head_dim), dtype=torch.float32, device=device, ) with torch.no_grad(): for layer_id in range(adapter.model.config.num_hidden_layers): if adapter.model_kv_cache.layer_sequence_length(layer_id) <= 0: continue adapter.model_kv_cache.decode_layer_torch( layer_id, zero_query, adapter.q_head_to_kv_head, query_scale=1.0, trace=None, ) _synchronize_device(device) def _begin_cuda_memory_region(device: Any) -> dict[str, int] | None: if _device_type(device) != "cuda" or not torch.cuda.is_available(): return None torch.cuda.synchronize(device=device) torch.cuda.reset_peak_memory_stats(device) stats = torch.cuda.memory_stats(device) return { "allocation_count": int(stats.get("allocation.all.allocated", 0)), "segment_count": int(stats.get("segment.all.allocated", 0)), } def _end_cuda_memory_region(device: Any, baseline: dict[str, int] | None) -> dict[str, int]: if baseline is None or _device_type(device) != "cuda" or not torch.cuda.is_available(): return {} torch.cuda.synchronize(device=device) stats = torch.cuda.memory_stats(device) return { "cuda_peak_memory_allocated_bytes": int(torch.cuda.max_memory_allocated(device)), "cuda_peak_memory_reserved_bytes": int(torch.cuda.max_memory_reserved(device)), "cuda_allocation_count": int(stats.get("allocation.all.allocated", 0)) - baseline["allocation_count"], "cuda_segment_allocation_count": int(stats.get("segment.all.allocated", 0)) - baseline["segment_count"], } def _default_attention_mask(input_ids) -> Any: return torch.ones_like(input_ids, dtype=torch.long) def _clone_attention_mask(attention_mask) -> Any: if attention_mask is None: return None return attention_mask.clone() def _normalize_input_ids(input_ids, *, device) -> Any: tensor = torch.as_tensor(input_ids, dtype=torch.long, device=device) if tensor.ndim != 2 or tensor.shape[0] != 1: raise ValueError("Phase 5 Llama harness requires input_ids with shape [1, seq_len]") return tensor def _ensure_attention_mask(input_ids, attention_mask, *, device) -> Any: if attention_mask is None: return _default_attention_mask(input_ids).to(device=device) mask = torch.as_tensor(attention_mask, dtype=torch.long, device=device) if mask.shape != input_ids.shape: raise ValueError("attention_mask must match input_ids shape") return mask def _tensor_to_float32_numpy(value: Any) -> np.ndarray: if torch.is_tensor(value): return value.detach().to(dtype=torch.float32).cpu().numpy() return np.asarray(value, dtype=np.float32) def _can_skip_decode_attention_mask(attention_mask) -> bool: if attention_mask is None: return True return bool(torch.all(attention_mask != 0).item()) def extract_past_key_values_arrays(past_key_values) -> list[tuple[np.ndarray, np.ndarray]]: layers = getattr(past_key_values, "layers", None) if layers is None: raise ValueError("past_key_values must expose a .layers cache structure") extracted: list[tuple[np.ndarray, np.ndarray]] = [] for layer in layers: keys = layer.keys.detach().to(dtype=torch.float32).cpu().numpy() values = layer.values.detach().to(dtype=torch.float32).cpu().numpy() if keys.shape[0] != 1 or values.shape[0] != 1: raise ValueError("Phase 5 Llama harness requires batch=1 past_key_values") extracted.append((keys, values)) return extracted def extract_past_key_values_tensors(past_key_values) -> list[tuple[Any, Any]]: layers = getattr(past_key_values, "layers", None) if layers is None: raise ValueError("past_key_values must expose a .layers cache structure") extracted: list[tuple[Any, Any]] = [] for layer in layers: keys = layer.keys.detach().to(dtype=torch.float32) values = layer.values.detach().to(dtype=torch.float32) if keys.shape[0] != 1 or values.shape[0] != 1: raise ValueError("Phase 5 Llama harness requires batch=1 past_key_values") extracted.append((keys, values)) return extracted def _prefill_layer_nbytes(prefill_layers: Sequence[tuple[Any, Any]]) -> int: total = 0 for layer_keys, layer_values in prefill_layers: if torch.is_tensor(layer_keys): total += int(layer_keys.numel() * layer_keys.element_size()) else: keys = np.asarray(layer_keys) total += int(keys.nbytes) if torch.is_tensor(layer_values): total += int(layer_values.numel() * layer_values.element_size()) else: values = np.asarray(layer_values) total += int(values.nbytes) return total def _dense_kv_bytes_after_decode( prefill_layers: Sequence[tuple[Any, Any]], *, generated_token_count: int, ) -> int: if not prefill_layers: return 0 layer_keys, _ = prefill_layers[0] if torch.is_tensor(layer_keys): seq_len = int(layer_keys.shape[2]) kv_heads = int(layer_keys.shape[1]) head_dim = int(layer_keys.shape[3]) dtype_bytes = int(layer_keys.element_size()) else: keys = np.asarray(layer_keys) seq_len = int(keys.shape[2]) kv_heads = int(keys.shape[1]) head_dim = int(keys.shape[3]) dtype_bytes = int(keys.dtype.itemsize) total_tokens = seq_len + max(generated_token_count - 1, 0) layer_count = len(prefill_layers) return int(layer_count * 2 * kv_heads * total_tokens * head_dim * dtype_bytes) class DotCacheLlamaAttention(nn.Module): def __init__(self, base_attention: nn.Module, adapter: "LlamaDotCacheModelAdapter") -> None: super().__init__() self.base_attention = base_attention self.adapter = adapter self.layer_idx = int(base_attention.layer_idx) self.config = base_attention.config def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values=None, cache_position: torch.LongTensor | None = None, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: if self.adapter.mode == "dense" and not self.adapter.capture_enabled: return self.base_attention( hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, **kwargs, ) if self.adapter.mode == "dense": return self._forward_dense_with_capture( hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, **kwargs, ) return self._forward_dotcache( hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, **kwargs, ) def _project_qkv( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if position_embeddings is None: raise ValueError("position_embeddings are required for the Llama attention path") input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.base_attention.head_dim) query_states = self.base_attention.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.base_attention.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.base_attention.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings return llama_mod.apply_rotary_pos_emb(query_states, key_states, cos, sin), value_states def _forward_dense_with_capture( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values=None, cache_position: torch.LongTensor | None = None, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: input_shape = hidden_states.shape[:-1] (query_states, key_states), value_states = self._project_qkv(hidden_states, position_embeddings) fresh_key_states = key_states fresh_value_states = value_states if past_key_values is not None: cos, sin = position_embeddings cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) attention_interface = llama_mod.ALL_ATTENTION_FUNCTIONS.get_interface( self.base_attention.config._attn_implementation, llama_mod.eager_attention_forward, ) attn_output, attn_weights = attention_interface( self.base_attention, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.base_attention.attention_dropout, scaling=self.base_attention.scaling, **kwargs, ) reshaped_output = attn_output.reshape(*input_shape, -1).contiguous() projected_output = self.base_attention.o_proj(reshaped_output) if self.adapter.capture_enabled and tuple(hidden_states.shape[:2]) == (1, 1): token_index = self.adapter.current_token_index(cache_position) self.adapter.record_replay( LlamaReplayRecord( step_index=self.adapter.capture_step_index, layer_id=self.layer_idx, token_index=token_index, query_states=query_states[0, :, 0, :].detach().to(dtype=torch.float32).cpu().numpy(), key_states=fresh_key_states[0, :, 0, :].detach().to(dtype=torch.float32).cpu().numpy(), value_states=fresh_value_states[0, :, 0, :].detach().to(dtype=torch.float32).cpu().numpy(), context_states=attn_output[0, 0].detach().to(dtype=torch.float32).cpu().numpy(), output_states=projected_output[0, 0].detach().to(dtype=torch.float32).cpu().numpy(), ) ) return projected_output, attn_weights def _forward_dotcache( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values=None, cache_position: torch.LongTensor | None = None, **kwargs: Any, ) -> tuple[torch.Tensor, None]: del attention_mask, kwargs if past_key_values is not None: raise ValueError("DotCache decode mode manages its own KV cache and requires past_key_values=None") if tuple(hidden_states.shape[:2]) != (1, 1): raise ValueError("DotCache decode mode only supports batch=1 and query_len=1") token_index = self.adapter.current_token_index(cache_position) ((query_states, key_states), value_states), qkv_ms = _timed_call( lambda: self._project_qkv(hidden_states, position_embeddings), device=hidden_states.device, ) query_step = query_states[0, :, 0, :].detach().to(dtype=torch.float32) key_step = key_states[0].detach().to(dtype=torch.float32) value_step = value_states[0].detach().to(dtype=torch.float32) self.adapter.record_layer_runtime(self.layer_idx, qkv_projection_ms=qkv_ms) _, append_ms = _timed_call( lambda: self.adapter.model_kv_cache.append_step_torch( self.layer_idx, key_step, value_step, token_index, trace=self.adapter.active_trace, ) if _torch_backend_matches_device(self.adapter.backend, hidden_states.device.type) else self.adapter.model_kv_cache.append_step( self.layer_idx, key_step.cpu().numpy(), value_step.cpu().numpy(), token_index, trace=self.adapter.active_trace, ), device=hidden_states.device, ) self.adapter.append_runtime_ms_total += append_ms context_states, decode_ms = _timed_call( lambda: self.adapter.model_kv_cache.decode_layer_torch( self.layer_idx, query_step, self.adapter.q_head_to_kv_head, query_scale=float(self.base_attention.scaling), trace=self.adapter.active_trace, ) if _torch_backend_matches_device(self.adapter.backend, hidden_states.device.type) else self.adapter.model_kv_cache.decode_layer( self.layer_idx, query_step.detach().cpu().numpy(), self.adapter.q_head_to_kv_head, query_scale=float(self.base_attention.scaling), trace=self.adapter.active_trace, ), device=hidden_states.device, ) self.adapter.decode_runtime_ms_total += decode_ms def _project_output(): local_context_states = context_states if not torch.is_tensor(local_context_states): local_context_states = torch.as_tensor(local_context_states, dtype=torch.float32, device=hidden_states.device) context_tensor = local_context_states.to(dtype=hidden_states.dtype, device=hidden_states.device).unsqueeze(0) return self.base_attention.o_proj(context_tensor.reshape(1, 1, -1).contiguous()) projected_output, output_projection_ms = _timed_call(_project_output, device=hidden_states.device) self.adapter.record_layer_runtime( self.layer_idx, append_ms=append_ms, decode_ms=decode_ms, output_projection_ms=output_projection_ms, ) if self.adapter.capture_enabled: self.adapter.record_replay( LlamaReplayRecord( step_index=self.adapter.capture_step_index, layer_id=self.layer_idx, token_index=token_index, query_states=query_step.detach().cpu().numpy(), key_states=key_step[:, 0, :].detach().cpu().numpy(), value_states=value_step[:, 0, :].detach().cpu().numpy(), context_states=context_states.detach().cpu().numpy().astype(np.float32, copy=False), output_states=projected_output[0, 0].detach().to(dtype=torch.float32).cpu().numpy(), ) ) return projected_output, None class LlamaDotCacheModelAdapter: def __init__( self, model, dotcache_config: DotCacheConfig, *, backend: str = "auto", cache: PreparedPageCache | None = None, ) -> None: _require_transformers() self.model = model self.dotcache_config = dotcache_config self.backend = backend self.cache = cache if cache is not None else PreparedPageCache() self.model_kv_cache = ModelPagedKVCache( config=dotcache_config, num_hidden_layers=model.config.num_hidden_layers, num_attention_heads=model.config.num_attention_heads, num_key_value_heads=model.config.num_key_value_heads, backend=backend, cache=self.cache, ) self.q_head_to_kv_head = self.model_kv_cache.default_q_head_to_kv_head.copy() self.mode: AttentionMode = "dense" self.capture_enabled = False self.capture_step_index = -1 self.active_trace: ExecutionTrace | None = None self._pending_records: list[LlamaReplayRecord] = [] self._wrappers: list[DotCacheLlamaAttention] = [] self.append_runtime_ms_total = 0.0 self.decode_runtime_ms_total = 0.0 self.qkv_projection_ms_total = 0.0 self.output_projection_ms_total = 0.0 self.layer_runtime_profiles = [LlamaLayerRuntimeProfile(layer_id=layer_id) for layer_id in range(model.config.num_hidden_layers)] self._current_token_index_override: int | None = None self._install_wrappers() @property def device(self): return next(self.model.parameters()).device def _install_wrappers(self) -> None: for layer in self.model.model.layers[: self.model.config.num_hidden_layers]: wrapper = DotCacheLlamaAttention(layer.self_attn, self) layer.self_attn = wrapper self._wrappers.append(wrapper) def set_mode(self, mode: AttentionMode) -> None: self.mode = mode def set_capture(self, enabled: bool) -> None: self.capture_enabled = bool(enabled) def begin_capture_step(self, step_index: int) -> None: self.capture_step_index = int(step_index) self._pending_records = [] self.capture_enabled = True def end_capture_step(self) -> list[LlamaReplayRecord]: records = list(self._pending_records) self._pending_records = [] self.capture_enabled = False self.capture_step_index = -1 return records def record_replay(self, record: LlamaReplayRecord) -> None: if self.capture_step_index < 0: return self._pending_records.append(record) def current_token_index(self, cache_position) -> int: if self._current_token_index_override is not None: return self._current_token_index_override if cache_position is None: raise ValueError("cache_position is required for the Phase 5 Llama path") token_positions = cache_position.reshape(-1) if token_positions.numel() != 1: raise ValueError("Phase 5 Llama path requires a single cache_position per decode step") return int(token_positions.item()) def set_current_token_index(self, token_index: int | None) -> None: self._current_token_index_override = None if token_index is None else int(token_index) def clear(self) -> None: self.model_kv_cache.clear() self._pending_records = [] self.capture_enabled = False self.capture_step_index = -1 self.active_trace = None self._current_token_index_override = None self.reset_runtime_metrics() def reconfigure(self, dotcache_config: DotCacheConfig, *, backend: str | None = None) -> None: self.dotcache_config = dotcache_config if backend is not None: self.backend = backend self.model_kv_cache = ModelPagedKVCache( 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=self.model.config.num_key_value_heads, backend=self.backend, cache=self.cache, ) self.q_head_to_kv_head = self.model_kv_cache.default_q_head_to_kv_head.copy() self.clear() def reset_runtime_metrics(self) -> None: self.append_runtime_ms_total = 0.0 self.decode_runtime_ms_total = 0.0 self.qkv_projection_ms_total = 0.0 self.output_projection_ms_total = 0.0 for profile in self.layer_runtime_profiles: profile.reset() def record_layer_runtime( self, layer_id: int, *, qkv_projection_ms: float = 0.0, append_ms: float = 0.0, decode_ms: float = 0.0, output_projection_ms: float = 0.0, ) -> None: profile = self.layer_runtime_profiles[layer_id] if qkv_projection_ms > 0.0: profile.call_count += 1 profile.qkv_projection_ms_total += qkv_projection_ms self.qkv_projection_ms_total += qkv_projection_ms if append_ms > 0.0: profile.append_ms_total += append_ms if decode_ms > 0.0: profile.decode_ms_total += decode_ms if output_projection_ms > 0.0: profile.output_projection_ms_total += output_projection_ms self.output_projection_ms_total += output_projection_ms def runtime_profile_summary(self, *, model_forward_ms_total: float) -> dict[str, Any]: per_layer = [profile.to_dict() for profile in self.layer_runtime_profiles if profile.call_count > 0] accounted_ms_total = ( self.qkv_projection_ms_total + self.append_runtime_ms_total + self.decode_runtime_ms_total + self.output_projection_ms_total ) return { "model_forward_ms_total": float(model_forward_ms_total), "qkv_projection_ms_total": float(self.qkv_projection_ms_total), "append_runtime_ms_total": float(self.append_runtime_ms_total), "decode_runtime_ms_total": float(self.decode_runtime_ms_total), "output_projection_ms_total": float(self.output_projection_ms_total), "other_overhead_ms_total": float(max(model_forward_ms_total - accounted_ms_total, 0.0)), "per_layer": per_layer, } def load_prefill_cache(self, past_key_values, *, trace: ExecutionTrace | None = None) -> None: if _torch_backend_matches_device(self.backend, self.device.type): self.load_prefill_cache_tensors(extract_past_key_values_tensors(past_key_values), trace=trace) else: self.load_prefill_cache_arrays(extract_past_key_values_arrays(past_key_values), trace=trace) def load_prefill_cache_arrays( self, prefill_layers: Sequence[tuple[np.ndarray, np.ndarray]], *, trace: ExecutionTrace | None = None, ) -> None: if len(prefill_layers) != self.model.config.num_hidden_layers: raise ValueError("prefill_layers must align with model.config.num_hidden_layers") self.model_kv_cache.clear() for layer_idx, (layer_keys, layer_values) in enumerate(prefill_layers): self.model_kv_cache.ingest_prefill_cache(layer_idx, layer_keys, layer_values, trace=trace) self.model_kv_cache.prepare_static_pages(trace=trace) def load_prefill_cache_tensors( self, prefill_layers: Sequence[tuple[Any, Any]], *, trace: ExecutionTrace | None = None, ) -> None: if len(prefill_layers) != self.model.config.num_hidden_layers: raise ValueError("prefill_layers must align with model.config.num_hidden_layers") self.model_kv_cache.clear() for layer_idx, (layer_keys, layer_values) in enumerate(prefill_layers): self.model_kv_cache.ingest_prefill_cache_torch(layer_idx, layer_keys, layer_values, trace=trace) self.model_kv_cache.prepare_static_pages(trace=trace) @dataclass(slots=True) class LlamaDotCacheHarness: model: Any tokenizer: Any | None adapter: LlamaDotCacheModelAdapter @classmethod def from_pretrained( cls, model_id: str, dotcache_config: DotCacheConfig, *, backend: str = "auto", device: str | None = None, torch_dtype: str = "float16", ) -> "LlamaDotCacheHarness": _require_transformers() dtype = getattr(torch, torch_dtype) resolved_device = _default_model_device() if device is None else device auth_kwargs = resolve_hf_auth_kwargs() model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, **auth_kwargs) model.to(resolved_device) model.eval() tokenizer = AutoTokenizer.from_pretrained(model_id, **auth_kwargs) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id adapter = LlamaDotCacheModelAdapter(model, dotcache_config, backend=backend) return cls(model=model, tokenizer=tokenizer, adapter=adapter) def tokenize_prompt(self, prompt: str) -> tuple[Any, Any]: if self.tokenizer is None: raise ValueError("tokenizer is unavailable for text prompt input") encoded = self.tokenizer(prompt, return_tensors="pt") input_ids = encoded["input_ids"].to(self.adapter.device) attention_mask = encoded["attention_mask"].to(self.adapter.device) return input_ids, attention_mask def run_replay( self, *, prompt: str | None = None, input_ids=None, attention_mask=None, decode_steps: int = 4, ) -> dict[str, float | int]: return run_llama_replay_harness( self.model, self.adapter, prompt=prompt, input_ids=input_ids, attention_mask=attention_mask, decode_steps=decode_steps, tokenizer=self.tokenizer, ) def capture_page_traces( self, *, output_dir: str | Path, tokens_per_page: int, kinds: tuple[str, ...] = ("K", "V"), prompt: str | None = None, input_ids=None, attention_mask=None, decode_steps: int = 4, ) -> dict[str, Any]: return run_llama_page_trace_capture_harness( self.model, self.adapter, output_dir=output_dir, tokens_per_page=tokens_per_page, kinds=kinds, prompt=prompt, input_ids=input_ids, attention_mask=attention_mask, tokenizer=self.tokenizer, decode_steps=decode_steps, ) def generate_greedy( self, *, prompt: str | None = None, input_ids=None, attention_mask=None, max_new_tokens: int = 8, profile: bool = False, ) -> dict[str, Any]: return run_llama_generation_harness( self.model, self.adapter, prompt=prompt, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, tokenizer=self.tokenizer, profile=profile, ) def evaluate_loss( self, *, prompt: str | None = None, input_ids=None, attention_mask=None, prefix_length: int, eval_steps: int, ) -> dict[str, Any]: return run_llama_loss_harness( self.model, self.adapter, prompt=prompt, input_ids=input_ids, attention_mask=attention_mask, prefix_length=prefix_length, eval_steps=eval_steps, tokenizer=self.tokenizer, ) def _prefill_prompt( model, adapter: LlamaDotCacheModelAdapter, input_ids, attention_mask, ): adapter.set_mode("dense") adapter.set_capture(False) outputs = _run_inference(lambda: model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)) if _torch_backend_matches_device(adapter.backend, input_ids.device.type): prefill_layers = extract_past_key_values_tensors(outputs.past_key_values) else: prefill_layers = extract_past_key_values_arrays(outputs.past_key_values) first_generated_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) return outputs, prefill_layers, first_generated_token def _run_dense_greedy_decode( model, adapter: LlamaDotCacheModelAdapter, *, input_ids, attention_mask, max_new_tokens: int, capture: bool, ) -> dict[str, Any]: (prefill_outputs, prefill_layers, first_generated_token), prefill_ms = _timed_call( lambda: _prefill_prompt(model, adapter, input_ids, attention_mask), device=input_ids.device, ) if max_new_tokens <= 0: return { "prefill_layers": prefill_layers, "generated_ids": [], "decode_inputs": [], "step_logits": [], "capture_records": [], "prefill_outputs": prefill_outputs, "prefill_ms": prefill_ms, } generated_ids = [int(first_generated_token.item())] if max_new_tokens == 1: return { "prefill_layers": prefill_layers, "generated_ids": generated_ids, "decode_inputs": [], "step_logits": [], "capture_records": [], "prefill_outputs": prefill_outputs, "prefill_ms": prefill_ms, } adapter.set_mode("dense") adapter.set_capture(False) past_key_values = prefill_outputs.past_key_values current_input_ids = first_generated_token current_attention_mask = torch.cat( [attention_mask, torch.ones((1, 1), dtype=attention_mask.dtype, device=attention_mask.device)], dim=1, ) cache_position = torch.tensor([input_ids.shape[1]], dtype=torch.long, device=input_ids.device) decode_inputs: list[Any] = [] step_logits: list[np.ndarray] = [] capture_records: list[list[LlamaReplayRecord]] = [] dense_decode_ms_total = 0.0 for step_index in range(max_new_tokens - 1): decode_inputs.append(current_input_ids.detach().clone()) if capture: adapter.begin_capture_step(step_index) adapter.set_current_token_index(int(input_ids.shape[1] + step_index)) try: outputs, step_ms = _timed_call( lambda: _run_inference( lambda: model( input_ids=current_input_ids, attention_mask=current_attention_mask, past_key_values=past_key_values, use_cache=True, cache_position=cache_position, position_ids=cache_position.unsqueeze(0), ) ), device=input_ids.device, ) dense_decode_ms_total += step_ms finally: adapter.set_current_token_index(None) if capture: capture_records.append(adapter.end_capture_step()) past_key_values = outputs.past_key_values logits = outputs.logits[:, -1, :].detach().to(dtype=torch.float32).cpu().numpy() step_logits.append(logits) current_input_ids = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) generated_ids.append(int(current_input_ids.item())) current_attention_mask = torch.cat( [current_attention_mask, torch.ones((1, 1), dtype=current_attention_mask.dtype, device=current_attention_mask.device)], dim=1, ) cache_position = cache_position + 1 return { "prefill_layers": prefill_layers, "generated_ids": generated_ids, "decode_inputs": decode_inputs, "step_logits": step_logits, "capture_records": capture_records, "prefill_outputs": prefill_outputs, "prefill_ms": prefill_ms, "dense_decode_ms_total": dense_decode_ms_total, } def _run_dotcache_decode_inputs( model, adapter: LlamaDotCacheModelAdapter, *, input_ids, attention_mask, prefill_layers: Sequence[tuple[Any, Any]], decode_inputs: Sequence[Any], profile_backend: bool = False, ) -> dict[str, Any]: if prefill_layers and torch.is_tensor(prefill_layers[0][0]): adapter.load_prefill_cache_tensors(prefill_layers) else: adapter.load_prefill_cache_arrays(prefill_layers) _prewarm_torch_decode_layers(adapter, device=input_ids.device) adapter.set_mode("dotcache") adapter.reset_runtime_metrics() use_attention_mask = not _can_skip_decode_attention_mask(attention_mask) current_attention_mask = attention_mask if use_attention_mask else None step_logits: list[np.ndarray] = [] decode_ms_total = 0.0 trace_total = ExecutionTrace(capture_timings=profile_backend) for offset, decode_input in enumerate(decode_inputs): cache_position = torch.tensor([input_ids.shape[1] + offset], dtype=torch.long, device=input_ids.device) step_trace = ExecutionTrace(capture_timings=profile_backend) adapter.active_trace = step_trace adapter.set_current_token_index(int(input_ids.shape[1] + offset)) try: outputs, step_ms = _timed_call( lambda: _run_inference( lambda: model( input_ids=decode_input, attention_mask=current_attention_mask, use_cache=False, cache_position=cache_position, position_ids=cache_position.unsqueeze(0), ) ), device=input_ids.device, ) finally: adapter.active_trace = None adapter.set_current_token_index(None) decode_ms_total += step_ms trace_total.merge(step_trace) step_logits.append(outputs.logits[:, -1, :].detach().to(dtype=torch.float32).cpu().numpy()) return { "decode_ms_total": decode_ms_total, "append_runtime_ms_total": adapter.append_runtime_ms_total, "decode_runtime_ms_total": adapter.decode_runtime_ms_total, "step_logits": step_logits, "trace": trace_total, } def _run_dense_decode_inputs( model, adapter: LlamaDotCacheModelAdapter, *, input_ids, attention_mask, prefill_outputs, decode_inputs: Sequence[Any], ) -> dict[str, Any]: adapter.set_mode("dense") adapter.set_capture(False) past_key_values = prefill_outputs.past_key_values current_attention_mask = attention_mask cache_position = torch.tensor([input_ids.shape[1]], dtype=torch.long, device=input_ids.device) step_logits: list[np.ndarray] = [] decode_ms_total = 0.0 for decode_input in decode_inputs: start = time.perf_counter() outputs = _run_inference( lambda: model( input_ids=decode_input, attention_mask=current_attention_mask, past_key_values=past_key_values, use_cache=True, cache_position=cache_position, position_ids=cache_position.unsqueeze(0), ) ) decode_ms_total += (time.perf_counter() - start) * 1000.0 step_logits.append(outputs.logits[:, -1, :].detach().to(dtype=torch.float32).cpu().numpy()) past_key_values = outputs.past_key_values if current_attention_mask is not None: current_attention_mask = torch.cat( [current_attention_mask, torch.ones((1, 1), dtype=current_attention_mask.dtype, device=current_attention_mask.device)], dim=1, ) cache_position = cache_position + 1 return { "decode_ms_total": decode_ms_total, "step_logits": step_logits, } def _run_dotcache_greedy_decode( model, adapter: LlamaDotCacheModelAdapter, *, input_ids, attention_mask, prefill_layers: Sequence[tuple[Any, Any]], first_generated_token, max_new_tokens: int, profile_backend: bool = False, ) -> dict[str, Any]: if prefill_layers and torch.is_tensor(prefill_layers[0][0]): adapter.load_prefill_cache_tensors(prefill_layers) else: adapter.load_prefill_cache_arrays(prefill_layers) _prewarm_torch_decode_layers(adapter, device=input_ids.device) adapter.set_mode("dotcache") adapter.reset_runtime_metrics() generated_ids = [int(first_generated_token.item())] if max_new_tokens <= 1: return { "generated_ids": generated_ids, "decode_ms_total": 0.0, "append_runtime_ms_total": 0.0, "decode_runtime_ms_total": 0.0, "step_count": 0, "trace": ExecutionTrace(capture_timings=profile_backend), } current_input_ids = first_generated_token use_attention_mask = not _can_skip_decode_attention_mask(attention_mask) current_attention_mask = ( torch.cat( [attention_mask, torch.ones((1, 1), dtype=attention_mask.dtype, device=attention_mask.device)], dim=1, ) if use_attention_mask else None ) cache_position = torch.tensor([input_ids.shape[1]], dtype=torch.long, device=input_ids.device) current_token_index = int(input_ids.shape[1]) step_count = 0 decode_ms_total = 0.0 trace_total = ExecutionTrace(capture_timings=profile_backend) for _ in range(max_new_tokens - 1): step_trace = ExecutionTrace(capture_timings=profile_backend) adapter.active_trace = step_trace adapter.set_current_token_index(current_token_index) try: outputs, step_ms = _timed_call( lambda: _run_inference( lambda: model( input_ids=current_input_ids, attention_mask=current_attention_mask, use_cache=False, cache_position=cache_position, position_ids=cache_position.unsqueeze(0), ) ), device=input_ids.device, ) finally: adapter.active_trace = None adapter.set_current_token_index(None) decode_ms_total += step_ms trace_total.merge(step_trace) current_input_ids = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) generated_ids.append(int(current_input_ids.item())) step_count += 1 if current_attention_mask is not None: current_attention_mask = torch.cat( [current_attention_mask, torch.ones((1, 1), dtype=current_attention_mask.dtype, device=current_attention_mask.device)], dim=1, ) cache_position = cache_position + 1 current_token_index += 1 return { "generated_ids": generated_ids, "decode_ms_total": decode_ms_total, "append_runtime_ms_total": adapter.append_runtime_ms_total, "decode_runtime_ms_total": adapter.decode_runtime_ms_total, "step_count": step_count, "trace": trace_total, } def _aggregate_query_states_by_kv_head( query_states: np.ndarray, q_head_to_kv_head: np.ndarray, *, num_key_value_heads: int, ) -> np.ndarray: queries = np.asarray(query_states, dtype=np.float32) mapping = np.asarray(q_head_to_kv_head, dtype=np.int32) if queries.ndim != 2: raise ValueError("query_states must have shape [query_heads, head_dim]") if mapping.ndim != 1 or mapping.shape[0] != queries.shape[0]: raise ValueError("q_head_to_kv_head must have shape [query_heads]") kv_queries = np.zeros((int(num_key_value_heads), int(queries.shape[1])), dtype=np.float32) counts = np.zeros((int(num_key_value_heads),), dtype=np.int32) for q_head_id, kv_head_id in enumerate(mapping.tolist()): if kv_head_id < 0 or kv_head_id >= int(num_key_value_heads): raise ValueError("q_head_to_kv_head contains an out-of-range kv head") kv_queries[kv_head_id] += queries[q_head_id] counts[kv_head_id] += 1 counts = np.maximum(counts, 1) kv_queries /= counts[:, None].astype(np.float32) return kv_queries def _build_page_traces_from_streams( streams: dict[tuple[int, int, str], list[tuple[int, np.ndarray, np.ndarray | None]]], *, max_token_index: int, tokens_per_page: int, source: str, stage: str, ) -> list[PageTraceRecord]: page_traces: list[PageTraceRecord] = [] for (layer_id, kv_head_id, kind), entries in sorted(streams.items()): entries.sort(key=lambda item: item[0]) for offset in range(0, len(entries), int(tokens_per_page)): chunk = entries[offset : offset + int(tokens_per_page)] token_indices = [token_index for token_index, _, _ in chunk] values = np.stack([value for _, value, _ in chunk], axis=0).astype(np.float32, copy=False) queries = [query_vector for _, _, query_vector in chunk if query_vector is not None] query = None if queries: query = np.mean(np.stack(queries, axis=0), axis=0, dtype=np.float32).astype(np.float32, copy=False) page_traces.append( PageTraceRecord( source=source, kind=kind, # type: ignore[arg-type] layer_id=layer_id, kv_head_id=kv_head_id, token_start=int(token_indices[0]), token_age=max(max_token_index - int(token_indices[-1]), 0), values=values, query=query, notes=[ f"stage={stage}", "query_aggregation=mean_mapped_q_heads" if query is not None else "query_aggregation=none", f"token_indices={token_indices[0]}..{token_indices[-1]}", ], ) ) return page_traces def build_llama_page_trace_records( per_step_records: list[list[LlamaReplayRecord]], *, q_head_to_kv_head: np.ndarray, tokens_per_page: int, kinds: tuple[str, ...] = ("K", "V"), source: str = "llama_dense_capture", ) -> list[PageTraceRecord]: if int(tokens_per_page) <= 0: raise ValueError("tokens_per_page must be positive") normalized_kinds = tuple(str(kind).upper() for kind in kinds) invalid_kinds = [kind for kind in normalized_kinds if kind not in {"K", "V"}] if invalid_kinds: raise ValueError(f"unsupported capture kinds: {invalid_kinds}") streams: dict[tuple[int, int, str], list[tuple[int, np.ndarray, np.ndarray | None]]] = {} max_token_index = -1 for step_records in per_step_records: for record in step_records: max_token_index = max(max_token_index, int(record.token_index)) kv_head_count = int(record.key_states.shape[0]) kv_queries = _aggregate_query_states_by_kv_head( record.query_states, q_head_to_kv_head, num_key_value_heads=kv_head_count, ) for kv_head_id in range(kv_head_count): if "K" in normalized_kinds: streams.setdefault((int(record.layer_id), kv_head_id, "K"), []).append( ( int(record.token_index), np.asarray(record.key_states[kv_head_id], dtype=np.float32), np.asarray(kv_queries[kv_head_id], dtype=np.float32), ) ) if "V" in normalized_kinds: streams.setdefault((int(record.layer_id), kv_head_id, "V"), []).append( ( int(record.token_index), np.asarray(record.value_states[kv_head_id], dtype=np.float32), np.asarray(kv_queries[kv_head_id], dtype=np.float32), ) ) return _build_page_traces_from_streams( streams, max_token_index=max_token_index, tokens_per_page=tokens_per_page, source=source, stage="decode", ) def build_llama_prefill_page_trace_records( prefill_layers: Sequence[tuple[Any, Any]], *, tokens_per_page: int, kinds: tuple[str, ...] = ("K", "V"), source: str = "llama_dense_capture", max_token_index: int | None = None, ) -> list[PageTraceRecord]: if int(tokens_per_page) <= 0: raise ValueError("tokens_per_page must be positive") normalized_kinds = tuple(str(kind).upper() for kind in kinds) invalid_kinds = [kind for kind in normalized_kinds if kind not in {"K", "V"}] if invalid_kinds: raise ValueError(f"unsupported capture kinds: {invalid_kinds}") streams: dict[tuple[int, int, str], list[tuple[int, np.ndarray, np.ndarray | None]]] = {} resolved_max_token_index = -1 if max_token_index is None else int(max_token_index) for layer_id, (layer_keys, layer_values) in enumerate(prefill_layers): key_array = _tensor_to_float32_numpy(layer_keys) value_array = _tensor_to_float32_numpy(layer_values) if key_array.ndim != 4 or value_array.ndim != 4 or key_array.shape[0] != 1 or value_array.shape[0] != 1: raise ValueError("prefill layers must have shape [1, kv_heads, seq_len, head_dim]") if key_array.shape[:3] != value_array.shape[:3]: raise ValueError("prefill key and value tensors must align on batch, kv_head, and seq_len") _, kv_head_count, seq_len, _ = key_array.shape resolved_max_token_index = max(resolved_max_token_index, int(seq_len) - 1) for kv_head_id in range(int(kv_head_count)): for token_index in range(int(seq_len)): if "K" in normalized_kinds: streams.setdefault((int(layer_id), kv_head_id, "K"), []).append( ( int(token_index), np.asarray(key_array[0, kv_head_id, token_index], dtype=np.float32), None, ) ) if "V" in normalized_kinds: streams.setdefault((int(layer_id), kv_head_id, "V"), []).append( ( int(token_index), np.asarray(value_array[0, kv_head_id, token_index], dtype=np.float32), None, ) ) return _build_page_traces_from_streams( streams, max_token_index=resolved_max_token_index, tokens_per_page=tokens_per_page, source=source, stage="prefill", ) def export_llama_page_traces( per_step_records: list[list[LlamaReplayRecord]], *, q_head_to_kv_head: np.ndarray, output_dir: str | Path, tokens_per_page: int, kinds: tuple[str, ...] = ("K", "V"), source: str = "llama_dense_capture", prefill_layers: Sequence[tuple[Any, Any]] | None = None, prefill_token_count: int | None = None, ) -> dict[str, Any]: output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) page_traces = build_llama_page_trace_records( per_step_records, q_head_to_kv_head=q_head_to_kv_head, tokens_per_page=tokens_per_page, kinds=kinds, source=source, ) if prefill_layers: prefill_length = max(int(prefill_token_count or 0), 0) max_token_index = max(prefill_length - 1 + len(per_step_records), 0) page_traces = build_llama_prefill_page_trace_records( prefill_layers, tokens_per_page=tokens_per_page, kinds=kinds, source=source, max_token_index=max_token_index, ) + page_traces trace_paths: list[str] = [] counts_by_kind: dict[str, int] = {} counts_by_layer: dict[str, int] = {} counts_by_stage: dict[str, int] = {} for index, trace in enumerate(page_traces): stage = "unknown" for note in trace.notes: if note.startswith("stage="): stage = note.split("=", 1)[1] break trace_name = ( f"{stage}_layer{trace.layer_id:02d}_kv{trace.kv_head_id:02d}_{trace.kind.lower()}_" f"t{trace.token_start:06d}_n{trace.token_count:03d}_{index:04d}.npz" ) target = output_path / trace_name save_page_trace(trace, target) trace_paths.append(str(target)) counts_by_kind[trace.kind] = counts_by_kind.get(trace.kind, 0) + 1 counts_by_stage[stage] = counts_by_stage.get(stage, 0) + 1 layer_key = str(trace.layer_id) counts_by_layer[layer_key] = counts_by_layer.get(layer_key, 0) + 1 manifest = { "output_dir": str(output_path), "page_trace_count": len(page_traces), "page_trace_paths": trace_paths, "page_trace_counts_by_kind": dict(sorted(counts_by_kind.items())), "page_trace_counts_by_stage": dict(sorted(counts_by_stage.items())), "page_trace_counts_by_layer": dict(sorted(counts_by_layer.items())), "tokens_per_page": int(tokens_per_page), "kinds": list(kinds), "source": source, } (output_path / "manifest.json").write_text(json.dumps(manifest, sort_keys=True, indent=2) + "\n", encoding="utf-8") return manifest def run_llama_page_trace_capture_harness( model, adapter: LlamaDotCacheModelAdapter, *, output_dir: str | Path, tokens_per_page: int, kinds: tuple[str, ...] = ("K", "V"), prompt: str | None = None, input_ids=None, attention_mask=None, tokenizer=None, decode_steps: int = 4, ) -> dict[str, Any]: _require_transformers() if prompt is not None: if tokenizer is None: raise ValueError("tokenizer is required when prompt text is provided") encoded = tokenizer(prompt, return_tensors="pt") input_ids = encoded["input_ids"] attention_mask = encoded["attention_mask"] input_ids = _normalize_input_ids(input_ids, device=adapter.device) attention_mask = _ensure_attention_mask(input_ids, attention_mask, device=adapter.device) dense_result = _run_dense_greedy_decode( model, adapter, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=decode_steps + 1, capture=True, ) result: dict[str, Any] = { "runtime_mode": "dense_llama_page_trace_capture", "prompt_length": int(input_ids.shape[1]), "decode_steps": max(0, int(decode_steps)), "prefill_ms": float(dense_result["prefill_ms"]), "dense_decode_ms_per_step": float( dense_result.get("dense_decode_ms_total", 0.0) / max(max(int(decode_steps), 0), 1) ), "capture_record_count": int(sum(len(step_records) for step_records in dense_result["capture_records"])), "capture_step_count": int(len(dense_result["capture_records"])), "capture_layer_count": int( len( { int(record.layer_id) for step_records in dense_result["capture_records"] for record in step_records } ) ), } result.update( export_llama_page_traces( dense_result["capture_records"], q_head_to_kv_head=adapter.q_head_to_kv_head, output_dir=output_dir, tokens_per_page=tokens_per_page, kinds=kinds, prefill_layers=dense_result["prefill_layers"], prefill_token_count=int(input_ids.shape[1]), ) ) return result def run_llama_replay_harness( model, adapter: LlamaDotCacheModelAdapter, *, prompt: str | None = None, input_ids=None, attention_mask=None, decode_steps: int = 4, tokenizer=None, ) -> dict[str, float | int]: _require_transformers() if prompt is not None: if tokenizer is None: raise ValueError("tokenizer is required when prompt text is provided") encoded = tokenizer(prompt, return_tensors="pt") input_ids = encoded["input_ids"] attention_mask = encoded["attention_mask"] input_ids = _normalize_input_ids(input_ids, device=adapter.device) attention_mask = _ensure_attention_mask(input_ids, attention_mask, device=adapter.device) dense_result = _run_dense_greedy_decode( model, adapter, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=decode_steps + 1, capture=True, ) replay_cache = ModelPagedKVCache( config=adapter.dotcache_config, num_hidden_layers=model.config.num_hidden_layers, num_attention_heads=model.config.num_attention_heads, num_key_value_heads=model.config.num_key_value_heads, backend=adapter.backend, cache=PreparedPageCache(), ) for layer_idx, (layer_keys, layer_values) in enumerate(dense_result["prefill_layers"]): if torch.is_tensor(layer_keys): replay_cache.ingest_prefill_cache_torch(layer_idx, layer_keys, layer_values) else: replay_cache.ingest_prefill_cache(layer_idx, layer_keys, layer_values) replay_context_max_abs = 0.0 replay_context_max_rel = 0.0 for step_records in dense_result["capture_records"]: for record in step_records: replay_cache.append_step( record.layer_id, record.key_states[:, None, :], record.value_states[:, None, :], record.token_index, ) replay_context = replay_cache.decode_layer(record.layer_id, record.query_states, adapter.q_head_to_kv_head) delta = np.abs(replay_context - record.context_states) denom = np.maximum(np.abs(record.context_states), 1e-8) replay_context_max_abs = max(replay_context_max_abs, float(np.max(delta))) replay_context_max_rel = max(replay_context_max_rel, float(np.max(delta / denom))) dotcache_teacher_forced = _run_dotcache_decode_inputs( model, adapter, input_ids=input_ids, attention_mask=attention_mask, prefill_layers=dense_result["prefill_layers"], decode_inputs=dense_result["decode_inputs"], ) dense_logits = np.stack(dense_result["step_logits"], axis=0) if dense_result["step_logits"] else np.zeros((0, 1)) dotcache_logits = ( np.stack(dotcache_teacher_forced["step_logits"], axis=0) if dotcache_teacher_forced["step_logits"] else np.zeros((0, 1)) ) if dense_logits.size == 0: max_abs_logit_drift = 0.0 max_rel_logit_drift = 0.0 else: logit_delta = np.abs(dotcache_logits - dense_logits) logit_denom = np.maximum(np.abs(dense_logits), 1e-8) max_abs_logit_drift = float(np.max(logit_delta)) max_rel_logit_drift = float(np.max(logit_delta / logit_denom)) return { "decode_steps": max(0, decode_steps), "replay_context_max_abs_error": replay_context_max_abs, "replay_context_max_rel_error": replay_context_max_rel, "teacher_forced_logit_max_abs_error": max_abs_logit_drift, "teacher_forced_logit_max_rel_error": max_rel_logit_drift, } def run_llama_generation_harness( model, adapter: LlamaDotCacheModelAdapter, *, prompt: str | None = None, input_ids=None, attention_mask=None, max_new_tokens: int = 8, tokenizer=None, profile: bool = False, ) -> dict[str, Any]: _require_transformers() if prompt is not None: if tokenizer is None: raise ValueError("tokenizer is required when prompt text is provided") encoded = tokenizer(prompt, return_tensors="pt") input_ids = encoded["input_ids"] attention_mask = encoded["attention_mask"] input_ids = _normalize_input_ids(input_ids, device=adapter.device) attention_mask = _ensure_attention_mask(input_ids, attention_mask, device=adapter.device) dense_result = _run_dense_greedy_decode( model, adapter, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, capture=False, ) prefill_trace = ExecutionTrace(capture_timings=profile) prefill_cuda_baseline = _begin_cuda_memory_region(input_ids.device) if profile else None _, prefill_ingest_ms = _timed_call( lambda: adapter.load_prefill_cache_tensors(dense_result["prefill_layers"], trace=prefill_trace) if dense_result["prefill_layers"] and torch.is_tensor(dense_result["prefill_layers"][0][0]) else adapter.load_prefill_cache_arrays(dense_result["prefill_layers"], trace=prefill_trace), device=input_ids.device, ) prefill_cuda_stats = _end_cuda_memory_region(input_ids.device, prefill_cuda_baseline) if profile else {} if max_new_tokens <= 1: generated_ids = dense_result["generated_ids"] decode_ms_per_step = 0.0 append_ms_per_step = 0.0 decode_trace = ExecutionTrace(capture_timings=profile) step_count = 0 append_runtime_ms_per_step = 0.0 decode_runtime_ms_per_step = 0.0 dotcache_profile = adapter.runtime_profile_summary(model_forward_ms_total=0.0) if profile else None dotcache_cuda_stats: dict[str, int] = {} else: dotcache_cuda_baseline = _begin_cuda_memory_region(input_ids.device) if profile else None dotcache_result = _run_dotcache_greedy_decode( model, adapter, input_ids=input_ids, attention_mask=attention_mask, prefill_layers=dense_result["prefill_layers"], first_generated_token=torch.as_tensor([[dense_result["generated_ids"][0]]], dtype=torch.long, device=input_ids.device), max_new_tokens=max_new_tokens, profile_backend=profile, ) step_count = int(dotcache_result["step_count"]) decode_trace = dotcache_result["trace"] generated_ids = dotcache_result["generated_ids"] decode_ms_per_step = dotcache_result["decode_ms_total"] / max(step_count, 1) append_runtime_ms_per_step = dotcache_result["append_runtime_ms_total"] / max(step_count, 1) decode_runtime_ms_per_step = dotcache_result["decode_runtime_ms_total"] / max(step_count, 1) append_ms_per_step = append_runtime_ms_per_step dotcache_profile = adapter.runtime_profile_summary(model_forward_ms_total=float(dotcache_result["decode_ms_total"])) if profile else None dotcache_cuda_stats = _end_cuda_memory_region(input_ids.device, dotcache_cuda_baseline) if profile else {} dense_generated_ids = dense_result["generated_ids"] dense_step_count = max(max_new_tokens - 1, 0) dense_decode_ms_per_step = float(dense_result["dense_decode_ms_total"] / max(dense_step_count, 1)) if dense_step_count > 0 else 0.0 dense_prefill_kv_cache_bytes = _prefill_layer_nbytes(dense_result["prefill_layers"]) dense_final_kv_cache_bytes = _dense_kv_bytes_after_decode( dense_result["prefill_layers"], generated_token_count=len(dense_generated_ids), ) agreement_prefix = sum( int(lhs == rhs) for lhs, rhs in zip(generated_ids, dense_generated_ids, strict=False) ) agreement_rate = agreement_prefix / max(min(len(generated_ids), len(dense_generated_ids)), 1) teacher_forced = _run_dotcache_decode_inputs( model, adapter, input_ids=input_ids, attention_mask=attention_mask, prefill_layers=dense_result["prefill_layers"], decode_inputs=dense_result["decode_inputs"], ) if dense_result["step_logits"]: dense_logits = np.stack(dense_result["step_logits"], axis=0) forced_logits = np.stack(teacher_forced["step_logits"], axis=0) logit_delta = np.abs(forced_logits - dense_logits) logit_denom = np.maximum(np.abs(dense_logits), 1e-8) max_abs_logit_drift = float(np.max(logit_delta)) max_rel_logit_drift = float(np.max(logit_delta / logit_denom)) else: max_abs_logit_drift = 0.0 max_rel_logit_drift = 0.0 resident_byte_summary = adapter.model_kv_cache.resident_byte_summary() result: dict[str, Any] = { "prompt_length": int(input_ids.shape[1]), "decode_steps": max(max_new_tokens - 1, 0), "prefill_ms": float(dense_result["prefill_ms"]), "dense_decode_ms_per_step": dense_decode_ms_per_step, "dense_prefill_kv_cache_bytes": dense_prefill_kv_cache_bytes, "dense_final_kv_cache_bytes": dense_final_kv_cache_bytes, "dense_generated_ids": dense_generated_ids, "dotcache_generated_ids": generated_ids, "greedy_token_agreement_rate": float(agreement_rate), "prefill_cache_ingest_host_to_device_bytes": prefill_trace.host_to_device_bytes, "prefill_cache_ingest_ms": float(prefill_ingest_ms), "decode_ms_per_step": float(decode_ms_per_step), "append_ms_per_step": float(append_ms_per_step), "append_runtime_ms_per_step": float(append_runtime_ms_per_step), "decode_runtime_ms_per_step": float(decode_runtime_ms_per_step), "resident_bytes": int(resident_byte_summary["resident_bytes"]), "kv_resident_bytes": int(resident_byte_summary["kv_resident_bytes"]), "prepared_page_cache_resident_bytes": int(resident_byte_summary["prepared_page_cache_resident_bytes"]), "direct_page_resident_bytes": int(resident_byte_summary["direct_page_resident_bytes"]), "tail_resident_bytes": int(resident_byte_summary["tail_resident_bytes"]), "prepared_chunk_cache_budget_bytes": int(resident_byte_summary["prepared_chunk_cache_budget_bytes"]), "prepared_chunk_resident_bytes": int(resident_byte_summary["prepared_chunk_resident_bytes"]), "dotcache_vs_dense_kv_bytes_ratio": float( resident_byte_summary["kv_resident_bytes"] / max(dense_final_kv_cache_bytes, 1) ), "dotcache_vs_dense_total_resident_bytes_ratio": float( resident_byte_summary["resident_bytes"] / max(dense_final_kv_cache_bytes, 1) ), "dotcache_vs_dense_decode_speedup": float(dense_decode_ms_per_step / max(decode_ms_per_step, 1e-8)) if decode_ms_per_step > 0.0 else 0.0, "decode_host_to_device_bytes_per_step": decode_trace.host_to_device_bytes / max(step_count, 1), "prefill_prepare_ms": float(prefill_trace.prepare_ms_total), "decode_prepare_ms_per_step": float(decode_trace.prepare_ms_total / max(step_count, 1)), "decode_score_ms_per_step": float(decode_trace.score_ms_total / max(step_count, 1)), "decode_softmax_ms_per_step": float(decode_trace.softmax_ms_total / max(step_count, 1)), "decode_mix_ms_per_step": float(decode_trace.mix_ms_total / max(step_count, 1)), "decode_unpack_ms_per_step": float(decode_trace.unpack_ms_total / max(step_count, 1)), "decode_fwht_ms_per_step": float(decode_trace.fwht_ms_total / max(step_count, 1)), "teacher_forced_logit_max_abs_error": max_abs_logit_drift, "teacher_forced_logit_max_rel_error": max_rel_logit_drift, } result.update(adapter.model_kv_cache.page_mode_summary()) if tokenizer is not None: result["dense_text"] = tokenizer.decode(dense_generated_ids, skip_special_tokens=True) result["dotcache_text"] = tokenizer.decode(generated_ids, skip_special_tokens=True) if profile: result["profile"] = { "device_type": input_ids.device.type, "prefill_cache_ingest": { "ms_total": float(prefill_ingest_ms), "host_to_device_bytes": int(prefill_trace.host_to_device_bytes), "trace": prefill_trace.to_dict(), **prefill_cuda_stats, }, "dotcache_decode": { **({} if dotcache_profile is None else dotcache_profile), "step_count": int(step_count), "host_to_device_bytes_total": int(decode_trace.host_to_device_bytes), "host_to_device_bytes_per_step": float(decode_trace.host_to_device_bytes / max(step_count, 1)), "trace": decode_trace.to_dict(), **dotcache_cuda_stats, }, } return result def run_llama_loss_harness( model, adapter: LlamaDotCacheModelAdapter, *, prompt: str | None = None, input_ids=None, attention_mask=None, prefix_length: int, eval_steps: int, tokenizer=None, ) -> dict[str, Any]: _require_transformers() if prompt is not None: if tokenizer is None: raise ValueError("tokenizer is required when prompt text is provided") encoded = tokenizer(prompt, return_tensors="pt") input_ids = encoded["input_ids"] attention_mask = encoded["attention_mask"] input_ids = _normalize_input_ids(input_ids, device=adapter.device) attention_mask = _ensure_attention_mask(input_ids, attention_mask, device=adapter.device) if prefix_length <= 0 or prefix_length >= int(input_ids.shape[1]): raise ValueError("prefix_length must be in [1, sequence_length)") available_eval_steps = int(input_ids.shape[1]) - prefix_length if eval_steps <= 0 or eval_steps > available_eval_steps: raise ValueError("eval_steps must be positive and fit inside the provided sequence after prefix_length") prefix_input_ids = input_ids[:, :prefix_length] prefix_attention_mask = attention_mask[:, :prefix_length] continuation_ids = input_ids[:, prefix_length : prefix_length + eval_steps] decode_inputs = [continuation_ids[:, index : index + 1] for index in range(max(eval_steps - 1, 0))] prefill_start = time.perf_counter() prefill_outputs, prefill_layers, _ = _prefill_prompt(model, adapter, prefix_input_ids, prefix_attention_mask) prefill_ms = (time.perf_counter() - prefill_start) * 1000.0 dense_prefill_logits = prefill_outputs.logits[:, -1, :].detach().to(dtype=torch.float32).cpu().numpy() dense_decode = _run_dense_decode_inputs( model, adapter, input_ids=prefix_input_ids, attention_mask=prefix_attention_mask, prefill_outputs=prefill_outputs, decode_inputs=decode_inputs, ) dotcache_decode = _run_dotcache_decode_inputs( model, adapter, input_ids=prefix_input_ids, attention_mask=prefix_attention_mask, prefill_layers=prefill_layers, decode_inputs=decode_inputs, ) dense_logits_list = [dense_prefill_logits, *dense_decode["step_logits"]] dotcache_logits_list = [dense_prefill_logits, *dotcache_decode["step_logits"]] dense_logits = np.concatenate(dense_logits_list, axis=0).astype(np.float32, copy=False) dotcache_logits = np.concatenate(dotcache_logits_list, axis=0).astype(np.float32, copy=False) target_tokens = continuation_ids[0, : dense_logits.shape[0]].detach().cpu().numpy().astype(np.int64, copy=False) def _loss_metrics(logits: np.ndarray) -> tuple[float, float, np.ndarray]: max_logits = np.max(logits, axis=-1, keepdims=True) stabilized = logits - max_logits log_probs = stabilized - np.log(np.sum(np.exp(stabilized), axis=-1, keepdims=True)) token_losses = -log_probs[np.arange(target_tokens.shape[0]), target_tokens] mean_loss = float(np.mean(token_losses)) perplexity = float(np.exp(min(mean_loss, 50.0))) predictions = np.argmax(logits, axis=-1).astype(np.int64, copy=False) return mean_loss, perplexity, predictions dense_loss, dense_perplexity, dense_predictions = _loss_metrics(dense_logits) dotcache_loss, dotcache_perplexity, dotcache_predictions = _loss_metrics(dotcache_logits) token_agreement = float(np.mean((dense_predictions == dotcache_predictions).astype(np.float32))) target_agreement = float(np.mean((dotcache_predictions == target_tokens).astype(np.float32))) logit_delta = np.abs(dotcache_logits - dense_logits) logit_denom = np.maximum(np.abs(dense_logits), 1e-8) result: dict[str, Any] = { "sequence_length": int(input_ids.shape[1]), "prefix_length": int(prefix_length), "eval_steps": int(eval_steps), "prefill_ms": float(prefill_ms), "dense_decode_ms_per_step": float(dense_decode["decode_ms_total"] / max(len(decode_inputs), 1)), "dotcache_decode_ms_per_step": float(dotcache_decode["decode_ms_total"] / max(len(decode_inputs), 1)), "dotcache_append_runtime_ms_per_step": float(dotcache_decode["append_runtime_ms_total"] / max(len(decode_inputs), 1)), "dotcache_decode_runtime_ms_per_step": float(dotcache_decode["decode_runtime_ms_total"] / max(len(decode_inputs), 1)), "dense_teacher_forced_loss": dense_loss, "dense_teacher_forced_perplexity": dense_perplexity, "dotcache_teacher_forced_loss": dotcache_loss, "dotcache_teacher_forced_perplexity": dotcache_perplexity, "teacher_forced_loss_delta": float(dotcache_loss - dense_loss), "teacher_forced_perplexity_ratio": float(dotcache_perplexity / max(dense_perplexity, 1e-8)), "teacher_forced_token_agreement_rate": token_agreement, "teacher_forced_target_match_rate": target_agreement, "teacher_forced_logit_max_abs_error": float(np.max(logit_delta)), "teacher_forced_logit_max_rel_error": float(np.max(logit_delta / logit_denom)), "prefill_cache_ingest_host_to_device_bytes": dotcache_decode["trace"].host_to_device_bytes, } result.update(adapter.model_kv_cache.page_mode_summary()) return result