Zamba2-1.2B-instruct — LiteRT-LM

Zyphra/Zamba2-1.2B-instruct converted to the LiteRT-LM (.litertlm) format for on-device inference with Google's 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

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 plus a hybrid-cache patch (reproduction script + patch: hf-to-litertlm zamba2_work/):

• 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.