chunked_model.py

"""
chunked_model.py — Memory-efficient layer-by-layer inference engine.

Drop-in replacement for airllm.  Zero airllm dependency.

How it works
────────────
• Model weights are split into small shard files (~chunk_mb MB each) once on
  first load, then reused on every subsequent run.
• During inference each layer's weights are loaded from its shard, the layer
  forward pass is executed, then the weights are released from RAM.
• An in-memory layer cache avoids re-reading the same shard twice in one
  generation call (huge speedup for multi-token generation).
• KV cache: after the prefill pass only one new token is forwarded per step,
  so the per-token cost is one layer-by-layer pass over 1 token — much faster
  than airllm's O(n²) approach.

Supported architectures
───────────────────────
  Qwen2ForCausalLM · Qwen3ForCausalLM · Qwen3_5ForConditionalGeneration
  LlamaForCausalLM · MistralForCausalLM · MixtralForCausalLM
  GemmaForCausalLM · Gemma2ForCausalLM · Phi3ForCausalLM

Usage
─────
  from chunked_model import ChunkedModel
  model = ChunkedModel("./model", chunk_mb=75)
  out   = model.generate(input_ids, max_new_tokens=200, temperature=0.7)
"""

from __future__ import annotations
import gc, json, os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn.functional as F
from safetensors import safe_open
from safetensors.torch import save_file
from transformers import AutoConfig


# ─────────────────────────────────────────────────────────────────────────────
#  Math primitives
# ─────────────────────────────────────────────────────────────────────────────

def _rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    # Compute norm in float32 for stability, then cast back to original dtype
    x_f32 = x.float()
    variance = x_f32.pow(2).mean(-1, keepdim=True)
    normed = x_f32 * torch.rsqrt(variance + eps)
    return (w.float() * normed).to(x.dtype)


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    h = x.shape[-1] // 2
    return torch.cat([-x[..., h:], x[..., :h]], dim=-1)


