Add G4 NPU AOT build (Gemma3-1B-IT int4-hadamard, Tensor-G4) — runs in Box/EdgeTPU-app/litert_lm

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Gemma3-1B-IT_mix4blk8_jit_ekv1280_Google_Tensor_G4.litertlm filter=lfs diff=lfs merge=lfs -text
+Gemma3-1B-IT_q4_block32_jit_ekv1280_Google_Tensor_G4.litertlm filter=lfs diff=lfs merge=lfs -text
+Gemma3-1B-IT_int4_aot_ekv1280_Google_Tensor_G4.litertlm filter=lfs diff=lfs merge=lfs -text
+Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm filter=lfs diff=lfs merge=lfs -text
+Gemma3-1B-IT_attn8head8_jit_ekv4096_Google_Tensor_G4.litertlm filter=lfs diff=lfs merge=lfs -text

Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c2962c8b18ff600ab37a1b15dfcccd6d1e5d8c155ad1ea46fe5fd0ab50dd6ee
+size 874631968

MANIFEST.md

@@ -0,0 +1,11 @@
+# Manifest
+
+| File | Size (bytes) | SHA-256 |
+|---|---|---|
+| `Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm` | 874,631,968 | `1c2962c8b18ff600ab37a1b15dfcccd6d1e5d8c155ad1ea46fe5fd0ab50dd6ee` |
+
+Verify after download:
+
+```bash
+sha256sum Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm
+```

RANK2_FC_SURGERY.md

@@ -0,0 +1,47 @@
+# Gemma-3-1B on the Tensor G4 NPU — the rank-2 FC surgery
+
+The compiler-bug workaround that makes Gemma-3/4 AoT-compile to the Tensor G4 edgetpu, plus
+working-model notes.
+
+---
+
+## The rank-2 FULLY_CONNECTED surgery — the key enabler (affects G4 AND G5)
+
+The DarwiNN AoT compiler **`OP_FAILED`s (`file.cc:4780`) on rank-3 `FULLY_CONNECTED` ops.**
+Gemma-3 (and Gemma-4) emit these ~360× across the model, so the AoT compile fails outright
+without a fix.
+
+**Fix:** wrap every rank-3 FC as `RESHAPE → (rank-2) FC → RESHAPE` — same math, rank the
+compiler accepts. Script: `recipe/fc_rank2_surgery.py` (`surgery(in_tflite, out_tflite)`), run
+after quantization, before `aot_compile`. ~360 FCs rewritten for gemma-3-1b. The G5 uses the same
+compiler family, so the surgery carries over to G5 as well.
+
+---
+
+## Profile
+
+**AoT (int4 per-channel + Hadamard):** ~100% NPU, prefill-monster (~660 tok/s prefill),
+power-efficient — best for long-input / RAG. Decode is dispatch-bound rather than
+compute-bound, so it sits around 16 tok/s regardless of prefill length.
+
+Behavior note: tone/"feel" differences vs a reference are **runtime config** (system prompt +
+sampling), not the quant — a faithful int4 copy matches the reference once system prompt and
+top-k / top-p / temperature match.
+
+---
+
+## On quantization support tables
+
+Don't trust a fixed "what quantization compiles" matrix, including any you find in older copies
+of this repo. Per-quant support is **compiler-build-dependent**: a result from one toolchain
+build does not generalize to another, and rows that read "unsupported" have repeatedly turned out
+to be one build's behaviour rather than a hardware or format limit.
+
+Retest against your own target compiler rather than trusting a table.
+
+---
+
+## Runtime requirement (brief)
+
+int4 AoT DGCs require a compatible GoogleTensor dispatch library to load on-device; a stock app's
+dispatch may reject them as a firmware mismatch.

README.md

