AOT recipe — Gemma3-1B-IT int4 + Hadamard → Google Tensor NPU

The exact recipe behind every .litertlm in this repo. A from-scratch G4 rebuild made this way is bit-identical to the published G4 build (see "Does your build match?" below).

Executable version with guardrails: recipe/gemma3_int4hr_recipe.py.

Toolchain (pin these)

Component	Version
`ai-edge-litert`	2.2.0
Google Tensor compiler (`liblitert_plugin_compiler.so`, beta SDK)	BuildID `87ea807a…`
`litert-torch`	with `generative.export_hf` (split-cache NPU export)
`ai-edge-quantizer`	with `recipe.dynamic_wi4c_hr_afp32`

The compiler .so must be on LD_LIBRARY_PATH; the plugin loads it by soname and ignores GOOGLE_TENSOR_COMPILER_LIB.

Steps

1. Source weights

google/gemma-3-1b-it-qat-int4-unquantized — Google's QAT checkpoint, not plain gemma-3-1b-it.

2. Export (float)

from litert_torch.generative.export_hf import export as E
E.export(model=SRC, output_dir=OUT, task="text_generation",
         prefill_lengths=[128], cache_length=4096,
         externalize_embedder=True, split_cache=True,
         enable_gpu_dynamic_cache=False, enable_gpu_dynamic_prefill=False,
         keep_temporary_files=True, quantization_recipe=None)

Outputs: model.tflite (float prefill_decode, ~4.0 GB), embedder_quantized.tflite, auxiliary.tflite (rope / mask / cache-update), tokenizer.

⛔ Trap: export_hf also writes model_quantized.tflite (a default int8 copy, ~1.0 GB), and the bundle it produces contains that int8 copy. Quantizing it again with int4 = double quantization: it still compiles, with the same size and layout, but the output is word salad. Always quantize the float model.tflite.

3. Quantize — int4 per-channel + Hadamard, once

from ai_edge_quantizer import quantizer, recipe as R
q = quantizer.Quantizer("model.tflite")
q.load_quantization_recipe(R.dynamic_wi4c_hr_afp32())
q.quantize().export_model("prefill_decode_int4hr.tflite")

4. Rank-2 FC surgery (G4, G3)

recipe/fc_rank2_surgery.py → surgery(in, out): wraps every rank-3 FULLY_CONNECTED as RESHAPE → rank-2 FC → RESHAPE. Same math. See RANK2_FC_SURGERY.md.

5. AOT compile — large-model-support ONLY

from ai_edge_litert.aot import aot_compile as aot_lib
from ai_edge_litert.aot.vendors import import_vendor
tgt = import_vendor.import_vendor("GOOGLE").create(
    {"backend_id": "GOOGLE", "soc_model": "Tensor_G4"}).target   # or Tensor_G3 / G5 / G6
res = aot_lib.aot_compile("prefill_decode_int4hr_r2.tflite", output_dir=OUT, target=[tgt],
                          google_tensor_enable_large_model_support=True)
res.models[0].save("prefill_decode_G4.tflite", export_only=True)

LMS = True is required. Without it the compile fails with INTERNAL.
Do NOT pass google_tensor_sharding_intensity. The verified builds use no sharding flag.
Expected result: 2 subgraphs (prefill_128, decode), 2 DISPATCH_OPs, 0 builtin ops, 560,536,512 B.

6. Bundle

Put the compiled prefill_decode into a .litertlm together with the embedder, aux, tokenizer and metadata sections, using litert-lm-builder toml --path model.toml output --path <out>.litertlm. Check the result with litert-lm-peek --litertlm_file <out>.litertlm.

Does your build match?

With the pinned toolchain, the compiled G4 prefill_decode (step 5 output, rank-2) should be 560,536,512 B, md5 f1102df2fdc182d7738bd495ad8f9cf0, bit-identical to the one in the published G4 bundle. A different md5 means something in your pipeline drifted (most often: quantizing export_hf's default int8 copy instead of the float model.tflite). On device, the rebuild and the published bundle also give word-for-word identical replies to the same prompt (greedy decoding).

Per-SoC notes

SoC	Rank-2	Result
G4	yes	✅ compiled + verified on device
G3	yes	✅ compiled, ⚠️ untested on hardware
G5	tried with and without	❌ not yet compiling with this recipe: INTERNAL deep in codegen (6–9 min), with and without rank-2. Under investigation (LMS off, ekv1280, int8 per-channel).
G6	tried with and without	❌ same as G5: INTERNAL after 7–9 min, with and without rank-2. Under investigation.

G3 and G4 compile with the recipe above as-is. G5/G6 need a different setting that hasn't been found yet. This table will be updated when one compiles.