scripts/decode_v12_fast.py

"""KV-cache speculative inference for CYPHER V12 encoder-decoder.

Why this exists:
  Current generate() in cypher_omega_v3_arch.py recomputes the full decoder
  forward at each step (no caching). For N output tokens, cross-attention K,V
  are recomputed N times even though encoder_hidden is constant. Self-attention
  recomputes K,V for all N-1 previous decoder tokens at step N.

  This module precomputes cross-attn K,V once, and caches self-attn K,V
  incrementally — expected speedup 5-15x at 100 tokens, scaling with length.

Correctness strategy:
  We extract weights from the trained nn.MultiheadAttention modules (in_proj_weight
  shape [3*H, H] split into Q/K/V projections, bias=False per arch). We
  reimplement the forward with cache. Smoke test verifies output equivalence
  vs non-cached baseline before any benchmark.

Usage:
  fast = CachedDecoder(model, tokenizer)
  ids, stats = fast.generate(prompt_ids, attn_mask, max_new_tokens=100,
                             temperature=0.7, top_k=40)
"""
from __future__ import annotations
import math
import time
import torch
import torch.nn.functional as F
from typing import Optional


class CachedDecoder:
    """KV-cache wrapper around CypherEncoderDecoderV3.

    Caches:
      - encoder_hidden (once)
      - cross-attn K, V per decoder block (once, derived from encoder_hidden)
      - self-attn K, V per decoder block (incrementally appended)
    """

    def __init__(self, model, tokenizer):
        self.model = model
        self.tok = tokenizer
        self.device = next(model.parameters()).device
        self.n_layers = len(model.decoder_blocks)
        self.hidden = model.hidden_size
        self.num_heads = 16  # from arch
        self.head_dim = self.hidden // self.num_heads
        # Pre-extract weights for fast forward (avoid attribute lookup per step)
        self._extract_weights()

    def _extract_weights(self):
        """Pre-extract block weights into flat lists indexed by layer."""
        self.w_self_in = []   # nn.MultiheadAttention in_proj_weight: [3H, H]
        self.w_self_out = []  # nn.MultiheadAttention out_proj.weight: [H, H]
        self.w_cross_q = []
        self.w_cross_k = []
        self.w_cross_v = []
        self.w_cross_out = []
        self.norm1 = []
        self.norm2 = []
        self.norm3 = []
        self.ffn = []
        for blk in self.model.decoder_blocks:
            sa = blk.self_attn
            ca = blk.cross_attn
            self.w_self_in.append(sa.in_proj_weight)
            self.w_self_out.append(sa.out_proj.weight)
            self.w_cross_q.append(ca.q_proj.weight)
            self.w_cross_k.append(ca.k_proj.weight)
            self.w_cross_v.append(ca.v_proj.weight)
            self.w_cross_out.append(ca.out_proj.weight)
            self.norm1.append(blk.norm1)
            self.norm2.append(blk.norm2)
            self.norm3.append(blk.norm3)
            self.ffn.append(blk.ffn)

    def _prefill_cross_kv(self, encoder_hidden):
        """Compute cross-attn K, V once per block. Returns list of (K, V) per layer."""
        B, Te, H = encoder_hidden.shape
        cross_cache = []
        for i in range(self.n_layers):
            K = F.linear(encoder_hidden, self.w_cross_k[i])  # [B, Te, H]
            V = F.linear(encoder_hidden, self.w_cross_v[i])
            # Reshape to [B, num_heads, Te, head_dim]
            K = K.view(B, Te, self.num_heads, self.head_dim).transpose(1, 2)
            V = V.view(B, Te, self.num_heads, self.head_dim).transpose(1, 2)
            cross_cache.append((K, V))
        return cross_cache

    def _self_attn_step(self, layer_idx, x_t, self_cache):
        """1-step self-attention with cache. x_t: [B, 1, H].

        self_cache[layer_idx] = (K_cum, V_cum) shape [B, h, Tcur, dk] or None.
        """
        B, _, H = x_t.shape
        h = self.num_heads
        dk = self.head_dim
        # Project Q, K, V from packed in_proj_weight [3H, H]
        # split: 0..H = Q, H..2H = K, 2H..3H = V
        Q = F.linear(x_t, self.w_self_in[layer_idx][:H])     # [B, 1, H]
        K_new = F.linear(x_t, self.w_self_in[layer_idx][H:2*H])
        V_new = F.linear(x_t, self.w_self_in[layer_idx][2*H:3*H])
        Q = Q.view(B, 1, h, dk).transpose(1, 2)   # [B, h, 1, dk]
        K_new = K_new.view(B, 1, h, dk).transpose(1, 2)
        V_new = V_new.view(B, 1, h, dk).transpose(1, 2)

        # Append to cache
        if self_cache[layer_idx] is None:
            K_cum, V_cum = K_new, V_new
        else:
            K_prev, V_prev = self_cache[layer_idx]
            K_cum = torch.cat([K_prev, K_new], dim=2)  # [B, h, Tcur+1, dk]
            V_cum = torch.cat([V_prev, V_new], dim=2)
        self_cache[layer_idx] = (K_cum, V_cum)

        # Causal attention (Q[1] attends to all K[Tcur+1])
        # scaled dot-product (no mask needed: latest Q attends to all past+current K)
        attn = F.scaled_dot_product_attention(Q, K_cum, V_cum, is_causal=False)
        # → [B, h, 1, dk]
        attn = attn.transpose(1, 2).contiguous().view(B, 1, H)
        out = F.linear(attn, self.w_self_out[layer_idx])
        return out

    def _cross_attn_step(self, layer_idx, x_t, cross_cache, encoder_pad_mask):
        """1-step cross-attention with cached K, V. x_t: [B, 1, H]."""
        B, _, H = x_t.shape
        h = self.num_heads
        dk = self.head_dim
        Q = F.linear(x_t, self.w_cross_q[layer_idx])      # [B, 1, H]
        Q = Q.view(B, 1, h, dk).transpose(1, 2)            # [B, h, 1, dk]
        K_cached, V_cached = cross_cache[layer_idx]        # [B, h, Te, dk]

        # Manual scaled dot-product with pad mask
        scores = torch.matmul(Q, K_cached.transpose(-2, -1)) / math.sqrt(dk)  # [B, h, 1, Te]
        if encoder_pad_mask is not None:
            mask = encoder_pad_mask.unsqueeze(1).unsqueeze(2)  # [B, 1, 1, Te]
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V_cached)                 # [B, h, 1, dk]
        out = out.transpose(1, 2).contiguous().view(B, 1, H)
        out = F.linear(out, self.w_cross_out[layer_idx])
        return out

    def _step_forward(self, token_id, position, encoder_pad_mask,
                      cross_cache, self_cache):
        """1-token forward: embed → 12 decoder blocks (cached) → lm_head.

        token_id: [B] long tensor
        position: int (position index for pos embedding)
        """
        B = token_id.size(0)
        # Embed
        tok = token_id.view(B, 1)
        tok_emb = self.model.encoder.token_embedding(tok)
        pos = torch.tensor([[position]], device=self.device).expand(B, 1)
        pos_emb = self.model.dec_pos_embedding(pos)
        x = self.model.dec_embed_norm(tok_emb + pos_emb)   # [B, 1, H]
        # No dropout in eval mode

        for i in range(self.n_layers):
            # 1) self-attention with cache (residual + norm1)
            sa_out = self._self_attn_step(i, x, self_cache)
            x = self.norm1[i](x + sa_out)
            # 2) cross-attention with cached K, V (residual + norm2)
            ca_out = self._cross_attn_step(i, x, cross_cache, encoder_pad_mask)
            x = self.norm2[i](x + ca_out)
            # 3) FFN (residual + norm3)
            ff_out = self.ffn[i](x)
            x = self.norm3[i](x + ff_out)

        x = self.model.decoder_norm(x)
        logits = self.model.encoder.lm_head(x)             # [B, 1, V]
        return logits[:, -1, :]                            # [B, V]

    @torch.no_grad()
    def generate(self, encoder_input_ids, encoder_attn_mask,
                 bos_token_id=1, eos_token_id=4, sep_token_id=2,
                 max_new_tokens=100, temperature=0.7, top_k=40):
        """Fast cached generation. Greedy + top-k sampling."""
        self.model.eval()
        B = encoder_input_ids.size(0)

        # 1) Encode (once)
        encoder_hidden, encoder_pad_mask = self.model.encode(
            encoder_input_ids, encoder_attn_mask
        )

        # 2) Prefill cross-attn K, V (once)
        cross_cache = self._prefill_cross_kv(encoder_hidden)
        self_cache = [None] * self.n_layers

        # 3) Decode loop
        generated = []
        current_token = torch.full((B,), bos_token_id, dtype=torch.long, device=self.device)
        for step in range(max_new_tokens):
            logits = self._step_forward(current_token, step, encoder_pad_mask,
                                        cross_cache, self_cache)
            # NaN guard
            if torch.isnan(logits).any():
                logits = torch.nan_to_num(logits, nan=-1e9)
            logits = logits / max(temperature, 1e-6)
            # Top-k
            if top_k and top_k > 0 and top_k < logits.size(-1):
                v, _ = torch.topk(logits, top_k)
                logits = logits.masked_fill(logits < v[:, [-1]], float("-inf"))
            logits[:, 0] = float("-inf")  # never predict PAD
            probs = F.softmax(logits, dim=-1)
            sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)  # [B]
            generated.append(sampled.unsqueeze(1))
            current_token = sampled
            if B == 1 and sampled.item() in (eos_token_id, sep_token_id):
                break

        return torch.cat(generated, dim=1) if generated else torch.empty(B, 0,
                                                                         dtype=torch.long, device=self.device)