@@ -0,0 +1,161 @@
+---
+license: gemma
+base_model: google/gemma-3-1b-it
+base_model_relation: quantized
+pipeline_tag: text-generation
+library_name: litert-lm
+tags:
+- litert-lm
+- litertlm
+- gemma3
+- chat
+- text-generation
+- npu
+- edge-tpu
+- tensor-g4
+- pixel
+- on-device
+- quantized
+- int4
+extra_gated_heading: Access Gemma3-1B-IT (Tensor G4 NPU) on Hugging Face
+extra_gated_prompt: >-
+  These are quantized, repackaged derivatives of google/gemma-3-1b-it and are
+  governed by the Gemma license. To access them you must review and agree to
+  Google's Gemma Terms of Use. Please ensure you are logged in to Hugging Face
+  and click below.
+extra_gated_button_content: Acknowledge license
+---
+
+# Gemma3-1B-IT — Tensor G4 NPU (AOT, int4 per-channel + Hadamard)
+
+> [!NOTE]
+> A quantized, repackaged variant of
+> [`google/gemma-3-1b-it`](https://huggingface.co/google/gemma-3-1b-it) that runs on the
+> **Google Tensor G4 Edge TPU ("rio") NPU** — the SoC in the Pixel 9 / 9 Pro / 9 Pro XL / 9 Pro Fold.
+> Packaged as a `.litertlm` bundle for the
+> [LiteRT-LM](https://github.com/google-ai-edge/LiteRT-LM) runtime.
+
+This is, to our knowledge, the **first publicly available int4 AOT Tensor-G4 NPU build of Gemma 3 1B** —
+and it is **~half the size of the only other Tensor-NPU Gemma 3 1B bundle** published
+(`litert-community/Gemma3-1B-IT_q8_ekv1280_Google_Tensor_G5.litertlm` is **1,678 MB**; this build is
+**875 MB = 0.52×**), because it is int4-per-channel rather than int8.
+
+`litert-community/Gemma3-1B-IT` ships Tensor **G5** and **G6** builds. There is no G4 build there.
+
+---
+
+## The build
+
+| | |
+|---|---|
+| File | `Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm` |
+| Mode | **AOT** — precompiled DarwiNN DGC |
+| Quantization | int4 **per-channel + Hadamard rotation** (all FC / BMM / embed) |
+| NPU share | **~100%** |
+| Embedded KV cache | 4096 |
+| Size | 874,631,968 B (875 MB) |
+
+26-layer Gemma 3 1B-IT, embedded KV cache, embedded tokenizer + chat template, `.litertlm` v1.
+Built from Google's mobile-QAT checkpoint
+[`google/gemma-3-1b-it-qat-int4-unquantized`](https://huggingface.co/google/gemma-3-1b-it-qat-int4-unquantized).
+
+---
+
+## Performance (Pixel 9 Pro XL, Tensor G4)
+
+Measured on this exact bundle, 2026-09-10:
+
+| prefill | decode | prefill tok/s | decode tok/s |
+|---|---|---|---|
+| 128 | 64 | 608.83 | 15.93 |
+| 512 | 128 | 659.22 | 16.17 |
+| 1024 | 64 | 663.99 | 15.89 |
+
+**Prefill scales; decode is pinned.** Prefill rises 608 → 659 → 664 tok/s as the batch grows
+128 → 1024, while decode holds at 15.89–16.17 tok/s regardless (±0.9% across an 8× change in
+prefill length). That is a clean confirmation that decode is **dispatch-bound, not compute-bound** —
+more work per dispatch makes prefill faster, but decode always does exactly one token per dispatch,
+so nothing moves it.
+
+The rise from 608 to 664 is fixed per-dispatch overhead amortizing: at 128 tokens you pay it once
+per 128, at 1024 once per 1024. Extrapolating, the per-dispatch prefill floor is roughly **~670 tok/s**
+for this model.
+
+**AOT profile = prefill-monster / decode-laggard.** The Edge TPU crushes batched matmuls (prefill) but
+a single-token decode underutilizes it — so this build is the fastest option on the phone for
+prefill-heavy work (RAG, long-context, summarization, classification), and a strong speculative-decode
+(MTP) candidate. Speculative decode is the real next perf lever and is not wired into this build.
+
+> Behaviour/"feel" differences vs a reference are **runtime config** (system prompt + sampling), not the
+> quant — a faithful int4 copy matches the reference once system prompt and top-k / top-p / temperature match.
+
+---
+
+## Quantization
+
+### The rank-2 FC surgery — the key enabler (affects G4 **and** G5)
+
+The DarwiNN AOT compiler `OP_FAILED`s on **rank-3 `FULLY_CONNECTED`** ops, which Gemma-3 emits ~360× — so
+the AOT compile fails outright without a fix. **`recipe/fc_rank2_surgery.py`** wraps every rank-3 FC as
+`RESHAPE → (rank-2) FC → RESHAPE` (same math, a rank the compiler accepts). This is what lets a full int4
+Gemma 3 actually compile to a single G4 DGC. We are not aware of an upstream fix.
+
+Details in **[`RANK2_FC_SURGERY.md`](./RANK2_FC_SURGERY.md)**.
+
+### Hadamard rotation
+
+On top of int4 per-channel, this build applies a **Hadamard rotation** to the weights before quantization
+(a `DECOMPOSED_HADAMARD_ROTATION`-style outlier-suppression transform). Spreading outliers across the
+rotated basis tightens the per-channel int4 grid — most visibly on the 262k-vocab head, where it improves
+logit precision and coherence over plain per-channel int4, at a small size cost.
+
+### Per-quant AOT support is compiler-build-dependent
+
+There is **no fixed per-bit-width support matrix** for the DarwiNN AOT compiler — what a given toolchain
+build accepts varies by compiler version. The shipping `gemma4-E2B` G4-NPU bundle is a **2-subgraph** model
+carrying **int2, int4, and int8** weight tensors (verified directly), so mixed int2/int4/int8 does
+AOT-compile. The verified enabler for compiling a full Gemma to a G4 DGC is the **rank-2 FC surgery**,
+not any particular bit width.
+
+Granularity is a build choice, not a compiler wall:
+
+- **int4 per-channel + Hadamard** (one scale per output channel, rotated) → this build.
+- **int4 per-channel** (one scale per output channel) → also coherent.
+- **int4 per-tensor** also AOT-compiles, but one scale over the 262k-vocab head → word salad.
+
+---
+
+## Running
+
+This is a **precompiled DarwiNN DGC**. It requires a **compatible GoogleTensor dispatch library** to load
+on-device; a stock app's dispatch may reject it as a firmware mismatch. Runtime/dispatch specifics are out
+of scope for this repo.
+
+| Surface | Support |
+|---|---|
+| LiteRT-LM CLI (`--backend=npu`) | ✅ with a compatible dispatch runtime |
+| Google AI Edge Gallery (NPU) | needs a compatible dispatch runtime |
+| Tensor G4 (Pixel 9 series) | ✅ |
+| Tensor G5 (Pixel 10) | ❌ DGC is G4-targeted |
+| Older Tensor (G2 / G3) | ❌ |
+| llama.cpp / GGUF, Web, iOS, desktop, non-Pixel | ❌ `.litertlm` is Pixel-NPU-only |
+
+---
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `Gemma3-1B-IT_int4hadamard_aot_ekv4096_G4.litertlm` | the model — per-channel int4 + Hadamard, 4096 ctx, ~100% NPU |
+| `RANK2_FC_SURGERY.md` | the rank-2 FC surgery writeup + why per-quant AOT support is compiler-build-dependent |
+| `recipe/fc_rank2_surgery.py` | rank-3 → rank-2 FC surgery (the DarwiNN AOT enabler) |
+| `MANIFEST.md` | checksum + size |
+
+---
+
+## License & attribution
+
+Base model **Gemma 3 1B-IT** © Google, under the [Gemma Terms of Use](https://ai.google.dev/gemma/terms).
+This quantized/repackaged derivative inherits the Gemma license. Runtime: **LiteRT-LM** / **LiteRT**
+(Apache-2.0). Quantization, rank-2 FC surgery, per-channel int4 + Hadamard AOT compile, and G4 NPU
+packaging by **@xThr45hx**.

recipe/fc_rank2_surgery.py

@@ -0,0 +1,162 @@
+"""fc_rank2_surgery.py — post-export tflite graph rewrite: wrap every rank-3
+FULLY_CONNECTED so its activation is rank-2, satisfying darwinn G4
+(MASTER_NPU_SPEC: FC rank-2 only; rank-3 [1,T,K] => OP_FAILED line 4780).
+
+Per rank-3 FC  (in:[1,T,K] -> out:[1,T,N]):
+    in[1,T,K] --RESHAPE--> in2[T*?,K] --FC--> out2[T*?,N] --RESHAPE--> out[1,T,N]
+Reshape is numerically a no-op so quant params copy unchanged (reshape input/
+output share scale+zp). Done with the flatbuffers object API (ModelT).
+
+Local-only: needs ai_edge_litert schema + flatbuffers (no torch, no box).
+Usage: python fc_rank2_surgery.py IN.tflite OUT.tflite
+"""
+import sys, copy
+import numpy as np
+from ai_edge_litert import schema_py_generated as schema
+import flatbuffers
+
+FC = schema.BuiltinOperator.FULLY_CONNECTED
+RESHAPE = schema.BuiltinOperator.RESHAPE
+INT32 = schema.TensorType.INT32
+
+
+def opcode_index_for(model, builtin):
+    """Return index of an OperatorCode for `builtin`, adding one if missing."""
+    for i, oc in enumerate(model.operatorCodes):
+        bc = oc.builtinCode if oc.builtinCode else oc.deprecatedBuiltinCode
+        if bc == builtin:
+            return i
+    oc = schema.OperatorCodeT()
+    oc.builtinCode = builtin
+    oc.deprecatedBuiltinCode = min(builtin, 127)
+    oc.version = 1
+    model.operatorCodes.append(oc)
+    return len(model.operatorCodes) - 1
+
+
+def add_buffer(model, data_bytes=None):
+    b = schema.BufferT()
+    if data_bytes is not None:
+        b.data = list(data_bytes)
+    model.buffers.append(b)
+    return len(model.buffers) - 1
+
+
+def clone_quant(q):
+    if q is None:
+        return None
+    nq = schema.QuantizationParametersT()
+    nq.scale = list(q.scale) if q.scale is not None else None
+    nq.zeroPoint = list(q.zeroPoint) if q.zeroPoint is not None else None
+    nq.quantizedDimension = q.quantizedDimension
+    nq.min = list(q.min) if q.min is not None else None
+    nq.max = list(q.max) if q.max is not None else None
+    return nq
+
+
+def new_tensor(sg, name, shape, ttype, quant):
+    t = schema.TensorT()
+    t.name = name.encode() if isinstance(name, str) else name
+    t.shape = list(shape)
+    t.type = ttype
+    t.buffer = 0  # intermediate (no data)
+    t.quantization = quant
+    sg.tensors.append(t)
+    return len(sg.tensors) - 1
+
+
+def surgery(in_path, out_path):
+    buf = open(in_path, "rb").read()
+    root = schema.Model.GetRootAs(buf, 0)        # for external-buffer offsets
+    model = schema.ModelT.InitFromPackedBuf(buf, 0)
+    # Pull EXTERNAL buffers (offset>1, data appended after the flatbuffer) inline
+    # so Pack preserves the ~1GB of weights. Model is <2GB so inline re-pack is legal.
+    pulled = 0
+    for i in range(root.BuffersLength()):
+        fb = root.Buffers(i)
+        off, size = fb.Offset(), fb.Size()
+        if off and off > 1 and size:
+            model.buffers[i].data = np.frombuffer(buf[off:off + size], dtype=np.uint8)
+            model.buffers[i].offset = 0
+            model.buffers[i].size = 0
+            pulled += 1
+    print(f"  pulled {pulled} external buffers inline")
+    reshape_oc = opcode_index_for(model, RESHAPE)
+
+    total_wrapped = 0
+    total_shared = 0
+    for sgi, sg in enumerate(model.subgraphs):
+        new_ops = []
+        wrapped = 0
+        shared = 0
+        in_reshape_cache = {}  # (orig_in0, flat, K) -> in2 tensor idx (share qkv/gate-up input reshapes)
+        for op in sg.operators:
+            oc = model.operatorCodes[op.opcodeIndex]
+            bc = oc.builtinCode if oc.builtinCode else oc.deprecatedBuiltinCode
+            if bc != FC:
+                new_ops.append(op)
+                continue
+            in_t = sg.tensors[op.inputs[0]]
+            shp = list(in_t.shape) if in_t.shape is not None else []
+            if len(shp) <= 2:
+                new_ops.append(op)
+                continue
+            K = int(shp[-1])
+            flat = int(np.prod(shp[:-1]))  # T (and any leading dims) collapsed
+            out_idx = int(op.outputs[0])
+            out_t = sg.tensors[out_idx]
+            out_shape = list(out_t.shape)
+            N = out_shape[-1]
+
+            # --- pre-reshape: in[..,K] -> in2[flat, K]  (SHARED across FCs with same input) ---
+            orig_in0 = int(op.inputs[0])
+            cache_key = (orig_in0, flat, K)
+            if cache_key in in_reshape_cache:
+                in2 = in_reshape_cache[cache_key]   # reuse — e.g. q/k/v share one reshape
+                shared += 1
+            else:
+                in2 = new_tensor(sg, f"fc_r2_in_{sgi}_{wrapped}", [flat, K],
+                                 in_t.type, clone_quant(in_t.quantization))
+                shp_buf = add_buffer(model, np.array([flat, K], dtype=np.int32).tobytes())
+                shp_t = new_tensor(sg, f"fc_r2_in_shape_{sgi}_{wrapped}", [2], INT32, None)
+                sg.tensors[shp_t].buffer = shp_buf
+                r1 = schema.OperatorT(); r1.opcodeIndex = reshape_oc
+                r1.inputs = [orig_in0, shp_t]; r1.outputs = [in2]
+                ro = schema.ReshapeOptionsT(); ro.newShape = [flat, K]
+                r1.builtinOptionsType = schema.BuiltinOptions.ReshapeOptions
+                r1.builtinOptions = ro
+                new_ops.append(r1)
+                in_reshape_cache[cache_key] = in2
+
+            # --- FC now rank-2: in2[flat,K] -> fcout[flat,N] ---
+            fcout = new_tensor(sg, f"fc_r2_out_{sgi}_{wrapped}", [flat, N],
+                               out_t.type, clone_quant(out_t.quantization))
+            inl = list(int(x) for x in op.inputs); inl[0] = in2; op.inputs = inl
+            outl = list(int(x) for x in op.outputs); outl[0] = fcout; op.outputs = outl
+            new_ops.append(op)
+
+            # --- post-reshape: fcout[flat,N] -> out[orig shape] ---
+            oshp_buf = add_buffer(model, np.array(out_shape, dtype=np.int32).tobytes())
+            oshp_t = new_tensor(sg, f"fc_r2_out_shape_{sgi}_{wrapped}", [len(out_shape)], INT32, None)
+            sg.tensors[oshp_t].buffer = oshp_buf
+            r2 = schema.OperatorT(); r2.opcodeIndex = reshape_oc
+            r2.inputs = [fcout, oshp_t]; r2.outputs = [out_idx]
+            ro2 = schema.ReshapeOptionsT(); ro2.newShape = list(out_shape)
+            r2.builtinOptionsType = schema.BuiltinOptions.ReshapeOptions
+            r2.builtinOptions = ro2
+            new_ops.append(r2)
+            wrapped += 1
+        sg.operators = new_ops
+        total_wrapped += wrapped
+        total_shared += shared
+        if wrapped:
+            print(f"  subgraph {sgi}: wrapped {wrapped} rank-3 FC ({shared} input-reshapes shared)")
+
+    b = flatbuffers.Builder(0)
+    b.Finish(model.Pack(b), b"TFL3")
+    open(out_path, "wb").write(b.Output())
+    print(f"wrapped {total_wrapped} FC total -> {out_path}")
+
+
+if __name__ == "__main__":
+    surgery(sys.argv[1], sys.argv[2])