"""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")