def smoke_equivalence(model, tokenizer, prompt, max_new=20, temperature=1e-9):
    """Verify cached generate() matches non-cached for identical (deterministic) sampling.

    Uses near-zero temperature for argmax-like greedy. Both paths should sample
    the same token IDs (modulo numerical noise at temperature=0).
    """
    device = next(model.parameters()).device
    p_ids = tokenizer._tok.encode(prompt).ids
    enc = torch.tensor([[tokenizer.CLS] + p_ids + [tokenizer.SEP]], device=device)
    enc_attn = torch.ones_like(enc)

    torch.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    base = model.generate(
        enc, enc_attn, bos_token_id=tokenizer.CLS, eos_token_id=tokenizer.EOS,
        sep_token_id=tokenizer.SEP, max_new_tokens=max_new,
        temperature=temperature, top_k=1,  # greedy
    )

    torch.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    fast = CachedDecoder(model, tokenizer)
    cached, _ = (fast.generate(
        enc, enc_attn, bos_token_id=tokenizer.CLS, eos_token_id=tokenizer.EOS,
        sep_token_id=tokenizer.SEP, max_new_tokens=max_new,
        temperature=temperature, top_k=1,
    ), None)

    base_list = base[0].tolist()
    cached_list = cached[0].tolist()
    n = min(len(base_list), len(cached_list))
    matches = sum(1 for i in range(n) if base_list[i] == cached_list[i])
    return matches, n, base_list, cached_list


