"""
Inference Engine - Core forward pass loop with adaptive offloading.

Supports both GPU and CPU-only operation:

GPU mode:
  - 3-level async prefetch pipeline
  - BF16 on Ampere+ (sm80+), FP16 on older GPUs
  - Smart layer caching (hot layers stay in VRAM)

CPU mmap mode (4 GB RAM, full accuracy):
  - OS page cache manages weight RAM automatically
  - FP16 weights stored in SafeTensors mmap
  - BF16 compute on modern CPUs (AVX-512 BF16), else FP32
  - Explicit GC after each layer to release Python references
  - No quantization needed — full model accuracy

The engine implements the full autoregressive generation loop:

    For each token to generate:
        1. Embed input tokens (chunked if prompt > prefill_chunk_size)
        2. For each layer:
           a. Prefetch next layer (GPU async / OS read-ahead for CPU)
           b. Load current layer weights (cache hit, CPU buffer, or disk)
           c. If MoE: run router, load selected experts only
           d. Compute attention + MLP/MoE forward pass with NTK-Aware RoPE
           e. Update paged KV cache (in-place, no O(n²) torch.cat)
           f. Evict layer weights (GPU) / release Python refs (CPU)
        3. Apply LM head
        4. Sample next token (temperature / top-k / top-p / min-p / rep-penalty)

Key improvements vs v0.1:
  - Paged KV cache: pre-allocated, in-place writes — O(1) per token vs O(n²)
  - KV INT8 compression: halves KV cache VRAM, doubles max context
  - NTK-Aware RoPE: extends context 4–8× beyond training length
  - Min-P sampling: better than pure top-p for open-ended generation
  - Repetition penalty: prevents looping in long outputs
  - Chunked prefill: prevents OOM on long prompts
  - stream_generate(): yields tokens one-by-one (essential for slow inference)
"""

from __future__ import annotations

import gc
import logging
from typing import Callable, Iterator, Optional

import torch
import torch.nn.functional as F

from dokodemo_ai.engine.cache import AdaptiveLayerCache
from dokodemo_ai.engine.scheduler import ExpertScheduler
from dokodemo_ai.engine.streaming import TensorStreamer
from dokodemo_ai.graph.compiler import ComponentInfo, LayerInfo, ModelGraph
from dokodemo_ai.utils.profiler import Profiler

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Paged KV Cache with optional INT8 compression
# ---------------------------------------------------------------------------


