# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE). # # Re-authoring of the EmbeddingGemma 2 text encoder architecture as described by # modeling_embedding_gemma2.py / configuration_embedding_gemma2.py in Hugging Face transformers 5.19 # (Copyright 2026 The HuggingFace Team, Apache-2.0). Model weights: Google DeepMind (Gemma Team), Apache-2.0, # https://huggingface.co/google/embeddinggemma-2 — see ../NOTICE for the full attribution and modification list. """EmbeddingGemma 2 text encoder re-authored for Core AI export (standard [B, S, D] layout, batch 1). Everything the sentence-transformers pipeline does is inside one graph: input_ids [1, S] int32 + attention_mask [1, S] (1 = token, 0 = right padding) -> scaled token embedding (sqrt(512) folded into the fp16 table) -> projection-only per-layer embeddings (PLE) -> 24 bidirectional encoder layers (5 sliding-window-512 : 1 full, GQA 4 heads / 2 kv (sliding, hd 256) or 1 kv (full, hd 512)) -> final RMSNorm -> 512->768 projection -> masked mean pooling -> L2 normalization -> embedding [1, 768] Precision modes (PRECISION_MODE, set before construction): "ane" every op fp16 (Neural Engine eligible). RMSNorm and the final L2 norm divide the row by its max-abs value before squaring (the model's activations reach ~2000, whose square overflows fp16; this rescaling idea comes from the README of the FluidInference Core ML port of this model and was re-derived here). Softmax in fp16 (scores are bounded because q/k are RMS-normalized with scale 1.0). "gpu" fp16 weights/matmuls, fp32 islands for every RMSNorm, the softmax, the pooling and the L2 norm (the HF reference computes the norms and the softmax in fp32). "fp32" the whole model in fp32 (correctness gate for the re-authoring itself; never exported). """ from __future__ import annotations import json import math from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F PRECISION_MODE = "ane" # "ane" | "gpu" | "fp32" GQA_MODE = "grouped5d" # "expand" (HF repeat_kv: expand+reshape) | "grouped5d" (no repeat, 5-D batched matmul) | "repeat_interleave" # Global (full-attention, MQA: 1 kv head, head_dim 512) layers 5/11/17/23. "grouped5d" = the GQA_MODE path. # "mqa4d" (the released bundle) = plain 4-D matmul with k/v broadcast over the heads. With grouped5d the Neural # Engine compiler rejects the MQA matmul ("Unsupported mps.matmul op for this ANE architecture"), every global layer # becomes a GPU region, and MPSGraph then places the short buckets entirely on the GPU; mqa4d keeps one GPU region # (the token gather) and everything else on the ANE. Mathematically identical (fp32 max|diff| 3e-7). GLOBAL_ATTN_MODE = __import__("os").environ.get("EG2_GLOBAL_ATTN", "mqa4d") # repeat_interleave made the Neural Engine compiler's fused SDPA fail on the MQA layers (1 kv head, 4 q heads): # "'mps_spi.sdpa' op failed: grouping for the value tensor does not match the one available on the key tensor". MASK_VALUE = -10000.0 # finite (ANE softmax mishandles -inf); fp16-exact NORM_FLOOR = 1e-3 # floor for the row max-abs in "ane" RMSNorm (sqrt(eps) = 1e-3 for eps 1e-6) class Cfg: def __init__(self, model_dir: str): c = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] self.hidden = c["hidden_size"] # 512 self.inter = c["intermediate_size"] # 2048 self.n_layers = c["num_hidden_layers"] # 24 self.n_heads = c["num_attention_heads"] # 4 self.layer_types = c["layer_types"] self.sliding_window = c["sliding_window"] # 512 (inclusive distance) self.eps = c["rms_norm_eps"] self.vocab = c["vocab_size"] self.embed_dim = c["embedding_dim"] # 768 self.ple_dim = c["hidden_size_per_layer_input"] # 512 self.act = c["hidden_activation"] assert self.act == "gelu_pytorch_tanh", self.act self.rope_theta = {k: v["rope_theta"] for k, v in c["rope_parameters"].items()} per = c.get("per_layer_config", {}) self.head_dim = [per.get(f"{i:02d}", {}).get("head_dim", c["head_dim"]) for i in range(self.n_layers)] self.kv_heads = [per.get(f"{i:02d}", {}).get("num_key_value_heads", c["num_key_value_heads"]) for i in range(self.n_layers)] def _f16(x: float) -> torch.Tensor: return torch.tensor(x, dtype=torch.float16) def _recip(m: torch.Tensor) -> torch.Tensor: """1/m for m > 0 without a divide: the Neural Engine compiler rejects `mps.divide` (falls back to the GPU); rsqrt is native. rsqrt(m)*rsqrt(m) stays in fp16 range for m in [6e-5, 65504].""" r = torch.rsqrt(m) return r * r class RMSNorm(nn.Module): """Gemma RMSNorm: x * rsqrt(mean(x^2) + eps) * w (w stored as the raw weight, no +1).""" def __init__(self, dim: int, eps: float, with_scale: bool = True): super().__init__() self.with_scale = with_scale if with_scale: self.weight = nn.Parameter(torch.ones(dim), requires_grad=False) self.eps = eps self.register_buffer("eps16", _f16(eps), persistent=False) self.register_buffer("sqrt_eps16", _f16(math.sqrt(eps)), persistent=False) self.register_buffer("floor16", _f16(NORM_FLOOR), persistent=False) def forward(self, x: torch.Tensor) -> torch.Tensor: if PRECISION_MODE == "ane": inv_m = _recip(torch.maximum(x.abs().amax(-1, keepdim=True), self.floor16)) u = x * inv_m r = self.sqrt_eps16 * inv_m inv = torch.rsqrt((u * u).mean(-1, keepdim=True) + r * r) y = u * inv return y * self.weight if self.with_scale else y xf = x.float() y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) if self.with_scale: y = y * self.weight.float() return y.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) class Attention(nn.Module): def __init__(self, cfg: Cfg, i: int): super().__init__() self.hd = cfg.head_dim[i] self.n_heads = cfg.n_heads self.n_kv = cfg.kv_heads[i] self.groups = self.n_heads // self.n_kv self.q_proj = nn.Linear(cfg.hidden, cfg.n_heads * self.hd, bias=False) self.k_proj = nn.Linear(cfg.hidden, self.n_kv * self.hd, bias=False) self.v_proj = nn.Linear(cfg.hidden, self.n_kv * self.hd, bias=False) self.o_proj = nn.Linear(cfg.n_heads * self.hd, cfg.hidden, bias=False) self.q_norm = RMSNorm(self.hd, cfg.eps) self.k_norm = RMSNorm(self.hd, cfg.eps) self.v_norm = RMSNorm(self.hd, cfg.eps, with_scale=False) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: B, S, _ = x.shape q = self.q_proj(x).view(B, S, self.n_heads, self.hd).transpose(1, 2) k = self.k_proj(x).view(B, S, self.n_kv, self.hd).transpose(1, 2) v = self.v_proj(x).view(B, S, self.n_kv, self.hd).transpose(1, 2) q, k, v = self.q_norm(q), self.k_norm(k), self.v_norm(v) q = q * cos + rotate_half(q) * sin k = k * cos + rotate_half(k) * sin if self.n_kv == 1 and GLOBAL_ATTN_MODE == "mqa4d": # MQA (global layers, 1 kv head): plain 4-D matmul, k/v broadcast over the 4 query heads, no repeat scores = torch.matmul(q, k.transpose(-1, -2)) + bias p = torch.softmax(scores, dim=-1) if PRECISION_MODE == "ane" else torch.softmax(scores.float(), dim=-1).to(q.dtype) o = torch.matmul(p, v) elif GQA_MODE == "grouped5d" and self.groups > 1: qg = q.view(B, self.n_kv, self.groups, S, self.hd) scores = torch.matmul(qg, k.transpose(-1, -2)[:, :, None]) + bias[:, None] # [B, n_kv, groups, S, S] p = torch.softmax(scores, dim=-1) if PRECISION_MODE == "ane" else torch.softmax(scores.float(), dim=-1).to(q.dtype) o = torch.matmul(p, v[:, :, None]).view(B, self.n_heads, S, self.hd) else: if self.groups > 1: if GQA_MODE == "repeat_interleave": k = k.repeat_interleave(self.groups, dim=1) v = v.repeat_interleave(self.groups, dim=1) else: k = k[:, :, None].expand(B, self.n_kv, self.groups, S, self.hd).reshape(B, self.n_heads, S, self.hd) v = v[:, :, None].expand(B, self.n_kv, self.groups, S, self.hd).reshape(B, self.n_heads, S, self.hd) scores = torch.matmul(q, k.transpose(-1, -2)) + bias # scale 1.0 (the model's q/k norms carry it) p = torch.softmax(scores, dim=-1) if PRECISION_MODE == "ane" else torch.softmax(scores.float(), dim=-1).to(q.dtype) o = torch.matmul(p, v) o = o.transpose(1, 2).reshape(B, S, self.n_heads * self.hd) return self.o_proj(o) class MLP(nn.Module): def __init__(self, cfg: Cfg): super().__init__() self.gate_proj = nn.Linear(cfg.hidden, cfg.inter, bias=False) self.up_proj = nn.Linear(cfg.hidden, cfg.inter, bias=False) self.down_proj = nn.Linear(cfg.inter, cfg.hidden, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x)) class PLEBlock(nn.Module): def __init__(self, cfg: Cfg): super().__init__() self.per_layer_input_gate = nn.Linear(cfg.hidden, cfg.ple_dim, bias=False) self.per_layer_projection = nn.Linear(cfg.ple_dim, cfg.hidden, bias=False) self.post_per_layer_input_norm = RMSNorm(cfg.hidden, cfg.eps) def forward(self, x: torch.Tensor, ple: torch.Tensor) -> torch.Tensor: h = F.gelu(self.per_layer_input_gate(x), approximate="tanh") * ple return x + self.post_per_layer_input_norm(self.per_layer_projection(h)) class Layer(nn.Module): def __init__(self, cfg: Cfg, i: int): super().__init__() self.self_attn = Attention(cfg, i) self.mlp = MLP(cfg) self.input_layernorm = RMSNorm(cfg.hidden, cfg.eps) self.post_attention_layernorm = RMSNorm(cfg.hidden, cfg.eps) self.pre_feedforward_layernorm = RMSNorm(cfg.hidden, cfg.eps) self.post_feedforward_layernorm = RMSNorm(cfg.hidden, cfg.eps) self.ple_block = PLEBlock(cfg) self.register_buffer("layer_scalar", torch.ones(1)) def forward(self, x, ple, cos, sin, bias): x = x + self.post_attention_layernorm(self.self_attn(self.input_layernorm(x), cos, sin, bias)) x = x + self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(x))) x = self.ple_block(x, ple) return x * self.layer_scalar class EmbeddingGemma2Text(nn.Module): def __init__(self, cfg: Cfg, max_seq: int = 8192): super().__init__() self.cfg = cfg self.max_seq = max_seq self.embed_tokens = nn.Embedding(cfg.vocab, cfg.hidden) # sqrt(hidden) folded in at load self.per_layer_model_projection = nn.Linear(cfg.hidden, cfg.n_layers * cfg.ple_dim, bias=False) # scale folded self.per_layer_projection_norm = RMSNorm(cfg.ple_dim, cfg.eps) self.layers = nn.ModuleList([Layer(cfg, i) for i in range(cfg.n_layers)]) self.norm = RMSNorm(cfg.hidden, cfg.eps) self.embedding_projection = nn.Linear(cfg.hidden, cfg.embed_dim, bias=False) # RoPE tables [max_seq, hd] per (layer type); head dims differ: sliding 256 @ theta 1e4, full 512 @ theta 1e6 for lt in ("sliding_attention", "full_attention"): hd = cfg.head_dim[cfg.layer_types.index(lt)] inv = 1.0 / (cfg.rope_theta[lt] ** (torch.arange(0, hd, 2, dtype=torch.float32) / hd)) fr = torch.arange(max_seq, dtype=torch.float32)[:, None] * inv[None] emb = torch.cat((fr, fr), dim=-1) self.register_buffer(f"cos_{lt}", emb.cos(), persistent=False) self.register_buffer(f"sin_{lt}", emb.sin(), persistent=False) self.register_buffer("mask_value", _f16(MASK_VALUE), persistent=False) self.register_buffer("one", _f16(1.0), persistent=False) self.register_buffer("zero", _f16(0.0), persistent=False) self.register_buffer("floor16", _f16(NORM_FLOOR), persistent=False) self.register_buffer("window", torch.tensor(cfg.sliding_window, dtype=torch.int32), persistent=False) def _biases(self, attention_mask: torch.Tensor, S: int): pad = (self.one - attention_mask) * self.mask_value # [1, S] 0 for tokens, -1e4 for padding keys full = pad[:, None, None, :] # [1, 1, 1, S] broadcast over heads and queries pos = torch.arange(S, device=attention_mask.device, dtype=torch.int32) dist = (pos[:, None] - pos[None, :]).abs() band = torch.where(dist <= self.window, self.zero, self.mask_value) # [S, S] sliding = full + band[None, None] return full, sliding def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: S = input_ids.shape[1] x = self.embed_tokens(input_ids) # [1, S, 512], already scaled ple = self.per_layer_projection_norm( self.per_layer_model_projection(x).view(1, S, self.cfg.n_layers, self.cfg.ple_dim)) full, sliding = self._biases(attention_mask, S) cs = {} for lt in ("sliding_attention", "full_attention"): cs[lt] = (getattr(self, f"cos_{lt}")[:S].to(x.dtype)[None, None], getattr(self, f"sin_{lt}")[:S].to(x.dtype)[None, None]) for i, layer in enumerate(self.layers): lt = self.cfg.layer_types[i] x = layer(x, ple[:, :, i, :], cs[lt][0], cs[lt][1], sliding if lt == "sliding_attention" else full) x = self.embedding_projection(self.norm(x)) # [1, S, 768] # masked mean pooling: weights sum to 1, so no large partial sums in fp16 if PRECISION_MODE == "ane": w = attention_mask * _recip(attention_mask.sum(-1, keepdim=True)) # [1, S] (count exact in fp16 up to 2048) pooled = torch.matmul(w[:, None, :], x)[:, 0, :] # [1, 768] u = pooled * _recip(torch.maximum(pooled.abs().amax(-1, keepdim=True), self.floor16)) return u * torch.rsqrt((u * u).sum(-1, keepdim=True)) xf = x.float() wf = attention_mask.float() pooled = (xf * wf[:, :, None]).sum(1) / wf.sum(-1, keepdim=True) return (pooled / pooled.norm(dim=-1, keepdim=True)).to(x.dtype) def load_model(model_dir: str, dtype: torch.dtype, max_seq: int = 8192) -> EmbeddingGemma2Text: from safetensors import safe_open cfg = Cfg(model_dir) m = EmbeddingGemma2Text(cfg, max_seq) sd = {} with safe_open(str(Path(model_dir) / "model.safetensors"), "pt") as f: for k in f.keys(): if not k.startswith("language_model."): continue t = f.get_tensor(k).float() k = k[len("language_model."):] if k == "embed_tokens.weight": t = t * math.sqrt(cfg.hidden) elif k == "ple.per_layer_model_projection.weight": k = "per_layer_model_projection.weight" t = t * cfg.hidden ** -0.5 elif k == "ple.per_layer_projection_norm.weight": k = "per_layer_projection_norm.weight" sd[k] = t.to(dtype) missing, unexpected = m.load_state_dict(sd, strict=False) missing = [k for k in missing if not any(s in k for s in ("cos_", "sin_", "eps16", "sqrt_eps16", "floor16", "mask_value", "one", "zero", "window"))] assert not missing and not unexpected, (missing, unexpected) return m.to(dtype).eval()