diff --git a/.github/workflows/smoke-reference.yml b/.github/workflows/smoke-reference.yml index a2c7480b..3dac4285 100644 --- a/.github/workflows/smoke-reference.yml +++ b/.github/workflows/smoke-reference.yml @@ -39,6 +39,10 @@ on: description: "Direct URL to Apertus-8B-Instruct-2509-Q4_K_S.gguf (~4.6 GB). Enables the Apertus golden-token parity gate (QK-norm + xIELU + ungated FFN). Leave blank to skip." required: false default: "" + gemma3n_gguf_url: + description: "Direct URL to gemma-3n-E2B-it-Q4_K_M.gguf (~3.0 GB). Enables the Gemma 3n golden-token parity gate on the DSL lane (AltUp + Laurel + sparsity + PLE + shared KV; needs a large-memory runner: 20g test heap). Leave blank to skip." + required: false + default: "" gemma4_safetensors_dir_url: description: "Direct URL to a tar.gz containing the Gemma-4 E2B SafeTensors checkpoint directory. Leave blank to skip the kgemma test." required: false @@ -110,6 +114,18 @@ jobs: echo "APERTUS_GGUF_PATH=$RUNNER_TEMP/models/apertus/Apertus-8B-Instruct-2509-Q4_K_S.gguf" >> "$GITHUB_ENV" # The 8B parity gate needs more than the module's 6g default test heap. echo "APERTUS_HEAP_ARG=-PapertusTestMaxHeap=12g" >> "$GITHUB_ENV" + + - name: Stage Gemma 3n E2B GGUF + if: inputs.gemma3n_gguf_url != '' + env: + URL: ${{ inputs.gemma3n_gguf_url }} + run: | + set -euo pipefail + mkdir -p "$RUNNER_TEMP/models/gemma3n" + curl -fsSL "$URL" -o "$RUNNER_TEMP/models/gemma3n/gemma-3n-E2B-it-Q4_K_M.gguf" + echo "GEMMA3N_E2B_GGUF=$RUNNER_TEMP/models/gemma3n/gemma-3n-E2B-it-Q4_K_M.gguf" >> "$GITHUB_ENV" + # The E2B parity gate self-skips below 16 GB test heap. + echo "GEMMA3N_HEAP_ARG=-PgemmaTestMaxHeap=20g" >> "$GITHUB_ENV" if: inputs.gemma4_safetensors_dir_url != '' env: URL: ${{ inputs.gemma4_safetensors_dir_url }} @@ -149,6 +165,7 @@ jobs: -Dorg.gradle.configuration-cache=true \ -PsmokeReference -PincludeIntegration \ ${APERTUS_HEAP_ARG:-} \ + ${GEMMA3N_HEAP_ARG:-} \ test - name: Disk space (after run) diff --git a/CHANGELOG.md b/CHANGELOG.md index f8fee006..3ea8fff3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,50 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — Gemma 3n StableHLO/IREE export harness + hybrid-AI design note + +- **`exportGemma3n`** (`Gemma3nExportHarness`, SmolLM2/FunctionGemma redecode pattern): + traces `gemma3nNetwork()` to StableHLO with external bf16 params and an in-graph argMax + tail. Mobile-honest contract: **`per_layer_inputs` is a graph INPUT** computed on the + CPU from the packed PLE table at runtime (PLE's design point — those parameters stay + off the accelerator), so the parameter archive carries the trunk + token embedding + only. `PerLayerEmbedding` gained a traceable `indexSelect` path while recording; + `GEMMA3N_LAYERS` truncates the trunk for pipeline verification on smaller hosts. Full + E2B emission is **blocked on engine SKaiNET#1247** (trace memory co-residency + an HLO + converter operand-linkage defect) — the harness hard-fails on both signatures instead + of shipping a silently-unservable module. +- New antora explanation page `explanation/gemma3n.adoc` (why Gemma 3n's mobile-first + architecture and why SKaiNET fits it) and pre-PRD design note + `docs/specs/matformer-hybrid-on-device-ai.md` (MatFormer elasticity in SKaiNET + + hybrid on-device/cloud routing: draft-first, escalate-on-evidence). + +### Added — Gemma 3n runs on the DSL path, parity-gated (#377) + +- **`gemma3nNetwork()` + `Gemma3nModel`** — the full Gemma 3n text architecture declared + in the DSL, faithful to HF `modeling_gemma3n.py`: **AltUp** (four parallel hidden + streams with the tanh modality router; `Gemma3nAltUpBlock` per layer, + `Gemma3nAltUpGlobals` for the magnitude-renormed stream init/merge), **Laurel**, + **Gaussian-top-k activation sparsity** on the first ten layers (driven by the GGUF's + precomputed per-layer std multipliers; `-inf` = off), **PLE feeding the non-active + streams** (reusing the gemma-4 lane's `PerLayerEmbedding` — the math is identical), + per-type **shared KV** for the last ten layers, hybrid sliding/global attention with + dual RoPE bases, q/k-norm + parameterless v-norm, attention scale 1.0. All math goes + through `ctx.ops`, so the model is traceable for the StableHLO → IREE mobile path. +- **The hand-rolled `Gemma3nRuntime` was never faithful to real checkpoints**: it loaded + the PLE tensors but never applied them, had no Laurel, ignored the AltUp router, and + its `E2B_DEFAULT` config claimed AltUp/sparsity were E4B-only — the real E2B GGUF has + `altup.num_inputs=4` and first-10-layer sparsity. The GGUF CLI paths (kgemma, unified + skainet-cli) now route gemma3n through the DSL lane; SafeTensors stays on the legacy + runtime until the DSL grows that leg. +- **`Gemma3nGoldenTokenParityTest`** (#346 gate, the last ungated generative family): + full 32-step greedy text equality vs mainline llama.cpp b10621 on + `gemma-3n-E2B-it-Q4_K_M.gguf`, on the exact CLI path — engine loading stays + packed/MAPPED (the PLE table row-dequants on demand). Wired into the smoke-reference + tier (`gemma3n_gguf_url` + 20g heap arg); `smoke-models.json` gains a Gemma3n-E2B row. + Metadata parsing now reads the real llama.cpp GGUF keys (`sliding_window_pattern` + booleans, per-layer `activation_sparsity_scale`, `rope.freq_base` fallback, + `rms_norm_eps`, per-layer `feed_forward_length`). + ### Fixed — Qwen tool calling follows the official Qwen3 chat template - **`QwenChatTemplate` rewritten against the official Qwen3 `chat_template`** (verified diff --git a/README.md b/README.md index 4a1ea519..ae1a443f 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ Use the version shown in this README as the source of truth for first-run snippe > The list below describes the project's **intended** scope. Maturity varies > widely per item and many paths are unverified — see the project-status note above. -- **Multi-model support.** Llama / Mistral, Qwen 2 / 2.5 / 3, Gemma 3 / 4, Apertus (Swiss AI) and BitNet b1.58 are each **verified token-for-token against a reference implementation** (mainline llama.cpp; bitnet.cpp + the HF BF16 reference for BitNet) by model-gated golden-token parity tests on the DSL path the CLIs ship. BERT (vs sentence-transformers) and T5/GTR (real-weights round-trip) are verified on the embedding side. Gemma 3n and Voxtral are the remaining unverified families — see the status table below. +- **Multi-model support.** Llama / Mistral, Qwen 2 / 2.5 / 3, Gemma 3 / 4, Apertus (Swiss AI) and BitNet b1.58 are each **verified token-for-token against a reference implementation** (mainline llama.cpp; bitnet.cpp + the HF BF16 reference for BitNet) by model-gated golden-token parity tests on the DSL path the CLIs ship. BERT (vs sentence-transformers) and T5/GTR (real-weights round-trip) are verified on the embedding side. Voxtral is the remaining unverified family — see the status table below. - **Native CPU performance.** Auto-discovers SKaiNET's priority-100 FFM (Foreign Function & Memory) native kernel provider when present (4–6× faster Q4_K matmul, 1.5–1.8× faster FP32 SGEMM vs the priority-50 Panama Vector path; Linux x86_64 / macOS ARM64 / Windows x86_64 in the published JAR — no manual setup). On **Android**, the runtime facades ship the engine's JNI NEON backend the same way — native kernels out of the box, ~6.4× measured on SmolLM2-135M Q8_0 (see the [supported-targets matrix](#supported-targets)). - **Tool calling (experimental).** Family-specific chat templates and tool-call parsers (Llama 3, Qwen, Gemma, Apertus, ChatML/Hermes) and a Java surface (`KLlamaJava`, `JavaTools.definition`, `JavaAgentLoop`) exist, but tool calling is **not reliable yet** — it may fail to trigger or parse even when plain generation works. - **GGUF + SafeTensors loading.** Streaming reader for any model size; `NATIVE_OPTIMIZED` quant policy keeps weights in their packed SIMD-friendly form. @@ -89,7 +89,7 @@ Honest status — see the project-status note at the top of this README. | **Llama / Mistral** | Verified: `LlamaGoldenTokenParityTest` asserts **full greedy text equality against mainline llama.cpp** on Llama-3.2-1B-Instruct Q8_0, on the DSL path the CLIs ship; smoke rows exercise the same path. | | **Qwen 2 / 2.5 / 3** | Verified: `QwenGoldenTokenParityTest` asserts **full greedy text equality against mainline llama.cpp** for both family variants — Qwen2.5-0.5B-Instruct Q8_0 (attention projection biases) and Qwen3-1.7B Q8_0 (QK-norm); smoke rows for both. | | **Gemma 3 / 4** | Verified: `Gemma4ChatGoldenTokenTest` asserts golden-token parity against llama.cpp on Gemma-4 E2B GGUF; gemma3 (FunctionGemma) and gemma4 checkpoints run the same DSL lane (`gemmaNetwork()`), GGUF and SafeTensors. **Gemma 2 has no supported path** (the CLI refuses it loudly). | -| **Gemma 3n** | Own family module (`llm-inference/gemma3n` + `kgemma3n`), hand-rolled runtime (AltUp / per-layer embeddings / activation sparsity), dense-FP32 loading. Works, but **no parity gate yet** — the remainder is tracked in #377. | +| **Gemma 3n** | Verified: `Gemma3nGoldenTokenParityTest` asserts **full greedy text equality against mainline llama.cpp** on gemma-3n-E2B-it Q4_K_M — the DSL lane (`gemma3nNetwork()`: AltUp × 4 streams, Laurel, activation sparsity, PLE, shared KV) with packed/MAPPED loading. Exports StableHLO for the IREE mobile path (`exportGemma3n`); see `docs/…/explanation/gemma3n.adoc`. | | **Apertus** | Verified: `ApertusGoldenTokenParityTest` asserts **full greedy text equality against mainline llama.cpp** on Apertus-8B-Instruct Q4_K_S — QK-norm, xIELU per-layer activations and the ungated FFN exercised end-to-end. | | **BitNet b1.58** | Packed I2_S path end-to-end on the eager JVM path: 2-bit ternary weights (0.25 B/weight), fused `BITNET_PLANES` lm_head, two-stage candidate decode. Greedy decode verified **token-for-token against bitnet.cpp and the HF BF16 reference** on 2B4T; model-gated parity + smoke tests. See `docs/modules/ROOT/pages/explanation/bitnet.adoc`. | | **BERT** | Sentence embeddings on the DSL path (`bertNetwork()` + `BertEncoderRuntime`, eager or traced/fused) — verified against sentence-transformers on MongoDB/mdbr-leaf. One-call `BertEmbeddingModel.fromHuggingFace(...)` with built-in Hub download; MEAN or CLS pooling and retrieval prefixes cover LEAF, BGE and E5-style models. No text generation, no tool calling. | diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index 09fb5271..6ac9d568 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -35,5 +35,6 @@ * xref:explanation/tokenizer-internals.adoc[Tokenizer Internals] * xref:explanation/weight-quantization.adoc[Weight Quantization and Numeric Representation] * xref:explanation/bitnet.adoc[BitNet b1.58 — Ternary Inference End to End] +* xref:explanation/gemma3n.adoc[Gemma 3n — Mobile-First Architecture] * xref:explanation/embeddings.adoc[How Embeddings Work] * xref:explanation/android-eager-vs-compiled.adoc[Eager vs. Compiled on Android] diff --git a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc b/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc index 043e90b0..e4b172c6 100644 --- a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc +++ b/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc @@ -82,7 +82,6 @@ If the architecture uses standard building blocks (MHA, RMSNorm, FFN), the DSL a Some architectures have components the DSL cannot express: * *Qwen3.5 DeltaNet* -- hybrid DeltaNet (linear attention + SSM) layers with causal 1D convolution -* *Gemma3n* -- variable FFN dimensions per layer (MatFormer), per-layer embeddings * *Voxtral* -- ODE flow matching for audio codec These use `DecoderRuntime` directly. @@ -126,9 +125,9 @@ The goal is to extend the DSL to support these patterns over time. |`voxtralBackboneNetwork()` |Partial DSL -|Gemma3n -|_none_ -|Hand-coded only +|Gemma 3n +|`gemma3nNetwork()` +|DSL lane for GGUF (AltUp/Laurel/sparsity/PLE as DSL modules, parity-gated); legacy runtime remains for SafeTensors only |Qwen3.5 |_none_ diff --git a/docs/modules/ROOT/pages/explanation/gemma3n.adoc b/docs/modules/ROOT/pages/explanation/gemma3n.adoc new file mode 100644 index 00000000..881f9dc6 --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/gemma3n.adoc @@ -0,0 +1,93 @@ += Gemma 3n — Mobile-First Architecture, and Why SKaiNET Fits It +:description: What makes Gemma 3n different, and why SKaiNET's DSL → DAG → StableHLO / eager design serves it well. + +Gemma 3n is Google's mobile-first generation of open models +(https://developers.googleblog.com/en/introducing-gemma-3n-developer-guide/[developer guide]): +it is engineered so that a model with 5B raw parameters (E2B) or 8B (E4B) runs in the +memory footprint of a much smaller one (~2 GB / ~3 GB respectively). Every one of its +architectural tricks is a memory-or-latency trade tailored to phones — and each maps +naturally onto a SKaiNET mechanism. + +== What Gemma 3n actually is + +*MatFormer (Matryoshka Transformer)*:: +A nested transformer built for elastic inference: the E4B model contains a fully +functional E2B sub-model, co-trained. Beyond the two pre-extracted sizes, *Mix-n-Match* +slices custom-sized models between them by adjusting per-layer feed-forward widths +(8192 → 16384) and skipping layers. This is why the GGUF stores +`feed_forward_length` as a *per-layer array* — and why `gemma3nNetwork()` builds each +layer's FFN width independently. + +*Per-Layer Embeddings (PLE)*:: +A large per-layer embedding table (262 144 × layers·256) that improves quality without +growing the accelerator-resident core: only the ~2B (E2B) / ~4B (E4B) trunk weights need +fast memory — the PLE parameters can stay on CPU and be *gathered per token*. + +*KV-cache sharing*:: +The last block of layers reuses the K/V of the last non-shared layer of the same +attention type — Google reports a *2× prefill improvement* over Gemma 3 4B. + +*Hybrid attention + AltUp + activation sparsity*:: +4-out-of-5 layers use a 512-token sliding window (10k RoPE base); every 5th sees the +full context (1M base). AltUp keeps four parallel hidden streams but routes only one +through the expensive layers; the first ten layers apply Gaussian-top-k activation +sparsity (95%) to their FFN gates. + +== Why SKaiNET is a natural fit + +*One declarative definition, two execution paths.*:: +`gemma3nNetwork()` declares the architecture once, through the transformer DSL. The same +module tree runs **eagerly** on the JVM (`OptimizedLLMRuntime` — the path the +golden-token parity gate certifies token-for-token against mainline llama.cpp) and +**traces to a compute graph** for StableHLO emission and `iree-compile` to a mobile +`vmfb` (`exportGemma3n`). No hand-written second implementation to drift. + +*PLE lands exactly where Google designed it to.*:: +SKaiNET's loader keeps the PLE table *packed* (its stored quantization, ~2 GB instead of +8 GB dense) and gathers rows on demand on the CPU — for eager decode via a row-dequant +wrapper, and for the compiled path by making `per_layer_inputs` a *graph input* computed +host-side, so the accelerator archive carries only the trunk. That is the PLE +memory-split from the Gemma 3n paper, realized by the engine's `WeightForm` machinery +rather than bespoke code. + +*Packed, memory-mapped weights.*:: +The engine's `StreamingGgufParametersLoader` serves quantized tensors zero-copy from +file-backed pages (`MAPPED` residency). A phone-class memory budget is the entire point +of Gemma 3n; the loading path honors it instead of inflating everything to FP32. + +*Per-layer heterogeneity is free in the DSL.*:: +MatFormer's variable FFN widths, the 4+1 sliding/global pattern, dual RoPE bases and +per-type shared KV caches (`OwnerReadOnlyKVCache`) are all per-layer decisions in a +plain Kotlin loop — the DSL builds a different stage per layer, no special casing. + +*Kotlin Multiplatform reach.*:: +The same codebase targets JVM/desktop, Android (eager via the JNI NEON backend, or +compiled via `llm-runtime/iree-android`), and Kotlin/Native — matching Gemma 3n's +"everywhere on-device" distribution story. + +*Faithfulness is gated, not claimed.*:: +`Gemma3nGoldenTokenParityTest` asserts full 32-step greedy text equality against +mainline llama.cpp on a real E2B checkpoint — AltUp router, Laurel, sparsity, PLE and +KV-sharing all exercised. See the +xref:../index.adoc#_supported_model_families[verified-model matrix]. + +== Where the pieces live + +|=== +|Concern |Code + +|DSL definition +|`llm-inference/gemma3n` — `gemma3nNetwork()`, `Gemma3nModel`, `Gemma3nAltUpBlock`, + `Gemma3nLaurelBlock`, `Gemma3nSparseGeGluFFN`, `Gemma3nPerLayerApply` + +|PLE machinery (shared with Gemma 4) +|`llm-inference/gemma` — `PerLayerEmbedding` (packed row-dequant + traceable + `indexSelect` path) + +|Compiled export +|`Gemma3nExportHarness` / `exportGemma3n` gradle task → `gemma3n-gen.mlir` + + `gemma3n.safetensors` (bf16) + `manifest.json` + +|Parity gate +|`Gemma3nGoldenTokenParityTest` (model-gated, smoke-reference tier) +|=== diff --git a/docs/modules/ROOT/pages/index.adoc b/docs/modules/ROOT/pages/index.adoc index 13855331..ae22108f 100644 --- a/docs/modules/ROOT/pages/index.adoc +++ b/docs/modules/ROOT/pages/index.adoc @@ -43,9 +43,9 @@ headers record the exact oracle build and commands. |Gemma 3n |Gemma 3n E2B/E4B -|Not yet — hand-rolled runtime (AltUp, per-layer embeddings); gate tracked in issue #377 +|mainline llama.cpp — full greedy text equality (gemma-3n-E2B-it Q4_K_M; AltUp, Laurel, sparsity, PLE, shared KV) |No -|Hand-coded +|`gemma3nNetwork()` |Apertus |Apertus 8B diff --git a/docs/specs/matformer-hybrid-on-device-ai.md b/docs/specs/matformer-hybrid-on-device-ai.md new file mode 100644 index 00000000..bd5c3c91 --- /dev/null +++ b/docs/specs/matformer-hybrid-on-device-ai.md @@ -0,0 +1,173 @@ +# MatFormer & Hybrid On-Device AI — analysis and architecture proposal + +Status: **pre-PRD design note** (2026-09-02). Scope: (1) what MatFormer is and how far our +gemma3n implementation already carries it, (2) a clean SKaiNET-native design for elastic +MatFormer inference, (3) an architecture for **hybrid on-device/cloud serving** — simple +queries answered locally and instantly, complex ones escalated to the cloud, with the +routing decision fast enough that the user never perceives a seam. + +--- + +## 1. MatFormer — what it actually is + +**MatFormer (Matryoshka Transformer)** trains one transformer so that *prefixes of every +FFN's hidden dimension are themselves complete, usable models*. During training, each +layer's FFN is optimized at g nested widths (e.g. 4096 ⊂ 8192 ⊂ 12288 ⊂ 16384); the loss +is the average over the nested submodels. Result: one weight file, a *spectrum* of +deployable models. + +- Paper: Devvrit, Kudugunta, et al., **"MatFormer: Nested Transformer for Elastic + Inference"**, [arXiv:2310.07707](https://arxiv.org/abs/2310.07707). +- Lineage: Kusupati et al., **"Matryoshka Representation Learning"**, + [arXiv:2205.13147](https://arxiv.org/abs/2205.13147) (same nesting idea for embeddings). +- Production instance: **Gemma 3n** + ([developer guide](https://developers.googleblog.com/en/introducing-gemma-3n-developer-guide/)): + E4B (8B raw) co-trains a nested **E2B** (5B raw) submodel. **Mix-n-Match** slices + intermediate sizes by choosing each layer's FFN width in [8192, 16384] (and optionally + skipping layers). Google ships a + [MatFormer Lab colab](https://ai.google.dev/gemma/docs/gemma-3n#matformer) for picking + slice configs benchmarked on MMLU. +- Reference implementations: HF `transformers` `Gemma3nTextMLP` (per-layer + `intermediate_size[layer_idx]` — the sliced widths arrive as plain config), llama.cpp + `gemma3n` (same), Google AI Edge / MediaPipe LLM Inference for the on-device runtime. + +**Key operational insight:** at inference time MatFormer is *not* exotic. A slice is just +per-layer FFN widths + (optionally) fewer layers, reading **prefix sub-ranges of the same +weight tensors**. All elasticity lives in the loader/config, not in new math. + +## 2. Where our implementation stands (post #377 DSL lane) + +Already in place, verified token-for-token vs llama.cpp: + +- `Gemma3nModelMetadata.feedForwardLengths: List` — per-layer FFN widths parsed from + the GGUF (`feed_forward_length` per-layer array). +- `gemma3nNetwork()` builds **each layer's FFN width independently** — a Mix-n-Match + config is *already representable*; only E2B's uniform 8192 has been exercised. +- PLE, AltUp, Laurel, sparsity, shared KV as DSL modules; packed/MAPPED loading; + StableHLO export (`exportGemma3n`) with PLE-as-host-input for the mobile path. + +Missing for real MatFormer elasticity: + +1. **Slice-aware weight views.** An E4B file with an E2B (or custom) slice config must + bind `ffn_gate/up[0:width]` and `ffn_down[:, 0:width]` — prefix *views* of the stored + tensors, no copies. SKaiNET tensors already support `narrow`; the loader needs a + "slice plan" hook. +2. **Slice config surface.** `Gemma3nSliceConfig(perLayerFfn: List, skipLayers: + Set)`, loadable from a JSON sidecar (the MatFormer Lab output format). +3. **Elastic runtime switch.** Two `Module` builds (fast/full) over the *same* weight + map — memory cost is the weight file once (mmap) + two small module trees; switching + is picking which module the runtime steps. On the compiled path: one vmfb per slice + (IREE graphs are static), same `.irpa` parameter archive shared between them. + +### Clean SKaiNET design (fits both execution philosophies) + +- **Eager (DSL → modules):** `gemma3nNetwork(metadata, slice)` where `slice` rewrites + `feedForwardLengths` and the layer list; `Gemma3nWeightLoader` gains + `sliceView(tensor, layer)` returning narrow views. One weight map, N module trees. +- **Compiled (DSL → DAG → StableHLO):** trace each slice once; emit + `gemma3n-e2b.vmfb`, `gemma3n-e3b.vmfb`, … all referencing **one parameter scope** — + IREE loads params by name, and prefix-sliced tensors export as their own named params + (only the sliced FFNs duplicate bytes in the archive; everything else is shared). + Alternative (later): emit the full-width graph and pass runtime width as a dynamic dim + — rejected for now; IREE static shapes are the proven lane. + +## 3. Hybrid on-device / cloud serving + +### The product requirement + +Simple queries ("set a timer", "summarize this note", quick factual Q&A) answer +**offline, instantly, privately**. Complex ones (long reasoning, tool orchestration, +fresh knowledge) go to a **cloud model**. The user sees ONE assistant: no mode switch, no +spinner-then-restart. That forces the routing decision to be either (a) made in ≲50 ms, +or (b) made *behind* an already-streaming local answer. + +### Design principle: draft-first, escalate-on-evidence + +Never block on the router. **The local model always starts.** Routing signals accumulate +in three layers, cheapest first: + +| Layer | Signal | Cost | Acts | +|---|---|---|---| +| L0 | static features: prompt length, tool availability offline, connectivity, battery/thermal, PII policy ("must stay local") | < 1 ms | pre-pick lane; hard-pin local for private content | +| L1 | tiny router: embedding of the prompt → complexity classifier (MRL-truncated embedding + logistic head, or a distilled BERT — we have `bertNetwork()` + LEAF/BGE embedders on-device) | 5–20 ms | route obvious cases before first token | +| L2 | local model's own uncertainty while drafting: first-k token logit margin / entropy, `` self-probe, draft perplexity trend | free (by-product of decoding) | escalate mid-stream | +| L3 | cloud verification (optional, speculative-cascade mode): cloud model verifies/extends the local draft à la speculative decoding across tiers | network RTT | quality backstop | + +The seam is hidden by **stream discipline**: hold the first ~150–250 ms of local tokens +in a presentation buffer. If L0–L2 escalate within the buffer window, the user never saw +the local draft; if the local answer is already flowing and L2 fires, the handoff sends +`(prompt, draft-so-far)` to the cloud with an instruction to continue/repair — visible at +worst as a brief pause, never as a restart. + +### Data flow + +```mermaid +graph TD + U[User query] --> L0{L0 static gate
<1 ms} + L0 -->|must-local / offline| LOCAL + L0 -->|obviously heavy| CLOUD + L0 -->|uncertain| L1[L1 tiny router
on-device embedding + head] + L1 -->|simple| LOCAL[Local decode
gemma3n slice via SKaiNET
eager JVM/NEON or IREE vmfb] + L1 -->|complex| CLOUD[Cloud model] + LOCAL -->|tokens + confidence| BUF[Stream buffer 150-250 ms] + BUF -->|confident| UI[UI stream] + LOCAL -.->|L2: entropy spike /
margin collapse| HAND[Handoff: prompt + draft] + HAND --> CLOUD + CLOUD --> UI + CLOUD -.->|L3 speculative cascade:
verify local draft tokens| BUF +``` + +### Where MatFormer earns its keep here + +Elastic width is the **third axis of the routing decision**: instead of binary +local/cloud, the device picks E2B-width for battery/latency, a wider Mix-n-Match slice +when plugged in or when L1 says "medium difficulty" — same weights, no extra download. +Google's guide names exactly this ("elastic execution … dynamically switch E4B/E2B based +on the task and device load") as a future capability; the slice-view design in §2 is our +path to shipping it. + +### SKaiNET implementation sketch + +- **`HybridSession`** (new, `llm-agent`): wraps two `InferenceRuntime`s — local + (`OptimizedLLMRuntime` over gemma3n) and remote (an `InferenceRuntime` adapter over an + HTTP streaming API). Owns the L0–L2 state machine and the stream buffer. `ChatSession` + API stays the user surface, so tool calling and templates work in both lanes. +- **Router**: `bertNetwork()`-based embedder (already verified vs sentence-transformers) + + a trained head; ships as a tiny side model. L2 signals read from the logits SKaiNET + already returns each step (`sampleFromTensor` exposes the tensor — margin/entropy is + a 10-line addition). +- **Handoff protocol**: cloud request carries `{messages, localDraft, draftLogprobs?}`; + the cloud either continues (prefix-cache friendly) or rewrites. For L3, the same + payload makes the cloud a *verifier* (speculative cascade). +- **On Android**: local lane = IREE vmfb (compiled slice) or eager NEON; router+buffer + logic is common KMP code shared with iOS/desktop. + +### Papers / prior art to anchor the PRD + +- Routing: **RouteLLM** (Ong et al., [arXiv:2406.18665](https://arxiv.org/abs/2406.18665); + [github.com/lm-sys/RouteLLM](https://github.com/lm-sys/RouteLLM)) — trained binary + routers, 2×+ cost cuts at matched quality. **Hybrid LLM** (Ding et al., ICLR 2024, + [arXiv:2404.14618](https://arxiv.org/abs/2404.14618)) — quality-aware small/large + routing. **FrugalGPT** (Chen et al., [arXiv:2305.05176](https://arxiv.org/abs/2305.05176)) — + cascades with stop-signals. +- Draft/verify across tiers: **speculative decoding** (Leviathan et al., + [arXiv:2211.17192](https://arxiv.org/abs/2211.17192)); **speculative cascades** + (Narasimhan et al., [arXiv:2405.19261](https://arxiv.org/abs/2405.19261)) — the formal + blend of cascades + speculative execution we call L3. +- Confidence signals: token-entropy/margin early-exit literature (e.g. CALM, Schuster et + al., [arXiv:2207.07061](https://arxiv.org/abs/2207.07061)). +- Elasticity: MatFormer + MRL above; Gemma 3n guide for the production framing. +- Product precedent: Apple Intelligence's on-device/Private-Cloud-Compute split — the UX + bar for "seamless" (no user-visible routing). + +### Open questions for the PRD + +1. Escalation budget: max acceptable silent hold (proposal: 250 ms) and mid-stream + handoff UX (pause vs. visible "thinking harder"). +2. Router training data: which task taxonomy, and do we log (opt-in) local drafts + + outcomes to train it. +3. Privacy contract: which content classes are hard-pinned local (L0) regardless of + quality cost. +4. Cloud protocol: continue-from-draft vs. fresh generation; prefix-cache assumptions. +5. MatFormer slice policy: fixed two slices (fast/full) first, or continuous + Mix-n-Match from device telemetry. diff --git a/llm-apps/skainet-cli/build.gradle.kts b/llm-apps/skainet-cli/build.gradle.kts index 814b68af..d0638747 100644 --- a/llm-apps/skainet-cli/build.gradle.kts +++ b/llm-apps/skainet-cli/build.gradle.kts @@ -25,6 +25,7 @@ dependencies { implementation(project(":llm-inference:qwen")) implementation(project(":llm-inference:bitnet")) implementation(project(":llm-inference:gemma")) + implementation(project(":llm-inference:gemma3n")) implementation(project(":llm-inference:apertus")) // SKaiNET core libraries diff --git a/llm-apps/skainet-cli/src/main/kotlin/sk/ainet/apps/skainet/cli/Main.kt b/llm-apps/skainet-cli/src/main/kotlin/sk/ainet/apps/skainet/cli/Main.kt index f080686d..0fabe729 100644 --- a/llm-apps/skainet-cli/src/main/kotlin/sk/ainet/apps/skainet/cli/Main.kt +++ b/llm-apps/skainet-cli/src/main/kotlin/sk/ainet/apps/skainet/cli/Main.kt @@ -242,26 +242,34 @@ fun main(args: Array) { val runtime: InferenceRuntime = if (modelInfo.family == ModelFamily.GEMMA) { // ModelFamily.GEMMA claims every gemma* architecture, but this DSL lane serves - // gemma3/gemma4 only (#376): 3n needs the hand-rolled runtime (AltUp/PLE/activation - // sparsity — kgemma CLI, split tracked in #377), and gemma2 has no supported path. - when (modelInfo.architecture) { - "gemma3n" -> error( - "Gemma 3n is not supported by the unified CLI's DSL lane — use the kgemma " + - "CLI (:llm-runtime:kgemma), which carries its hand-rolled runtime (#377).", - ) - "gemma2", "gemma" -> error( - "Architecture '${modelInfo.architecture}' has no supported path — the Gemma " + - "lane serves gemma3/gemma4 checkpoints (#376).", - ) - } - println("Loading Gemma GGUF model from $modelPath via gemmaNetwork() + OptimizedLLMRuntime (engine loader, keep-packed, mapped)...") - if (cliArgs.contextLength != null) { - println(" --context flag currently ignored on the Gemma path; uses model default capped to 4096.") + // gemma3/gemma4/gemma3n (#376, #377): gemma3n runs its own DSL lane + // (gemma3nNetwork() — AltUp/Laurel/sparsity/PLE, parity-gated vs llama.cpp); + // gemma2 has no supported path. + if (modelInfo.architecture == "gemma3n") { + println("Loading Gemma 3n GGUF model from $modelPath via gemma3nNetwork() + OptimizedLLMRuntime (engine loader, keep-packed, mapped)...") + val model3n = kotlinx.coroutines.runBlocking { + sk.ainet.models.gemma3n.Gemma3nNetworkLoader.fromGguf( + ctx, + { JvmRandomAccessSource.open(modelPath.toString()) }, + ) + } + OptimizedLLMRuntime(model3n, ctx, OptimizedLLMMode.DIRECT, FP32::class) + } else { + when (modelInfo.architecture) { + "gemma2", "gemma" -> error( + "Architecture '${modelInfo.architecture}' has no supported path — the Gemma " + + "lane serves gemma3/gemma4/gemma3n checkpoints (#376, #377).", + ) + } + println("Loading Gemma GGUF model from $modelPath via gemmaNetwork() + OptimizedLLMRuntime (engine loader, keep-packed, mapped)...") + if (cliArgs.contextLength != null) { + println(" --context flag currently ignored on the Gemma path; uses model default capped to 4096.") + } + val model = GemmaNetworkLoader.fromGguf( + randomAccessProvider = { JvmRandomAccessSource.open(modelPath.toString()) } + ).load(ctx) + OptimizedLLMRuntime(model, ctx, OptimizedLLMMode.DIRECT, FP32::class) } - val model = GemmaNetworkLoader.fromGguf( - randomAccessProvider = { JvmRandomAccessSource.open(modelPath.toString()) } - ).load(ctx) - OptimizedLLMRuntime(model, ctx, OptimizedLLMMode.DIRECT, FP32::class) } else if (modelInfo.family == ModelFamily.APERTUS) { println("Loading Apertus GGUF model from $modelPath via apertusNetwork() + OptimizedLLMRuntime (engine loader, keep-packed, mapped)...") if (cliArgs.contextLength != null) { diff --git a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/PerLayerEmbedding.kt b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/PerLayerEmbedding.kt index c3efc3aa..bebd0492 100644 --- a/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/PerLayerEmbedding.kt +++ b/llm-inference/gemma/src/commonMain/kotlin/sk/ainet/models/gemma/PerLayerEmbedding.kt @@ -126,6 +126,34 @@ public class PerLayerEmbedding( ): Tensor { val ops = ctx.ops + // Graph-recording path (StableHLO/IREE export): every step must go through + // ctx.ops to be traceable — the manual buffer gather below would bake the + // fixture tokens' rows in as frozen constants. indexSelect keeps the + // token-identity lookup a real graph op on the (dense) PLE table. + if (ctx.isRecording) { + val idsShapeR = tokenIds.shape + require(idsShapeR.rank == 2) { + "$name.compute: tokenIds must be [batch, seq], got shape=${tokenIds.shape}" + } + val batchR = idsShapeR[0] + val seqR = idsShapeR[1] + @Suppress("UNCHECKED_CAST") + val idsFlatT = ops.reshape(tokenIds as Tensor, Shape(batchR * seqR)) + @Suppress("UNCHECKED_CAST") + var rawR = ops.indexSelect( + embedTokensWeight, idsFlatT as Tensor, dim = 0, + ) // [B*S, perLayerTotal] + rawR = ops.mulScalar(rawR, embedScale) + val rawReshapedR = ops.reshape(rawR, Shape(batchR, seqR, numLayers, perLayerDim)) + + val flatEmbedsR = ops.reshape(inputsEmbeds, Shape(batchR * seqR, hiddenSize)) + var projR = linearProject(ops, flatEmbedsR, modelProjWeight) + projR = ops.reshape(projR, Shape(batchR, seqR, numLayers, perLayerDim)) + projR = ops.mulScalar(projR, projScale) + projR = projectionNorm.forward(projR, ctx) + return ops.mulScalar(ops.add(projR, rawReshapedR), inputScale) + } + // (1) Token-identity: gather rows of embedTokensWeight by tokenIds. // We do the gather via a manual buffer build; `ops.gather` would also // work but needs Int32 handling that varies by backend. Since the diff --git a/llm-inference/gemma3n/api/jvm/gemma3n.api b/llm-inference/gemma3n/api/jvm/gemma3n.api index 610462e0..9ce2f85b 100644 --- a/llm-inference/gemma3n/api/jvm/gemma3n.api +++ b/llm-inference/gemma3n/api/jvm/gemma3n.api @@ -49,6 +49,28 @@ public abstract interface class sk/ainet/models/gemma3n/AttentionBackend { public abstract fun reset ()V } +public final class sk/ainet/models/gemma3n/Gemma3nAltUpBlock : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (IIIFLkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (IIIFLkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun correct (Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Ljava/util/List; + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public fun getParams ()Ljava/util/List; + public final fun getRouterNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun predict (Ljava/util/List;Lsk/ainet/context/ExecutionContext;)Ljava/util/List; + public final fun scaleCorrectedOutput (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; +} + +public final class sk/ainet/models/gemma3n/Gemma3nAltUpGlobals : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (IILkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (IILkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public fun getParams ()Ljava/util/List; + public final fun initStreams (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Ljava/util/List; + public final fun mergeStreams (Ljava/util/List;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; +} + public final class sk/ainet/models/gemma3n/Gemma3nAttentionBackend : sk/ainet/models/gemma3n/AttentionBackend { public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/gemma3n/Gemma3nRuntimeWeights;Lkotlin/reflect/KClass;Lsk/ainet/models/gemma3n/Gemma3nConfig;Lsk/ainet/models/gemma3n/Gemma3nKvCache;)V public synthetic fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/gemma3n/Gemma3nRuntimeWeights;Lkotlin/reflect/KClass;Lsk/ainet/models/gemma3n/Gemma3nConfig;Lsk/ainet/models/gemma3n/Gemma3nKvCache;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -123,6 +145,47 @@ public final class sk/ainet/models/gemma3n/Gemma3nConfigParser { public final fun parseFromJson (Ljava/lang/String;)Lsk/ainet/models/gemma3n/Gemma3nModelMetadata; } +public final class sk/ainet/models/gemma3n/Gemma3nExportCliKt { + public static final fun main ()V + public static synthetic fun main ([Ljava/lang/String;)V +} + +public final class sk/ainet/models/gemma3n/Gemma3nExportHarness { + public static final field FN_REDECODE Ljava/lang/String; + public static final field INSTANCE Lsk/ainet/models/gemma3n/Gemma3nExportHarness; + public static final field PARAMETER_SCOPE Ljava/lang/String; + public final fun export (Ljava/lang/String;Ljava/lang/String;IZLjava/lang/Integer;)Lsk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult; + public static synthetic fun export$default (Lsk/ainet/models/gemma3n/Gemma3nExportHarness;Ljava/lang/String;Ljava/lang/String;IZLjava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult; +} + +public final class sk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult { + public fun (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;IJII)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()Ljava/lang/String; + public final fun component4 ()I + public final fun component5 ()J + public final fun component6 ()I + public final fun component7 ()I + public final fun copy (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;IJII)Lsk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult; + public static synthetic fun copy$default (Lsk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;IJIIILjava/lang/Object;)Lsk/ainet/models/gemma3n/Gemma3nExportHarness$RedecodeResult; + public fun equals (Ljava/lang/Object;)Z + public final fun getExternalParamCount ()I + public final fun getManifestPath ()Ljava/lang/String; + public final fun getMlirPath ()Ljava/lang/String; + public final fun getSafetensorsPath ()Ljava/lang/String; + public final fun getSeq ()I + public final fun getVocabSize ()I + public final fun getWeightMiB ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/models/gemma3n/Gemma3nGGUFNameResolver : sk/ainet/io/weights/WeightNameResolver { + public fun ()V + public fun resolve (Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; +} + public abstract interface class sk/ainet/models/gemma3n/Gemma3nKvCache : sk/ainet/apps/llm/KvCache { } @@ -130,6 +193,16 @@ public final class sk/ainet/models/gemma3n/Gemma3nKvCacheKt { public static final fun createOptimalGemma3nKvCache (Lsk/ainet/models/gemma3n/Gemma3nConfig;I)Lsk/ainet/models/gemma3n/Gemma3nKvCache; } +public final class sk/ainet/models/gemma3n/Gemma3nLaurelBlock : sk/ainet/lang/nn/Module { + public fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun getLinearLeft ()Lsk/ainet/lang/nn/transformer/VoidDense; + public final fun getLinearRight ()Lsk/ainet/lang/nn/transformer/VoidDense; + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public final fun getPostNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; +} + public final class sk/ainet/models/gemma3n/Gemma3nLayerWeights { public fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/gemma3n/AltUpLayerWeights;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public synthetic fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/gemma3n/AltUpLayerWeights;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -184,14 +257,45 @@ public final class sk/ainet/models/gemma3n/Gemma3nLayerWeights { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/models/gemma3n/Gemma3nModel : sk/ainet/lang/nn/Module { + public fun (Lsk/ainet/lang/nn/layers/EmbeddingAdapter;Lsk/ainet/models/gemma/PerLayerEmbedding;Lsk/ainet/models/gemma3n/Gemma3nAltUpGlobals;Ljava/util/List;Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/lang/nn/transformer/VoidDense;Lkotlin/reflect/KClass;IFLjava/lang/String;)V + public synthetic fun (Lsk/ainet/lang/nn/layers/EmbeddingAdapter;Lsk/ainet/models/gemma/PerLayerEmbedding;Lsk/ainet/models/gemma3n/Gemma3nAltUpGlobals;Ljava/util/List;Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/lang/nn/transformer/VoidDense;Lkotlin/reflect/KClass;IFLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun getActiveIdx ()I + public final fun getAltupGlobals ()Lsk/ainet/models/gemma3n/Gemma3nAltUpGlobals; + public final fun getBlocks ()Ljava/util/List; + public final fun getDtype ()Lkotlin/reflect/KClass; + public final fun getEmbedScale ()F + public final fun getExternalPerLayerInputs ()Lsk/ainet/lang/tensor/Tensor; + public final fun getLmHead ()Lsk/ainet/lang/nn/transformer/VoidDense; + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public final fun getOutputNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun getPle ()Lsk/ainet/models/gemma/PerLayerEmbedding; + public final fun getTokenEmbedding ()Lsk/ainet/lang/nn/layers/EmbeddingAdapter; + public final fun setExternalPerLayerInputs (Lsk/ainet/lang/tensor/Tensor;)V +} + +public final class sk/ainet/models/gemma3n/Gemma3nModel$BlockRefs { + public fun (Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/lang/nn/transformer/MultiHeadAttention;Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/models/gemma3n/Gemma3nSparseGeGluFFN;Lsk/ainet/lang/nn/normalization/RMSNormalization;Lsk/ainet/models/gemma3n/Gemma3nLaurelBlock;Lsk/ainet/models/gemma3n/Gemma3nAltUpBlock;Lsk/ainet/models/gemma3n/Gemma3nPerLayerApply;)V + public final fun getAltup ()Lsk/ainet/models/gemma3n/Gemma3nAltUpBlock; + public final fun getAttnNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun getFfn ()Lsk/ainet/models/gemma3n/Gemma3nSparseGeGluFFN; + public final fun getFfnNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun getLaurel ()Lsk/ainet/models/gemma3n/Gemma3nLaurelBlock; + public final fun getMha ()Lsk/ainet/lang/nn/transformer/MultiHeadAttention; + public final fun getPerLayer ()Lsk/ainet/models/gemma3n/Gemma3nPerLayerApply; + public final fun getPostAttnNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun getPostFfwNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; +} + public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata { public static final field Companion Lsk/ainet/models/gemma3n/Gemma3nModelMetadata$Companion; public static final field DEFAULT_KV_SHARED_LAYERS I public static final field DEFAULT_ROPE_BASE_GLOBAL F public static final field DEFAULT_ROPE_BASE_LOCAL F public static final field DEFAULT_SLIDING_WINDOW I - public fun (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;F)V - public synthetic fun (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FFLjava/util/List;)V + public synthetic fun (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FFLjava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component10 ()I public final fun component11 ()I @@ -204,6 +308,8 @@ public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata { public final fun component18 ()Ljava/util/List; public final fun component19 ()F public final fun component2 ()I + public final fun component20 ()F + public final fun component21 ()Ljava/util/List; public final fun component3 ()I public final fun component4 ()I public final fun component5 ()I @@ -211,11 +317,12 @@ public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata { public final fun component7 ()I public final fun component8 ()Ljava/util/List; public final fun component9 ()I - public final fun copy (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;F)Lsk/ainet/models/gemma3n/Gemma3nModelMetadata; - public static synthetic fun copy$default (Lsk/ainet/models/gemma3n/Gemma3nModelMetadata;Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FILjava/lang/Object;)Lsk/ainet/models/gemma3n/Gemma3nModelMetadata; + public final fun copy (Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FFLjava/util/List;)Lsk/ainet/models/gemma3n/Gemma3nModelMetadata; + public static synthetic fun copy$default (Lsk/ainet/models/gemma3n/Gemma3nModelMetadata;Ljava/lang/String;IIIIIILjava/util/List;IIIFFILjava/util/List;IILjava/util/List;FFLjava/util/List;ILjava/lang/Object;)Lsk/ainet/models/gemma3n/Gemma3nModelMetadata; public fun equals (Ljava/lang/Object;)Z public final fun getActivationSparsityPattern ()Ljava/util/List; public final fun getActivationSparsityScale ()F + public final fun getActivationSparsityScales ()Ljava/util/List; public final fun getAltupActiveIdx ()I public final fun getArchitecture ()Ljava/lang/String; public final fun getBlockCount ()I @@ -233,6 +340,7 @@ public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata { public final fun getLayerType (I)Lsk/ainet/models/gemma/LayerType; public final fun getNumAltupInputs ()I public final fun getPerLayerEmbeddingLength ()I + public final fun getRmsNormEps ()F public final fun getRopeBase (I)F public final fun getRopeBaseGlobal ()F public final fun getRopeBaseLocal ()F @@ -240,6 +348,7 @@ public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata { public final fun getVocabSize ()I public fun hashCode ()I public final fun isKvShared (I)Z + public final fun sparsityScaleFor (I)Ljava/lang/Float; public fun toString ()Ljava/lang/String; } @@ -247,6 +356,29 @@ public final class sk/ainet/models/gemma3n/Gemma3nModelMetadata$Companion { public final fun getDEFAULT_LAYER_PATTERN ()Ljava/util/List; } +public final class sk/ainet/models/gemma3n/Gemma3nNetworkDefKt { + public static final field LAUREL_RANK I + public static final fun gemma3nNetwork (Lsk/ainet/models/gemma3n/Gemma3nModelMetadata;Lkotlin/reflect/KClass;III)Lsk/ainet/lang/nn/Module; + public static synthetic fun gemma3nNetwork$default (Lsk/ainet/models/gemma3n/Gemma3nModelMetadata;Lkotlin/reflect/KClass;IIIILjava/lang/Object;)Lsk/ainet/lang/nn/Module; +} + +public final class sk/ainet/models/gemma3n/Gemma3nNetworkLoader { + public static final field INSTANCE Lsk/ainet/models/gemma3n/Gemma3nNetworkLoader; + public final fun fromWeights (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/gemma3n/Gemma3nWeights;Lkotlin/reflect/KClass;Ljava/lang/Integer;ZLjava/lang/Integer;)Lsk/ainet/lang/nn/Module; + public static synthetic fun fromWeights$default (Lsk/ainet/models/gemma3n/Gemma3nNetworkLoader;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/gemma3n/Gemma3nWeights;Lkotlin/reflect/KClass;Ljava/lang/Integer;ZLjava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; +} + +public final class sk/ainet/models/gemma3n/Gemma3nPerLayerApply : sk/ainet/lang/nn/Module { + public fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun computeDelta (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; + public final fun getInpGate ()Lsk/ainet/lang/nn/transformer/VoidDense; + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public final fun getPostNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization; + public final fun getProj ()Lsk/ainet/lang/nn/transformer/VoidDense; +} + public final class sk/ainet/models/gemma3n/Gemma3nRuntime : sk/ainet/apps/llm/DecoderRuntime { public static final field BOS_TOKEN I public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/gemma3n/Gemma3nRuntimeWeights;Lsk/ainet/models/gemma3n/AttentionBackend;Lkotlin/reflect/KClass;Lsk/ainet/models/gemma3n/Gemma3nConfig;FLkotlin/random/Random;)V @@ -314,6 +446,17 @@ public final class sk/ainet/models/gemma3n/Gemma3nSafeTensorsWeightLoaderKt { public static final fun loadGemma3nRuntimeWeightsFromSafeTensors (Lsk/ainet/context/ExecutionContext;Ljava/lang/String;Lkotlin/reflect/KClass;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; } +public final class sk/ainet/models/gemma3n/Gemma3nSparseGeGluFFN : sk/ainet/lang/nn/Module { + public fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (IIFLkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun getDown ()Lsk/ainet/lang/nn/transformer/VoidDense; + public final fun getGate ()Lsk/ainet/lang/nn/transformer/VoidDense; + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public final fun getSparsityEnabled ()Z + public final fun getUp ()Lsk/ainet/lang/nn/transformer/VoidDense; +} + public final class sk/ainet/models/gemma3n/Gemma3nTensorNames { public static final field ALTUP_PROJ Ljava/lang/String; public static final field ALTUP_UNEMBD_PROJ Ljava/lang/String; diff --git a/llm-inference/gemma3n/build.gradle.kts b/llm-inference/gemma3n/build.gradle.kts index 8e0a75b2..e74f9e86 100644 --- a/llm-inference/gemma3n/build.gradle.kts +++ b/llm-inference/gemma3n/build.gradle.kts @@ -67,6 +67,18 @@ kotlin { implementation(libs.skainet.backend.api) } + val jvmMain by getting { + dependencies { + implementation(project.dependencies.platform(project(":llm-bom"))) + // The StableHLO export harness (Gemma3nExportHarness) is host tooling: + // trace gemma3nNetwork() to a ComputeGraph and lower to StableHLO. JVM-only + // — compile-hlo/-dag publish no JS variant (same note as jvmTest below). + implementation(libs.skainet.compile.dag) + implementation(libs.skainet.compile.hlo) + implementation(libs.skainet.backend.cpu) + } + } + val jvmTest by getting { dependencies { implementation(project.dependencies.platform(project(":llm-bom"))) @@ -149,3 +161,21 @@ tasks.matching { it.name == "jsBrowserTest" || it.name == "wasmJsBrowserTest" }. ?.failOnNoDiscoveredTests = false enabled = includeBrowserTests } + +// Gemma 3n compiled-export entry point (SmolLM2/FunctionGemma pattern): +// GEMMA3N_GGUF=…gemma-3n-E2B-it-Q4_K_M.gguf GEMMA3N_OUT_DIR=build/gemma3n-export \ +// ./gradlew :llm-inference:gemma3n:exportGemma3n -PexportMaxHeap=30g +tasks.register("exportGemma3n") { + group = "bridge" + description = "Export Gemma 3n -> StableHLO MLIR (redecode, argMax tail) + safetensors + manifest from the GGUF." + val jvmMainComp = kotlin.jvm().compilations.getByName("main") + dependsOn(jvmMainComp.compileTaskProvider) + classpath = jvmMainComp.output.allOutputs + jvmMainComp.runtimeDependencyFiles + mainClass.set("sk.ainet.models.gemma3n.Gemma3nExportCliKt") + // Dense trunk + trace zeros + graph constant copies peak ~44 GB on E2B (the + // all-zero trace pages compress well under macOS) — override to fit the host. + maxHeapSize = (findProperty("exportMaxHeap") as? String) ?: "46g" + listOf("GEMMA3N_GGUF", "GEMMA3N_OUT_DIR", "GEN_SEQ", "GEMMA3N_DTYPE", "GEMMA3N_LAYERS").forEach { k -> + System.getenv(k)?.let { environment(k, it) } + } +} diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nDslModules.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nDslModules.kt new file mode 100644 index 00000000..933f41f7 --- /dev/null +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nDslModules.kt @@ -0,0 +1,315 @@ +package sk.ainet.models.gemma3n + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.normalization.RMSNormalization +import sk.ainet.lang.nn.topology.ModuleParameter +import sk.ainet.lang.nn.topology.ModuleParameters +import sk.ainet.lang.nn.transformer.VoidDense +import sk.ainet.lang.nn.transformer.linearProject +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.DType +import kotlin.reflect.KClass + +/* + * DSL modules for the Gemma 3n-specific machinery (the #377 DSL migration): AltUp, Laurel, + * activation-sparsity FFN and the per-layer-input application. All math goes through + * `ctx.ops` so the modules are traceable for the StableHLO → IREE export path. The + * reference implementation is HF `transformers` `modeling_gemma3n.py` (verified against + * the installed 5.x source); working tensors are rank-2 `[seq, hidden]`. + */ + +@Suppress("UNCHECKED_CAST") +internal fun voidParam(name: String, shape: Shape, dtype: KClass?): ModuleParameter = + ModuleParameter.WeightParameter( + name, + VoidOpsTensor( + object : TensorData { + override val shape: Shape = shape + override fun get(vararg indices: Int): V = 0.0f as V + override fun set(vararg indices: Int, value: V) {} + }, + (dtype ?: Any::class) as KClass, + ), + ) + +/** + * Per-layer AltUp (Alternating Updates) block — HF `Gemma3nTextAltUp`. + * + * Maintains `numInputs` parallel hidden streams; only the active one runs the expensive + * transformer sub-layers, the rest are predicted/corrected via a learned router: + * + * ``` + * modalities(x) = tanh( modality_router( router_norm(x) * 1/hidden ) ) + * predict: coefs = prediction_coefs(modalities) # [S, n²] + * pred_i = h_i + Σ_j coefs[:, i·n+j] ⊙ h_j + * correct: coefs = correction_coefs(modalities(activated)) + 1 + * innovation = activated − pred_active + * corr_i = pred_i + coefs[:, i] ⊙ innovation + * scale_corrected_output(x) = x * correct_output_scale # [hidden] + * ``` + */ +public class Gemma3nAltUpBlock( + private val hiddenSize: Int, + private val numInputs: Int, + private val activeIdx: Int, + rmsEps: Float, + private val dtype: KClass? = null, + override val name: String = "altup", +) : Module(), ModuleParameters { + + /** `router_norm` — scale-full RMSNorm over hidden, plain (non-unit-offset) weight. */ + public val routerNorm: RMSNormalization = RMSNormalization( + intArrayOf(hiddenSize), rmsEps.toDouble(), unitOffset = false, name = "$name.altup_router_norm", dtype = dtype, + ) + + override val params: List> = listOf( + voidParam("$name.altup_router.weight", Shape(numInputs, hiddenSize), dtype), + voidParam("$name.altup_predict_coef.weight", Shape(numInputs * numInputs, numInputs), dtype), + voidParam("$name.altup_correct_coef.weight", Shape(numInputs, numInputs), dtype), + voidParam("$name.altup_correct_scale.weight", Shape(hiddenSize), dtype), + ) + + override val modules: List> = listOf(routerNorm) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor = input + + private fun modalities(x: Tensor, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + val normed = routerNorm.forward(x, ctx) + val scaled = ops.mulScalar(normed, 1.0f / hiddenSize) + return ops.tanh(linearProject(ops, scaled, params[0].value)) // [S, n] + } + + /** One learned scalar column `[.., 1]` broadcast-multiplied over `[.., H]` — + * rank-general (the trunk runs rank-2 `[S, H]` eagerly, rank-3 `[B, S, H]` under + * the export trace). */ + private fun scaleBy(coefs: Tensor, col: Int, x: Tensor, ctx: ExecutionContext): Tensor = + ctx.ops.multiply(x, ctx.ops.narrow(coefs, dim = coefs.rank - 1, start = col, length = 1)) + + public fun predict(streams: List>, ctx: ExecutionContext): List> { + val ops = ctx.ops + val m = modalities(streams[activeIdx], ctx) + val coefs = linearProject(ops, m, params[1].value) // [S, n²] + return List(numInputs) { i -> + var pred = streams[i] + for (j in 0 until numInputs) { + pred = ops.add(pred, scaleBy(coefs, i * numInputs + j, streams[j], ctx)) + } + pred + } + } + + public fun correct( + predictions: List>, + activated: Tensor, + ctx: ExecutionContext, + ): List> { + val ops = ctx.ops + val m = modalities(activated, ctx) + val coefs = ops.addScalar(linearProject(ops, m, params[2].value), 1.0f) // [S, n] + val innovation = ops.subtract(activated, predictions[activeIdx]) + return List(numInputs) { i -> + ops.add(predictions[i], scaleBy(coefs, i, innovation, ctx)) + } + } + + /** `x * correct_output_scale` (element-wise over hidden). */ + public fun scaleCorrectedOutput(x: Tensor, ctx: ExecutionContext): Tensor = + ctx.ops.multiply(x, params[3].value) +} + +/** + * Model-level AltUp stream projections — HF `altup_projections` / `altup_unembed_projections`. + * The GGUF stores each set as ONE 3D tensor (`altup_proj.weight`, `altup_unembd_proj.weight`, + * logical `[numInputs-1, hidden, hidden]`); slices are narrowed out at forward time. + * + * Both directions renormalize the projected stream to the active stream's per-token RMS + * magnitude (HF: `target_magnitude / max(rms(proj), 1e-5)`). + */ +public class Gemma3nAltUpGlobals( + private val hiddenSize: Int, + private val numInputs: Int, + private val dtype: KClass? = null, + override val name: String = "altup_globals", +) : Module(), ModuleParameters { + + override val params: List> = listOf( + voidParam("$name.altup_proj.weight", Shape(numInputs - 1, hiddenSize, hiddenSize), dtype), + voidParam("$name.altup_unembd_proj.weight", Shape(numInputs - 1, hiddenSize, hiddenSize), dtype), + ) + + override val modules: List> = emptyList() + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor = input + + /** Per-token RMS magnitude `[S, 1]`: `sqrt(mean(x², dim=-1))`. */ + private fun magnitude(x: Tensor, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + val meanSq = ops.mean(ops.multiply(x, x), dim = -1) // [S] + return ops.unsqueeze(ops.sqrt(meanSq), dim = -1) // [S, 1] + } + + private fun sliceOf(param: ModuleParameter, k: Int, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + // The GGUF stores the stack as ne=[hidden, hidden, numExtra] (ggml: ne2 slowest), and + // the engine surfaces the raw buffer under that ne-ordered shape — so the slice index + // is SLOWEST in memory. Reinterpret row-major as [numExtra, hidden, hidden] first, + // then narrow the leading dim; each slice's buffer is the converter's [out, in] + // row-major matrix. + val stacked = ops.reshape(param.value, Shape(numInputs - 1, hiddenSize, hiddenSize)) + val sliced = ops.narrow(stacked, dim = 0, start = k, length = 1) + return ops.reshape(sliced, Shape(hiddenSize, hiddenSize)) + } + + private fun projectRenormed( + x0mag: Tensor, + stream: Tensor, + param: ModuleParameter, + k: Int, + ctx: ExecutionContext, + ): Tensor { + val ops = ctx.ops + val proj = linearProject(ops, stream, sliceOf(param, k, ctx)) + val newMagSq = ops.mean(ops.multiply(proj, proj), dim = -1) // [S] + val newMag = ops.unsqueeze(ops.sqrt(ops.clamp(newMagSq, 1e-5f, Float.MAX_VALUE)), dim = -1) + return ops.multiply(proj, ops.divide(x0mag, newMag)) + } + + /** HF stream init: `[h0] + [renorm(altup_projections[k](h0))]`. */ + public fun initStreams(h0: Tensor, ctx: ExecutionContext): List> { + val mag = magnitude(h0, ctx) + return listOf(h0) + List(numInputs - 1) { k -> projectRenormed(mag, h0, params[0], k, ctx) } + } + + /** HF finalize: mean of `[h0] + [renorm(altup_unembed_projections[k](h_k+1))]`. */ + public fun mergeStreams(streams: List>, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + val mag = magnitude(streams[0], ctx) + var acc = streams[0] + for (k in 0 until numInputs - 1) { + acc = ops.add(acc, projectRenormed(mag, streams[k + 1], params[1], k, ctx)) + } + return ops.mulScalar(acc, 1.0f / numInputs) + } +} + +/** + * Laurel (Learned Augmented Residual Layer) — HF `Gemma3nTextLaurelBlock`: + * `x + post_laurel_norm(linear_right(linear_left(x)))`. + */ +public class Gemma3nLaurelBlock( + hiddenSize: Int, + laurelRank: Int, + rmsEps: Float, + dtype: KClass? = null, + override val name: String = "laurel", +) : Module() { + + public val linearLeft: VoidDense = VoidDense("$name.laurel_l", laurelRank, hiddenSize, dtype) + public val linearRight: VoidDense = VoidDense("$name.laurel_r", hiddenSize, laurelRank, dtype) + public val postNorm: RMSNormalization = RMSNormalization( + intArrayOf(hiddenSize), rmsEps.toDouble(), unitOffset = false, name = "$name.laurel_post_norm", dtype = dtype, + ) + + override val modules: List> = listOf(linearLeft, linearRight, postNorm) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { + val low = linearLeft.forward(input, ctx) + val back = linearRight.forward(low, ctx) + return ctx.ops.add(input, postNorm.forward(back, ctx)) + } +} + +/** + * Gemma 3n FFN — gelu-gated (`down(gelu(gate(x)) * up(x))`) with optional Gaussian-top-k + * activation sparsity on the gate projection (HF `Gemma3nTextMLP._gaussian_topk`): + * + * ``` + * cutoff = mean(gate, -1) + std_pop(gate, -1) * stdMultiplier + * gate = relu(gate - cutoff) + * ``` + * + * `stdMultiplier` comes precomputed per layer from the GGUF (`activation_sparsity_scale`, + * `Φ⁻¹(0.95) ≈ 1.6449` on sparse layers, `-inf` on the rest — non-finite disables the + * whole branch at build time). Std is population (unbiased=False): `sqrt(E[x²] − E[x]²)`. + */ +public class Gemma3nSparseGeGluFFN( + hiddenSize: Int, + ffnDim: Int, + private val stdMultiplier: Float, + dtype: KClass? = null, + override val name: String = "ffn", +) : Module() { + + // Param names follow the llama/HF convention the engine resolver maps to + // `blk.N.ffn_{gate,up,down}.weight`. + public val gate: VoidDense = VoidDense("$name.gate_proj", ffnDim, hiddenSize, dtype) + public val up: VoidDense = VoidDense("$name.up_proj", ffnDim, hiddenSize, dtype) + public val down: VoidDense = VoidDense("$name.down_proj", hiddenSize, ffnDim, dtype) + + public val sparsityEnabled: Boolean = stdMultiplier.isFinite() && stdMultiplier > 0f + + override val modules: List> = listOf(gate, up, down) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + var g = gate.forward(input, ctx) + if (sparsityEnabled) { + val mean = ops.mean(g, dim = -1) // [S] + val meanSq = ops.mean(ops.multiply(g, g), dim = -1) // [S] + val varPop = ops.subtract(meanSq, ops.multiply(mean, mean)) + val std = ops.sqrt(ops.clamp(varPop, 0f, Float.MAX_VALUE)) + val cutoff = ops.unsqueeze( + ops.add(mean, ops.mulScalar(std, stdMultiplier)), dim = -1, + ) // [S, 1] + g = ops.relu(ops.subtract(g, cutoff)) + } + val activated = ops.gelu(g) + val upOut = up.forward(input, ctx) + return down.forward(ops.multiply(activated, upOut), ctx) + } +} + +/** + * Per-layer-input application — the tail of HF `Gemma3nTextDecoderLayer.forward`. + * Takes the (scaled) corrected active stream and this layer's `per_layer_input` slice, + * returns the DELTA that gets added to the non-active streams: + * `post_norm( proj( gelu(inp_gate(x)) ⊙ per_layer_input ) )`. + * + * The gemma-4 lane's `PerLayerInputBlockHook` applies the same transform but adds it to + * the main residual (gemma-4 has no AltUp streams); gemma3n adds it to streams `1..n-1`, + * so this module returns the delta and `Gemma3nModel` does the stream adds. + */ +public class Gemma3nPerLayerApply( + hiddenSize: Int, + perLayerDim: Int, + rmsEps: Float, + dtype: KClass? = null, + override val name: String = "per_layer_input", +) : Module() { + + public val inpGate: VoidDense = VoidDense("$name.inp_gate", perLayerDim, hiddenSize, dtype) + public val proj: VoidDense = VoidDense("$name.proj", hiddenSize, perLayerDim, dtype) + public val postNorm: RMSNormalization = RMSNormalization( + intArrayOf(hiddenSize), rmsEps.toDouble(), unitOffset = false, name = "$name.post_norm", dtype = dtype, + ) + + override val modules: List> = listOf(inpGate, proj, postNorm) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor = input + + public fun computeDelta( + activeCorrected: Tensor, + perLayerInput: Tensor, + ctx: ExecutionContext, + ): Tensor { + val ops = ctx.ops + val gated = ops.gelu(inpGate.forward(activeCorrected, ctx)) + val mixed = ops.multiply(gated, perLayerInput) + return postNorm.forward(proj.forward(mixed, ctx), ctx) + } +} diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nGGUFNameResolver.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nGGUFNameResolver.kt new file mode 100644 index 00000000..65f1dc3f --- /dev/null +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nGGUFNameResolver.kt @@ -0,0 +1,40 @@ +package sk.ainet.models.gemma3n + +import sk.ainet.io.weights.WeightNameResolver +import sk.ainet.models.gemma.GemmaGGUFNameResolver + +/** + * Resolves DSL module paths to GGUF tensor names for the Gemma 3n family: the 3n-specific + * rules (AltUp per-layer + global tensors, Laurel) matched first, everything the gemma-4 + * lane already handles (sandwich norms, PLE names, llama-standard set) delegated to + * [GemmaGGUFNameResolver]. + */ +public class Gemma3nGGUFNameResolver : WeightNameResolver { + + private val gemma = GemmaGGUFNameResolver() + + override fun resolve(modulePath: String, paramName: String): String? { + val blockPrefix = modulePath.split("/").drop(1).firstOrNull { it.startsWith("blk.") } + + // Per-layer AltUp + Laurel params are named after their GGUF tensors already — + // ".altup_router.weight" etc. — so the rule is: take the tensor-suffix + // and prefix the block. + for (suffix in BLOCK_SUFFIXES) { + if (paramName.endsWith(".$suffix.weight") || paramName == "$suffix.weight") { + return if (blockPrefix != null) "$blockPrefix.$suffix.weight" else null + } + } + // Model-level AltUp stream projections (3D tensors, no block prefix). + if (paramName.endsWith(".altup_proj.weight")) return "altup_proj.weight" + if (paramName.endsWith(".altup_unembd_proj.weight")) return "altup_unembd_proj.weight" + + return gemma.resolve(modulePath, paramName) + } + + private companion object { + val BLOCK_SUFFIXES = listOf( + "altup_router_norm", "altup_router", "altup_predict_coef", "altup_correct_coef", + "altup_correct_scale", "laurel_l", "laurel_r", "laurel_post_norm", + ) + } +} diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModel.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModel.kt new file mode 100644 index 00000000..e984b74a --- /dev/null +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModel.kt @@ -0,0 +1,164 @@ +package sk.ainet.models.gemma3n + +import sk.ainet.apps.llm.HybridTransformerBlock +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.layers.EmbeddingAdapter +import sk.ainet.lang.nn.normalization.RMSNormalization +import sk.ainet.lang.nn.transformer.MultiHeadAttention +import sk.ainet.lang.nn.transformer.VoidDense +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType +import sk.ainet.models.gemma.PerLayerEmbedding +import kotlin.reflect.KClass + +/** + * Top-level Gemma 3n model — the DSL replacement for the hand-rolled `Gemma3nRuntime` + * (#377), faithful to HF `Gemma3nTextModel.forward` / `Gemma3nTextDecoderLayer.forward`. + * + * Follows the `GemmaModel` wrapper pattern (gemma-4 PLE precedent): the module tree is + * regular DSL modules (so `WeightMapper` binds every weight by name), but the forward + * orchestration is bespoke because AltUp threads `numInputs` parallel hidden streams + * through every layer — inexpressible as a plain Sequential: + * + * ``` + * h0 = embed(ids) * sqrt(hidden) + * ple = PerLayerEmbedding.compute(ids, h0) # [B, S, L, pleDim] + * streams = altupGlobals.initStreams(h0) # magnitude-renormed projections + * per layer: + * preds = altup.predict(streams) + * active = preds[activeIdx]; an = attn_norm(active) + * laurel = laurel(an) # an + norm(right(left(an))) + * attn = post_attention_norm( MHA(an) ) + * attnLaurel = ((active + attn) + laurel) / √2 + * ffw = post_ffw_norm( ffn( ffn_norm(attnLaurel) ) ) # sparsity on first layers + * streams = altup.correct(preds, attnLaurel + ffw) + * delta = perLayerApply( scale(streams[active]), ple[:, :, layer] ) + * streams[1:] += delta + * merged = altupGlobals.mergeStreams(streams) # renormed mean + * logits = lm_head( output_norm(merged) ) # tied embeddings, no softcap + * ``` + */ +public class Gemma3nModel( + public val tokenEmbedding: EmbeddingAdapter, + public val ple: PerLayerEmbedding, + public val altupGlobals: Gemma3nAltUpGlobals, + public val blocks: List>, + public val outputNorm: RMSNormalization, + public val lmHead: VoidDense, + public val dtype: KClass, + public val activeIdx: Int, + public val embedScale: Float, + override val name: String = "Gemma3nModel", +) : Module() { + + /** + * Export/runtime injection point: when set, [onForward] uses this tensor + * (`[batch, seq, numLayers, pleDim]`) as the per-layer inputs and skips [ple].compute — + * the compiled StableHLO graph takes per_layer_inputs as a second INPUT, computed on + * the CPU from the packed PLE table at runtime (PLE's whole design point: those + * parameters stay off the accelerator). Eager decode leaves this null. + */ + public var externalPerLayerInputs: Tensor? = null + + override val modules: List> = buildList { + add(tokenEmbedding) + add(ple) + add(altupGlobals) + addAll(blocks) + add(outputNorm) + add(lmHead) + } + + /** Typed handles into one block's module list (bound by construction in `gemma3nNetwork`). */ + public class BlockRefs( + public val attnNorm: RMSNormalization, + public val mha: MultiHeadAttention, + public val postAttnNorm: RMSNormalization, + public val ffnNorm: RMSNormalization, + public val ffn: Gemma3nSparseGeGluFFN, + public val postFfwNorm: RMSNormalization, + public val laurel: Gemma3nLaurelBlock, + public val altup: Gemma3nAltUpBlock, + public val perLayer: Gemma3nPerLayerApply, + ) + + @Suppress("UNCHECKED_CAST") + private fun refsFor(block: HybridTransformerBlock): BlockRefs { + val mods = block.modules + fun norm(id: String): RMSNormalization = + mods.filterIsInstance>().firstOrNull { it.name == id } + ?: error("Gemma3nModel: block ${block.name} has no RMSNorm '$id'") + return BlockRefs( + attnNorm = norm("attn_norm"), + mha = mods.filterIsInstance>().first(), + postAttnNorm = norm("post_attention_norm"), + ffnNorm = norm("ffn_norm"), + ffn = mods.filterIsInstance>().first(), + postFfwNorm = norm("post_ffw_norm"), + laurel = mods.filterIsInstance>().first(), + altup = mods.filterIsInstance>().first(), + perLayer = mods.filterIsInstance>().first(), + ) + } + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { + val ops = ctx.ops + val invSqrt2 = 0.70710678f + + // 1 — scaled embedding. + val rawEmbeds = tokenEmbedding.forward(input, ctx) + val h0 = if (embedScale != 1f) ops.mulScalar(rawEmbeds, embedScale) else rawEmbeds + + // 2 — per-layer inputs [B, S, L, pleDim] (identical math to gemma-4: token-identity + // gather * sqrt(pleDim), context projection * hidden^-0.5, norm, sum * 1/sqrt2). + val perLayerInputs = externalPerLayerInputs ?: run { + val ids2d = if (input.rank == 1) ops.unsqueeze(input, 0) else input + val embeds3d = if (h0.rank == 2) ops.unsqueeze(h0, 0) else h0 + ple.compute(ids2d, embeds3d, ctx, dtype) + } + + // 3 — AltUp stream init. + var streams = altupGlobals.initStreams(h0, ctx) + + // 4 — per-layer flow. + for ((layerIdx, block) in blocks.withIndex()) { + val r = refsFor(block) + val preds = r.altup.predict(streams, ctx) + val active = preds[activeIdx] + val an = r.attnNorm.forward(active, ctx) + val laurel = r.laurel.forward(an, ctx) + val attn = r.postAttnNorm.forward(r.mha.forward(an, ctx), ctx) + val attnGated = ops.add(active, attn) + val attnLaurel = ops.mulScalar(ops.add(attnGated, laurel), invSqrt2) + val ffw = r.postFfwNorm.forward(r.ffn.forward(r.ffnNorm.forward(attnLaurel, ctx), ctx), ctx) + val corrected = r.altup.correct(preds, ops.add(attnLaurel, ffw), ctx) + + // Per-layer input: transform the (scaled) corrected active stream and add the + // delta to the NON-active streams (HF: `corrected_predictions[1:] += ...`). + val scaledActive = r.altup.scaleCorrectedOutput(corrected[activeIdx], ctx) + val pliSlice = perLayerSlice(perLayerInputs, layerIdx, h0.rank, ctx) + val delta = r.perLayer.computeDelta(scaledActive, pliSlice, ctx) + streams = List(corrected.size) { i -> + if (i == 0) corrected[0] else ops.add(corrected[i], delta) + } + } + + // 5 — merge streams, final norm, tied lm_head (gemma3n has no final softcap). + val merged = altupGlobals.mergeStreams(streams, ctx) + return lmHead.forward(outputNorm.forward(merged, ctx), ctx) + } + + /** `per_layer_inputs[..., layerIdx, :]` matched to the trunk's working rank. */ + private fun perLayerSlice( + perLayerInputs: Tensor, + layerIdx: Int, + workingRank: Int, + ctx: ExecutionContext, + ): Tensor { + val ops = ctx.ops + // [B, S, L, pleDim] → narrow L → [B, S, 1, pleDim] → squeeze → [B, S, pleDim] + val slice = ops.squeeze(ops.narrow(perLayerInputs, dim = 2, start = layerIdx, length = 1), dim = 2) + return if (workingRank == 2) ops.squeeze(slice, dim = 0) else slice + } +} diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModelMetadata.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModelMetadata.kt index 85045df1..d3225dd8 100644 --- a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModelMetadata.kt +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nModelMetadata.kt @@ -34,8 +34,19 @@ public data class Gemma3nModelMetadata( /** Per-layer activation sparsity rates. Empty means no sparsity. */ val activationSparsityPattern: List = emptyList(), /** Activation sparsity scale factor (from GGUF: gemma3n.activation_sparsity_scale). */ - val activationSparsityScale: Float = 0f + val activationSparsityScale: Float = 0f, + /** RMSNorm epsilon (`gemma3n.attention.layer_norm_rms_epsilon`; real E2B/E4B: 1e-6). */ + val rmsNormEps: Float = 1e-6f, + /** + * Per-layer activation-sparsity std multipliers (`gemma3n.activation_sparsity_scale`). + * Real checkpoints store `Φ⁻¹(target_sparsity) ≈ 1.6449` on sparse layers and `-inf` + * on the rest — non-finite (or empty list) disables sparsity for that layer. + */ + val activationSparsityScales: List = emptyList(), ) { + /** The std multiplier for one layer, or `null` when sparsity is off there. */ + public fun sparsityScaleFor(layerIdx: Int): Float? = + activationSparsityScales.getOrNull(layerIdx)?.takeIf { it.isFinite() && it > 0f } /** * Returns the layer type at the given layer index. * Pattern repeats: ["sliding", "sliding", "sliding", "sliding", "full"] diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkDef.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkDef.kt new file mode 100644 index 00000000..deea0178 --- /dev/null +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkDef.kt @@ -0,0 +1,195 @@ +package sk.ainet.models.gemma3n + +import sk.ainet.apps.llm.HybridTransformerBlock +import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.dsl.NeuralNetworkDslImpl +import sk.ainet.lang.nn.dsl.StageImpl +import sk.ainet.lang.nn.dsl.embedding +import sk.ainet.lang.nn.dsl.multiHeadAttention +import sk.ainet.lang.nn.dsl.rmsNorm +import sk.ainet.lang.nn.layers.EmbeddingAdapter +import sk.ainet.lang.nn.normalization.RMSNormalization +import sk.ainet.lang.nn.transformer.OwnerReadOnlyKVCache +import sk.ainet.lang.nn.transformer.PositionalKVCache +import sk.ainet.lang.nn.transformer.RoPEMode +import sk.ainet.lang.nn.transformer.VoidDense +import sk.ainet.lang.types.DType +import sk.ainet.models.gemma.LayerType +import sk.ainet.models.gemma.PerLayerEmbedding +import kotlin.reflect.KClass + +/** + * Gemma 3n architecture defined via the network DSL — the #377 DSL migration, replacing the + * hand-rolled `Gemma3nRuntime`. Faithful to HF `Gemma3nTextModel` (see [Gemma3nModel] for + * the per-layer flow) and buildable into a compute-graph tape for the StableHLO → IREE + * mobile path. + * + * What gemma3n adds over the gemma-4 lane's `gemmaNetwork()` (which already carries hybrid + * sliding/global attention, dual RoPE bases, per-layer FFN dims, per-type shared KV, PLE, + * q/k-norm, parameterless v-norm and attention scale 1.0): + * [Gemma3nAltUpBlock] (4 parallel streams + router), [Gemma3nLaurelBlock], + * [Gemma3nSparseGeGluFFN] (Gaussian-top-k on the first layers) and the PLE delta going to + * the non-active AltUp streams instead of the residual. + */ +public fun gemma3nNetwork( + metadata: Gemma3nModelMetadata, + dtype: KClass, + maxInferenceLen: Int = minOf(metadata.contextLength, 4096), + /** HF `laurel_rank` (64 on real checkpoints; not in the GGUF — the loader derives it + * from `blk.0.laurel_l`'s shape). */ + laurelRank: Int = LAUREL_RANK, + /** Number of layers the PLE tensors cover — normally [Gemma3nModelMetadata.blockCount], + * but a layer-truncated export build keeps the FULL table so the stored + * `per_layer_*` tensors still bind shape-exact. */ + pleNumLayers: Int = metadata.blockCount, +): Module { + val dim = metadata.embeddingLength + val nHeads = metadata.headCount + val nKVHeads = metadata.kvHeadCount + val nLayers = metadata.blockCount + val headDim = metadata.headDim + val seqLen = maxInferenceLen + val vocabSize = metadata.vocabSize + val eps = metadata.rmsNormEps + + val nnCtx = DefaultNeuralNetworkExecutionContext() + val dslImpl = NeuralNetworkDslImpl(nnCtx, dtype) + dslImpl.embedding(vocabSize, dim, id = "token_embd") + + // KV sharing: same owner-per-attention-type scheme as gemma-4 (HF: a shared layer + // reuses the K/V of the LAST non-shared layer of the same type). + val firstSharedLayer = nLayers - metadata.kvSharedLayers + val typeOwners = mutableMapOf>() + val typeOwnerLayerIdx = mutableMapOf() + if (metadata.kvSharedLayers > 0) { + for (l in 0 until firstSharedLayer) { + typeOwnerLayerIdx[metadata.getLayerType(l)] = l + } + } + + for (layer in 0 until nLayers) { + val layerType = metadata.getLayerType(layer) + val isGlobal = layerType == LayerType.GLOBAL + val ropeBase = metadata.getRopeBase(layer) + val slidingWindow = if (isGlobal) null else metadata.slidingWindow + val ffnDim = metadata.feedForwardLengths.getOrElse(layer) { metadata.feedForwardLengths.last() } + val isInSharedGroup = metadata.kvSharedLayers > 0 && layer >= firstSharedLayer + + val stage = StageImpl(nnCtx, "blk.$layer", dtype) + stage.rmsNorm(dim, eps, id = "attn_norm", unitOffset = false) + stage.multiHeadAttention( + dim = dim, + nHeads = nHeads, + nKVHeads = nKVHeads, + causal = true, + // HF Gemma3nTextAttention: per-head RMSNorm (with scale) on Q and K before RoPE, + // parameterless per-head RMSNorm on V, attention scaling fixed to 1.0. + qkNorm = true, + qkNormUnitOffset = false, + qkNormEps = eps, + attentionScale = 1.0f, + vNormNoScale = true, + id = "attn", + slidingWindow = slidingWindow, + ) { + rope( + headDim = headDim, + maxSeqLen = seqLen, + mode = RoPEMode.SPLIT_HALF, + base = ropeBase, + ) + if (!isInSharedGroup) { + val own = PositionalKVCache( + maxSeqLen = seqLen, + nKVHeads = nKVHeads, + headDim = headDim, + name = "blk.$layer.attn.kv_cache", + ) + kvCache(own) + if (typeOwnerLayerIdx[layerType] == layer) typeOwners[layerType] = own + } else { + val ownerCache = typeOwners[layerType] + ?: error( + "gemma3n: kv-shared layer $layer (type=$layerType) has no non-shared " + + "owner of the same type before firstSharedLayer=$firstSharedLayer", + ) + kvCache(OwnerReadOnlyKVCache(delegate = ownerCache, name = "blk.$layer.attn.kv_cache")) + } + } + stage.rmsNorm(dim, eps, id = "post_attention_norm", unitOffset = false) + stage.rmsNorm(dim, eps, id = "ffn_norm", unitOffset = false) // pre_feedforward_layernorm + stage.modules += Gemma3nSparseGeGluFFN( + hiddenSize = dim, + ffnDim = ffnDim, + stdMultiplier = metadata.sparsityScaleFor(layer) ?: Float.NEGATIVE_INFINITY, + dtype = dtype, + name = "ffn", + ) + stage.rmsNorm(dim, eps, id = "post_ffw_norm", unitOffset = false) + stage.modules += Gemma3nLaurelBlock( + hiddenSize = dim, + laurelRank = laurelRank, + rmsEps = eps, + dtype = dtype, + name = "laurel", + ) + stage.modules += Gemma3nAltUpBlock( + hiddenSize = dim, + numInputs = metadata.numAltupInputs, + activeIdx = metadata.altupActiveIdx, + rmsEps = eps, + dtype = dtype, + name = "altup", + ) + stage.modules += Gemma3nPerLayerApply( + hiddenSize = dim, + perLayerDim = metadata.perLayerEmbeddingLength, + rmsEps = eps, + dtype = dtype, + name = "per_layer_input", + ) + dslImpl.modules += HybridTransformerBlock(stage.modules.toList(), name = "blk.$layer") + } + + dslImpl.rmsNorm(dim, eps, id = "output_norm", unitOffset = false) + // Void placeholder — the 262k vocab head would eagerly allocate ~2 GB of zeros otherwise; + // gemma3n ties the head to token_embd, bound by the loader. + dslImpl.modules += VoidDense("output", vocabSize, dim, dtype = dtype) + + @Suppress("UNCHECKED_CAST") + val tokenEmbedding = dslImpl.modules[0] as EmbeddingAdapter + val blocks = dslImpl.modules.filterIsInstance>() + val outputNorm = dslImpl.modules[dslImpl.modules.size - 2] as RMSNormalization + @Suppress("UNCHECKED_CAST") + val lmHead = dslImpl.modules[dslImpl.modules.size - 1] as VoidDense + + val ple = PerLayerEmbedding( + vocabSize = vocabSize, + hiddenSize = dim, + numLayers = pleNumLayers, + perLayerDim = metadata.perLayerEmbeddingLength, + rmsEps = eps, + ) + + return Gemma3nModel( + tokenEmbedding = tokenEmbedding, + ple = ple, + altupGlobals = Gemma3nAltUpGlobals( + hiddenSize = dim, + numInputs = metadata.numAltupInputs, + dtype = dtype, + ), + blocks = blocks, + outputNorm = outputNorm, + lmHead = lmHead, + dtype = dtype, + activeIdx = metadata.altupActiveIdx, + // HF: embed_scale = hidden_size ** 0.5 (bf16-rounded in the reference; we bisect + // against llama.cpp if the rounding ever matters at parity tolerance). + embedScale = kotlin.math.sqrt(dim.toFloat()), + ) +} + +/** HF `laurel_rank` — constant 64 across gemma-3n checkpoints (not stored in the GGUF). */ +public const val LAUREL_RANK: Int = 64 diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkLoader.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkLoader.kt new file mode 100644 index 00000000..e3c9e6e7 --- /dev/null +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nNetworkLoader.kt @@ -0,0 +1,88 @@ +package sk.ainet.models.gemma3n + +import sk.ainet.context.ExecutionContext +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.weights.MappingConfig +import sk.ainet.io.weights.WeightMapper +import sk.ainet.io.weights.WeightTensor +import sk.ainet.lang.nn.Module +import sk.ainet.lang.types.DType +import kotlin.reflect.KClass + +/** + * End-to-end loader for the Gemma 3n DSL path (#377): loads a real gemma3n GGUF through the + * engine-delegated [Gemma3nWeightLoader] (packed/MAPPED by default; the PLE table stays + * packed with row-dequant), builds [gemma3nNetwork] and binds every weight via + * [WeightMapper] + [Gemma3nGGUFNameResolver]. + */ +public object Gemma3nNetworkLoader { + + public suspend inline fun fromGguf( + ctx: ExecutionContext, + noinline randomAccessProvider: () -> RandomAccessSource, + maxInferenceLen: Int? = null, + debug: Boolean = false, + ): Module { + val weights = Gemma3nWeightLoader(randomAccessProvider).loadToMapStreaming(ctx) + return fromWeights(ctx, weights, T::class, maxInferenceLen, debug) + } + + public fun fromWeights( + ctx: ExecutionContext, + weights: Gemma3nWeights, + dtype: KClass, + maxInferenceLen: Int? = null, + debug: Boolean = false, + /** PLE table coverage — see [gemma3nNetwork]'s `pleNumLayers` (export truncation). */ + pleNumLayers: Int? = null, + ): Module { + val md = weights.metadata + // laurel_rank is not a GGUF field — read it off the checkpoint's own tensor. + val laurelRank = weights.tensors["blk.0.laurel_l.weight"]?.shape?.get(0) ?: LAUREL_RANK + + val model = gemma3nNetwork( + md, + dtype, + maxInferenceLen = maxInferenceLen ?: minOf(md.contextLength, 4096), + laurelRank = laurelRank, + pleNumLayers = pleNumLayers ?: md.blockCount, + ) + + val weightTensors = weights.tensors.map { (name, tensor) -> + WeightTensor(name = name, shape = tensor.shape.dimensions.toList(), tensor = tensor) + } + val config = MappingConfig( + usePathBasedMatching = false, + fallbackToShapeMatching = false, + debug = debug, + nameResolver = Gemma3nGGUFNameResolver(), + ) + val result = WeightMapper.applyWeights(model, weightTensors, config) + + // gemma3n has no bias tensors; every non-bias DSL param must bind, and no loaded + // tensor may silently go unused (the qwen-bias lesson, transformers#352). + val unmappedNonBias = result.missingParams.filter { !it.contains(".bias") } + require(unmappedNonBias.isEmpty()) { + buildString { + appendLine("gemma3n: failed to map ${unmappedNonBias.size} weight parameters:") + unmappedNonBias.take(20).forEach { appendLine(" - $it") } + if (result.unusedTensors.isNotEmpty()) { + appendLine("Unused tensors (${result.unusedTensors.size}):") + result.unusedTensors.take(20).forEach { appendLine(" - $it") } + } + }.trim() + } + require(result.unusedTensors.isEmpty()) { + "gemma3n: tensors present in the GGUF but never bound: ${result.unusedTensors.take(20)}" + } + return model + } + + public inline fun fromWeights( + ctx: ExecutionContext, + weights: Gemma3nWeights, + maxInferenceLen: Int? = null, + debug: Boolean = false, + pleNumLayers: Int? = null, + ): Module = fromWeights(ctx, weights, T::class, maxInferenceLen, debug, pleNumLayers) +} diff --git a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nWeightLoader.kt b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nWeightLoader.kt index 315442a6..78fbb9d2 100644 --- a/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nWeightLoader.kt +++ b/llm-inference/gemma3n/src/commonMain/kotlin/sk/ainet/models/gemma3n/Gemma3nWeightLoader.kt @@ -260,7 +260,12 @@ public class Gemma3nWeightLoader private constructor( weightForm = defaultForm, weightFormFor = { name -> when (name) { + // Embedding.gather needs element access — always dense. Gemma3nTensorNames.TOKEN_EMBEDDINGS -> GEMMA_DEQUANTIZE_ALL + // The PLE table ALWAYS stays packed (row-dequant wrapper) — that is + // PLE's design point: per-layer embeddings live off-accelerator, on + // the CPU side, even for the compiled mobile path (where + // per_layer_inputs is a graph INPUT, not a baked parameter). Gemma3nTensorNames.PER_LAYER_TOKEN_EMBD -> WeightForm(shape = WeightShapeOrientation.OUT_IN) else -> null @@ -386,9 +391,14 @@ public class Gemma3nWeightLoader private constructor( // Gemma 3n specific val slidingWindow = fields["$prefix.attention.sliding_window"]?.toIntValue() ?: Gemma3nModelMetadata.DEFAULT_SLIDING_WINDOW + // Real llama.cpp gemma3n GGUFs declare only `rope.freq_base` (the global/full base, + // 1M) — the sliding/local base is the SWA default 10k (llama.cpp + // `rope_freq_base_train_swa`). Legacy keys tried first for synthetic fixtures. val ropeBaseLocal = fields["$prefix.rope.freq_base_local"]?.toFloatValue() + ?: fields["$prefix.rope.freq_base_swa"]?.toFloatValue() ?: Gemma3nModelMetadata.DEFAULT_ROPE_BASE_LOCAL val ropeBaseGlobal = fields["$prefix.rope.freq_base_global"]?.toFloatValue() + ?: fields["$prefix.rope.freq_base"]?.toFloatValue() ?: Gemma3nModelMetadata.DEFAULT_ROPE_BASE_GLOBAL val kvSharedLayers = fields["$prefix.kv_shared_layers"]?.toIntValue() ?: fields["$prefix.attention.shared_kv_layers"]?.toIntValue() @@ -408,6 +418,15 @@ public class Gemma3nWeightLoader private constructor( // Activation sparsity val activationSparsityPattern = extractStreamingActivationSparsityPattern(fields, prefix) val activationSparsityScale = fields["$prefix.activation_sparsity_scale"]?.toFloatValue() ?: 0f + // Real GGUFs store the per-layer std multipliers directly (1.6449 on sparse layers, + // -inf on the rest); the DSL path consumes this list. + val activationSparsityScales = when (val v = fields["$prefix.activation_sparsity_scale"]) { + is List<*> -> v.mapNotNull { (it as? Number)?.toFloat() } + is FloatArray -> v.toList() + is DoubleArray -> v.map { it.toFloat() } + else -> emptyList() + } + val rmsNormEps = fields["$prefix.attention.layer_norm_rms_epsilon"]?.toFloatValue() ?: 1e-6f return Gemma3nModelMetadata( architecture = arch, @@ -428,7 +447,9 @@ public class Gemma3nWeightLoader private constructor( numAltupInputs = numAltupInputs, altupActiveIdx = altupActiveIdx, activationSparsityPattern = activationSparsityPattern, - activationSparsityScale = activationSparsityScale + activationSparsityScale = activationSparsityScale, + rmsNormEps = rmsNormEps, + activationSparsityScales = activationSparsityScales, ) } @@ -483,6 +504,11 @@ public class Gemma3nWeightLoader private constructor( if (perLayerValue != null && perLayerValue is List<*>) { return perLayerValue.mapNotNull { (it as? Number)?.toInt() } } + // Real llama.cpp GGUFs store the per-layer array under the SINGULAR key. + val singular = fields["$prefix.feed_forward_length"] + if (singular is List<*> && singular.size > 1) { + return singular.mapNotNull { (it as? Number)?.toInt() } + } // Fall back to single FFN length val ffnLength = fields["$prefix.feed_forward_length"]?.toIntValue() ?: (embeddingLength * 4) @@ -508,6 +534,11 @@ public class Gemma3nWeightLoader private constructor( if (patternValue != null && patternValue is List<*>) { return patternValue.mapNotNull { it as? String } } + // Real llama.cpp GGUFs store per-layer booleans: true = sliding, false = full. + val swaPattern = fields["$prefix.attention.sliding_window_pattern"] + if (swaPattern is List<*> && swaPattern.isNotEmpty() && swaPattern.first() is Boolean) { + return swaPattern.map { if (it == true) "sliding" else "full" } + } return Gemma3nModelMetadata.DEFAULT_LAYER_PATTERN } diff --git a/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportCli.kt b/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportCli.kt new file mode 100644 index 00000000..22d2ab98 --- /dev/null +++ b/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportCli.kt @@ -0,0 +1,27 @@ +package sk.ainet.models.gemma3n + +/** + * CLI entry for [Gemma3nExportHarness] — the `exportGemma3n` gradle task + * (SmolLM2/FunctionGemma pattern). + * + * Env: + * - `GEMMA3N_GGUF` path to gemma-3n GGUF (required) + * - `GEMMA3N_OUT_DIR` output directory (default `build/gemma3n-export`) + * - `GEN_SEQ` fixed sequence length of the redecode graph (default 24) + * - `GEMMA3N_DTYPE` `bf16` (default) or `f32` external params + */ +public fun main() { + val gguf = System.getenv("GEMMA3N_GGUF") + ?: error("GEMMA3N_GGUF must point at a gemma-3n GGUF checkpoint") + val outDir = System.getenv("GEMMA3N_OUT_DIR") ?: "build/gemma3n-export" + val seq = System.getenv("GEN_SEQ")?.toIntOrNull() ?: 24 + val bf16 = (System.getenv("GEMMA3N_DTYPE") ?: "bf16").lowercase() != "f32" + val layers = System.getenv("GEMMA3N_LAYERS")?.toIntOrNull() + + val r = Gemma3nExportHarness.export(gguf = gguf, outDir = outDir, seq = seq, bf16 = bf16, layers = layers) + println("gemma3n export complete:") + println(" mlir = ${r.mlirPath}") + println(" safetensors = ${r.safetensorsPath} (${r.weightMiB} MiB, ${r.externalParamCount} params)") + println(" manifest = ${r.manifestPath}") + println(" seq=${r.seq} vocab=${r.vocabSize} fn=@${Gemma3nExportHarness.FN_REDECODE}") +} diff --git a/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportHarness.kt b/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportHarness.kt new file mode 100644 index 00000000..1110d552 --- /dev/null +++ b/llm-inference/gemma3n/src/jvmMain/kotlin/sk/ainet/models/gemma3n/Gemma3nExportHarness.kt @@ -0,0 +1,312 @@ +package sk.ainet.models.gemma3n + +import kotlinx.coroutines.runBlocking +import sk.ainet.compile.hlo.ConstantMaterializationPolicy +import sk.ainet.compile.hlo.ExternalParameterRef +import sk.ainet.compile.hlo.StableHloConverterFactory +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.transformer.MultiHeadAttention +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.types.FP32 +import sk.ainet.tape.Execution +import java.io.BufferedOutputStream +import java.io.File +import java.io.FileOutputStream +import java.nio.ByteBuffer +import java.nio.ByteOrder + +/** + * Gemma 3n compiled-export harness — the StableHLO → IREE mobile path for the #377 DSL + * lane (SmolLM2/FunctionGemma redecode pattern): author `gemma3nNetwork()` from the real + * GGUF, trace ONE fixed `[1, seq]` prefill pass ending in the DSL argMax tail, and emit + * portable StableHLO with every weight lifted to an EXTERNAL parameter + * (`scope = "model"`). `func @gemma3n` returns `tensor` — the + * `llm-runtime/iree-android` `IreeRedecodeSession` contract, so a compiled + * `(vmfb, irpa, "gemma3n")` triple runs on-device unchanged. + * + * Everything gemma3n-specific traces through `ctx.ops`: the four AltUp streams with the + * tanh router, Laurel, Gaussian-top-k sparsity, and the PLE token gather + * (`PerLayerEmbedding` switches to an `indexSelect` graph op while recording — the eager + * packed row-dequant path is host-side and would bake constants). KV caches are stripped + * before tracing (a single fixed-seq pass needs none; `OwnerReadOnlyKVCache` followers + * would otherwise record the stateful eager copy path). + * + * **The PLE table is NOT a graph parameter.** `per_layer_inputs` (`[1, seq, L, pleDim]`) + * is the graph's SECOND INPUT, computed on the CPU from the packed PLE table at runtime — + * PLE's design point per Google's Gemma 3n guide: per-layer embeddings live + * off-accelerator. That keeps the 262k×7680 table (2 GB packed, 8 GB dense) out of the + * parameter archive entirely: the archive carries the trunk + token embedding only + * (~4 GB bf16 for E2B; int8 narrowing is follow-up scope, FunctionGemma + * `rewriteGlobalsToInt8` precedent). The host load must be dense FP32 (packed tensors + * cannot become graph constants), and trace zeros + graph copies sit alongside it — + * E2B peaks ~44 GB transient heap; run on a large-memory host with + * `-PexportMaxHeap=46g` (the all-zero trace pages compress well under macOS). + * + * Writes `/gemma3n-gen.mlir` + `/gemma3n.safetensors` + `manifest.json`. + * vmfb compilation happens outside Kotlin (iree-compile; the repo pins a Torq fork for + * `iree-run-module` — see `llm-runtime/gemma-iree`). + */ +@OptIn(ExperimentalMemoryApi::class) +public object Gemma3nExportHarness { + + public const val FN_REDECODE: String = "gemma3n" + public const val PARAMETER_SCOPE: String = "model" + + public data class RedecodeResult( + val mlirPath: String, + val safetensorsPath: String, + val manifestPath: String, + val externalParamCount: Int, + val weightMiB: Long, + val seq: Int, + val vocabSize: Int, + ) + + public fun export( + gguf: String, + outDir: String, + seq: Int = 24, + bf16: Boolean = true, + /** + * Truncate the exported trunk to the first N layers. Full-E2B export needs a + * ≥64 GB host (the engine's trace/export pipeline keeps dense weights, per-op + * zero buffers and graph constant copies co-resident — engine issue filed); + * a truncated export verifies the whole pipeline end-to-end on smaller hosts + * and is what the model-gated smoke uses. + */ + layers: Int? = null, + ): RedecodeResult = runBlocking { + val ctx = DirectCpuExecutionContext.create() + // DENSE load is REQUIRED for export: packed tensors cannot be extracted as graph + // constants — the tracer silently turns them into opaque function ARGUMENTS + // (measured: 190+ weight args, zero dot ops — an unservable module). The sanity + // check below makes that failure loud if it ever regresses. + val weights = Gemma3nWeightLoader( + randomAccessProvider = { JvmRandomAccessSource.open(gguf) }, + weightForm = WeightForm( + encoding = EncodingRequest.DequantizeTo(FP32), + shape = WeightShapeOrientation.OUT_IN, + ), + ).loadToMapStreaming(ctx) + + val fullLayers = weights.metadata.blockCount + val n = layers?.coerceIn(1, fullLayers) ?: fullLayers + val weightsN = if (n == fullLayers) weights else { + val md = weights.metadata + val firstShared = fullLayers - md.kvSharedLayers + Gemma3nWeights( + md.copy( + blockCount = n, + // Shared-KV followers only exist past firstShared; a truncated trunk + // below that point has no sharing. + kvSharedLayers = (n - firstShared).coerceAtLeast(0), + feedForwardLengths = md.feedForwardLengths.take(n), + layerPattern = md.layerPattern.take(n).ifEmpty { md.layerPattern }, + activationSparsityScales = md.activationSparsityScales.take(n), + ), + weights.tensors.filterKeys { key -> + !key.startsWith("blk.") || + (key.removePrefix("blk.").substringBefore('.').toIntOrNull() ?: 0) < n + }, + ) + } + val model = Gemma3nNetworkLoader.fromWeights( + ctx, weightsN, maxInferenceLen = seq, pleNumLayers = fullLayers, + ) + + fun stripKvCache(m: Module<*, *>) { + if (m is MultiHeadAttention<*, *>) m.kvCache = null + m.modules.forEach { stripKvCache(it) } + } + stripKvCache(model) + + val input = VoidOpsTensor( + object : TensorData { + override val shape = Shape(1, seq) + override fun get(vararg indices: Int): Float = 0.0f + override fun set(vararg indices: Int, value: Float) {} + }, + FP32::class, + ) + // Second graph input: per_layer_inputs, computed host-side (CPU gather over the + // packed PLE table) by the on-device session, exactly like the eager path does. + val md0 = weights.metadata + val pliInput = VoidOpsTensor( + object : TensorData { + override val shape = Shape(1, seq, fullLayers, md0.perLayerEmbeddingLength) + override fun get(vararg indices: Int): Float = 0.0f + override fun set(vararg indices: Int, value: Float) {} + }, + FP32::class, + ) + (model as Gemma3nModel).externalPerLayerInputs = pliInput + + val tapeCtx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + val tape = tapeCtx.record { + val ct = currentTape ?: error("no tape") + Execution.tapeStack.pushTape(ct) + try { + val ectx = this as ExecutionContext + val logits = model.forward(input, ectx) // [1, seq, vocab] f32 + val idx = ectx.ops.argMax(logits, dim = -1) // [1, seq] i32 + ectx.ops.squeeze(idx, 0) // [seq] i32 — the redecode runtime contract + } finally { + Execution.tapeStack.popTape() + } + }.first + + val graph = (tape as DefaultExecutionTape).toComputeGraph( + synthesizeExternalInputs = true, embedConstants = true, + ) + val module = StableHloConverterFactory + .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = PARAMETER_SCOPE)) + .convert(graph, FN_REDECODE) + + val out = File(outDir).apply { mkdirs() } + val ext = module.externalParameters + + // Sanity: the redecode contract has exactly TWO function inputs (tokens, + // per_layer_inputs). More means weights leaked into the signature (the + // packed-tensor failure mode) — the module would be unservable. + val sig = Regex("func\\.func @$FN_REDECODE\\(([^)]*)\\)").find(module.content)?.groupValues?.get(1) + ?: error("gemma3n export: emitted module has no @$FN_REDECODE function") + val argCount = Regex("%arg\\d+").findAll(sig).count() + require(argCount == 2) { + "gemma3n export: expected 2 graph inputs (tokens, per_layer_inputs) but the " + + "function signature has $argCount — weights leaked into the signature " + + "(packed tensors cannot become graph constants; the loader must dequantize)." + } + // Emission completeness: the engine converter currently reports operand-linkage + // failures as MLIR comments and still exits 0, leaving a compute-free module + // (SKaiNET#1247). Make that a hard error here. + val failures = Regex("// Conversion failed for node ([^:]+):").findAll(module.content) + .map { it.groupValues[1] }.take(5).toList() + require(failures.isEmpty()) { + "gemma3n export: the StableHLO converter failed on ${failures.size}+ nodes " + + "(first: $failures) — the emitted module is not servable. See SKaiNET#1247." + } + require(!module.content.contains("// Warning: No output values found")) { + "gemma3n export: emitted function has no outputs — unservable module (SKaiNET#1247)." + } + + val mlir = if (bf16) rewriteGlobalsToBf16(module.content) else module.content + val mlirFile = File(out, "gemma3n-gen.mlir").apply { writeText(mlir) } + + val stFile = File(out, "gemma3n.safetensors") + writeSafetensors(ext, stFile, bf16) + + val manifestFile = File(out, "manifest.json").apply { + writeText( + """ + { + "family": "gemma3n", + "function": "$FN_REDECODE", + "parameterScope": "$PARAMETER_SCOPE", + "seq": $seq, + "layers": $n, + "layersTotal": $fullLayers, + "vocabSize": ${weights.metadata.vocabSize}, + "dtype": "${if (bf16) "bf16" else "f32"}", + "inputs": ["tokens[1,$seq]i32", "per_layer_inputs[1,$seq,${weights.metadata.blockCount},${weights.metadata.perLayerEmbeddingLength}]f32"], + "perLayerInputsOnHost": true, + "mlir": "${mlirFile.name}", + "safetensors": "${stFile.name}" + } + """.trimIndent() + "\n", + ) + } + + val totalF32 = ext.sumOf { it.source.sizeInBytes } + RedecodeResult( + mlirPath = mlirFile.absolutePath, + safetensorsPath = stFile.absolutePath, + manifestPath = manifestFile.absolutePath, + externalParamCount = ext.size, + weightMiB = (if (bf16) totalF32 / 2 else totalF32) / (1024 * 1024), + seq = seq, + vocabSize = weights.metadata.vocabSize, + ) + } + + /** + * FunctionGemma/SmolLM2 safetensors contract, with one difference: bf16 conversion is + * CHUNKED — gemma3n's PLE table is 2.01B elements, and the single-`ByteArray` + * conversion buffer the smaller models use would exceed the JVM array limit. + */ + private fun writeSafetensors(ext: List, stFile: File, bf16: Boolean) { + val dtype = if (bf16) "BF16" else "F32" + val bpe = if (bf16) 2L else 4L + var off = 0L + val hdr = StringBuilder("{") + ext.forEachIndexed { i, e -> + val count = e.source.sizeInBytes / 4 + val len = count * bpe + if (i > 0) hdr.append(",") + hdr.append("\"${e.key}\":{\"dtype\":\"$dtype\",\"shape\":[$count],\"data_offsets\":[$off,${off + len}]}") + off += len + } + hdr.append("}") + val headerBytes = hdr.toString().encodeToByteArray() + BufferedOutputStream(FileOutputStream(stFile), 1 shl 20).use { os -> + os.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array()) + os.write(headerBytes) + val chunkElems = 1 shl 24 // 16M floats per conversion chunk (32 MiB bf16 out) + for (e in ext) { + val src = e.source as BufferHandle.Owned + if (bf16) { + val data = src.data + val n = (src.sizeInBytes / 4).toInt() + var done = 0 + while (done < n) { + val take = minOf(chunkElems, n - done) + val obuf = ByteArray(take * 2) + for (j in 0 until take) { + val o = src.offset + (done + j) * 4 + val fb = (data[o].toInt() and 0xFF) or + ((data[o + 1].toInt() and 0xFF) shl 8) or + ((data[o + 2].toInt() and 0xFF) shl 16) or + ((data[o + 3].toInt() and 0xFF) shl 24) + val bf = Bf16TensorData.floatToBf16Bits(Float.fromBits(fb)) + obuf[j * 2] = (bf and 0xFF).toByte() + obuf[j * 2 + 1] = ((bf ushr 8) and 0xFF).toByte() + } + os.write(obuf) + done += take + } + } else { + os.write(src.data, src.offset, src.sizeInBytes.toInt()) + } + } + } + } + + /** Copied from `FunctionGemmaExportHarness.rewriteGlobalsToBf16` (identical contract). */ + private fun rewriteGlobalsToBf16(mlir: String): String { + var m = mlir + m = Regex("""(util\.global private @\w+ = #flow\.parameter\.named<"[^"]*"::"[^"]*"> : tensor<[0-9x]*x)f32>""") + .replace(m) { it.groupValues[1] + "bf16>" } + m = Regex("""(%\w+) = util\.global\.load @(\w+) : tensor<([0-9x]*)xf32>""") + .replace(m) { r -> + val ssa = r.groupValues[1] + val g = r.groupValues[2] + val shape = r.groupValues[3] + "${ssa}_h = util.global.load @$g : tensor<${shape}xbf16>\n" + + " $ssa = stablehlo.convert ${ssa}_h : (tensor<${shape}xbf16>) -> tensor<${shape}xf32>" + } + return m + } +} diff --git a/llm-inference/gemma3n/src/jvmTest/kotlin/sk/ainet/models/gemma3n/Gemma3nGoldenTokenParityTest.kt b/llm-inference/gemma3n/src/jvmTest/kotlin/sk/ainet/models/gemma3n/Gemma3nGoldenTokenParityTest.kt new file mode 100644 index 00000000..bf992f9f --- /dev/null +++ b/llm-inference/gemma3n/src/jvmTest/kotlin/sk/ainet/models/gemma3n/Gemma3nGoldenTokenParityTest.kt @@ -0,0 +1,100 @@ +package sk.ainet.models.gemma3n + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlinx.coroutines.runBlocking +import sk.ainet.apps.llm.OptimizedLLMMode +import sk.ainet.apps.llm.OptimizedLLMRuntime +import sk.ainet.apps.llm.sampleFromTensor +import sk.ainet.apps.llm.tokenizer.TokenizerFactory +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.gguf.StreamingGGUFReader +import sk.ainet.lang.types.FP32 + +/** + * The #346 maturity-gate parity probe for the Gemma 3n family (#377 DSL migration), against + * **mainline llama.cpp**: full 32-step greedy text equality on the DSL path + * ([Gemma3nWeightLoader] engine loading, packed/MAPPED → [Gemma3nNetworkLoader.fromWeights] + * → [OptimizedLLMRuntime]). This is the first time the family has a reference gate at all — + * and it exercises everything gemma3n adds over gemma-4: AltUp's four parallel streams with + * the tanh router, Laurel, Gaussian-top-k activation sparsity on the first ten layers, PLE + * feeding the non-active streams, and per-type shared KV for the last ten layers. + * + * Model-gated: runs only when `GEMMA3N_E2B_GGUF` points at `gemma-3n-E2B-it-Q4_K_M.gguf` + * AND the test JVM has ≥ 16 GB heap (`-PgemmaTestMaxHeap=20g`); skips quietly otherwise. + * The fixture header records the exact oracle build and commands. + */ +@org.junit.jupiter.api.Tag("smoke-reference") +@org.junit.jupiter.api.Tag("integration") +class Gemma3nGoldenTokenParityTest { + + private data class Fixture( + val prompt: String, + val steps: Int, + val promptTokens: List, + val oracleText: String, + ) + + private fun loadFixture(): Fixture { + val raw = checkNotNull(javaClass.getResourceAsStream("/gemma3n-e2b/golden-greedy-e2b.txt")) { + "fixture /gemma3n-e2b/golden-greedy-e2b.txt missing from test resources" + }.bufferedReader().readLines() + val map = raw.filter { it.isNotBlank() && !it.startsWith("#") } + .associate { it.substringBefore('=') to it.substringAfter('=') } + return Fixture( + prompt = map.getValue("prompt"), + steps = map.getValue("steps").toInt(), + promptTokens = map.getValue("prompt_tokens").split(',').map { it.trim().toInt() }, + oracleText = map.getValue("oracle_text"), + ) + } + + @Test + fun greedyDecodeMatchesMainlineLlamaCpp() { + val modelPath = System.getenv("GEMMA3N_E2B_GGUF") + if (modelPath.isNullOrBlank()) { + println("PARITY skipped: GEMMA3N_E2B_GGUF not set") + return + } + val maxHeapGb = Runtime.getRuntime().maxMemory() / (1024L * 1024L * 1024L) + if (maxHeapGb < 16) { + println("PARITY skipped: heap=$maxHeapGb GB < 16 GB; rerun with -PgemmaTestMaxHeap=20g") + return + } + val fixture = loadFixture() + val ctx = DirectCpuExecutionContext() + + // 1 — prompt tokenization parity (a failure here names the tokenizer, not the model). + val fields = StreamingGGUFReader.open(JvmRandomAccessSource.open(modelPath)).use { it.fields } + val tokenizer = TokenizerFactory.fromGgufFields(fields) + // gemma adds a BOS token (id 2); mirror llama.cpp's add_special encoding. + val raw = tokenizer.encode(fixture.prompt) + val encoded = if (raw.isNotEmpty() && raw[0] == tokenizer.bosTokenId) raw + else intArrayOf(tokenizer.bosTokenId) + raw + assertEquals( + fixture.promptTokens, encoded.toList(), + "prompt tokenization must match mainline llama.cpp", + ) + + // 2 — greedy continuation on the DSL path. + val weights = runBlocking { + Gemma3nWeightLoader( + randomAccessProvider = { JvmRandomAccessSource.open(modelPath) }, + ).loadToMapStreaming(ctx) + } + val model = Gemma3nNetworkLoader.fromWeights(ctx, weights) + val runtime = OptimizedLLMRuntime(model, ctx, OptimizedLLMMode.DIRECT, FP32::class) + for (i in 0 until fixture.promptTokens.size - 1) runtime.forward(fixture.promptTokens[i]) + var token = fixture.promptTokens.last() + val text = StringBuilder() + repeat(fixture.steps) { + token = sampleFromTensor(runtime.forward(token), 0f) + text.append(tokenizer.decode(token)) + } + assertEquals( + fixture.oracleText, text.toString(), + "greedy decode must equal mainline llama.cpp token-for-token", + ) + } +} diff --git a/llm-inference/gemma3n/src/jvmTest/resources/gemma3n-e2b/golden-greedy-e2b.txt b/llm-inference/gemma3n/src/jvmTest/resources/gemma3n-e2b/golden-greedy-e2b.txt new file mode 100644 index 00000000..955523b0 --- /dev/null +++ b/llm-inference/gemma3n/src/jvmTest/resources/gemma3n-e2b/golden-greedy-e2b.txt @@ -0,0 +1,19 @@ +# gemma-3n-E2B-it Q4_K_M greedy parity fixture (#377 DSL migration) +# +# Oracle: MAINLINE llama.cpp (brew, version 0.3.0 build 10621 commit c1d0e7a00). Asserts FULL +# cross-implementation greedy text equality on the new gemma3n DSL path +# (Gemma3nWeightLoader engine loading, packed/MAPPED -> gemma3nNetwork() -> +# OptimizedLLMRuntime) — AltUp (4 streams + router), Laurel, first-10-layer activation +# sparsity, PLE into the non-active streams, per-type shared KV (last 10 layers), hybrid +# sliding/global attention with dual RoPE bases: all exercised end-to-end. +# The prompt is chosen for greedy decisiveness (min top-1/top-2 gap 1.77 nats across all 32 +# steps; "The capital of France is" hits a 0.12-nat tie at step 3 that Q4_K +# cross-implementation noise legitimately flips). NOTE: gemma adds a BOS token (id 2) — +# prompt_tokens is the oracle's add_special encoding, BOS included. +# llama-server -m gemma-3n-E2B-it-Q4_K_M.gguf -ngl 0 -t 4 --port 8816 +# curl :8816/tokenize -d '{"content": "The first ten prime numbers are", "add_special": true}' +# curl :8816/completion -d '{"prompt": "The first ten prime numbers are", "n_predict": 32, "temperature": 0, "top_k": 1}' +prompt=The first ten prime numbers are +steps=32 +prompt_tokens=2,818,1171,3595,8355,4945,659 +oracle_text= 2, 3, 5, 7, 11, 13, 17, 19, 23, diff --git a/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/cli/Main.kt b/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/cli/Main.kt index 874e4dde..e7fbcfab 100644 --- a/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/cli/Main.kt +++ b/llm-runtime/kgemma/src/jvmMain/kotlin/sk/ainet/apps/kgemma/cli/Main.kt @@ -234,19 +234,32 @@ fun main(args: Array) { } } GemmaVariant.GEMMA3N -> { - val ingestion = Gemma3nIngestion( - ctx = ctx, - dtype = FP32::class, - config = Gemma3nLoadConfig() - ) when (format) { ModelFormat.GGUF -> { - println("Loading Gemma 3n GGUF model from $modelPath (streaming mode)...") - ingestion.loadRuntimeStreaming { - JvmRandomAccessSource.open(modelPath.toString()) + // The DSL lane (#377): engine loading (packed/MAPPED), + // gemma3nNetwork() with AltUp/Laurel/sparsity/PLE, verified + // token-for-token vs mainline llama.cpp by + // Gemma3nGoldenTokenParityTest. The hand-rolled + // Gemma3nRuntime never applied PLE and predates the gate. + println("Loading Gemma 3n GGUF model from $modelPath via gemma3nNetwork() + OptimizedLLMRuntime (engine loader, keep-packed, mapped)...") + val model = kotlinx.coroutines.runBlocking { + sk.ainet.models.gemma3n.Gemma3nNetworkLoader.fromGguf( + ctx, + { JvmRandomAccessSource.open(modelPath.toString()) }, + ) } + sk.ainet.apps.llm.OptimizedLLMRuntime( + model, ctx, sk.ainet.apps.llm.OptimizedLLMMode.DIRECT, FP32::class, + ) } ModelFormat.SAFETENSORS -> { + // SafeTensors stays on the legacy hand-rolled runtime until the DSL + // lane grows a SafeTensors leg (tracked with the #377 remainder). + val ingestion = Gemma3nIngestion( + ctx = ctx, + dtype = FP32::class, + config = Gemma3nLoadConfig() + ) val modelDir = if (modelPath.isDirectory()) modelPath else modelPath.parent ?: modelPath val indexPath = modelDir.resolve("model.safetensors.index.json") val safetensorsPath = if (indexPath.exists()) indexPath.toString() diff --git a/tests/smoke/smoke-models.json b/tests/smoke/smoke-models.json index bed73a85..9d578431 100644 --- a/tests/smoke/smoke-models.json +++ b/tests/smoke/smoke-models.json @@ -63,6 +63,14 @@ "format": "gguf", "steps": 16 }, + { + "name": "Gemma3n-E2B-GGUF", + "runner": "skainet", + "model": "gemma-3n-E2B-it-Q4_K_M.gguf", + "format": "gguf", + "steps": 12, + "prompt": "The first ten prime numbers are" + }, { "name": "Gemma4-E4B-GGUF", "runner": "kgemma",