README.md

---
license: apache-2.0
base_model: Zyphra/Zamba2-1.2B-instruct
tags:
  - litert
  - litert-lm
  - litertlm
  - on-device
  - edge
  - hybrid
  - mamba2
  - zamba2
pipeline_tag: text-generation
library_name: litert-lm
---

# Zamba2-1.2B-instruct — LiteRT-LM

[Zyphra/Zamba2-1.2B-instruct](https://huggingface.co/Zyphra/Zamba2-1.2B-instruct) converted to the **LiteRT-LM** (`.litertlm`) format for on-device inference with Google's [LiteRT-LM](https://github.com/google-ai-edge/litert-lm) runtime. **Requires litert-lm ≥ 0.15.**

Zamba2 is Zyphra's shared-attention hybrid: a Mamba2 selective-scan backbone (32 layers) with a **single shared transformer block** applied at 6 interleaved positions — one set of attention+MLP weights reused at every position, specialized by per-position LoRA adapters, attending over the concatenation of the running hidden state and the original embeddings. To our knowledge this is the first Zamba2 conversion to a mobile runtime.

| File | Recipe | Size |
|---|---|---|
| `Zamba2-1.2B-instruct_int8.litertlm` | int8 dynamic on linears + embedding (convs and the scan stay float); fp32 activations declared for GPU | 1.33 GB |

## Correctness

- **Logits parity vs PyTorch**: the float export matches the HF model teacher-forced across 48 decode positions — top-1 and top-5 identical at every position, mean per-position logit correlation 1.0000, mean KL ≈ 0.
- **8-question sanity gate**: **8/8 on every lane** — GPU and CPU, litert-lm 0.15.0 and 0.16.0. No degeneration, no greedy flips.
- **Prompt-length robustness**: hermetic prefill-chunk sweep (fresh engine per length) — CPU fills 12–51 and GPU fills 12–31 all clean.
- **iPhone 17 Pro (Metal)**: the 8-item composite quality probe answers **8/8 on CPU**; on GPU 6/8, where both misses (8×7 and the rhyme) are questions the **HF fp32 reference itself answers incorrectly** on this composite ("63"; a broken echo) — the GPU path tracks the reference model's own behavior, and Mac and iPhone GPU produce identical text. The individual 8-question gate is 8/8 on both backends.

## Usage

```bash
litert-lm run ./Zamba2-1.2B-instruct_int8.litertlm --prompt "What is the capital of France? Answer in one word."

# GPU
litert-lm run ./Zamba2-1.2B-instruct_int8.litertlm --backend gpu --cache no --prompt "..."
```

Multi-length prefill signatures (1–1024) are exported so the runtime picks tight chunks. The bundle carries the tokenizer and the stock ChatML Zamba2 chat template.

## Performance

`litert-lm benchmark` (litert-lm 0.16.0), Apple M4 Max, `-p 256 -d 256 --runs 3 --cache no`, quiet machine:

| Backend | Prefill (256) | Decode | TTFT |
|---|---|---|---|
| GPU | 1033 tok/s | 74.0 tok/s | 0.26 s |
| CPU | 450 tok/s | 22.7 tok/s | 0.61 s |

On device (cold start, single runs, composite prompt, quality harness):

| Device | Backend | Prefill | Decode | TTFT | Peak memory |
|---|---|---|---|---|---|
| iPhone 17 Pro | **GPU (Metal)** | 96.6 tok/s | **12.8 tok/s** | 1.63 s | 5.43 GB |
| iPhone 17 Pro | CPU | 87.3 tok/s | 7.4 tok/s | 1.73 s | 1.82 GB |

Honest notes:
- **GPU runs with fp32 activations** (declared in the bundle) — expect a corresponding memory multiple over CPU. The 5.43 GB GPU peak includes the six shared-attention positions' wide KV caches (32 KV heads × 128 head dim at 4096 context) held in fp32; a 12 GB phone runs the full 12-signature ladder with no memory pressure.
- On the long composite probe the two backends split two model-edge questions differently (see Correctness) — deterministically, and identically on Mac and iPhone. For maximum fidelity to the HF reference use GPU; for the best composite score use CPU.

## Conversion notes

Converted with [`litert-torch`](https://github.com/google-ai-edge/litert-torch) plus a hybrid-cache patch (reproduction script + patch: [hf-to-litertlm `zamba2_work/`](https://github.com/john-rocky/hf-to-litertlm)):

- **Folded selective scan**: the Mamba2 scan is re-expressed as batched matmuls with chunk and head axes folded into the batch axis (all tensors rank ≤ 4, no `BROADCAST_TO`, no int64 index math) — this is what makes the graph fully delegable on GPU.
- **Min-only dt clamp handling**: Zamba2 clamps `softplus(dt)` at `time_step_min` with no upper clamp; padded prefill positions are forced to exact identity steps AFTER the clamp (without this, every runtime pad token decays the recurrent state).
- **Shared block + adapters**: the tied transformer block traces once per position with its own LoRA adapter statically selected; tied weights are stored once.
- **Composite hybrid cache layer**: the 6 shared-attention positions hold KV + conv + recurrent state at ONE layer index (the runtime binds states by tensor name, so co-residency is just packaging).
- **Prefill-pad guard**: the runtime runs partially-filled prefill chunks; pad positions are made exact identity steps for the SSM and the stored conv window is gathered at the last valid column.
- **Streaming detokenization**: the tokenizer's `Strip` decoder is removed from the bundle — Zamba2's metaspace (SP-BPE) tokenizer otherwise loses every interior space under the runtime's per-token streaming decode; the only behavior change is a sequence-initial space, which the runtime trims.
- **Quantization**: post-hoc dynamic int8 over linears + embedding only; convs and the scan stay float.

## License and changes

Distributed under **Apache 2.0** (inherited from the base model). **Changes from the original work:** weights converted from safetensors to LiteRT flatbuffers and quantized as described above; tokenizer and chat template repackaged unmodified. This repository is a community conversion and is not affiliated with Zyphra.