def _apply_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    position_ids: torch.Tensor,
    head_dim: int,
    rope_theta: float = 1_000_000.0,
    partial_factor: float = 1.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Apply Rotary Position Embedding to query and key tensors.

    partial_factor < 1.0: only apply RoPE to the first int(head_dim*partial_factor)
    dimensions (used by Qwen3.5 which uses partial_rotary_factor=0.25).
    """
    device = q.device
    rot_dim = int(head_dim * partial_factor)
    if rot_dim % 2 != 0:
        rot_dim -= 1

    inv_freq = 1.0 / (
        rope_theta ** (
            torch.arange(0, rot_dim, 2, device=device, dtype=torch.float32) / rot_dim
        )
    )
    freqs = torch.einsum("bi,j->bij", position_ids.float(), inv_freq)
    emb   = torch.cat([freqs, freqs], dim=-1)    # (B, T, rot_dim)
    cos   = emb.cos().unsqueeze(1)               # (B, 1, T, rot_dim)
    sin   = emb.sin().unsqueeze(1)

    # Cast cos/sin to q's dtype so RoPE doesn't upcast bfloat16 tensors
    cos = cos.to(q.dtype)
    sin = sin.to(q.dtype)

    if partial_factor < 1.0:
        q_rot, q_pass = q[..., :rot_dim], q[..., rot_dim:]
        k_rot, k_pass = k[..., :rot_dim], k[..., rot_dim:]
        q_rot = q_rot * cos + _rotate_half(q_rot) * sin
        k_rot = k_rot * cos + _rotate_half(k_rot) * sin
        q = torch.cat([q_rot, q_pass], dim=-1)
        k = torch.cat([k_rot, k_pass], dim=-1)
    else:
        q = q * cos + _rotate_half(q) * sin
        k = k * cos + _rotate_half(k) * sin
    return q, k


def _swiglu(
    x: torch.Tensor,
    gate_w: torch.Tensor,
    up_w: torch.Tensor,
    down_w: torch.Tensor,
    chunk_rows: int = 0,
) -> torch.Tensor:
    """SwiGLU feed-forward: down( silu(gate(x)) ⊙ up(x) ).

    chunk_rows > 0: compute the intermediate dimension in slices so that
    only `chunk_rows` worth of gate/up/down are live at once.  Use this for
    very large models where the MLP weights exceed the chunk budget.
    """
    interm = gate_w.shape[0]
    if chunk_rows <= 0 or chunk_rows >= interm:
        return F.linear(
            F.silu(F.linear(x, gate_w)) * F.linear(x, up_w),
            down_w,
        )
    # Chunked path — accumulate into output buffer row-slice by row-slice
    out = torch.zeros(*x.shape[:-1], down_w.shape[0], dtype=x.dtype, device=x.device)
    for start in range(0, interm, chunk_rows):
        end  = min(start + chunk_rows, interm)
        gate = F.silu(F.linear(x, gate_w[start:end]))
        up   = F.linear(x, up_w[start:end])
        out += F.linear(gate * up, down_w[:, start:end])
        del gate, up
    return out


# ─────────────────────────────────────────────────────────────────────────────
#  Shard manager
# ─────────────────────────────────────────────────────────────────────────────

class _ShardManager:
    """
    One-time split of model safetensors into fixed-size shard files.
    Builds index.json mapping every weight key to its shard filename.
    On subsequent runs the existing shards + index are reused as-is.
    """

    _INDEX = "index.json"

    def __init__(self, model_path: Path, chunk_mb: int, dtype: torch.dtype):
        self.model_path  = model_path
        self.chunk_bytes = chunk_mb * 1024 * 1024
        self.dtype       = dtype
        self.shard_dir   = model_path / f"_chunks_{chunk_mb}mb"
        self.index: Dict[str, str] = {}

    # ── Public API ────────────────────────────────────────────────────────────

    def prepare(self, num_layers: int) -> None:
        idx = self.shard_dir / self._INDEX
        if idx.exists():
            with open(idx) as f:
                self.index = json.load(f)
            n_shards = len(set(self.index.values()))
            print(f"[chunked] Reusing {n_shards} shards ({self.shard_dir.name})",
                  flush=True)
        else:
            self._build()

    def set_layer_prefix(self, prefix: str) -> None:
        """Configure the key prefix used by load_layer (e.g. 'model.language_model.')."""
        self._layer_prefix = prefix

    def load_layer(self, i: int) -> Dict[str, torch.Tensor]:
        pfx = getattr(self, "_layer_prefix", "model.")
        key = f"{pfx}layers.{i}."
        return self._load_iter(k for k in self.index if k.startswith(key))

    def load_keys(self, *keys: str) -> Dict[str, torch.Tensor]:
        return self._load_iter(k for k in keys if k in self.index)

    # ── Internal ──────────────────────────────────────────────────────────────

    def _load_iter(self, keys) -> Dict[str, torch.Tensor]:
        shard_map: Dict[str, List[str]] = {}
        for k in keys:
            sf = self.index.get(k)
            if sf:
                shard_map.setdefault(sf, []).append(k)
        out: Dict[str, torch.Tensor] = {}
        for sf, ks in shard_map.items():
            with safe_open(str(self.shard_dir / sf), framework="pt", device="cpu") as f:
                for k in ks:
                    # Cast to target dtype at load time (shards keep original dtype)
                    out[k] = f.get_tensor(k).to(self.dtype)
        return out

    # Key prefixes to SKIP during shard building (vision encoder, MTP, etc.)
    _SKIP_PREFIXES = ("model.visual", "mtp.", "model.embed.visual")

    def _read_source(self) -> Dict[str, torch.Tensor]:
        """Read all text-model weights from safetensors, keeping original dtype.

        Vision-encoder and MTP weights are skipped — they are not needed for
        text-only inference and would double RAM usage during shard building.
        dtype conversion (to float32) happens at inference time, not here.
        """
        weights: Dict[str, torch.Tensor] = {}
        for st in sorted(self.model_path.glob("*.safetensors")):
            with safe_open(str(st), framework="pt", device="cpu") as f:
                for k in f.keys():
                    if any(k.startswith(pfx) for pfx in self._SKIP_PREFIXES):
                        continue
                    weights[k] = f.get_tensor(k)   # keep original dtype (bf16/fp16)
        if not weights:
            raise RuntimeError(f"No .safetensors files found in {self.model_path}")
        return weights

    @staticmethod
    def _group_of(key: str) -> str:
        # Support both model.layers.{i}. and model.language_model.layers.{i}.
        if ".layers." in key and not key.startswith("model.visual") \
                and not key.startswith("mtp"):
            parts = key.split(".layers.")
            if len(parts) >= 2:
                idx = parts[1].split(".")[0]
                if idx.isdigit():
                    return "layer_" + idx.zfill(5)
        if "embed_tokens" in key:
            return "00_embed"
        return "zz_head"

    def _build(self) -> None:
        chunk_mb = self.chunk_bytes // (1024 * 1024)
        self.shard_dir.mkdir(parents=True, exist_ok=True)
        print(f"[chunked] Building shards (~{chunk_mb} MB each) — one-time setup ...",
              flush=True)

        weights = self._read_source()

        # Group weights by logical layer
        groups: Dict[str, Dict[str, torch.Tensor]] = {}
        for k, t in weights.items():
            g = self._group_of(k)
            groups.setdefault(g, {})[k] = t

        index: Dict[str, str] = {}
        shard_n = 0

        def _flush(bucket: Dict[str, torch.Tensor]) -> str:
            nonlocal shard_n
            fname = f"s{shard_n:05d}.safetensors"
            save_file(bucket, str(self.shard_dir / fname))
            shard_n += 1
            return fname

        for g_name in sorted(groups):
            group = groups[g_name]
            g_bytes = sum(t.numel() * t.element_size() for t in group.values())

            if g_bytes <= self.chunk_bytes:
                fname = _flush(group)
                for k in group:
                    index[k] = fname
            else:
                bucket: Dict[str, torch.Tensor] = {}
                bucket_bytes = 0
                for k, t in group.items():
                    tb = t.numel() * t.element_size()
                    if tb > self.chunk_bytes and not bucket:
                        # Single oversized tensor — save alone; chunked inference
                        # handles it at runtime via _swiglu(chunk_rows=...)
                        fname = _flush({k: t})
                        index[k] = fname
                        continue
                    if bucket and bucket_bytes + tb > self.chunk_bytes:
                        fname = _flush(bucket)
                        for bk in bucket:
                            index[bk] = fname
                        bucket, bucket_bytes = {}, 0
                    bucket[k] = t
                    bucket_bytes += tb
                if bucket:
                    fname = _flush(bucket)
                    for bk in bucket:
                        index[bk] = fname

        with open(self.shard_dir / self._INDEX, "w") as f:
            json.dump(index, f, indent=2)
        self.index = index

        total_mb = sum(
            (self.shard_dir / sf).stat().st_size
            for sf in set(index.values())
        ) // (1024 * 1024)
        print(f"[chunked] {shard_n} shards created ({total_mb} MB) in {self.shard_dir.name}",
              flush=True)


# ─────────────────────────────────────────────────────────────────────────────
#  In-memory layer cache
# ─────────────────────────────────────────────────────────────────────────────

class _LayerCache:
    """
    Caches layer weight dicts in memory to avoid re-reading shards on every
    decode step.  max_layers=None means unlimited (keep everything).
    For small models (≤ ~4 GB weights) this is effectively a full warm cache
    after the first generate() call.
    """

    def __init__(self, max_layers: Optional[int] = None):
        self._cache: Dict[str, Dict[str, torch.Tensor]] = {}
        self._order: List[str] = []   # LRU insertion order
        self.max   = max_layers

    def get(self, key: str) -> Optional[Dict[str, torch.Tensor]]:
        if key in self._cache:
            self._order.remove(key)
            self._order.append(key)
            return self._cache[key]
        return None

    def put(self, key: str, weights: Dict[str, torch.Tensor]) -> None:
        if self.max is not None and len(self._cache) >= self.max and key not in self._cache:
            evict = self._order.pop(0)
            del self._cache[evict]
        self._cache[key] = weights
        if key in self._order:
            self._order.remove(key)
        self._order.append(key)

    def clear(self) -> None:
        self._cache.clear()
        self._order.clear()
        gc.collect()

    def __len__(self) -> int:
        return len(self._cache)


# ─────────────────────────────────────────────────────────────────────────────
#  ChunkedModel — the inference engine
# ─────────────────────────────────────────────────────────────────────────────

class ChunkedModel:
    """
    Memory-efficient transformer inference — loads one shard at a time.

    Parameters
    ----------
    model_path : str | Path
        Directory containing config.json + model.safetensors.
    chunk_mb : int
        Target shard size in MB (default 75).  Smaller = less peak RAM,
        more disk I/O on cold start.
    dtype : torch.dtype
        float32 (default, safest) or float16 (faster, needs modern CPU).
    cache_layers : int | None
        How many layer weight dicts to keep in RAM between steps.
        None (default) = keep everything — fast generation, uses more RAM.
        Set to e.g. 4 for 70B+ models to cap working-set memory.
    mlp_chunk_rows : int
        > 0 to split MLP projections into row-slices at runtime.
        Auto-computed when 0 (default).
    """

    SUPPORTED = {
        "Qwen2ForCausalLM", "Qwen3ForCausalLM",
        "Qwen3_5ForConditionalGeneration",
        "LlamaForCausalLM",
        "MistralForCausalLM", "MixtralForCausalLM",
        "GemmaForCausalLM",  "Gemma2ForCausalLM",
        "Phi3ForCausalLM",
    }

    def __init__(
        self,
        model_path: str,
        chunk_mb: int = 75,
        dtype: Optional[torch.dtype] = None,   # None → auto-detect from config
        cache_layers: Optional[int] = None,
        mlp_chunk_rows: int = 0,
    ):
        self.model_path = Path(model_path)

        if not self.model_path.is_dir():
            raise RuntimeError(
                f"Model directory not found: {self.model_path}\n"
                "Run 'bash install.sh' to download the model."
            )

        print(f"[chunked] Loading config from {self.model_path.name}", flush=True)
        cfg = AutoConfig.from_pretrained(str(self.model_path), local_files_only=True)
        self.cfg = cfg

        arch = (cfg.architectures or ["Unknown"])[0]
        if arch not in self.SUPPORTED:
            print(f"[chunked] WARNING: {arch} not in supported list, attempting generic path")

        # Auto-detect dtype from model config — bfloat16 halves RAM vs float32
        if dtype is None:
            raw = getattr(cfg, "torch_dtype", None) or getattr(cfg, "dtype", None)
            raw = str(raw) if raw is not None else ""
            if "bfloat16" in raw or raw is torch.bfloat16:
                dtype = torch.bfloat16
            elif "float16" in raw or raw is torch.float16:
                dtype = torch.float16
            else:
                dtype = torch.float32
        self.dtype = dtype

        # ── Config extraction: try top-level then text_config fallback ──────────
        # Qwen3.5 / multimodal models nest text params under cfg.text_config
        def _get(attr, default=None):
            v = getattr(cfg, attr, None)
            if v is None and hasattr(cfg, "text_config"):
                v = getattr(cfg.text_config, attr, None)
            return v if v is not None else default

        self.arch    = arch
        self.L       = _get("num_hidden_layers")
        self.H       = _get("num_attention_heads")
        self.KVH     = _get("num_key_value_heads", self.H)
        self.D       = _get("hidden_size")
        self.Dh      = _get("head_dim", self.D // self.H)
        self.interm  = _get("intermediate_size")
        self.vocab   = _get("vocab_size")
        self.eps     = float(_get("rms_norm_eps", 1e-6))
        self.tied    = bool(_get("tie_word_embeddings", False))

        # rope_theta may be nested in rope_parameters dict (Qwen3.5)
        rope_params  = _get("rope_parameters") or {}
        self.theta   = float(
            rope_params.get("rope_theta", None)
            or _get("rope_theta", None)
            or 1_000_000.0
        )

        # partial RoPE: Qwen3.5 uses partial_rotary_factor=0.25
        self.rope_partial = float(_get("partial_rotary_factor", 1.0))

        # Qwen3 / Qwen3.5: per-head RMSNorm on Q and K
        self.qk_norm = "Qwen3" in arch

        # Hybrid layer types: Qwen3.5 has linear_attention + full_attention
        raw_layer_types = _get("layer_types", [])
        if raw_layer_types:
            self.layer_types = list(raw_layer_types)
        else:
            self.layer_types = ["full_attention"] * self.L

        # Linear attention dims (Qwen3.5 SSM layers)
        self.lin_kh  = _get("linear_num_key_heads",   self.H)   # K heads
        self.lin_vh  = _get("linear_num_value_heads",  self.H)   # V heads
        self.lin_kdh = _get("linear_key_head_dim",     self.Dh)  # K head dim
        # V head dim = out_proj_input / lin_vh  (auto from weights)
        self.lin_qdim = self.lin_kh * self.lin_kdh  # Q total dim (=K total dim)
        self.lin_kdim = self.lin_kh * self.lin_kdh  # K total dim
        # V total dim = out_proj rows = hidden_size … use out_proj shape at runtime

        # ── Weight key prefix detection ─────────────────────────────────────────
        # Standard models use "model." prefix; Qwen3.5 uses "model.language_model."
        self.prefix = "model.language_model." \
            if arch == "Qwen3_5ForConditionalGeneration" else "model."

        print(
            f"[chunked] {arch}  |  {self.L} layers  |  "
            f"hidden={self.D}  |  "
            f"heads={self.H}(kv={self.KVH})  |  "
            f"prefix={self.prefix.rstrip('.')}",
            flush=True,
        )

        has_linear = any(lt == "linear_attention" for lt in self.layer_types)
        if has_linear:
            n_lin  = sum(1 for lt in self.layer_types if lt == "linear_attention")
            n_full = self.L - n_lin
            print(f"[chunked] Hybrid layers: {n_full} full_attention + {n_lin} linear_attention (SSM approx)",
                  flush=True)

        # MLP runtime chunking (for huge layer weights)
        if mlp_chunk_rows > 0:
            self.mlp_chunk = mlp_chunk_rows
        else:
            row_b  = self.D * _dtype_bytes(dtype)
            rpc    = max(256, (chunk_mb * 1024 * 1024 // 4) // row_b)
            mlp_b  = 3 * self.interm * self.D * _dtype_bytes(dtype)
            self.mlp_chunk = rpc if mlp_b > chunk_mb * 1024 * 1024 else 0

        # Shard storage
        self._sm = _ShardManager(self.model_path, chunk_mb, dtype)
        self._sm.set_layer_prefix(self.prefix)
        self._sm.prepare(self.L)

        # In-memory cache (avoids re-reading shards during decode steps)
        # Default: unlimited — warm on first generate(), stays warm.
        self._cache = _LayerCache(max_layers=cache_layers)
        # Always-pinned: embed and head weights (small overhead, huge speedup)
        self._embed_w: Optional[torch.Tensor] = None
        self._norm_w:  Optional[torch.Tensor] = None
        self._head_w:  Optional[torch.Tensor] = None

        if self.mlp_chunk > 0:
            print(f"[chunked] MLP sub-chunking active: {self.mlp_chunk} rows/pass",
                  flush=True)
        print(
            f"[chunked] Ready — chunk={chunk_mb} MB  "
            f"rope_theta={self.theta:.0f}  "
            f"rope_partial={self.rope_partial}  "
            f"tied={self.tied}  "
            f"layer_cache={'unlimited' if cache_layers is None else cache_layers}",
            flush=True,
        )

    # ── Cached weight loaders ─────────────────────────────────────────────────

    def _embed(self) -> torch.Tensor:
        if self._embed_w is None:
            key = self.prefix + "embed_tokens.weight"
            w = self._sm.load_keys(key)
            self._embed_w = w[key]
        return self._embed_w

    def _norm(self) -> torch.Tensor:
        if self._norm_w is None:
            key = self.prefix + "norm.weight"
            w = self._sm.load_keys(key)
            self._norm_w = w[key]
        return self._norm_w

    def _lm_head(self) -> torch.Tensor:
        if self._head_w is None:
            lm_key = self.prefix + "lm_head.weight"
            if self.tied or lm_key not in self._sm.index:
                self._head_w = self._embed()
            else:
                w = self._sm.load_keys(lm_key)
                self._head_w = w[lm_key]
        return self._head_w

    def _layer(self, i: int) -> Dict[str, torch.Tensor]:
        key = f"layer_{i}"
        cached = self._cache.get(key)
        if cached is not None:
            return cached
        w = self._sm.load_layer(i)
        self._cache.put(key, w)
        return w

    # ── Linear-attention (SSM approximation) forward ──────────────────────────

    def _linear_attn_step(
        self,
        hidden: torch.Tensor,
        w: dict,
        p: str,
    ) -> torch.Tensor:
        """
        Approximate forward for Qwen3.5 linear_attention (Mamba/SSM) layers.

        The true computation is a state-space recurrence; here we approximate
        it as standard scaled-dot-product attention over the same QKV projections.
        Quality is slightly lower than the full SSM but produces coherent text.

        All dims are auto-detected from actual weight shapes (no config needed):
          in_proj_qkv : [qkv_total, D]  →  Q + K + V  (V = out_proj input dim)
          in_proj_z   : [gate_dim, D]   →  gate
          norm.weight : [vdh]            →  value head dim (e.g. 128)
          out_proj    : [D, v_dim]
        """
        B, T, D = hidden.shape

        out_proj_w = w[p + "linear_attn.out_proj.weight"]   # [D, v_dim]
        v_dim      = out_proj_w.shape[1]                    # e.g. 4096

        qkv = F.linear(hidden, w[p + "linear_attn.in_proj_qkv.weight"])  # [B,T,qkv_total]
        z   = F.linear(hidden, w[p + "linear_attn.in_proj_z.weight"])     # [B,T,gate_dim]

        # Auto-derive Q and K dims: total - V, split equally between Q and K
        qkv_total = qkv.shape[-1]               # e.g. 8192
        q_dim = k_dim = (qkv_total - v_dim) // 2  # e.g. 2048 each

        q = qkv[..., :q_dim]                    # [B,T,q_dim]
        k = qkv[..., q_dim:q_dim + k_dim]       # [B,T,k_dim]
        v = qkv[..., q_dim + k_dim:]            # [B,T,v_dim]

        # Head count from norm.weight (norm applied per value head)
        norm_key = p + "linear_attn.norm.weight"
        nw = w.get(norm_key)
        if nw is not None:
            vdh = int(nw.shape[0])              # value head dim (e.g. 128)
        else:
            vdh = max(1, v_dim // 32)           # fallback
        kh  = v_dim // vdh                     # number of heads (e.g. 32)
        kdh = q_dim // kh                      # query/key head dim (e.g. 64)

        q = q.view(B, T, kh, kdh).transpose(1, 2)   # [B, kh, T, kdh]
        k = k.view(B, T, kh, kdh).transpose(1, 2)   # [B, kh, T, kdh]
        v = v.view(B, T, kh, vdh).transpose(1, 2)   # [B, kh, T, vdh]

        is_causal = T > 1
        attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        attn = attn.transpose(1, 2).contiguous().view(B, T, v_dim)  # [B,T,v_dim]

        # Per-value-head RMS norm (norm.weight tiled across all heads)
        if nw is not None:
            nw_tiled = nw.repeat(kh) if kh > 1 else nw   # [vdh * kh] = [v_dim]
            attn = _rms_norm(attn, nw_tiled.to(attn.dtype), self.eps)

        # Gating: element-wise silu gate
        gated = attn * F.silu(z)   # [B,T,v_dim]

        return F.linear(gated, out_proj_w)   # [B,T,D]

    # ── Forward pass ──────────────────────────────────────────────────────────

    def _forward(
        self,
        input_ids: torch.Tensor,
        past_kv:   Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
    ):
        B, T   = input_ids.shape
        past_n = past_kv[0][0].shape[2] if past_kv is not None else 0
        pos_ids = torch.arange(past_n, past_n + T,
                               device=input_ids.device).unsqueeze(0)

        # Embedding
        hidden = F.embedding(input_ids, self._embed().to(self.dtype))

        new_kv: List[Tuple[torch.Tensor, torch.Tensor]] = []

        for i in range(self.L):
            w  = self._layer(i)
            p  = f"{self.prefix}layers.{i}."
            ltype = self.layer_types[i] if i < len(self.layer_types) else "full_attention"

            # ── Pre-attention norm ──────────────────────────────────────────
            res    = hidden
            hidden = _rms_norm(hidden, w[p + "input_layernorm.weight"], self.eps)

            # ── Attention block (full or linear) ────────────────────────────
            if ltype == "linear_attention":
                # SSM approximation — no KV cache for these layers
                attn_out = self._linear_attn_step(hidden, w, p)
                new_kv.append((
                    torch.zeros(B, self.KVH, 0, self.Dh),
                    torch.zeros(B, self.KVH, 0, self.Dh),
                ))
            else:
                # Full self-attention (Qwen3.5: q_proj outputs Q+gate combined)
                q_raw = F.linear(hidden, w[p + "self_attn.q_proj.weight"])
                k     = F.linear(hidden, w[p + "self_attn.k_proj.weight"])
                v     = F.linear(hidden, w[p + "self_attn.v_proj.weight"])

                # Detect gated Q: q_proj output is 2× expected → split Q and gate
                expected_q = self.H * self.Dh   # e.g. 16*256=4096
                q_attn_gate = q_raw.shape[-1] == expected_q * 2  # Qwen3.5 gated
                if q_attn_gate:
                    q, q_gate = q_raw.chunk(2, dim=-1)  # each [B,T,H*Dh]
                else:
                    q, q_gate = q_raw, None

                q = q.view(B, T, self.H,   self.Dh).transpose(1, 2)
                k = k.view(B, T, self.KVH, self.Dh).transpose(1, 2)
                v = v.view(B, T, self.KVH, self.Dh).transpose(1, 2)

                # Per-head RMSNorm on Q and K (Qwen3 / Qwen3.5)
                if self.qk_norm:
                    q = _rms_norm(q, w[p + "self_attn.q_norm.weight"], self.eps)
                    k = _rms_norm(k, w[p + "self_attn.k_norm.weight"], self.eps)

                q, k = _apply_rope(q, k, pos_ids, self.Dh, self.theta,
                                   self.rope_partial)

                if past_kv is not None:
                    pk, pv = past_kv[i]
                    if pk.shape[2] > 0:   # skip empty SSM placeholders
                        k = torch.cat([pk, k], dim=2)
                        v = torch.cat([pv, v], dim=2)
                new_kv.append((k.detach(), v.detach()))

                # GQA: broadcast KV heads to match Q heads
                if self.KVH != self.H:
                    k = k.repeat_interleave(self.H // self.KVH, dim=1)
                    v = v.repeat_interleave(self.H // self.KVH, dim=1)

                is_causal = T > 1 and past_kv is None
                attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
                attn = attn.transpose(1, 2).contiguous().view(B, T, self.H * self.Dh)

                # Output gate (Qwen3.5 attn_output_gate=True)
                if q_attn_gate:
                    attn = attn * F.silu(q_gate)

                attn_out = F.linear(attn, w[p + "self_attn.o_proj.weight"])

            hidden = res + attn_out

            # ── Post-attention norm + FFN ───────────────────────────────────
            res    = hidden
            hidden = _rms_norm(hidden, w[p + "post_attention_layernorm.weight"], self.eps)
            hidden = res + _swiglu(
                hidden,
                w[p + "mlp.gate_proj.weight"],
                w[p + "mlp.up_proj.weight"],
                w[p + "mlp.down_proj.weight"],
                self.mlp_chunk,
            )

        # Final norm + LM head
        hidden = _rms_norm(hidden, self._norm().to(self.dtype), self.eps)
        logits = F.linear(hidden, self._lm_head().to(self.dtype))
        return logits, new_kv

    # ── Text generation ───────────────────────────────────────────────────────

    def _sample_next(
        self,
        logits: torch.Tensor,
        generated: torch.Tensor,
        temperature: float,
        top_p: float,
        top_k: int,
        repetition_penalty: float,
        do_sample: bool,
    ) -> torch.Tensor:
        """Sample the next token from logits."""
        next_logits = logits[:, -1, :].float()

        if repetition_penalty != 1.0:
            for tok in generated[0].tolist():
                v = next_logits[0, tok]
                next_logits[0, tok] = v / repetition_penalty if v > 0 \
                                      else v * repetition_penalty

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

        if temperature > 0:
            next_logits = next_logits / max(temperature, 1e-6)
        if top_k > 0:
            vals, _ = torch.topk(next_logits, top_k)
            next_logits[next_logits < vals[:, -1:]] = float("-inf")
        if 0.0 < top_p < 1.0:
            srt_l, srt_i = torch.sort(next_logits, descending=True)
            cum = torch.cumsum(F.softmax(srt_l, dim=-1), dim=-1)
            srt_l[cum - F.softmax(srt_l, dim=-1) > top_p] = float("-inf")
            next_logits = torch.full_like(next_logits, float("-inf")).scatter_(
                1, srt_i, srt_l
            )
        probs = F.softmax(next_logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 200,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        repetition_penalty: float = 1.0,
        eos_token_id: Optional[int] = None,
        pad_token_id: int = 0,
        use_cache: bool = True,
        do_sample: bool = True,
        **_,
    ) -> torch.Tensor:
        """
        Generate tokens with optional KV cache.

        With use_cache=True (default) the first forward pass processes the
        entire prompt, and every subsequent step processes only the one new
        token — dramatically faster than full-context re-computation.
        The layer weight cache means shard files are read from disk only once
        per generate() call; later steps use the in-memory copies.
        """
        generated = input_ids.clone()
        past_kv   = None
        cur_ids   = input_ids

        with torch.inference_mode():
            for _ in range(max_new_tokens):
                logits, new_kv = self._forward(cur_ids,
                                               past_kv if use_cache else None)
                if use_cache:
                    past_kv = new_kv

                next_token = self._sample_next(
                    logits, generated, temperature, top_p, top_k,
                    repetition_penalty, do_sample,
                )

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

                if eos_token_id is not None and (next_token == eos_token_id).all():
                    break

        return generated

    def generate_stream(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 200,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        repetition_penalty: float = 1.0,
        eos_token_id: Optional[int] = None,
        pad_token_id: int = 0,
        use_cache: bool = True,
        do_sample: bool = True,
        **_,
    ):
        """
        Token-streaming variant — yields each new token ID as a 1-D tensor
        the moment it is sampled, so callers can decode and stream to clients
        without waiting for the full generation to finish.
        """
        generated = input_ids.clone()
        past_kv   = None
        cur_ids   = input_ids

        with torch.inference_mode():
            for _ in range(max_new_tokens):
                logits, new_kv = self._forward(cur_ids,
                                               past_kv if use_cache else None)
                if use_cache:
                    past_kv = new_kv

                next_token = self._sample_next(
                    logits, generated, temperature, top_p, top_k,
                    repetition_penalty, do_sample,
                )

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

                yield next_token

                if eos_token_id is not None and (next_token == eos_token_id).all():
                    break


# ─────────────────────────────────────────────────────────────────────────────
#  Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _dtype_bytes(dt: torch.dtype) -> int:
    return {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(dt, 4)


# ─────────────────────────────────────────────────────────────────────────────
#  Self-test
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys, time
    from transformers import AutoTokenizer

    path    = sys.argv[1] if len(sys.argv) > 1 else "./model"
    chunk   = int(sys.argv[2]) if len(sys.argv) > 2 else 75

    print(f"\n=== ChunkedModel self-test  model={path}  chunk={chunk} MB ===\n")
    tok = AutoTokenizer.from_pretrained(path, local_files_only=True)
    m   = ChunkedModel(path, chunk_mb=chunk)

    for prompt in [
        "Say hello in one sentence.",
        "What is 2 + 2? Answer in one word.",
    ]:
        msgs = [{"role": "user", "content": prompt}]
        try:
            text = tok.apply_chat_template(msgs, tokenize=False,
                                           add_generation_prompt=True,
                                           enable_thinking=False)
        except TypeError:
            text = tok.apply_chat_template(msgs, tokenize=False,
                                           add_generation_prompt=True)
        ids  = tok(text, return_tensors="pt")["input_ids"]
        t0   = time.time()
        out  = m.generate(ids, max_new_tokens=40, temperature=0.7,
                          eos_token_id=tok.eos_token_id)
        ans  = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True).strip()
        print(f"Q: {prompt}")
        print(f"A: {ans}")
        print(f"   ({time.time()-t0:.1f}s)\n")

    print("=== done ===")