anziank/grio-llama-3.2-1b-instruct-coreml-stateful

Llama 3.2 1B Instruct (Stateful KV) — CoreML Stateful KV-Cache (ctx≤512)

On-device CoreML .mlpackage converted from `meta-llama/Llama-3.2-1B-Instruct` with stateful KV cache (ct.StateType) for efficient token-by-token generation.

Designed for the coremlkv: runner in GrioKit — CoreMLTranslationProvider handles stateful generation via makeState() + prediction(from:using:state:).

Model details

Verified output (greedy, deterministic)

Smoke harness: scripts/coreml/path_b_convert.py --model llama-3.2-1b-stateful --smoke-only against the pre-compiled .mlmodelc on .cpuAndGPU.

Basic capability (`"The capital of France is"`): "The capital of France is Paris." — 8 decode tokens. Matches PyTorch FP16 reference.

Translation pack (10 cases, system prompt: `"You are a French translator. Translate the user message to French. Output only the translation."`):

Postprocess pack (8 cases, light-touch transcript copy-edit): Llama 1B beats Qwen 1.5B on the headline sight → site diagnostic (case 3: we need to make sure the site is stable before launch — Qwen kept sight). It also preserves fillers more faithfully on case 1 (well, i think we should probably start with the budget review and then move into hiring. — Qwen dropped well, probably and rephrased). Weaker on capitalization (such, life in pink stay lowercase; Qwen capitalized). Net: better at the substantive fix-vs-leave-alone judgment, worse on surface formatting.

Headline takeaway: Llama 3.2 1B is a usable on-device French translator for short literal sentences and is the better light-touch copy-editor than Qwen 1.5B. Idiom handling and "translator vs. chatter" instruction-following are weaker than Qwen — pick by use case.

Observed performance

Device: Apple M1 Pro, 16 GB, macOS 26.5 Tahoe. Prompt: "The capital of France is" (cold prefill row) and translation-pack averages (decode row).

Translation-pack warm prefill range on `.cpuAndGPU`: 168–588 ms (depends on prompt length, 30–80 tokens after chat template). Decode latency stable at ~130 ms/tok across all 10 translation cases and 8 postprocess cases.

Recommendation: ship with .cpuAndGPU for production. The 1.8× ANE decode speedup that helps the Qwen 1.5B sibling does NOT apply here — Llama 1B produces semantically wrong output on ANE in this build, likely due to the FP32 SDPA islands interacting with ANE's Llama-specific op scheduling. If you need ANE-class throughput on a 1–2B Apple-Silicon model, use qwen2.5-1.5b-stateful instead.

Runtime gotchas (please read before integrating)

1. 1.Llama 3.2 uses `<|eot_id|>` (id 128009) as the chat-turn terminator in addition to <|end_of_text|> (id 128001). Read generation_config.json for the complete stop list.
2. 2.Do not load on `.cpuAndNE`. ANE produces incorrect tokens for this model in this conversion (see perf table). Pin computeUnits: .cpuAndGPU when configuring the runner.
3. 3.*`causalMask` is a length signal, not the actual causal mask* — shape [1, 1, 1, end_step], all-zero float16. The stateful wrapper inside the model reads only end_step from its shape and builds the proper (1, 1, q_len, end_step) lower-triangular mask internally. Do not try to send a "real" causal mask through the public interface — its query dimension is pinned to 1 in the export.
4. 4.Reset state between independent generations. Call CoreMLTranslationProvider.shared.resetSession() (or model.resetState() on the stateful LanguageModel) between unrelated requests. Stale KV state from a previous sequence corrupts subsequent generations.
5. 5.`.mlmodelc` is included alongside `.mlpackage`. Use the .mlmodelc path for loading — it skips the multi-minute ANE compilation that happens on first .mlpackage load.

Conversion notes (for the CoreML community)

1. 1.Use `attn_implementation="sdpa"`. Llama 3.2 export is cleaner with SDPA; eager produces a brittler op graph.
2. 2.Use `torch.export.default_decompositions()` for run_decompositions. {} (all decompositions) can SIGSEGV in coremltools at this scale on some macOS / coremltools combinations.
3. 3.`create_causal_mask` in `transformers` ≥ 5 returns 4-D masks unchanged. Sending a [1, 1, 1, end_step] length-signal mask straight into attention layers therefore broadcasts across the query axis and gives BIDIRECTIONAL attention during batch prefill. Llama 3.2 1B (4:1 GQA, 8 KV heads) is robust to this contamination — earlier conversions without the proper internal mask still produced sensible output, which initially masked the bug. The current build's wrapper constructs the correct (1, 1, q_len, end_step) lower-triangular mask internally so every prefill is causally correct; Llama output is now strictly more correct than before, not just "good enough by accident".
4. 4.FP32 SDPA islands — the wrapper runs attention in FP32 inside an otherwise-FP16 model. For Llama 1B this is a no-op safety net (its activations fit FP16 comfortably) but it does interact with ANE op scheduling and is the most likely cause of the .cpuAndNE divergence noted above. Removing the FP32 island would require Qwen-specific handling — a tradeoff against ANE correctness for Llama. Current build prioritizes correctness across model families over ANE compatibility for Llama specifically.
5. 5.FP16 greedy decoding is not byte-deterministic across backends. Semantically equivalent to PyTorch FP16 reference for every prompt we have tested, except on .cpuAndNE for this model (see perf table).

Usage (Swift — GrioKit coremlkv: runner)

import GrioKit

// Load via CoreMLTranslationProvider using the coremlkv: prefix
let provider = CoreMLTranslationProvider.shared
try await provider.configure(modelID: "coremlkv:llama-3.2-1b-stateful", computeUnits: .cpuAndGPU)

// Translate a sentence (stateful: KV cache persists across tokens)
let result = try await provider.translate(
    "The quick brown fox jumps over the lazy dog.",
    sourceLanguage: "English",
    targetLanguage: "French"
)
print(result) // e.g. "Le rapide renard brun saute par-dessus le chien paresseux."

// Reset KV state between unrelated requests
await provider.resetSession()

Keep tokenizer.json, tokenizer_config.json, config.json, generation_config.json (and chat_template.jinja for Gemma) as siblings of the .mlpackage on disk. All are bundled in this repo.

Reproducibility

Conversion done with coremltools==9.0, torch==2.7.0, transformers==5.8.1 via path_b_convert.py --model llama-3.2-1b-stateful --stateful.

Stateful KV weights are ct.StateType buffers (keyCache, valueCache), shaped (num_layers, 1, num_kv_heads, 512, head_dim). Post-convert quantization (--quantize int8) applies linear_quantize_weights with linear_symmetric mode to the saved .mlpackage; the state buffers are runtime tensors and are unaffected by weight quantization.

License

Llama 3.2 Community License. Weights from meta-llama/Llama-3.2-1B-Instruct by Meta. Re-uploaded as a CoreML port; original model card terms apply.