""" Reference forward pass for one decoder block under EPIC-Quant. Gap 1 closed: real attention via torch.nn.functional.scaled_dot_product_attention (F.scaled_dot_product_attention). On CUDA this is flash-attention; on CPU it's the math/efficient kernel. Either way, this replaces the O(seq^2) einsum we had before. The attention mask is built as a 4D bool additive mask per the SDPA contract: True = attend, False = mask out. We build causal + sliding-window together for sliding layers; pure causal for global layers. QK-norm is applied per-head after the projection. We don't implement RoPE in this reference forward — the engine's job is to measure quant error, not RoPE. To get exact parity with the real model, hook in `transformers.models.gemma4.modeling_gemma4.Gemma4RotaryEmbedding` (or read rope_parameters from config and use HF's apply_multimodal_rotary_pos_emb). """ from __future__ import annotations import time import math from typing import Dict, List, Optional import torch import torch.nn as nn import torch.nn.functional as F from .loader import MmapSafetensors from .engine import (EPICQuantEngine, QuantPolicy, PLEPolicy, KVPolicy, quantize_intN, dequantize_intN) from .packed import quantize_packed, dequantize_packed, total_packed_size_bytes from .layers import (LayerDims, layer_param_keys, ple_columns_for_layer, get_layer_dims) def _rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: var = x.float().pow(2).mean(-1, keepdim=True) x = x.float() * torch.rsqrt(var + eps) return (x.to(w.dtype) * w) def _gelu_tanh(x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate="tanh") def build_attn_mask(seq_len: int, is_global: bool, sliding_window: int, device: torch.device) -> torch.Tensor: """Return a 4D additive attention mask [1, 1, S, S] for SDPA. Convention: 0.0 = attend, -inf = mask out. Causal + sliding window if not global, pure causal if global. """ idx = torch.arange(seq_len, device=device) # Causal: q can attend to k if k_idx <= q_idx causal = idx[None, :] <= idx[:, None] # [S, S] if is_global: keep = causal else: # Sliding: also require q_idx - k_idx < sliding_window within = (idx[:, None] - idx[None, :]) < sliding_window keep = causal & within # Convert to additive: keep -> 0, mask -> -inf add = torch.zeros(seq_len, seq_len, dtype=torch.bfloat16, device=device) add.masked_fill_(~keep, float("-inf")) return add[None, None, :, :] def real_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_global: bool, sliding_window: int, layer_scalar: torch.Tensor) -> torch.Tensor: """Real attention via F.scaled_dot_product_attention. Inputs are post-projection, in [B, H, S, D] layout, BF16. GQA is handled by repeating K/V along the head axis. """ B, H_q, S, D = q.shape H_kv = k.shape[1] repeat = H_q // H_kv if repeat > 1: k = k.repeat_interleave(repeat, dim=1) v = v.repeat_interleave(repeat, dim=1) add_mask = build_attn_mask(S, is_global, sliding_window, q.device) # SDPA wants either a bool mask (True=keep) or a float additive mask # of shape [B, H, S, S] or [1, 1, S, S]. We use additive for clarity. out = F.scaled_dot_product_attention(q, k, v, attn_mask=add_mask, dropout_p=0.0, is_causal=False) out = out.transpose(1, 2).contiguous().view(B, S, H_q * D) return out * layer_scalar.float().to(out.dtype) def forward_one_layer(engine: EPICQuantEngine, layer_idx: int, hidden_states: torch.Tensor, token_ids: torch.Tensor, sliding_window: int = 512) -> Dict: """Run one decoder block with EPIC-Quant. Returns output + per-step stats.""" sf = engine.sf dims = get_layer_dims(layer_idx, engine.layer_types) is_global = dims.is_global base = f"model.language_model.layers.{layer_idx}" t0 = time.perf_counter() stats = {"layer": layer_idx, "is_global": is_global} # 1. Input layernorm w_in = sf.get_tensor(f"{base}.input_layernorm.weight") x = _rms_norm(hidden_states, w_in) stats["input_norm_ms"] = (time.perf_counter() - t0) * 1000 t1 = time.perf_counter() # 2. PLE — gather per-layer slice, then gate + project ple_vecs = [] for t in token_ids.tolist(): ple_vecs.append(engine.ple_cache.lookup(t, layer_idx)) ple_stack = torch.stack(ple_vecs) # [S, ple_dim] BF16 w_gate = sf.get_tensor(f"{base}.per_layer_input_gate.weight") # [256, 2560] w_proj = sf.get_tensor(f"{base}.per_layer_projection.weight") # [2560, 256] w_post = sf.get_tensor(f"{base}.post_per_layer_input_norm.weight") qbits = engine.quant.bits_ple_per_layer pg, sg = quantize_packed(w_gate, qbits) pp, sp = quantize_packed(w_proj, qbits) w_gate_dq = dequantize_packed(pg, sg, w_gate.shape[0], w_gate.shape[1], qbits) w_proj_dq = dequantize_packed(pp, sp, w_proj.shape[0], w_proj.shape[1], qbits) ple_gate = ple_stack @ w_gate_dq # [S, 2560] ple_value = ple_stack @ w_proj_dq.T ple_contrib = (ple_gate * ple_value) # [S, 2560] x = x + ple_contrib x = _rms_norm(x, w_post) stats["ple_ms"] = (time.perf_counter() - t1) * 1000 stats["ple_gate_packed_bytes"] = pg.numel() + sg.numel() * 2 stats["ple_proj_packed_bytes"] = pp.numel() + sp.numel() * 2 stats["ple_gate_recon_l2"] = (w_gate_dq.float() - w_gate.float()).norm().item() / \ (w_gate.float().norm().item() + 1e-9) stats["ple_proj_recon_l2"] = (w_proj_dq.float() - w_proj.float()).norm().item() / \ (w_proj.float().norm().item() + 1e-9) t2 = time.perf_counter() # 3. Self-attn: pack q/k/v/o, then dequant + matmul, then real SDPA Wq = sf.get_tensor(f"{base}.self_attn.q_proj.weight") Wk = sf.get_tensor(f"{base}.self_attn.k_proj.weight") Wv = sf.get_tensor(f"{base}.self_attn.v_proj.weight") Wo = sf.get_tensor(f"{base}.self_attn.o_proj.weight") qbits = engine.quant.bits_global_attn if is_global else engine.quant.bits_sliding_attn pq, sq = quantize_packed(Wq, qbits) pk, sk = quantize_packed(Wk, qbits) pv, sv = quantize_packed(Wv, qbits) po, so = quantize_packed(Wo, qbits) stats["attn_q_packed_bytes"] = total_packed_size_bytes(dims.q_out, dims.hidden, qbits) stats["attn_k_packed_bytes"] = total_packed_size_bytes(dims.kv_out, dims.hidden, qbits) stats["attn_v_packed_bytes"] = total_packed_size_bytes(dims.kv_out, dims.hidden, qbits) stats["attn_o_packed_bytes"] = total_packed_size_bytes(dims.hidden, dims.q_out, qbits) Wq_dq = dequantize_packed(pq, sq, dims.q_out, dims.hidden, qbits) Wk_dq = dequantize_packed(pk, sk, dims.kv_out, dims.hidden, qbits) Wv_dq = dequantize_packed(pv, sv, dims.kv_out, dims.hidden, qbits) Wo_dq = dequantize_packed(po, so, dims.hidden, dims.q_out, qbits) stats["attn_q_recon_l2"] = (Wq_dq.float() - Wq.float()).norm().item() / (Wq.float().norm().item() + 1e-9) stats["attn_k_recon_l2"] = (Wk_dq.float() - Wk.float()).norm().item() / (Wk.float().norm().item() + 1e-9) stats["attn_v_recon_l2"] = (Wv_dq.float() - Wv.float()).norm().item() / (Wv.float().norm().item() + 1e-9) stats["attn_o_recon_l2"] = (Wo_dq.float() - Wo.float()).norm().item() / (Wo.float().norm().item() + 1e-9) # Project B, S, H = hidden_states.shape q = (x @ Wq_dq.T) # [B, S, q_out] k = (x @ Wk_dq.T) # [B, S, kv_out] v = (x @ Wv_dq.T) q = q.view(B, S, dims.q_out // dims.head_dim, dims.head_dim).transpose(1, 2) k = k.view(B, S, dims.kv_out // dims.head_dim, dims.head_dim).transpose(1, 2) v = v.view(B, S, dims.kv_out // dims.head_dim, dims.head_dim).transpose(1, 2) layer_scalar = sf.get_tensor(f"{base}.layer_scalar") attn_out = real_attention(q, k, v, is_global, sliding_window, layer_scalar) x = attn_out @ Wo_dq.T stats["attn_ms"] = (time.perf_counter() - t2) * 1000 t3 = time.perf_counter() # 4. Post-attn norm + residual w_pa = sf.get_tensor(f"{base}.post_attention_layernorm.weight") hidden_states = hidden_states + _rms_norm(x, w_pa) x = _rms_norm(hidden_states, w_pa) stats["post_attn_ms"] = (time.perf_counter() - t3) * 1000 t4 = time.perf_counter() # 5. MLP — pack gate/up/down, dequant, matmul mbits = engine.quant.bits_global_mlp if is_global else engine.quant.bits_sliding_mlp Wg = sf.get_tensor(f"{base}.mlp.gate_proj.weight") Wu = sf.get_tensor(f"{base}.mlp.up_proj.weight") Wd = sf.get_tensor(f"{base}.mlp.down_proj.weight") pg, sg = quantize_packed(Wg, mbits) pu, su = quantize_packed(Wu, mbits) pd_, sd = quantize_packed(Wd, mbits) Wg_dq = dequantize_packed(pg, sg, Wg.shape[0], Wg.shape[1], mbits) Wu_dq = dequantize_packed(pu, su, Wu.shape[0], Wu.shape[1], mbits) Wd_dq = dequantize_packed(pd_, sd, Wd.shape[0], Wd.shape[1], mbits) stats["mlp_gate_packed_bytes"] = total_packed_size_bytes(Wg.shape[0], Wg.shape[1], mbits) stats["mlp_up_packed_bytes"] = total_packed_size_bytes(Wu.shape[0], Wu.shape[1], mbits) stats["mlp_down_packed_bytes"] = total_packed_size_bytes(Wd.shape[0], Wd.shape[1], mbits) stats["mlp_gate_recon_l2"] = (Wg_dq.float() - Wg.float()).norm().item() / (Wg.float().norm().item() + 1e-9) stats["mlp_up_recon_l2"] = (Wu_dq.float() - Wu.float()).norm().item() / (Wu.float().norm().item() + 1e-9) stats["mlp_down_recon_l2"] = (Wd_dq.float() - Wd.float()).norm().item() / (Wd.float().norm().item() + 1e-9) mlp_out = (_gelu_tanh(x @ Wg_dq.T) * (x @ Wu_dq.T)) @ Wd_dq.T w_pf = sf.get_tensor(f"{base}.pre_feedforward_layernorm.weight") w_pof = sf.get_tensor(f"{base}.post_feedforward_layernorm.weight") mlp_out = _rms_norm(mlp_out, w_pf) hidden_states = hidden_states + _rms_norm(mlp_out, w_pof) stats["mlp_ms"] = (time.perf_counter() - t4) * 1000 stats["total_ms"] = (time.perf_counter() - t0) * 1000 return {"hidden": hidden_states, "stats": stats}