if __name__ == "__main__":
    import sys
    sys.path.insert(0, "/workspace/CYPHER_V12/scripts")
    sys.path.insert(0, "/workspace/CYPHER_V12/scripts/cortex")
    from cortex.model_v2 import CypherCortexV2, CypherTokenizerV2, create_cypher_cortex_v2
    from cypher_omega_v3_arch import CypherEncoderDecoderV3

    device = torch.device("cuda:0")
    tok = CypherTokenizerV2("/workspace/CYPHER_V12/scripts/cortex/cypher_bpe_v2.json",
                            max_length=1024)
    encoder = create_cypher_cortex_v2(device=device, max_seq_len=4096,
                                       vocab_size=tok.vocab_size, mtp_heads=5)
    model = CypherEncoderDecoderV3(
        encoder_model=encoder, vocab_size=tok.vocab_size,
        hidden_size=1024, num_decoder_layers=12, num_heads=16, ff_size=4096,
        max_decoder_len=512, encoder_frozen=False, use_gradient_checkpointing=False,
    ).to(device)

    ckpt = torch.load("/workspace/CYPHER_V12/ckpts/cypher_omega_v12_FINAL.pt",
                       map_location=device, weights_only=False)
    state = ckpt.get("model", ckpt) if isinstance(ckpt, dict) else ckpt
    m, u = model.load_state_dict(state, strict=False)
    print(f"Ckpt loaded: missing={len(m)} unexpected={len(u)}")
    model.eval()

    # === Smoke test equivalence (greedy) ===
    print("\n=== SMOKE TEST EQUIVALENCE (greedy temperature=1e-9 top_k=1) ===")
    prompts = [
        "User: What is SQL injection?\nCYPHER-Ω: ",
        "User: Explain ransomware briefly.\nCYPHER-Ω: ",
    ]
    for p in prompts:
        matches, n, base, cached = smoke_equivalence(model, tok, p, max_new=20)
        print(f"  Prompt: {p[:40]!r}")
        print(f"    Match: {matches}/{n} tokens identical")
        if matches != n:
            print(f"    Base : {base[:10]}")
            print(f"    Cached: {cached[:10]}")

    # === Bench cached vs non-cached ===
    print("\n=== BENCH cached vs non-cached ===")
    bench_prompts = [
        "User: What is a SQL injection attack?\nCYPHER-Ω: ",
        "User: Explain ransomware in 3 sentences.\nCYPHER-Ω: ",
        "User: What does MITRE ATT&CK T1059 mean?\nCYPHER-Ω: ",
        "User: How to detect a phishing email?\nCYPHER-Ω: ",
        "User: Quelle est la différence entre IDS et IPS?\nCYPHER-Ω: ",
    ]
    fast = CachedDecoder(model, tok)

    # Warm up
    p = bench_prompts[0]
    p_ids = tok._tok.encode(p).ids
    enc = torch.tensor([[tok.CLS] + p_ids + [tok.SEP]], device=device)
    enc_attn = torch.ones_like(enc)
    _ = fast.generate(enc, enc_attn, bos_token_id=tok.CLS, eos_token_id=tok.EOS,
                       sep_token_id=tok.SEP, max_new_tokens=10,
                       temperature=0.7, top_k=40)

    # Bench
    for max_new in [50, 100, 200]:
        print(f"\n--- max_new_tokens={max_new} ---")
        base_total_tok = 0
        base_total_time = 0
        cached_total_tok = 0
        cached_total_time = 0
        for p in bench_prompts:
            p_ids = tok._tok.encode(p).ids
            enc = torch.tensor([[tok.CLS] + p_ids + [tok.SEP]], device=device)
            enc_attn = torch.ones_like(enc)

            torch.cuda.synchronize()
            t0 = time.perf_counter()
            base_out = model.generate(
                enc, enc_attn, bos_token_id=tok.CLS, eos_token_id=tok.EOS,
                sep_token_id=tok.SEP, max_new_tokens=max_new,
                temperature=0.7, top_k=40,
            )
            torch.cuda.synchronize()
            base_dt = time.perf_counter() - t0
            base_n = base_out.size(1)
            base_total_tok += base_n
            base_total_time += base_dt

            torch.cuda.synchronize()
            t0 = time.perf_counter()
            cached_out = fast.generate(
                enc, enc_attn, bos_token_id=tok.CLS, eos_token_id=tok.EOS,
                sep_token_id=tok.SEP, max_new_tokens=max_new,
                temperature=0.7, top_k=40,
            )
            torch.cuda.synchronize()
            cached_dt = time.perf_counter() - t0
            cached_n = cached_out.size(1)
            cached_total_tok += cached_n
            cached_total_time += cached_dt

        base_tps = base_total_tok / base_total_time
        cached_tps = cached_total_tok / cached_total_time
        speedup = cached_tps / base_tps
        print(f"  Baseline:  {base_total_tok} tok / {base_total_time:.2f}s = {base_tps:.1f} tok/s")
        print(f"  Cached:    {cached_total_tok} tok / {cached_total_time:.2f}s = {cached_tps:.1f} tok/s")
        print(f"  SPEEDUP:   {speedup:.2f}x")