| """Gemma 3 (text) for the ExecuTorch WebGPU playground: model, export and CPU validation. |
| |
| ExecuTorch 1.5.1's export_llm has no Gemma 3 model class, and its llama_transformer lacks the |
| pieces Gemma 3 needs (post-attention/post-FFN norms, embedding scaling, GELU-tanh MLP, per-layer |
| RoPE base, sliding-window layers). This module re-implements the Gemma 3 decoder with the ops the |
| public WebGPU runtime registers and lowers it through the same LLMEdgeManager pipeline export_llm |
| uses (same quantization transforms, edge config, to_executorch config), with WebGPUPartitioner. |
| |
| Program contract (same as the other playground LLMs): |
| forward(tokens int64 [1, 1..prefill_max], input_pos int64 [1]) -> logits float32 [1, vocab] |
| (last position), KV caches inside the program as fp32 mutable buffers. |
| |
| Attention: |
| * global layers: llama.sdpa_with_kv_cache (writes K/V at input_pos, causal over [0, pos+S)). |
| * sliding-window layers: the public runtime has no masked/windowed KV-cache op and its |
| stand-alone update_cache does not accept a dynamic position. So each local layer keeps a |
| position-indexed cache with PAD leading rows (position x lives at row x + PAD) and |
| 1. writes K/V with sdpa_with_kv_cache(k, k, v, cache, pos + PAD) (1-head dummy query; |
| its output is discarded, it only performs the cache write), |
| 2. reads a fixed-length window of L = window + prefill_max - 1 rows ending at the last |
| token of the chunk, and attends with bmm -> + mask -> softmax -> bmm. |
| The additive mask is band[prefill_max - S : prefill_max] (relative positions, constant) plus |
| padbias[pos + S - 1 : ... + L] (masks rows that hold no position yet). This reproduces |
| transformers' sliding mask exactly: key visible iff 0 <= q_pos - k_pos < sliding_window. |
| """ |
|
|
| import contextlib |
| import json |
| import re |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| NEG = -1.0e30 |
|
|
|
|
| def _load_custom_ops(): |
| import executorch.extension.llm.custom_ops.custom_ops |
|
|
|
|
| class RMSNorm(nn.Module): |
| """Gemma RMSNorm with the (1 + w) offset folded into `weight` (fp32), in the op order the |
| Vulkan rms_norm pattern matches: x * rsqrt(mean(x^2) + eps) * weight.""" |
|
|
| def __init__(self, dim, eps): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim), requires_grad=False) |
|
|
| def forward(self, x): |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return x * self.weight |
|
|
|
|
| def _rope(q, k, cos, sin): |
| |
| |
| from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern |
|
|
| return HfRotaryEmbeddingPattern()(q, k, cos, sin) |
|
|
|
|
| class Layer(nn.Module): |
| def __init__(self, cfg, sliding, ctx, pad): |
| super().__init__() |
| d, hd = cfg["hidden_size"], cfg["head_dim"] |
| self.nh, self.nkv, self.hd = ( |
| cfg["num_attention_heads"], |
| cfg["num_key_value_heads"], |
| hd, |
| ) |
| self.sliding = sliding |
| eps = cfg["rms_norm_eps"] |
| self.wq = nn.Linear(d, self.nh * hd, bias=False) |
| self.wk = nn.Linear(d, self.nkv * hd, bias=False) |
| self.wv = nn.Linear(d, self.nkv * hd, bias=False) |
| self.wo = nn.Linear(self.nh * hd, d, bias=False) |
| self.q_norm = RMSNorm(hd, eps) |
| self.k_norm = RMSNorm(hd, eps) |
| self.attention_norm = RMSNorm(d, eps) |
| self.post_attention_norm = RMSNorm(d, eps) |
| self.ffn_norm = RMSNorm(d, eps) |
| self.post_ffn_norm = RMSNorm(d, eps) |
| self.w1 = nn.Linear(d, cfg["intermediate_size"], bias=False) |
| self.w3 = nn.Linear(d, cfg["intermediate_size"], bias=False) |
| self.w2 = nn.Linear(cfg["intermediate_size"], d, bias=False) |
| self.rows = ctx + (pad if sliding else 0) |
| self.make_buffers() |
|
|
| def make_buffers(self): |
| shape = (1, self.rows, self.nkv, self.hd) |
| with torch.device("cpu"): |
| self.register_buffer("k_cache", torch.zeros(shape), persistent=False) |
| self.register_buffer("v_cache", torch.zeros(shape), persistent=False) |
|
|
| def attention(self, x, p, S, cos, sin, mask, pad, win): |
| q = self.wq(x).view(1, S, self.nh, self.hd) |
| k = self.wk(x).view(1, S, self.nkv, self.hd) |
| v = self.wv(x).view(1, S, self.nkv, self.hd) |
| q, k = _rope(self.q_norm(q), self.k_norm(k), cos, sin) |
| if not self.sliding: |
| out = torch.ops.llama.sdpa_with_kv_cache( |
| q, k, v, self.k_cache, self.v_cache, p, S, None, 0.0, True, None |
| ) |
| return out.view(1, S, self.nh * self.hd) |
| |
| torch.ops.llama.sdpa_with_kv_cache( |
| k, k, v, self.k_cache, self.v_cache, p + pad, S, None, 0.0, True, None |
| ) |
| start = p + S - 1 |
| kw = self.k_cache.narrow(1, start, win).view(1, win, self.hd).transpose(1, 2) |
| vw = self.v_cache.narrow(1, start, win).view(1, win, self.hd) |
| qf = q.reshape(1, S * self.nh, self.hd) |
| scores = torch.bmm(qf, kw) * (self.hd**-0.5) |
| scores = scores.view(1, S, self.nh, win) + mask |
| probs = torch.softmax(scores, dim=-1).view(1, S * self.nh, win) |
| return torch.bmm(probs, vw).view(1, S, self.nh * self.hd) |
|
|
| def forward(self, x, p, S, cos, sin, mask, pad, win): |
| h = x + self.post_attention_norm( |
| self.wo(self.attention(self.attention_norm(x), p, S, cos, sin, mask, pad, win)) |
| ) |
| x = self.ffn_norm(h) |
| ff = self.w2(F.gelu(self.w1(x), approximate="tanh") * self.w3(x)) |
| return h + self.post_ffn_norm(ff) |
|
|
|
|
| class Gemma3Text(nn.Module): |
| def __init__(self, cfg, max_context, prefill_max): |
| super().__init__() |
| assert cfg["num_key_value_heads"] == 1, ( |
| "window attention folds heads; needs 1 KV head" |
| ) |
| assert cfg["query_pre_attn_scalar"] == cfg["head_dim"], ( |
| "sdpa uses 1/sqrt(head_dim)" |
| ) |
| assert cfg.get("attn_logit_softcapping") is None |
| assert cfg.get("final_logit_softcapping") is None |
| assert cfg["hidden_activation"] == "gelu_pytorch_tanh" |
| self.cfg = cfg |
| self.ctx, self.prefill_max = max_context, prefill_max |
| self.window = cfg["sliding_window"] |
| self.win = self.window + prefill_max - 1 |
| self.pad = self.win - 1 |
| self.layer_types = layer_types(cfg) |
| d, V = cfg["hidden_size"], cfg["vocab_size"] |
| self.tok_embeddings = nn.Embedding(V, d) |
| |
| self.embed_scale = float(torch.tensor(d**0.5, dtype=torch.float32, device="cpu")) |
| self.layers = nn.ModuleList( |
| Layer(cfg, t == "sliding_attention", max_context, self.pad) |
| for t in self.layer_types |
| ) |
| self.norm = RMSNorm(d, cfg["rms_norm_eps"]) |
| self.output = nn.Linear(d, V, bias=False) |
| self.make_buffers() |
|
|
| def make_buffers(self): |
| """Constant/state buffers (rope tables, window masks, KV caches), always built on CPU.""" |
| with torch.device("cpu"): |
| cos_g, sin_g, cos_l, sin_l = rope_tables(self.cfg, self.ctx) |
| self.register_buffer("cos_global", cos_g, persistent=False) |
| self.register_buffer("sin_global", sin_g, persistent=False) |
| self.register_buffer("cos_local", cos_l, persistent=False) |
| self.register_buffer("sin_local", sin_l, persistent=False) |
| r = torch.arange(self.prefill_max)[:, None] |
| j = torch.arange(self.win)[None, :] |
| band = torch.where((j >= r) & (j <= r + self.window - 1), 0.0, NEG) |
| self.register_buffer("band", band.float(), persistent=False) |
| padbias = torch.zeros(self.pad + self.ctx) |
| padbias[: self.pad] = NEG |
| self.register_buffer("padbias", padbias, persistent=False) |
| for layer in self.layers: |
| layer.make_buffers() |
|
|
| def forward(self, tokens, input_pos): |
| S = tokens.shape[1] |
| p = input_pos[0].item() |
| torch._check_is_size(p) |
| torch._check(p + S <= self.ctx) |
| h = self.tok_embeddings(tokens) * self.embed_scale |
| cg, sg = self.cos_global.narrow(0, p, S), self.sin_global.narrow(0, p, S) |
| cl, sl = self.cos_local.narrow(0, p, S), self.sin_local.narrow(0, p, S) |
| mask = self.band.narrow(0, self.prefill_max - S, S) + self.padbias.narrow( |
| 0, p + S - 1, self.win |
| ) |
| mask = mask.view(1, S, 1, self.win) |
| for layer in self.layers: |
| if layer.sliding: |
| h = layer(h, p, S, cl, sl, mask, self.pad, self.win) |
| else: |
| h = layer(h, p, S, cg, sg, None, self.pad, self.win) |
| h = self.norm(h[:, -1, :]) |
| return self.output(h) |
|
|
|
|
| def layer_types(cfg): |
| if cfg.get("layer_types"): |
| return list(cfg["layer_types"]) |
| pat = cfg.get("sliding_window_pattern", 6) |
| return [ |
| "sliding_attention" if (i + 1) % pat else "full_attention" |
| for i in range(cfg["num_hidden_layers"]) |
| ] |
|
|
|
|
| def rope_tables(cfg, ctx): |
| """cos/sin tables computed by transformers' own Gemma3RotaryEmbedding (bit-identical to HF).""" |
| from transformers import Gemma3TextConfig |
| from transformers.models.gemma3.modeling_gemma3 import Gemma3RotaryEmbedding |
|
|
| hf_cfg = Gemma3TextConfig(**{k: v for k, v in cfg.items() if not k.startswith("_")}) |
| rope = Gemma3RotaryEmbedding(hf_cfg) |
| pos = torch.arange(ctx)[None] |
| x = torch.zeros(1, dtype=torch.float32) |
| cg, sg = rope(x, pos, "full_attention") |
| cl, sl = rope(x, pos, "sliding_attention") |
| return ( |
| cg[0].contiguous(), |
| sg[0].contiguous(), |
| cl[0].contiguous(), |
| sl[0].contiguous(), |
| ) |
|
|
|
|
| HF_TO_ET = { |
| "input_layernorm": "attention_norm", |
| "post_attention_layernorm": "post_attention_norm", |
| "pre_feedforward_layernorm": "ffn_norm", |
| "post_feedforward_layernorm": "post_ffn_norm", |
| "self_attn.q_proj": "wq", |
| "self_attn.k_proj": "wk", |
| "self_attn.v_proj": "wv", |
| "self_attn.o_proj": "wo", |
| "self_attn.q_norm": "q_norm", |
| "self_attn.k_norm": "k_norm", |
| "mlp.gate_proj": "w1", |
| "mlp.up_proj": "w3", |
| "mlp.down_proj": "w2", |
| } |
|
|
|
|
| def build(snapshot_dir, max_context, prefill_max, linear_dtype=torch.bfloat16): |
| """Gemma3Text with checkpoint weights. Linear/embedding weights stay in `linear_dtype` |
| (the checkpoint's bf16, so quantization sees the stored values, as export_llm does); norms |
| are (1 + w) in fp32 exactly as transformers computes them; output is tied to the embedding.""" |
| from safetensors.torch import load_file |
|
|
| snap = Path(snapshot_dir) |
| cfg = json.loads((snap / "config.json").read_text()) |
| sd = load_file(str(snap / "model.safetensors")) |
| with torch.device("meta"): |
| model = Gemma3Text(cfg, max_context, prefill_max) |
| model = model.to_empty(device="cpu") |
| model.make_buffers() |
| used = set() |
|
|
| def take(key): |
| used.add(key) |
| return sd[key] |
|
|
| def norm(dst, key): |
| dst.weight = nn.Parameter(1.0 + take(key).float(), requires_grad=False) |
|
|
| def lin(dst, key): |
| dst.weight = nn.Parameter(take(key).to(linear_dtype), requires_grad=False) |
|
|
| pre = "model." if "model.embed_tokens.weight" in sd else "" |
| emb = take(f"{pre}embed_tokens.weight").to(linear_dtype) |
| model.tok_embeddings.weight = nn.Parameter(emb, requires_grad=False) |
| model.output.weight = nn.Parameter(emb.clone(), requires_grad=False) |
| norm(model.norm, f"{pre}norm.weight") |
| for i, layer in enumerate(model.layers): |
| for hf, et in HF_TO_ET.items(): |
| key = f"{pre}layers.{i}.{hf}.weight" |
| dst = layer.get_submodule(et) |
| (norm if isinstance(dst, RMSNorm) else lin)(dst, key) |
| unused = set(sd) - used - {"lm_head.weight"} |
| assert not unused, f"unused checkpoint tensors: {sorted(unused)[:5]}" |
| return model.eval(), cfg |
|
|
|
|
| def reset_caches(model): |
| for layer in model.layers: |
| layer.k_cache.zero_() |
| layer.v_cache.zero_() |
|
|
|
|
| def quantize(model, group_size, embedding_group_size): |
| """Same transforms and order as export_llm (embedding first, then linears), quantizing the |
| checkpoint-dtype weights, then casting everything else to fp32.""" |
| from executorch.examples.models.llama.source_transformation.quantize import ( |
| get_quant_embedding_transform, |
| get_quant_weight_transform, |
| ) |
| from executorch.extension.llm.export.builder import DType |
|
|
| model = get_quant_embedding_transform( |
| f"4,{embedding_group_size}", False, quantize_with_hqq=True |
| )(model) |
| model = get_quant_weight_transform( |
| quantization_mode="4w", |
| group_size=group_size, |
| computation_dtype=DType.fp32, |
| checkpoint_dtype=DType.bf16, |
| quantize_with_hqq=True, |
| )(model) |
| return model.to(dtype=torch.float32) |
|
|
|
|
| def edge_manager(model, prefill_max, max_context, cfg, metadata_extra=None): |
| from executorch.extension.llm.export.builder import DType, LLMEdgeManager |
|
|
| metadata = { |
| "get_max_seq_len": prefill_max + 1, |
| "get_max_context_len": max_context, |
| "get_n_layers": cfg["num_hidden_layers"], |
| "get_vocab_size": cfg["vocab_size"], |
| "use_kv_cache": True, |
| "use_sdpa_with_kv_cache": True, |
| "enable_dynamic_shape": True, |
| "get_bos_id": cfg["bos_token_id"], |
| "get_eos_ids": list(cfg["eos_token_id"]), |
| **(metadata_extra or {}), |
| } |
| return LLMEdgeManager( |
| model=model, |
| modelname="gemma3", |
| max_seq_len=prefill_max + 1, |
| dtype=DType.fp32, |
| use_kv_cache=True, |
| example_inputs=( |
| torch.tensor([[2, 3, 4]], dtype=torch.long), |
| torch.tensor([0], dtype=torch.long), |
| ), |
| dynamic_shapes=({1: torch.export.Dim("token_dim", max=prefill_max)}, None), |
| enable_dynamic_shape=True, |
| metadata=metadata, |
| ) |
|
|
|
|
| def webgpu_partitioner(): |
| from executorch.backends.webgpu.partitioner import WebGPUPartitioner |
|
|
| |
| return WebGPUPartitioner( |
| compile_options={"require_dynamic_shapes": True, "force_fp16": False} |
| ) |
|
|
|
|
| _KV_CACHE = re.compile(r"(^|\.)(k|v)_cache$") |
|
|
|
|
| @contextlib.contextmanager |
| def live_kv_cache_reads(): |
| """Lowering-time fix so sliding-window layers read the live KV cache. |
| |
| The Vulkan/WebGPU lowering serializes mutable buffers (KV caches) as delegate constants. |
| llama.sdpa_with_kv_cache handles its own prepacking, so it binds the cache tensor itself and |
| its in-place writes persist across calls. Every other reader (here: the slice_copy that cuts |
| the attention window) gets an et_vk.prepack node, i.e. a separate GPU copy made once at |
| load, so it would only ever see the initial zeros. This wraps the backend's |
| insert_prepack_nodes pass and, for buffers named *.k_cache / *.v_cache only, removes those |
| prepack nodes and points their users at the cache tensor itself, so the slices bind the |
| same GPU buffer sdpa_with_kv_cache writes. The runtime is unchanged. |
| """ |
| import executorch.backends.vulkan.vulkan_preprocess as vp |
| from executorch.exir.dialects._ops import ops as exir_ops |
|
|
| orig = vp.insert_prepack_nodes |
| rewired = [] |
|
|
| def patched(program): |
| program = orig(program) |
| buffers = program.graph_signature.inputs_to_buffers |
| graph = program.graph_module.graph |
| for node in list(graph.nodes): |
| if node.target != exir_ops.edge.et_vk.prepack.default: |
| continue |
| src = node.args[0] |
| if src.op == "placeholder" and _KV_CACHE.search(buffers.get(src.name, "")): |
| node.replace_all_uses_with(src) |
| graph.erase_node(node) |
| rewired.append(buffers[src.name]) |
| program.graph_module.recompile() |
| return program |
|
|
| vp.insert_prepack_nodes = patched |
| try: |
| yield rewired |
| finally: |
| vp.insert_prepack_nodes = orig |
|
|
|
|
| def lower(model, prefill_max, max_context, cfg): |
| _load_custom_ops() |
| em = edge_manager(model, prefill_max, max_context, cfg).export() |
| em.run_canonical_optimizations() |
| em = em.export_to_edge() |
| with live_kv_cache_reads() as rewired: |
| em = em.to_backend([webgpu_partitioner()]) |
| n_local = sum(t == "sliding_attention" for t in layer_types(cfg)[: len(model.layers)]) |
| assert len(rewired) == 2 * n_local, f"rewired {len(rewired)} cache reads, expected {2 * n_local}: {rewired}" |
| em = em.to_executorch() |
| em.rewired_kv_caches = sorted(rewired) |
| return em |
|
|
