# 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`](./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)
```python
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
```python
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`](./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`](./RANK2_FC_SURGERY.md).

### 5. AOT compile — large-model-support ONLY
```python
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_OP`s, 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.