class KVCache:
    """
    Pre-allocated paged KV cache — eliminates the O(n²) torch.cat bottleneck.

    v0.1 used `torch.cat([old_k, new_k], dim=2)` every token.  For 1,000
    generated tokens that creates 1,000 tensors and copies ~500,000 KV
    rows total (1+2+…+1000).  This class pre-allocates a fixed buffer and
    fills it in-place — O(1) per token regardless of sequence length.

    Optional INT8 compression (kv_bits=8):
      - Stores K/V as INT8 with per-token-per-head FP16 scales.
      - Halves KV VRAM: doubles max context at the same memory cost.
      - Dequantizes on the fly before attention — negligible compute cost.
      - Example: Llama-70B at seq=4096 → 1.28 GB (FP16) → 640 MB (INT8).
    """

    def __init__(
        self,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
        max_seq_len: int = 2048,
        kv_bits: int = 16,
    ):
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.device = device
        self.dtype = dtype
        self.max_seq_len = max_seq_len
        self.kv_bits = kv_bits

        store_dtype = torch.int8 if kv_bits == 8 else dtype

        # Pre-allocated buffers: [num_layers, 1, num_kv_heads, max_seq_len, head_dim]
        self._k = torch.zeros(
            num_layers, 1, num_kv_heads, max_seq_len, head_dim,
            dtype=store_dtype, device=device,
        )
        self._v = torch.zeros(
            num_layers, 1, num_kv_heads, max_seq_len, head_dim,
            dtype=store_dtype, device=device,
        )

        if kv_bits == 8:
            # Per-token-per-head scale: [num_layers, 1, num_kv_heads, max_seq_len, 1]
            self._k_scale = torch.ones(
                num_layers, 1, num_kv_heads, max_seq_len, 1,
                dtype=torch.float16, device=device,
            )
            self._v_scale = torch.ones(
                num_layers, 1, num_kv_heads, max_seq_len, 1,
                dtype=torch.float16, device=device,
            )
        else:
            self._k_scale = None
            self._v_scale = None

        # Per-layer position counters (all advance together, tracked separately
        # so chunked prefill and layer-by-layer writes are correct)
        self._layer_pos: list[int] = [0] * num_layers

    # ------------------------------------------------------------------
    # Core API
    # ------------------------------------------------------------------

    def update(
        self,
        layer_idx: int,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Write new K/V into the pre-allocated buffer in-place.

        Returns the full K/V slice [0 : pos+seq] for this layer,
        ready to pass to scaled_dot_product_attention.
        """
        seq = key.shape[2]
        pos = self._layer_pos[layer_idx]
        end = min(pos + seq, self.max_seq_len)
        actual = end - pos

        if actual < seq:
            key = key[:, :, :actual, :]
            value = value[:, :, :actual, :]

        if self.kv_bits == 8:
            k_q, k_s = self._quantize(key)
            v_q, v_s = self._quantize(value)
            self._k[layer_idx, :, :, pos:end, :] = k_q
            self._v[layer_idx, :, :, pos:end, :] = v_q
            self._k_scale[layer_idx, :, :, pos:end, :] = k_s
            self._v_scale[layer_idx, :, :, pos:end, :] = v_s
            k_out = self._dequantize(
                self._k[layer_idx, :, :, :end, :],
                self._k_scale[layer_idx, :, :, :end, :],
            )
            v_out = self._dequantize(
                self._v[layer_idx, :, :, :end, :],
                self._v_scale[layer_idx, :, :, :end, :],
            )
        else:
            self._k[layer_idx, :, :, pos:end, :] = key.to(self.dtype)
            self._v[layer_idx, :, :, pos:end, :] = value.to(self.dtype)
            k_out = self._k[layer_idx, :, :, :end, :]
            v_out = self._v[layer_idx, :, :, :end, :]

        self._layer_pos[layer_idx] = end
        return k_out, v_out

    def get_seq_len(self, layer_idx: int = 0) -> int:
        return self._layer_pos[layer_idx]

    def clear(self) -> None:
        self._k.zero_()
        self._v.zero_()
        if self._k_scale is not None:
            self._k_scale.fill_(1.0)
            self._v_scale.fill_(1.0)
        self._layer_pos = [0] * self.num_layers

    # ------------------------------------------------------------------
    # INT8 helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _quantize(
        x: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Quantize x to INT8 with per-token-per-head absmax scaling."""
        # x: [1, num_kv_heads, seq, head_dim]
        scale = (
            x.float()
            .abs()
            .amax(dim=-1, keepdim=True)   # max over head_dim
            .clamp(min=1e-8)
            / 127.0
        )
        q = (x.float() / scale).round().clamp(-128, 127).to(torch.int8)
        return q, scale.to(torch.float16)

    @staticmethod
    def _dequantize(
        x: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        """Dequantize INT8 back to float using stored scales."""
        return (x.float() * scale.float()).to(scale.dtype)


# ---------------------------------------------------------------------------
# Inference Engine
# ---------------------------------------------------------------------------


class InferenceEngine:
    """
    Main inference engine for offloaded generation.

    Works on GPU (fast, async) and CPU (mmap, full accuracy, any model size).
    """

    def __init__(
        self,
        graph: ModelGraph,
        streamer: TensorStreamer,
        cache: AdaptiveLayerCache,
        scheduler: ExpertScheduler,
        device: torch.device,
        compute_dtype: torch.dtype = torch.float32,
        profiler: Optional[Profiler] = None,
    ):
        self.graph = graph
        self.streamer = streamer
        self.cache = cache
        self.scheduler = scheduler
        self.device = device
        self.compute_dtype = compute_dtype
        self.profiler = profiler or Profiler(enabled=False)
        self._is_cpu = device.type == "cpu"

        config = graph.config
        self.hidden_size = config.get(
            "hidden_size", config.get("d_model", config.get("n_embd", 4096))
        )
        self.num_heads = config.get(
            "num_attention_heads",
            config.get("n_head", config.get("num_heads", 32)),
        )
        self.num_kv_heads = config.get("num_key_value_heads", self.num_heads)
        self.head_dim = self.hidden_size // self.num_heads
        self.rms_norm_eps = config.get("rms_norm_eps", 1e-6)
        self.rope_theta = config.get("rope_theta", 10000.0)
        self.vocab_size = config.get("vocab_size", 32000)

        # Training context length — used for NTK-Aware RoPE scaling
        self._max_train_pos = config.get("max_position_embeddings", 4096)

        # RoPE cache: (max_seq_len, effective_theta) → (cos, sin)
        self._rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None
        self._rope_cache_len: int = 0

        # Persistent weights (kept resident)
        self._embedding_weight: Optional[torch.Tensor] = None
        self._head_weight: Optional[torch.Tensor] = None
        self._head_bias: Optional[torch.Tensor] = None
        self._final_norm_weight: Optional[torch.Tensor] = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def load_persistent_weights(self) -> None:
        """Load embedding and LM head weights (kept in memory permanently)."""
        with self.profiler.track("load_embedding"):
            emb_tensors = self.streamer.load_component(
                ComponentInfo(
                    name=self.graph.embedding.module_name,
                    kind="embedding",
                    params=self.graph.embedding.params,
                ),
                to_device=self.device,
            )
            for name, tensor in emb_tensors.items():
                if "weight" in name:
                    self._embedding_weight = tensor.to(self.compute_dtype)
                    break

        with self.profiler.track("load_head"):
            head_tensors = self.streamer.load_component(
                ComponentInfo(
                    name=self.graph.head.module_name,
                    kind="head",
                    params=self.graph.head.params,
                ),
                to_device=self.device,
            )
            for name, tensor in head_tensors.items():
                if "lm_head" in name and "weight" in name:
                    self._head_weight = tensor.to(self.compute_dtype)
                elif "lm_head" in name and "bias" in name:
                    self._head_bias = tensor.to(self.compute_dtype)
                elif "norm" in name and "weight" in name:
                    self._final_norm_weight = tensor.to(self.compute_dtype)

        if self._head_weight is None and self.graph.head.tied_to_embedding:
            self._head_weight = self._embedding_weight

        logger.info(
            "Persistent weights loaded. Compute dtype: %s", self.compute_dtype
        )

    @torch.inference_mode()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 0.9,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        do_sample: bool = True,
        prefill_chunk_size: int = 512,
        kv_cache_bits: int = 16,
        on_token: Optional[Callable[[int], None]] = None,
    ) -> torch.Tensor:
        """
        Generate tokens autoregressively.

        Args:
            input_ids: Input token IDs [batch_size, seq_len].
            max_new_tokens: Maximum number of new tokens to generate.
            temperature: Sampling temperature.
            top_k: Top-k sampling (0 = disabled).
            top_p: Nucleus sampling threshold (1.0 = disabled).
            min_p: Min-P threshold — removes tokens with
                   prob < min_p × max_prob (0.0 = disabled).
                   Better than top-p for creative generation.
            repetition_penalty: Penalise previously generated tokens
                   (1.0 = no penalty, >1.0 = penalise repetition).
            do_sample: If False, use greedy decoding.
            prefill_chunk_size: Process prefill in chunks of this size
                   to prevent OOM on long prompts (default 512).
            kv_cache_bits: KV cache precision (16 = FP16, 8 = INT8).
                   INT8 halves KV VRAM and doubles max context.
            on_token: Optional callback called with each new token ID
                   as it is generated (for sync streaming).

        Returns:
            Generated token IDs [batch_size, seq_len + n_generated].
        """
        self.profiler.start()

        if self._embedding_weight is None:
            self.load_persistent_weights()

        input_ids = input_ids.to(self.device)
        batch_size, seq_len = input_ids.shape
        max_seq_len = seq_len + max_new_tokens

        kv_dtype = self.compute_dtype if self._is_cpu else torch.float16

        kv_cache = KVCache(
            num_layers=self.graph.num_layers,
            num_kv_heads=self.num_kv_heads,
            head_dim=self.head_dim,
            device=self.device,
            dtype=kv_dtype,
            max_seq_len=max_seq_len,
            kv_bits=kv_cache_bits,
        )

        generated = input_ids

        # Prefill — chunked to prevent OOM on long prompts
        logger.info(
            "Prefill: %d tokens through %d layers on %s (chunk=%d)",
            seq_len, self.graph.num_layers, self.device, prefill_chunk_size,
        )
        with self.profiler.track("prefill"):
            if seq_len > prefill_chunk_size:
                hidden = self._chunked_prefill(
                    input_ids, kv_cache, prefill_chunk_size
                )
            else:
                hidden = self._forward_all_layers(
                    input_ids, kv_cache, is_prefill=True
                )
            logits = self._apply_head(hidden)

        # Decode — one token at a time
        for step in range(max_new_tokens):
            with self.profiler.track("sample"):
                next_token = self._sample(
                    logits[:, -1:, :],
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    min_p=min_p,
                    repetition_penalty=repetition_penalty,
                    do_sample=do_sample,
                    previous_tokens=generated,
                )

            generated = torch.cat([generated, next_token], dim=1)
            self.profiler.record_token()

            if on_token is not None:
                on_token(next_token.item())

            eos_id = self.graph.config.get("eos_token_id")
            if eos_id is not None:
                eos_ids = eos_id if isinstance(eos_id, list) else [eos_id]
                if next_token.item() in eos_ids:
                    break

            with self.profiler.track(f"decode_{step}"):
                hidden = self._forward_all_layers(
                    next_token, kv_cache, is_prefill=False
                )
                logits = self._apply_head(hidden)

        return generated

    @torch.inference_mode()
    def stream_generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 0.9,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        do_sample: bool = True,
        prefill_chunk_size: int = 512,
        kv_cache_bits: int = 16,
    ) -> Iterator[tuple[int, torch.Tensor]]:
        """
        Streaming generation — yields each token as it is produced.

        This is critical for disk-offloaded inference where each token
        can take 1–10 seconds. Without streaming, users wait minutes
        for the first output.

        Yields:
            (token_id: int, full_sequence: torch.Tensor)
            token_id  — the newly generated token (scalar int)
            full_sequence — all tokens so far [1, prompt_len + n_so_far]

        Example::

            for token_id, seq in model.stream_generate(input_ids):
                word = tokenizer.decode([token_id], skip_special_tokens=True)
                print(word, end="", flush=True)
        """
        if self._embedding_weight is None:
            self.load_persistent_weights()

        input_ids = input_ids.to(self.device)
        batch_size, seq_len = input_ids.shape
        max_seq_len = seq_len + max_new_tokens

        kv_dtype = self.compute_dtype if self._is_cpu else torch.float16
        kv_cache = KVCache(
            num_layers=self.graph.num_layers,
            num_kv_heads=self.num_kv_heads,
            head_dim=self.head_dim,
            device=self.device,
            dtype=kv_dtype,
            max_seq_len=max_seq_len,
            kv_bits=kv_cache_bits,
        )

        generated = input_ids

        # Prefill (not yielded — we only stream new tokens)
        if seq_len > prefill_chunk_size:
            hidden = self._chunked_prefill(
                input_ids, kv_cache, prefill_chunk_size
            )
        else:
            hidden = self._forward_all_layers(
                input_ids, kv_cache, is_prefill=True
            )
        logits = self._apply_head(hidden)

        for _ in range(max_new_tokens):
            next_token = self._sample(
                logits[:, -1:, :],
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                min_p=min_p,
                repetition_penalty=repetition_penalty,
                do_sample=do_sample,
                previous_tokens=generated,
            )

            generated = torch.cat([generated, next_token], dim=1)

            yield next_token.item(), generated.clone()

            eos_id = self.graph.config.get("eos_token_id")
            if eos_id is not None:
                eos_ids = eos_id if isinstance(eos_id, list) else [eos_id]
                if next_token.item() in eos_ids:
                    return

            hidden = self._forward_all_layers(
                next_token, kv_cache, is_prefill=False
            )
            logits = self._apply_head(hidden)

    # ------------------------------------------------------------------
    # Prefill helpers
    # ------------------------------------------------------------------

    def _chunked_prefill(
        self,
        input_ids: torch.Tensor,
        kv_cache: KVCache,
        chunk_size: int,
    ) -> torch.Tensor:
        """
        Process a long prompt in chunks to prevent attention OOM.

        The KV cache accumulates across chunks, so attention to earlier
        chunks is handled via the cached K/V.  PyTorch's is_causal=True
        SDPA correctly handles q_len < k_len by applying the causal mask
        only within the new query positions, allowing full attention to
        all cached positions — exactly the behaviour we need.
        """
        logger.info(
            "Chunked prefill: %d tokens → %d chunks of %d",
            input_ids.shape[1],
            (input_ids.shape[1] + chunk_size - 1) // chunk_size,
            chunk_size,
        )
        hidden = None
        chunks = input_ids.split(chunk_size, dim=1)
        for chunk in chunks:
            hidden = self._forward_all_layers(chunk, kv_cache, is_prefill=True)
        return hidden

    # ------------------------------------------------------------------
    # Forward pass
    # ------------------------------------------------------------------

    def _forward_all_layers(
        self,
        input_ids: torch.Tensor,
        kv_cache: KVCache,
        is_prefill: bool,
    ) -> torch.Tensor:
        """Forward pass through all transformer layers."""

        with self.profiler.track("embed"):
            emb_w = self._embedding_weight.to(self.compute_dtype)
            hidden = F.embedding(input_ids, emb_w)

        for i, layer_info in enumerate(self.graph.layers):
            if i + 1 < len(self.graph.layers):
                next_layer = self.graph.layers[i + 1]
                self.streamer.prefetch_layer(
                    next_layer, exclude_experts=next_layer.is_moe
                )

            with self.profiler.track(f"layer_{i}"):
                hidden = self._forward_layer(
                    layer_info, hidden, kv_cache, is_prefill
                )

            if self._is_cpu:
                gc.collect()

        if self._final_norm_weight is not None:
            with self.profiler.track("final_norm"):
                norm_w = self._final_norm_weight.to(self.compute_dtype)
                hidden = self._rms_norm(hidden, norm_w)

        return hidden

    def _forward_layer(
        self,
        layer: LayerInfo,
        hidden: torch.Tensor,
        kv_cache: KVCache,
        is_prefill: bool,
    ) -> torch.Tensor:
        """Forward pass through a single transformer layer."""

        with self.profiler.track(f"layer_{layer.index}.load"):
            layer_tensors = self._load_layer_weights(layer)

        attn_tensors: dict[str, torch.Tensor] = {}
        mlp_tensors:  dict[str, torch.Tensor] = {}
        norm_tensors: dict[str, torch.Tensor] = {}
        router_tensors: dict[str, torch.Tensor] = {}

        for name, tensor in layer_tensors.items():
            t = tensor.to(self.compute_dtype)
            if any(k in name for k in ["self_attn", "attention", "attn"]):
                attn_tensors[name] = t
            elif any(k in name for k in ["gate.", "router"]):
                router_tensors[name] = t
            elif any(k in name for k in ["mlp", "feed_forward", "ffn"]):
                mlp_tensors[name] = t
            elif any(k in name for k in ["norm", "layernorm"]):
                norm_tensors[name] = t

        del layer_tensors
        if self._is_cpu:
            gc.collect()

        norm_w = self._find_tensor(norm_tensors, ["input_layernorm", "pre"])
        normed = self._rms_norm(hidden, norm_w) if norm_w is not None else hidden

        with self.profiler.track(f"layer_{layer.index}.attention"):
            attn_out = self._compute_attention(
                normed, attn_tensors, kv_cache, layer.index, is_prefill
            )

        hidden = hidden + attn_out

        norm_w2 = self._find_tensor(
            norm_tensors,
            ["post_attention_layernorm", "post", "pre_feedforward"],
        )
        normed2 = (
            self._rms_norm(hidden, norm_w2) if norm_w2 is not None else hidden
        )

        if layer.is_moe:
            with self.profiler.track(f"layer_{layer.index}.moe"):
                mlp_out = self._compute_moe(layer, normed2, router_tensors)
        else:
            with self.profiler.track(f"layer_{layer.index}.mlp"):
                mlp_out = self._compute_mlp(normed2, mlp_tensors)

        hidden = hidden + mlp_out

        del attn_tensors, mlp_tensors, norm_tensors, router_tensors
        if self._is_cpu:
            gc.collect()

        return hidden

    # ------------------------------------------------------------------
    # Component computations
    # ------------------------------------------------------------------

    def _compute_attention(
        self,
        hidden: torch.Tensor,
        attn_tensors: dict[str, torch.Tensor],
        kv_cache: KVCache,
        layer_idx: int,
        is_prefill: bool,
    ) -> torch.Tensor:
        """Multi-head attention with paged KV cache and NTK-Aware RoPE."""
        batch, seq_len, _ = hidden.shape

        q_w = self._find_tensor(attn_tensors, ["q_proj"])
        k_w = self._find_tensor(attn_tensors, ["k_proj"])
        v_w = self._find_tensor(attn_tensors, ["v_proj"])
        o_w = self._find_tensor(attn_tensors, ["o_proj", "out_proj", "dense"])

        if q_w is not None and k_w is not None and v_w is not None:
            q = F.linear(hidden, q_w)
            k = F.linear(hidden, k_w)
            v = F.linear(hidden, v_w)
        else:
            qkv_w = self._find_tensor(
                attn_tensors, ["qkv", "query_key_value", "c_attn"]
            )
            if qkv_w is None:
                raise RuntimeError(
                    f"Cannot find Q/K/V projections in layer {layer_idx}. "
                    f"Keys: {list(attn_tensors.keys())}"
                )
            qkv = F.linear(hidden, qkv_w)
            q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # NTK-Aware RoPE: pass total context length for scaling decision
        offset = kv_cache.get_seq_len(layer_idx)
        full_len = offset + seq_len
        cos, sin = self._get_rotary_embeddings(full_len)
        q = self._apply_rope(q, cos, sin, offset)
        k = self._apply_rope(k, cos, sin, offset)

        # Paged KV cache — in-place write, no torch.cat
        k, v = kv_cache.update(layer_idx, k, v)

        if self.num_kv_heads < self.num_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)

        # is_causal=True with q_len < k_len correctly attends to all cached
        # tokens while applying causal mask only within the current chunk.
        attn_out = F.scaled_dot_product_attention(
            q, k, v, is_causal=is_prefill
        )

        attn_out = (
            attn_out.transpose(1, 2)
            .contiguous()
            .view(batch, seq_len, self.hidden_size)
        )
        if o_w is not None:
            attn_out = F.linear(attn_out, o_w)

        return attn_out

    def _compute_mlp(
        self,
        hidden: torch.Tensor,
        mlp_tensors: dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Feed-forward MLP (SwiGLU or standard GELU)."""
        gate_w = self._find_tensor(mlp_tensors, ["gate_proj", "w1"])
        up_w   = self._find_tensor(mlp_tensors, ["up_proj",   "w3"])
        down_w = self._find_tensor(mlp_tensors, ["down_proj", "w2"])

        if gate_w is not None and up_w is not None and down_w is not None:
            return F.linear(
                F.silu(F.linear(hidden, gate_w)) * F.linear(hidden, up_w),
                down_w,
            )

        fc1 = self._find_tensor(mlp_tensors, ["fc1", "c_fc", "dense_h_to_4h"])
        fc2 = self._find_tensor(mlp_tensors, ["fc2", "c_proj", "dense_4h_to_h"])

        if fc1 is not None and fc2 is not None:
            return F.linear(F.gelu(F.linear(hidden, fc1)), fc2)

        raise RuntimeError(
            f"Cannot identify MLP structure. Keys: {list(mlp_tensors.keys())}"
        )

    def _compute_moe(
        self,
        layer: LayerInfo,
        hidden: torch.Tensor,
        router_tensors: dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """MoE block with router-guided sparse expert loading."""
        expert_indices, routing_weights, expert_weight_dicts = (
            self.scheduler.route_and_load(layer, hidden, router_tensors)
        )

        expert_outputs = []
        for i, (expert_idx, expert_tensors) in enumerate(
            zip(expert_indices, expert_weight_dicts)
        ):
            expert_tensors_cast = {
                k: v.to(self.compute_dtype) for k, v in expert_tensors.items()
            }
            out = self._compute_single_expert(hidden, expert_tensors_cast)
            expert_outputs.append(out * routing_weights[i])
            del expert_tensors, expert_tensors_cast

        output = sum(expert_outputs)
        if self._is_cpu:
            gc.collect()
        return output

    def _compute_single_expert(
        self,
        hidden: torch.Tensor,
        expert_tensors: dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Single expert MLP forward pass."""
        gate_w = self._find_tensor(expert_tensors, ["gate_proj", "w1"])
        up_w   = self._find_tensor(expert_tensors, ["up_proj",   "w3"])
        down_w = self._find_tensor(expert_tensors, ["down_proj", "w2"])

        if gate_w is not None and up_w is not None and down_w is not None:
            return F.linear(
                F.silu(F.linear(hidden, gate_w)) * F.linear(hidden, up_w),
                down_w,
            )

        w1 = self._find_tensor(expert_tensors, ["w1", "fc1", "linear_1"])
        w2 = self._find_tensor(expert_tensors, ["w2", "fc2", "linear_2"])
        if w1 is not None and w2 is not None:
            return F.linear(F.silu(F.linear(hidden, w1)), w2)

        raise RuntimeError(
            f"Cannot identify expert structure. Keys: {list(expert_tensors.keys())}"
        )

    # ------------------------------------------------------------------
    # Head and sampling
    # ------------------------------------------------------------------

    def _apply_head(self, hidden: torch.Tensor) -> torch.Tensor:
        """Apply final LM head to get next-token logits."""
        with self.profiler.track("lm_head"):
            head_w = self._head_weight.to(self.compute_dtype)
            head_b = (
                self._head_bias.to(self.compute_dtype)
                if self._head_bias is not None
                else None
            )
            logits = F.linear(hidden, head_w, head_b)
        return logits

    def _sample(
        self,
        logits: torch.Tensor,
        temperature: float,
        top_k: int,
        top_p: float,
        do_sample: bool,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        previous_tokens: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Sample next token from logits.

        Applies (in order):
        1. Repetition penalty — discourages repeating tokens
        2. Temperature scaling
        3. Min-P filtering — removes tokens with prob < min_p × max_prob
        4. Top-K filtering
        5. Top-P (nucleus) filtering
        6. Softmax + multinomial sample (or greedy)
        """
        logits = logits[:, -1, :].clone()

        if not do_sample:
            return logits.argmax(dim=-1, keepdim=True)

        # 1. Repetition penalty
        if repetition_penalty != 1.0 and previous_tokens is not None:
            for token_id in set(previous_tokens[0].tolist()):
                if token_id < logits.shape[-1]:
                    if logits[0, token_id] > 0:
                        logits[0, token_id] /= repetition_penalty
                    else:
                        logits[0, token_id] *= repetition_penalty

        # 2. Temperature
        if temperature != 1.0:
            logits = logits / temperature

        # 3. Min-P filtering (before top-k/p so we don't mask needed tokens)
        if min_p > 0.0:
            probs_for_filter = F.softmax(logits, dim=-1)
            max_prob = probs_for_filter.max(dim=-1, keepdim=True)[0]
            logits = logits.masked_fill(
                probs_for_filter < min_p * max_prob, float("-inf")
            )

        # 4. Top-K
        if top_k > 0:
            k = min(top_k, logits.size(-1))
            threshold = torch.topk(logits, k)[0][..., -1:]
            logits = logits.masked_fill(logits < threshold, float("-inf"))

        # 5. Top-P (nucleus)
        if top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(
                F.softmax(sorted_logits, dim=-1), dim=-1
            )
            mask = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
            sorted_logits[mask] = float("-inf")
            logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)

        # 6. Sample
        probs = F.softmax(logits.float(), dim=-1)
        return torch.multinomial(probs, num_samples=1)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _load_layer_weights(self, layer: LayerInfo) -> dict[str, torch.Tensor]:
        """Load non-expert layer weights via cache + streamer."""
        all_tensors: dict[str, torch.Tensor] = {}

        for comp in layer.components:
            if comp.kind == "expert":
                continue
            tensors = self.cache.get_or_load(
                component=comp,
                layer_index=layer.index,
                loader=lambda c=comp: self.streamer.load_component(
                    c, to_device=self.device
                ),
            )
            all_tensors.update(tensors)

        return all_tensors

    def _rms_norm(
        self, hidden: torch.Tensor, weight: torch.Tensor
    ) -> torch.Tensor:
        """Root Mean Square Layer Normalization (float32 for numerical stability)."""
        h = hidden.float()
        rms = h.pow(2).mean(-1, keepdim=True)
        h = h * torch.rsqrt(rms + self.rms_norm_eps)
        return (h * weight.float()).to(hidden.dtype)

    def _get_rotary_embeddings(
        self, max_seq_len: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Compute or retrieve cached RoPE sin/cos tables with NTK-Aware scaling.

        NTK-Aware RoPE (YaRN-lite):
          When max_seq_len exceeds the training context length, the effective
          rope_theta base is scaled up so that the model can generalise to
          longer sequences.  The scaling formula:

              adjusted_theta = rope_theta × scale_factor^(d / (d − 2))

          where d = head_dim and scale_factor = max_seq_len / max_train_pos.

          This gives 4–8× longer usable context with no fine-tuning.
        """
        # Determine effective theta (possibly scaled for long contexts)
        if max_seq_len > self._max_train_pos:
            scale = max_seq_len / self._max_train_pos
            # NTK interpolation: shift frequencies to fit longer context
            effective_theta = self.rope_theta * (
                scale ** (self.head_dim / (self.head_dim - 2))
            )
        else:
            effective_theta = self.rope_theta

        # Return cached tables if sufficient length and same theta
        if (
            self._rope_cache is not None
            and self._rope_cache_len >= max_seq_len
            and not (max_seq_len > self._max_train_pos)  # invalidate on scale change
        ):
            cos, sin = self._rope_cache
            return cos[:max_seq_len], sin[:max_seq_len]

        inv_freq = 1.0 / (
            effective_theta
            ** (
                torch.arange(0, self.head_dim, 2, device=self.device).float()
                / self.head_dim
            )
        )
        t = torch.arange(max_seq_len, device=self.device).float()
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos().to(self.compute_dtype)
        sin = emb.sin().to(self.compute_dtype)
        self._rope_cache = (cos, sin)
        self._rope_cache_len = max_seq_len
        return cos, sin

    @staticmethod
    def _apply_rope(
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        offset: int = 0,
    ) -> torch.Tensor:
        """Apply rotary position embedding."""
        seq_len = x.shape[2]
        cos = cos[offset: offset + seq_len].unsqueeze(0).unsqueeze(0)
        sin = sin[offset: offset + seq_len].unsqueeze(0).unsqueeze(0)
        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
        return x * cos + torch.cat([-x2, x1], dim=-1) * sin

    @staticmethod
    def _find_tensor(
        tensors: dict[str, torch.Tensor],
        patterns: list[str],
    ) -> Optional[torch.Tensor]:
        """Find the first tensor whose name contains any of the given patterns."""
        for name, tensor in tensors.items():
            if any(p in name for p in patterns):
                return tensor
        return None
