Implement KerasHub → LiteRT-LM export with prefill/decode signatures by pctablet505 · Pull Request #2705 · keras-team/keras-hub

pctablet505 · 2026-04-25T12:35:10Z

Design Doc: KerasHub → LiteRT-LM Export

Export a KerasHub CausalLM directly to a .litertlm Task Bundle that the
on-device LiteRT-LM runtime can load and run — text, vision, and audio — with no
per-model native inference code.

model = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_1b")
model.export("model.litertlm", format="litertlm", prefill_seq_len=128)

1. Objective

Make CausalLM.export(filepath, format="litertlm") produce a complete
.litertlm Task Bundle that the LiteRT-LM runtime can load and execute on
Android/iOS. One Python call turns a Keras model into a deployable on-device LLM
artifact, covering:

Text-only CausalLM models (Gemma, Llama, Mistral, Qwen, Phi3, …).
Multimodal models that combine vision and/or audio with text (Gemma3,
Gemma3n, Gemma4, PaliGemma).

The bundle contains everything the runtime needs — the TFLite inference graph,
the tokenizer, and model metadata — so application developers do not need to
write native inference code.

2. Background & Motivation

2.1 The gap today

Keras already supports model.export(format="litert"), which produces a raw
.tflite graph with a single serving_default signature. That is only the math
kernel. Shipping an actual on-device chat experience on top of a bare .tflite
still requires native code for all of the following:

flowchart LR
    T[".tflite graph<br/>(from format=litert)"]
    subgraph Gap["Still hand-written per model"]
        direction TB
        G1["TFLite interpreter loading"]
        G2["KV-cache allocation & indexing"]
        G3["Tokenizer bundling & invocation"]
        G4["Prefill / decode splitting"]
        G5["Sampling & stop-token detection"]
        G6["Chat templating"]
    end
    X(["Substantial per-model native glue work"])
    T --> Gap --> X

Every team that wants to run a Keras LLM on a phone re-implements this glue. It
is model-specific and duplicated across the ecosystem.

2.2 What `.litertlm` provides

The .litertlm Task Bundle is the LiteRT-LM runtime's packaging format. It
encapsulates the inference graph plus the tokenizer and the metadata
(start/stop tokens, context window, model type, vision/audio config) that drive
cache management, prefill/decode dispatch, sampling, and chat templating inside
the runtime. Give the runtime a .litertlm file and it runs text generation —
no per-model native code.

This document describes the KerasHub exporter that produces that bundle.

2.3 Terminology

These terms are used throughout; skim them once and the rest reads cleanly.

Term	Meaning
LiteRT-LM	The on-device LLM runtime and its Android/iOS APIs.
`.litertlm`	The LiteRT-LM Task Bundle file format produced by this exporter.
`format="litertlm"`	The export-format keyword passed to `CausalLM.export()`.
`litert-torch`	Python package that traces PyTorch models and converts them to TFLite.
`litert-lm-builder`	Python package that packs TFLite models, tokenizer assets, and metadata into a `.litertlm`.
Signature	A named, static-shape entry point in a TFLite model (e.g. `prefill`, `decode`).
Prefill / Decode	Prefill consumes the whole prompt to seed the KV cache; decode emits one token at a time.
Bucketing	Tracing several fixed prefill lengths so the runtime dispatches to the smallest one that fits.
Backend constraint	A hint restricting a TFLite model to CPU, GPU, NPU, or a GPU variant.
AUX model	An upstream LiteRT-LM pattern that moves KV-cache management into a separate auxiliary graph.
KV cache	The per-layer key/value attention state carried between generation steps.
`LiteRTLMExportSpec`	A per-model-family configuration object (cache layout, model type, vision/audio behavior) resolved once per export; see §6.8.
EOS	End-of-sequence token; the primary stop token written to the bundle metadata.
TTFT	Time to first token — latency from prompt submission to the first generated token.
MoE	Mixture of experts (e.g. Mixtral, Qwen-MoE model families).
BPE	Byte-pair encoding, the tokenizer algorithm used by GPT-2/Llama/Qwen-style models.
MLIR	The compiler intermediate representation TFLite graphs are lowered through; surfaced here only in upstream converter error messages.

3. Goals & Non-Goals

Goals

Single-call export: CausalLM.export(path, format="litertlm", …) → a runnable
bundle.
Cover text-only and multimodal (vision/audio) CausalLM families.
Produce a device-portable graph (one bundle runs on CPU/GPU/NPU delegates).
Preserve token identity end-to-end (no lossy tokenizer conversion).
Support in-conversion quantization and document a post-export quantization
path.
Fail fast with explicit messages for unsupported models/configs, before
any expensive tracing.

Non-Goals

Non-transformer architectures (e.g. RWKV) and non-generative models (e.g.
OCR/PARSeq).
Exporting from the JAX or TensorFlow Keras backend (litert_torch consumes
PyTorch modules).
Custom on-device chat-template strings (the runtime owns template
formatting; the app owns content).
A new Keras-specific mobile SDK (LiteRT-LM is the target runtime).

4. Requirements & Constraints

The design is shaped by the contract the LiteRT-LM runtime imposes and by the
mechanics of torch.export. These are the forces every later decision answers
to.

#	Constraint	Source	Consequence in the design
R1	Runtime expects flat per-layer KV-cache tensors (`kv_cache_k_0…`, `kv_cache_v_0…`) and a 1-D `input_pos`	LiteRT-LM I/O contract	The Adapter translates KerasHub's stacked cache into flat tensors and back (§6.1).
R2	Runtime drives generation via two fixed entry points: prefill and decode	LiteRT-LM execution model	Two traced signatures with precise `input_pos` semantics (§6.2).
R3	Graph must be device-agnostic (same bundle on CPU/GPU/NPU)	Portability goal	Tracing is forced CPU-only so no device constants leak in (§6.1).
R4	`torch.export` needs static shapes and traceable ops	`litert_torch` / `torch.export`	Fixed prefill length per signature (bucketing); export-scoped traceability patches (§6.3).
R5	Tokenization must match training exactly	Correctness	Bundle the native SentencePiece or HF `tokenizer.json`; never lossy-convert (§6.5).
R6	Backend must be PyTorch at export time	`litert_torch` input	Hard guard: reject non-torch backends up front (§6.3, validation).
R7	Target hardware is ARM64 mobile	Deployment reality	Runtime notes assume `arm64-v8a`; emulators unsupported (§12).

5. Solution Overview

The exporter bridges two worlds. KerasHub generates text in Python using a
single stacked KV-cache tensor and dynamic shapes. LiteRT-LM expects a
small set of static-shape TFLite signatures over flat per-layer cache
tensors, plus a metadata/tokenizer bundle. The exporter closes that gap in four
stages:

flowchart LR
    M["CausalLM model<br/>(PyTorch backend)"]
    A["1 · Adapter<br/><b>KerasHubLiteRTAdapter</b><br/>stacked KV cache ⇄ flat tensors<br/>clean prefill/decode entry points"]
    S["2 · Signatures<br/><b>prefill</b> · <b>decode</b><br/>(+ optional vision encoder / adapter;<br/>audio always baked into prefill)<br/>static-shape traces"]
    C["3 · Convert<br/><b>litert_torch</b><br/>torch.export → TFLite<br/>per signature"]
    B["4 · Bundle<br/><b>LitertLmFileBuilder</b><br/>TFLite + tokenizer + metadata"]
    O[("model.litertlm")]
    RT["LiteRT-LM runtime<br/>Android / iOS"]

    M --> A --> S --> C --> B --> O
    O -.->|load & run| RT
    classDef hot fill:#1f4f8a,color:#fff
    class A,S hot

Three core components:

Adapter (KerasHubLiteRTAdapter, a torch.nn.Module). Reshapes the
stacked KV cache into the flat per-layer tensors the runtime binds
(delegating the exact stack/unstack translation to the resolved
LiteRTLMExportSpec so a future family with a different cache layout can
override it, §6.8), exposes
prefill/decode as clean module boundaries, and filters model kwargs so
one adapter serves text, vision, and audio models. Traced on CPU so the
graph stays portable. (Requirements R1,
R3.)
Signatures. prefill fills the cache from the whole prompt; decode
emits one token per step. Each is a static-shape trace; multiple prefill
buckets can be traced so the runtime skips padding work. (R2, R4.)
Bundle. litert_torch converts each signature to TFLite, then
LitertLmFileBuilder packages the TFLite model(s), the tokenizer asset, and
the LlmMetadata protobuf into the final .litertlm. (R5.)

5.1 How this PR would split (proposed, pending maintainer approval)

This PR is intentionally broad (full pipeline + multimodal + verification).
If the maintainers prefer a staged review, the natural split is six stacked
PRs — each independently reviewable and testable, each landing green:

flowchart TD
    P1["PR 1 · Export core (text-only)<br/>adapter, traceable_ops, ExportPlan,<br/>prefill/decode/bundle + unit tests"]
    P2["PR 2 · Family spec registry & metadata<br/>model_specs.py, LlmMetadata,<br/>chat stop tokens, function_gemma"]
    P3["PR 3 · Tokenizer handling<br/>HF auto-conversion + vocab validation"]
    P4["PR 4 · Multimodal vision<br/>vision spec fields, baked + separate<br/>encoder, Gemma3/PaliGemma + parity harness"]
    P5["PR 5 · Multimodal audio<br/>Gemma3n/Gemma4 audio path +<br/>gemma4 subtype metadata"]
    P6["PR 6 · Quantization & on-device numbers<br/>quant configs, bucketing,<br/>Pixel 9 benchmark table"]

    P1 --> P2 --> P3
    P1 --> P4 --> P5
    P3 --> P6
    P5 --> P6

PR	Scope	Why it stands alone
1	Text-only export end-to-end (`export.py`, `adapter.py`, `traceable_ops.py`) for the registered text families, with the smoke + text-parity tests.	The whole multimodal story is additive on top; text export is the smallest complete, useful feature.
2	`model_specs.py` registry, `LlmMetadata` population (incl. chat stop tokens and `function_gemma`).	Pure metadata layer; reviewable by proto field diff, no graph changes.
3	HF tokenizer auto-conversion + vocabulary sanity checks.	Self-contained tokenizer path with its own tests.
4	Vision: `vision_input_style`, baked-in and separate encoder (single-image runtime contract), Gemma3 + PaliGemma enablement, vision parity harness, on-device image-run evidence.	First multimodal increment; brings the runtime-contract findings with it.
5	Audio: `audio_input_style`, Gemma3n/Gemma4 audio wiring, `gemma4` subtype metadata (`skip_mel_spectrogram_extraction`).	Depends on PR 4's multimodal plumbing.
6	Quantization configs, prefill bucketing, the §10.4 on-device performance table.	Numbers only mean something once all bundles are final.

Already-split-out companion: PR #2848
(CI dependency fix: torch floor + torchvision, independent of this feature).

Known-issue follow-ups that stay out of every stage (tracked, not hidden):
the Gemma4 harness-interpreter xfail, the Gemma3n upstream converter bug,
the Falcon tracing bug, and the PaliGemma cache-sizing defect — each is a
labeled xfail with its root cause, and each gets its own issue/PR.

The one public entry point

model.export(
    "model.litertlm",
    format="litertlm",
    prefill_seq_len=128,          # int or list[int] (buckets)
    cache_length=None,            # optional; else max_sequence_length → preprocessor.sequence_length (warns)
    backend_constraint=None,       # optional: cpu | gpu | npu | gpu_artisan
    quant_config=quant_config,    # optional litert_torch QuantConfig
    separate_vision_encoder=False,
    hf_tokenizer_path=None,       # optional user tokenizer.json
    sampler_config=None,          # optional SamplerConfig → LlmMetadata.sampler_params
    llm_model_type=None,          # optional override, e.g. "function_gemma"
)

CausalLM.export() routes to export_to_litertlm() when format="litertlm";
all other formats fall through to the standard Keras exporter. The full kwargs
reference is in the Appendix.

6. Detailed Design

Code references below name files inside keras_hub/src/utils/litertlm/
(export.py, adapter.py, model_specs.py, traceable_ops.py) and point
at function names, not line numbers — line numbers rot with every push.

6.1 The adapter and KV-cache translation

KerasHub stores the KV cache as one stacked tensor
[batch, num_layers, 2, cache_length, num_kv_heads, head_dim] with a scalar
cache_update_index. LiteRT-LM instead binds flat per-layer tensors
(kv_cache_k_0, kv_cache_v_0, …, kv_cache_k_L, kv_cache_v_L) and a 1-D
input_pos. KerasHubLiteRTAdapter performs the translation on every call:

flowchart LR
    subgraph In["LiteRT-LM flat inputs"]
        K0[kv_cache_k_0]
        V0[kv_cache_v_0]
        KL[kv_cache_k_L]
        VL[kv_cache_v_L]
    end
    subgraph Ad["KerasHubLiteRTAdapter"]
        direction TB
        Stack["export_spec.stack_kv_cache →<br/>[B, L, 2, T, H, D]<br/>(or [B, L, 2, H, T, D] for gemma3n)"]
        Call["model.call_with_cache(<br/>tokens, cache, cache_update_index)"]
        Unstack["export_spec.unstack_kv_cache →<br/>per-layer tensors"]
        Stack --> Call --> Unstack
    end
    subgraph Out["LiteRT-LM flat outputs"]
        OK0[kv_cache_k_0]
        OKL[kv_cache_k_L]
    end
    In --> Stack
    Unstack --> Out

Two cache layouts, selected by backbone:

"standard" (default) — [batch, cache_length, num_kv_heads, head_dim]
(Gemma3, Llama, Mistral, …).
"gemma3n" — [batch, num_kv_heads, cache_length, head_dim] (Gemma3n).

CPU-only tracing (R3).
_cpu_default_device_scope() forces the default device to CPU for the whole
trace, so the exported graph carries no device-specific constants and the
bundle runs on any delegate.

Locking: torch's default device is process-global state, so a
module-level lock (_DEFAULT_DEVICE_LOCK) serializes concurrent export()
calls — two overlapping exports can never interleave their default-device
swaps.
Caveat: the lock does not protect unrelated GPU work. Because the
swap is process-global, any other thread doing ordinary GPU work while an
export holds the lock still observes torch.set_default_device("cpu") for the
whole conversion window.

Aliasing safety. Earlier revisions of this exporter cloned the stacked
cache and each per-layer output to stop torch.export aliasing returned KV
buffers with activation tensors. Those clones were removed because they are
unnecessary: Keras's cache-update ops (slice_update/scatter_update, used
by every model's call_with_cache) and torch.stack are purely functional —
they never alias their inputs — and torch.export's functionalization pass
materializes graph-output views into independent buffers on its own. Graph
inspection of a traced decode step (recorded as a comment on
stack_kv_cache/unstack_kv_cache in model_specs.py) confirms
unstack_kv_cache compiles to zero stack/cat nodes — view-based, no copy.
Note that this reasoning is not backed by an end-to-end, multi-step
runtime-generation regression test today; see
Testing Strategy for the coverage that does exist.

One adapter, many modalities. The adapter inspects model.call_with_cache
once and filters kwargs (img_embeddings, pixel_values, vision_mask,
audio_embeddings, …) to match the model's real signature — so the same code
path serves text-only, vision, and audio models.

KV-cache dtype. Exported KV tensors use the model's compute dtype. FP32 and
FP16 are known-good with the runtime; BF16 is untested — convert to FP16
before export for predictable results.

I/O cost note. A 28-layer model binds 56 flat KV input/output
tensors through the TFLite C++ API every decode step (plus torch.stack
copy overhead in stack_kv_cache; unstack_kv_cache is 0-copy). This is a
static graph-node count, not a profiled measurement. Cost and follow-up are
in the trade-offs table,
§10.1, and Future Work #3.

6.2 Prefill and decode signatures

The runtime drives generation with exactly two operations. Prefill seeds the
cache from the whole prompt and returns only the updated cache — no logits
(forward_prefill in adapter.py, return_logits=False). The first
logits come from the first decode step, which feeds the last prompt token
and returns logits[0]; every later decode step samples the token returned by
the step before it (forward_decode, return_logits=True).

sequenceDiagram
    participant R as LiteRT-LM Runtime
    participant P as PREFILL_DECODE model
    participant K as KV cache buffers

    Note over R,K: New turn, prompt of N tokens
    R->>R: input_pos = [s, s+1, …, s+N-1]

    Note over R,P: Prefill: seed the cache, no sampleable output
    R->>P: prefill(prompt[1,N], input_pos, kv_cache_k_0…v_L)
    P->>K: write updated KV caches
    P-->>R: updated kv_cache_k_0…v_L (KV caches only, no logits)

    Note over R,P: Decode step 1: last prompt token → first logits
    R->>P: decode(last_prompt_token, input_pos=[s+N-1], kv_cache_k_0…v_L)
    P->>K: write updated KV caches
    P-->>R: logits[0] + updated kv_cache_k_0…v_L

    loop until stop token (i = 1, 2, …)
        R->>R: sample token_id from the logits just returned
        R->>P: decode(token_id, input_pos=[s+N-1+i], kv_cache_k_0…v_L)
        P->>K: write updated KV caches
        P-->>R: logits[i] + updated kv_cache_k_0…v_L
    end

Signature	Input tokens	`input_pos`	Returns	Purpose
prefill	full prompt, `seq_len = N`	1-D int32; first element is the cache start index	updated KV caches only	fill the cache with the prompt
decode	single token, `seq_len = 1`	1-D int32, length 1; treated as a scalar cache index	logits + updated KV caches	one autoregressive step

Bucketing (R4). A fixed signature needs a fixed prompt length. To avoid
padding every prompt to the maximum, prefill_seq_len accepts a list[int]:
the exporter traces one prefill signature per bucket (prefill_32, prefill_64,
prefill_128) plus a single decode. At runtime the executor dispatches to the
smallest bucket that fits:

flowchart LR
    Prompt([Prompt of length N]) --> Vis{"Vision-capable<br/>model?"}
    Vis -->|yes| Fixed["No bucketing available:<br/>always run prefill at cache_length"]
    Vis -->|no| Pick{"Smallest bucket B<br/>with B ≥ N?"}
    Pick -->|found| Run["Run prefill_B, then decode loop"]
    Pick -->|none| Abort["Abort: prompt exceeds<br/>largest bucket"]

Because each bucket is an independent static-shape trace, export time scales
~linearly with bucket count. The exporter validates every requested
prefill_seq_len is a positive int ≤ cache_length; an over-length prompt
at runtime aborts (no graceful truncation yet — size buckets to your max prompt).

Bucketing applies to text-only signatures only; the vision-capable
restriction, its mechanism, and its consequence are stated once in
§6.6.

Observed impact (a single example measurement, not a guarantee — actual
latency varies by OS, runtime version, and thermals):

Prefill length	Single bucket (pad to max)	Buckets `[32,64,128]`	Improvement
32 tokens	~43 ms	~31 ms	~28%
64 tokens	~61 ms	~35 ms	~43%

(Pixel 9, Gemma3 270M, CPU backend.)

6.3 Export pipeline

export_to_litertlm() runs five phases top-to-bottom; sub-steps run
left-to-right.

flowchart TB
    subgraph P1["1 · Validate (fail fast)"]
        direction LR
        V1["path ends .litertlm"] --> V2["tokenizer supported"] --> V3["backend == torch"] --> V4["cache_structure supported"] --> V5["prefill_seq_len ≤ cache_length"]
    end
    subgraph P2["2 · Build sample inputs"]
        direction LR
        S1["one set per bucket"] --> S2["decode inputs"]
    end
    subgraph P3["3 · Trace signatures"]
        direction LR
        T1["prefill_&lt;bucket&gt;<br/>(vision/audio baked in by default)"] --> T2["decode"]
        T1 -.->|optional, vision only| T3["VISION_ENCODER / VISION_ADAPTER"]
    end
    subgraph P4["4 · Convert"]
        C1["litert_torch.convert — per signature"]
    end
    subgraph P5["5 · Package"]
        direction LR
        K1["materialize tokenizer"] --> K2["build LlmMetadata"] --> K3["LitertLmFileBuilder.build"]
    end
    P1 --> P2 --> P3 --> P4 --> P5

Implementation details that matter for review:

ExportPlan. export_to_litertlm() resolves the model-family spec and
every per-export-run setting once (phases 1-2), bundles them into one frozen
ExportPlan dataclass, and passes that single object into
_build_prefill_inputs, _trace_and_convert, and _assemble_bundle.
Per-signature conversion. litert_torch.convert runs the full
export/decomposition/FX-pass pipeline independently per signature; only the
final flatbuffer merge is shared. This is the root of the linear
bucket-count cost.
Traceability patches (R4). A few Keras torch ops are not export-traceable
as written — e.g. ops.slice coerces tensor start-indices to Python ints via
.tolist(), which fails when the index is dynamic (the decode input_pos).
traceable_ops.py holds a traceable replacement for each:
slice, dot_product_attention, one_hot, repeat, arange, take,
scatter_update, amax. Its traceable_ops_scope() combines the eight
individual patch scopes into a single contextlib.ExitStack-based context
manager, so _trace_and_convert opens one scope instead of nesting eight
with statements. Patches apply only during export and are restored in
finally; normal eager/training behavior is untouched outside the export
window. Like _cpu_default_device_scope()
(§6.1), these patches are
process-global for the scope's duration — each is a
unittest.mock.patch.object() override of a function on a shared Keras
torch-backend module — so a concurrent Keras-torch thread executing any of
the patched ops during an export would observe the patched behavior too.
The amax patch works around an upstream lowering gap — litert_torch's
layout pass checks 4-D aten.amax for an NHWC rewrite but has none
registered, so conversion aborted; the patch routes a single-integer-axis
reduction of a 4-D tensor through torch.max(dim=...) instead. This is what
unblocks GPT-OSS, whose attention-sink softmax hits exactly this shape
(upstream gap: litert-torch#1126).
JAX x64 and platform preservation. Importing litert_torch enables
jax_enable_x64 as a side effect, and LiteRT-LM's JAX bridge defaults to the
TPU platform if one is visible. The exporter saves and restores both flags
around tracing/conversion (_preserve_jax_x64_state(),
_preserve_jax_platforms_state()) — forcing jax_platforms="cpu" for the
duration so export does not contend with other processes for the TPU — so
JAX code elsewhere in the process sees its original configuration
afterward either way.
Conversion mode. The optional VISION_ENCODER / VISION_ADAPTER models
convert with lightweight_conversion=True; the main PREFILL_DECODE uses
lightweight_conversion=False.
backend_constraint normalization. The kwarg is lowercased before it
reaches LitertLmFileBuilder.add_tflite_model for every bundled TFLite
model, so a mixed-case input (e.g. "GPU") behaves identically to
"gpu".

6.4 Bundle packaging and metadata

LitertLmFileBuilder assembles the final archive:

flowchart LR
    subgraph Assets
        TFile["PREFILL_DECODE slot<br/>(file: model.tflite, or prefill_decode.tflite<br/>when vision is separated)<br/>audio always baked in"]
        VEnc["vision_encoder.tflite (opt)"]
        VAd["vision_adapter.tflite (opt)"]
        Eov["end_of_vision.tflite (opt,<br/>separate vision only)"]
        SFile["tokenizer (vocabulary.spm / tokenizer.json)"]
        MFile["llm_metadata.pb"]
    end
    subgraph Builder["LitertLmFileBuilder"]
        direction TB
        AddT["add_tflite_model …"] --> Build["build()"]
        AddS["add_sentencepiece / add_hf_tokenizer"] --> Build
        AddM["add_llm_metadata"] --> Build
    end
    TFile --> AddT
    VEnc --> AddT
    VAd --> AddT
    Eov --> AddT
    SFile --> AddS
    MFile --> AddM
    Build --> Out[("model.litertlm")]

A bundle contains (1) the TFLite model(s) — PREFILL_DECODE (one prefill
signature per bucket) plus optional VISION_ENCODER/VISION_ADAPTER and, on a
separate-vision export for families declaring an end_of_vision_token,
END_OF_VISION (§6.6); (2) a
tokenizer asset; (3) the LlmMetadata protobuf.

The bundle-internal slot is always named PREFILL_DECODE. On disk the graph file
is model.tflite in the default case, and prefill_decode.tflite only when
separate_vision_encoder=True (export.py) — the slot name, not the filename, is
what the runtime binds.

Atomic write. _assemble_bundle writes the final .litertlm to a temp
file in the same directory as the destination path, then os.replace()s it
into place on success — a crash mid-build (bundles can be large) never leaves
a truncated file at path. The temp file's permissions are explicitly reset
to 0666 & ~umask (matching what a plain open(path, "wb") would have
produced), since tempfile.mkstemp() otherwise creates it 0600.

Per-field status — every LlmMetadata field the exporter can touch.
The installed litert_lm_builder proto defines exactly eight top-level
LlmMetadata fields. "Set" means the exporter writes it; "Never set" means
no code path in the exporter touches it.

`LlmMetadata` field	Status	Detail
`start_token`	Set	`tokenizer.start_token_id`, when the tokenizer defines one.
`stop_tokens`	Set	Primary EOS (`tokenizer.end_token_id`) plus any family chat-turn token from `spec.get_chat_stop_token_ids()`; de-duplicated.
`max_num_tokens`	Set	Context window: `cache_length` kwarg → `backbone.max_sequence_length` → `preprocessor.sequence_length` (last fallback warns).
`llm_model_type`	Set	Family oneof selected from the resolved spec (`generic_model`, `gemma3`, `gemma3n`, `gemma4`, `qwen3`, `qwen2p5`, `function_gemma`). Drives the runtime chat template and any multimodal subtype.
`sampler_params`	Conditional	Written only when the caller passes a `SamplerConfig` (`sampler_config=` kwarg); omitted by default so the runtime picks its own policy. `GREEDY_SAMPLER_CONFIG` (`top_k=1`) exists for deterministic verification. Type derivation mirrors litert-torch's `export_hf` precedence with one deliberate deviation: `top_k` set → `TOP_K` (`k=1` for greedy — the proto's `GREEDY` type is never emitted because runtime 0.13.1 doesn't implement it, see §10.2), otherwise `TOP_P`.
`prompt_templates`	Never set	Deliberate non-goal — the runtime owns chat-template formatting by `LlmModelType` (Non-Goals); the exporter emits no template strings.
`jinja_prompt_template`	Never set	Same reason as `prompt_templates`: template formatting is a runtime responsibility, not an export output.
`channels`	Never set	A `Channel` is a `{channel_name, start, end}` reasoning-channel token span. No KerasHub family currently emits reasoning-channel tokens, so nothing is written.

Two fields live on a llm_model_type subtype (nested, not top-level), set
only for the families that own them:

Subtype field	Owning subtype	Status	Detail
Vision block (image H/W, start/end image tokens)	`gemma3`, `gemma3n`	Set for Gemma3/Gemma3n	Derived from the backbone `image_size`.
Vision block (patch W/H, `max_num_patches`, pooling kernel)	`gemma4`	Set for Gemma4	Derived from `image_size` and the encoder's `patch_size`/`pool_size`.
Audio start/end tokens	`gemma3n`, `gemma4`	Set for Gemma3n/Gemma4	Hardcoded `<
Function-calling block (`code_fence_start`, `code_fence_end`, `function_response_start`, `use_template_for_fc_format`)	`function_gemma`	Set for `function_gemma` only	Populated by `FunctionGemmaSpec`. The preset (`function_gemma_instruct_270m`) loads as a plain `Gemma3CausalLM`, so it is auto-detected by tokenizer: its vocabulary contains `<start_function_call>`/`<end_function_call>`, which plain Gemma3 presets lack (verified by a token→id→token round-trip, since some tokenizers map out-of-vocabulary strings to `<unk>` instead of raising). An explicit `llm_model_type="function_gemma"` kwarg remains as an override, but note the asymmetry: an explicit override sets the bare `function_gemma` oneof and skips this block's population (mirroring litert-torch, which skips its model-specific metadata builder whenever `litert_lm_model_type_override` is set) — only tokenizer auto-detection fills these fields. `constraint_mode` is left at its proto default.
`skip_mel_spectrogram_extraction`	`gemma4`	Set for Gemma4 only	Set to `False`, cited against ai-edge-litert's `audio_preprocessor_miniaudio.cc` (`False` = runtime performs log-mel extraction, matching Gemma4's host-side log-mel export contract). Also the proto3 default, so this is a contract statement and regression guard, not a byte change.

Two more attributes are commonly confused with LlmMetadata fields because
they travel alongside it in the same bundle, but they are not — they are
per-TFLite-section kwargs on LitertLmFileBuilder.add_tflite_model:

Attribute	Status	Detail
`backend_constraint`	Applied uniformly	Passed to every bundled TFLite section (PREFILL_DECODE and any vision sections), lowercased first. Accepted values `cpu`/`gpu`/`npu`/`gpu_artisan`, but only CPU is validated end-to-end on device today — GPU/NPU are accepted and stored, not yet verified to run (§10.2, §9).
`prefer_activation_type`	Never set	The builder API accepts it, but the exporter never passes it — no correctness stake today. See Future Work.

Chat-turn stop tokens, per family. _build_llm_metadata always adds
tokenizer.end_token_id. Some families additionally mark the end of a chat
turn with a second, distinct token, resolved by
LiteRTLMExportSpec.get_chat_stop_token_ids() (default: none) — one method
per family instead of one hardcoded probe:

Family	Extra stop token	Note
Gemma, Gemma3, Gemma3n, Gemma4, PaliGemma (`GemmaSpec` and its subclasses)	`<end_of_turn>`	shared SentencePiece convention across the whole Gemma family
Llama3 (`Llama3Spec`)	`<	eot_id
Qwen, QwenMoe (`Qwen2p5FamilySpec`)	`<	im_end
Qwen3, Qwen3Moe, Qwen3.5 (`Qwen3FamilySpec`/`Qwen3_5Spec`)	`<	im_end
Phi3 (`Phi3Spec`)	`<	end
Base Llama (v1/v2), and every other registered or unregistered family (Mistral, Mixtral, GPT2, Bloom, GPT-NeoX, GPT-OSS, OPT, Falcon, SmolLM3)	none	only the primary EOS is emitted; see Limitations

Duplicate ids (e.g. Qwen3's) are de-duplicated before writing
meta.stop_tokens.

cache_length fallback. Most backbones (Gemma, Llama, Mistral, Qwen) do
not define max_sequence_length, so without an explicit value cache_length
— and therefore max_num_tokens — would fall back to
preprocessor.sequence_length, a tokenization default that is not necessarily
the model's true context window. The exporter therefore accepts an explicit
cache_length kwarg (threaded through CausalLM.export() →
export_to_litertlm(), see the
Appendix) and raises a
UserWarning whenever it still has to fall back to the preprocessor value, so
a context-crippled bundle is never silent.

Model-type mapping. resolve_export_spec() (model_specs.py) walks a
lazy (module_path, class_name, spec_factory) registry
(_EXPORT_SPEC_REGISTRY) with isinstance checks, in registration order —
first match wins. There is no class-name heuristic: every entry is an
explicit isinstance check against an imported class, and an unregistered
model gets the base LiteRTLMExportSpec (model_type="generic_model"). See
§6.8 for the full spec-registry mechanism.

Current registry, exactly:

Gemma4CausalLM → Gemma4Spec (gemma4)
Gemma3nCausalLM → Gemma3nSpec (gemma3n)
Gemma3CausalLM → Gemma3Spec (gemma3)
PaliGemmaCausalLM → PaliGemmaSpec (generic_model — no dedicated LlmModelType vision subtype; registered only to share the Gemma-family chat-stop-token behavior, see [§10.3](#103-multimodal))
GemmaCausalLM → GemmaSpec (generic_model — registered only to share the same chat-stop-token behavior)
Qwen3MoeCausalLM / Qwen3CausalLM → Qwen3FamilySpec (qwen3)
Qwen3_5CausalLM → Qwen3_5Spec (qwen3 — same oneof as Qwen3; a distinct spec class only for the hybrid-cache rejection message, see [Limitations](#101-export-time))
QwenMoeCausalLM / QwenCausalLM → Qwen2p5FamilySpec (qwen2p5)
Llama3CausalLM → Llama3Spec (generic_model — registered ahead of the base Llama entry so its <|eot_id|> override wins; the proto has no dedicated "llama" oneof)
LlamaCausalLM (v1/v2) and every unregistered family → base LiteRTLMExportSpec (generic_model)
Phi3CausalLM → Phi3Spec (generic_model — no dedicated LlmModelType subtype; registered only for the <|end|> chat-stop-token override, see the chat-stop-token table above)

The mapping is scoped to the LlmModelType values the installed
litert_lm_builder proto defines: generic_model, gemma3n, function_gemma,
gemma3, qwen3, qwen2p5, gemma4, fast_vlm. function_gemma is not
class-detectable (the preset loads as a plain Gemma3CausalLM), so
auto-detection selects it by tokenizer content — the function-calling special
tokens described in the subtype table above — checked before the class
registry so the model does not fall through to Gemma3Spec; the explicit
llm_model_type="function_gemma" kwarg stays available as an override.
Nothing maps to fast_vlm.

6.5 Tokenizer handling

The runtime loads exactly two on-device asset types — a SentencePiece
.spm or an HF tokenizer.json. The exporter picks between them, and (for the
HF one) its source, by tokenizer class and whether the user supplied a file:

flowchart TD
    Start([Tokenizer bundling]) --> User{"hf_tokenizer_path given?"}
    User -->|yes| HFUser["bundle user tokenizer.json<br/>add_hf_tokenizer()"]
    User -->|no| SPM{"SentencePiece /<br/>vocabulary.spm?"}
    SPM -->|yes| SPMOut["materialize vocabulary.spm<br/>add_sentencepiece_tokenizer()"]
    SPM -->|no| BPE{"BytePairTokenizer<br/>subclass?"}
    BPE -->|yes| HFAuto["serialize wrapped HF tokenizer →<br/>tokenizer.json (hf_tokenizer_converter.py)"]
    BPE -->|no| Err["ValueError:<br/>unsupported tokenizer class"]
    HFUser --> Done([Tokenizer in .litertlm])
    SPMOut --> Done
    HFAuto --> Done

SentencePiece asset — SentencePieceTokenizer subclasses → materialize
vocabulary.spm, bundle via add_sentencepiece_tokenizer(). The runtime
parses the SentencePiece protobuf directly and supports constrained decoding.
HF tokenizer.json, user-provided — bundled verbatim via
add_hf_tokenizer() into the .litertlm's HF_Tokenizer_Zlib section.
Functional for ordinary encode/decode.
HF tokenizer.json, auto-serialized from BytePair — same asset as (2),
just generated instead of user-supplied. Every BytePairTokenizer subclass
already wraps an actual HF tokenizers.Tokenizer object internally
(self._tokenizer), which it uses for its own encode_batch/decode_batch.
hf_tokenizer_converter.py simply calls .to_str() on that same object and
bundles the result exactly as in (2). This is not a format conversion —
no new tokenizer is constructed and no token IDs are re-derived. It is
literally the object KerasHub already tokenizes with, serialized to JSON,
so token identity is byte-exact by construction, not by validation.

Tokenizer family	KerasHub class	Builder method	Asset
Gemma, Mistral, Phi3, Llama	`SentencePieceTokenizer`	`add_sentencepiece_tokenizer()`	`vocabulary.spm`
GPT2, Llama3, Qwen/Qwen3(+MoE), Bloom, Falcon, GPT-NeoX, GPT-OSS, OPT, SmolLM3	`BytePairTokenizer`	`add_hf_tokenizer()`	`tokenizer.json` (auto or user)

Why not lossy-convert (R5). Upstream tooling can convert BPE →
SentencePiece, but it reports ~1% token-ID mismatch (e.g. Llama3.2). The
exporter refuses that path and bundles HF JSON instead, preserving token
identity. The trade-off is narrower runtime support for HF tokenizers today:
constrained decoding is SentencePiece-only, and streaming load of an
HF_Tokenizer_Zlib section is not yet implemented.

hf_tokenizer_path vocabulary sanity check. A user-supplied tokenizer no
longer gets only a file-existence/extension check. export.py now compares
the tokenizer's implied vocabulary size (max token id + 1, read directly from
tokenizer.json, across model.vocab and added_tokens) against the model's
embedding-table size (backbone.vocabulary_size, or
token_embedding.input_dim as a fallback), and raises a ValueError when the
two differ by more than 300 tokens and by at least a 5x ratio in either
direction. This is a coarse heuristic aimed at catching an obviously wrong
tokenizer (e.g. one pointed at a different model family's tokenizer.json) —
it is not exact validation, and a modest mismatch (e.g. a handful of added
special tokens) still passes silently.

Why HF format, not a KerasHub-native one. Two independent constraints —
not a preference for one tokenizer ecosystem over another:

The runtime only accepts two on-device tokenizer formats.
LitertLmFileBuilder exposes exactly add_sentencepiece_tokenizer() and
add_hf_tokenizer() (SP_Tokenizer/HF_Tokenizer_Zlib section types) —
there is no third, KerasHub-specific format the LiteRT-LM runtime can load
on-device. Adding one would be a LiteRT-LM runtime change, outside this
PR's (and this repo's) control.
For BPE families, "the HF tokenizer" and "the KerasHub tokenizer" are
the same object, not two competing implementations.
BytePairTokenizer already builds a tokenizers.Tokenizer BPE instance
from KerasHub's own preset vocab/merges, and every eager encode/decode
call KerasHub makes today runs through that object. The exporter calls
.to_str() on it (hf_tokenizer_converter.py) — no new dependency, no
data fetched from HuggingFace's hub, no format conversion. It is KerasHub's
own vocabulary and merge rules, serialized through the container format the
object already produces. Native SentencePiece would be the only zero-HF
path, and it isn't an option for BPE families: they were never
SentencePiece-tokenized, and converting them to SPM is exactly the ~1%
token-mismatch lossy path this section rejects (R5).

6.6 Multimodal design

The mental model

A model is text-only, or it has vision and/or audio. The exporter decides how to
handle each modality at export time, reading everything it needs from the
model's resolved LiteRTLMExportSpec (model_specs.py) — it does not sniff
the model at trace time. From the spec it decides three things:

How each modality's input is shaped — the vision_input_style and
audio_input_style fields (defined below).
Whether the vision encoder may be split into its own bundle sections — the
supports_separate_vision spec flag, only relevant when the caller passes
separate_vision_encoder=True.
What bundle shape results — one fused PREFILL_DECODE, or that plus
separate VISION_ENCODER/VISION_ADAPTER sections.

Terms, defined once (used everywhere below)

Two spec fields carry all the per-family multimodal behavior.

vision_input_style — how a family delivers vision input to the graph.

Value	Meaning	Families
`raw_images`	Adapter calls the vision encoder inline on `[B, N, H, W, 3]` images.	Gemma3, PaliGemma
`patch_values`	Adapter calls the encoder with pre-processed `pixel_values` + `pixel_position_ids`.	Gemma4
`embedded_pixel_values`	Encoder runs inside the backbone; `call_with_cache` consumes raw `pixel_values` and the adapter never calls a separate encoder.	Gemma3n

audio_input_style — how a family delivers audio input.

Value	Meaning	Families
`None`	No audio encoder.	all text-only + vision-only families
`embedded_mel`	Encoder inside the backbone; `call_with_cache` consumes the pre-extracted mel spectrogram directly.	Gemma3n
`standalone_mel`	Adapter calls `backbone.audio_encoder(audio_mel, audio_mel_mask)` as an in-trace stage, then feeds the embeddings into `call_with_cache`.	Gemma4

The adapter dispatches on these spec fields in adapter.py and raises a typed
ValueError if the declared style and the supplied tensors disagree, rather
than silently returning nothing.

The adapter asks the spec; it never branches on the model class. The diagram
below is the whole decision, including the two paths that end in a hard
ValueError.

flowchart TD
    Start(["export(separate_vision_encoder=?)"]) --> Vis{"has_vision?<br/>(spec.get_vision_config)"}

    Vis -->|no| Aud0{"has_audio?"}
    Aud0 -->|no| TextOnly["Text-only:<br/>PREFILL_DECODE only"]
    Aud0 -->|yes| AudBake["Audio baked into PREFILL_DECODE<br/>(no separate-audio path exists)"]

    Vis -->|yes| Style{"vision_input_style?"}
    Style -->|"embedded_pixel_values<br/>(Gemma3n)"| EmbBuck
    Style -->|"raw_images / patch_values<br/>(Gemma3, PaliGemma, Gemma4)"| SepReq{"separate_vision_encoder<br/>requested?"}

    SepReq -->|no| Baked["Baked-in:<br/>encoder inside PREFILL_DECODE"]
    SepReq -->|yes| SupSep{"spec.supports_separate_vision?"}
    SupSep -->|yes| Separate["Separate:<br/>VISION_ENCODER + VISION_ADAPTER<br/>+ optional END_OF_VISION"]

    EmbBuck{"separate requested?"} -->|no| Baked
    EmbBuck -->|yes| RejA[["ValueError:<br/>Gemma3n runs its encoder<br/>inside the backbone"]]

    Baked --> Buck{"any prefill_seq_len<br/>!= cache_length?"}
    Separate --> Buck
    Buck -->|"yes, and<br/>allows_vision_bucketing=False"| RejB[["ValueError:<br/>vision families forbid<br/>bucketing"]]
    Buck -->|no| OK(["Bundle ready"])

    classDef reject fill:#7a1f1f,color:#fff
    class RejA,RejB reject

cache_structure is a sibling spec field describing the KV-cache tensor shape;
only "single_stacked" is supported today (see §10.1 and the
glossary). It is orthogonal to the vision/audio styles.

Per-family mode matrix

Family	`vision_input_style`	Baked-in vision?	Separate vision?	Audio (`audio_input_style`)	Vision metadata set
Gemma3	`raw_images`	yes	yes	none	`gemma3` subtype: start/end image tokens, image H/W
Gemma3n	`embedded_pixel_values`	yes	no → `ValueError`	`embedded_mel`	`gemma3n` subtype (same helper as Gemma3)
Gemma4	`patch_values`	yes	yes	`standalone_mel`	`gemma4` subtype: start/end image tokens, patch W/H, `max_num_patches`, pooling kernel
PaliGemma	`raw_images` (+ `flatten_image_batch`)	yes	yes	none	none — `generic_model`, documented no-op (§10.3)

Audio metadata: Gemma3n and Gemma4 set start/end audio tokens (<|audio> /
<audio|>); Gemma4 additionally sets skip_mel_spectrogram_extraction=False,
cited against ai-edge-litert's audio_preprocessor_miniaudio.cc.

Dataflow: baked-in mode (default)

The vision/audio encoders run inside the PREFILL_DECODE graph, so
image/audio inputs are processed alongside text tokens in one pass. Decode stays
text-only, because the multimodal KV-cache is already seeded after prefill.

flowchart LR
    IMG["raw images /<br/>pixel_values"] --> PD["PREFILL_DECODE<br/>(encoder runs inside)"]
    AUD["audio mel"] --> PD
    TOK["text tokens"] --> PD
    PD --> OUT["KV caches + logits"]

Baking the encoder into PREFILL_DECODE is a keras-hub-only design point.
Upstream export_hf always exports vision separately, so there is no upstream
contract for this mode to conform to — it is a disclosed design choice, not a
gap.

Dataflow: separate mode (`separate_vision_encoder=True`)

The vision tower is exported as its own bundle sections, and PREFILL_DECODE
consumes pre-computed mm_embeddings — a smaller main graph and no re-encoding
the same image across turns.

flowchart LR
    IMG["raw images /<br/>pixel_values"] --> VE["VISION_ENCODER"]
    VE -->|features| VA["VISION_ADAPTER"]
    VA -->|mm_embedding| PD["PREFILL_DECODE"]
    TOK["text tokens"] --> PD
    VA -.-> EOV["END_OF_VISION<br/>(if end_of_vision_token declared)"]
    PD --> OUT["KV caches + logits"]

Two notes on the section shapes:

VISION_ADAPTER is intentionally a pass-through rename (features →
mm_embedding; see KerasHubVisionAdapter in adapter.py). keras-hub
already projects features inside
the encoder, so the adapter model is a no-op. The two-slot
VISION_ENCODER/VISION_ADAPTER split exists only because the LiteRT-LM
bundle format defines both slots (an upstream contract) — this PR conforms to
it, it does not introduce it.
END_OF_VISION is emitted on a separate-vision export only for families
that declare an end_of_vision_token — Gemma3 (<end_of_image>) and Gemma4
(<image|>), but not PaliGemma or Gemma3n. It is added for parity with
upstream's END_OF_VISION section. An on-device A/B test (Pixel 9,
litertlm-android 0.13.1) ran the same Gemma3 separate-vision bundle with
and without the section: both run identically — the runtime does not
require it; it is upstream-parity only.

Restrictions

Bucketing is banned for all 4 vision-capable families — Gemma3, Gemma3n,
Gemma4, and PaliGemma — via the allows_vision_bucketing=False spec default
(enforced in _validate_export_args, whose error text names all four). The
constraint was first observed for Gemma3's image attention-mask needing
cache_length == input_length, but no per-family assessment has confirmed
the other three need it; relaxing it is a future, numerics-gated change (flip
allows_vision_bucketing=True on that family's spec — already spec-flagged).
Consequence: every prefill on a vision-capable bundle runs at the full
cache_length, even for a text-only prompt, so those families lose the
bucketing TTFT win from §6.2. This applies
in both baked-in and separate modes.
Separate vision is per-family, read straight from the
supports_separate_vision flag: supported for Gemma3/Gemma4/PaliGemma,
rejected for Gemma3n (its encoder lives in the backbone). The flag keeps this
matrix readable from the spec rather than re-derived from the input style.
Audio has no separate-export path at all, and this is upstream-gated, not
a keras-hub gap. The bundle format's AUDIO_* slots exist and the adapter
already calls backbone.audio_encoder standalone, so tracing is not the
blocker — what is missing is a published upstream reference contract (vision has
a documented VISION_ENCODER/VISION_ADAPTER contract; audio has only a
one-off example for the Moonshine speech-recognition model, not a documented
slot contract). See §10.3 and
Future Work #4.

6.7 Quantization

In-conversion — pass a litert_torch QuantConfig as quant_config; it is
forwarded to litert_torch.convert() for in-graph quantization. Recipes from
litert_torch.generative.quantize.quant_recipes:

Recipe	Parameters	Description
`full_dynamic_recipe()`	`mcfg`, `weight_dtype`, `granularity`	Dynamic-range weights (activations FP32). Good first try.
`full_weight_only_recipe()`	`mcfg`, `weight_dtype`, `granularity`	Weight-only; activations FP32. Best size/quality after validation.
`full_fp16_recipe()`	`mcfg` only	FP16 weights+activations. Rejects `weight_dtype`/`granularity`.

weight_dtype: INT8 (default), INT4, FP16, FP32. granularity:
CHANNELWISE (default), BLOCKWISE_{32,64,128,256}.

from litert_torch.generative.quantize.quant_recipes import full_weight_only_recipe
import litert_torch.generative.quantize.quant_attrs as quant_attrs

model.export(
    "model.litertlm", format="litertlm", prefill_seq_len=128,
    quant_config=full_weight_only_recipe(
        weight_dtype=quant_attrs.Dtype.INT4,
        granularity=quant_attrs.Granularity.BLOCKWISE_128,
    ),
)

Post-export — if quant_config is omitted, export is FP32/FP16. To shrink:
export → extract TFLite from the bundle → run ai-edge-quantizer (INT8/INT4,
dynamic/static, weight-only) → repackage with LitertLmFileBuilder. Observed:
Gemma3 270M ≈ 1.1 GB (FP32) → ≈ 275 MB (INT8 dynamic).

6.8 Model support and detection

Cache config is extracted from the model at export time, with fallbacks for
naming differences across backbones:

Field	Source	Fallback
`num_layers`	`backbone.num_layers`	`backbone.num_hidden_layers`
`cache_length`	`backbone.max_sequence_length`	`preprocessor.sequence_length`
`num_kv_heads`	`backbone.num_key_value_heads`	`num_query_heads` → `num_heads` → `num_attention_heads`; raises `ValueError` if none is found (no silent default)
`head_dim`	`backbone.head_dim`	`hidden_dim // num_query_heads` (then `num_heads`, `num_attention_heads`)
`dtype`	`model.compute_dtype` / `backbone.compute_dtype`	`float32`

Extension points. All model-family behavior lives on one class per family —
LiteRTLMExportSpec (model_specs.py), resolved once per export by
resolve_export_spec() (§6.4).

What the spec owns: KV-cache layout plus the
stack_kv_cache/unstack_kv_cache translation (the default matches today's
single-stacked behavior; no family overrides it yet, but the seam exists for
a genuinely different layout — e.g. Qwen3.5's hybrid cache,
Limitations), LlmModelType, vision/audio config
(vision_input_style/audio_input_style, §6.6),
chat stop tokens, and the multimodal adapter hooks (Gemma4's embedding
reshape, Gemma3n's forced padding mask).
Unaffected: HF tokenizer conversion — a separate mechanism that
introspects the tokenizer class, not the model family:

Behavior	Current mechanism	How to extend
Model-family behavior (cache layout, `LlmModelType`, vision/audio config, multimodal adapter hooks, KV-cache stack/unstack)	`LiteRTLMExportSpec` registry: lazy `isinstance` checks against `_EXPORT_SPEC_REGISTRY`, first match wins; an unregistered model gets the base `LiteRTLMExportSpec` (`generic_model`)	Add a `LiteRTLMExportSpec` subclass overriding the relevant class attributes/methods, then register `(module_path, class_name, subclass)` in `_EXPORT_SPEC_REGISTRY`
HF tokenizer conversion	serialization of the HF `tokenizers.Tokenizer` object `BytePairTokenizer` already wraps	extend `hf_tokenizer_converter.py` (a from-scratch converter for tokenizer classes that don't wrap an HF object, see Future Work)

Long-term: add a litertlm_config() protocol on CausalLM/Backbone so
new models declare these values explicitly instead of relying on introspection.

Supported families (each has a *_causal_lm_test.py):

Family	Class	Notes
Gemma	`GemmaCausalLM`	text-only baseline
Gemma3	`Gemma3CausalLM`	vision+text; baked-in & separate encoder; the multimodal family with passing numeric-parity export tests (Gemma4's parity test exists but is currently strict-`xfail` — see its row below), see Testing Strategy
Gemma3n	`Gemma3nCausalLM`	vision+audio+text; baked-in only. Export is currently blocked by a pre-existing upstream `litert-torch` conversion bug — see the export-blocked table below. (An earlier int64/int32 MLIR tracing failure was already fixed by an `int32` cast, Appendix; the converter defect is the remaining blocker.)
Gemma4	`Gemma4CausalLM`	vision+audio+text; baked-in & separate encoder; export produces a bundle, but the multimodal numeric-parity test is currently a strict `xfail` — root-caused to an upstream ai_edge_litert builtin-interpreter memory-planning bug, not an export defect (§10.3)
Llama / Llama3	`LlamaCausalLM` / `Llama3CausalLM`	Llama → `generic_model`; Llama3 HF JSON
Mistral / Mixtral	`MistralCausalLM` / `MixtralCausalLM`	Mixtral is MoE
Qwen / Qwen3 (+ MoE)	`QwenCausalLM` / `Qwen3CausalLM` / `*MoeCausalLM`	HF JSON
Phi3	`Phi3CausalLM`
PaliGemma	`PaliGemmaCausalLM`	`vit_encoder`, single-image input; smoke-tested only — multimodal numeric parity blocked by a known cache-sizing defect (§10.3)
GPT2, Bloom, GPT-NeoX, OPT, GPT-OSS, SmolLM3	`*CausalLM`	HF JSON (auto-converted); text-only, `generic_model`
Qwen3.5	`Qwen3_5CausalLM`	`LlmModelType` mapped (`Qwen3_5Spec` → `qwen3`); export fails fast with a documented `ValueError` — see Limitations for why

Export-blocked (tracked). Two families are in scope but format="litertlm"
export currently fails on a tracked bug (not a design rejection). Each is marked
@pytest.mark.xfail(strict=True) in its *_causal_lm_test.py with the cause
below; the label flips to a passing test once the underlying issue is fixed.

Family	Class	Whose bug	Cause (matches the test-file `xfail` reason)
Falcon	`FalconCausalLM`	keras-hub (fix planned as a separate issue + PR, not this one)	torch.export tracing failure: `torch.fx._ModuleNotInstalledAsSubmoduleError` in `FalconAttention` (module not registered as a submodule during tracing); under diagnosis
Gemma3n	`Gemma3nCausalLM`	upstream `litert-torch` converter	conversion fails on an aten.view shape/stride mismatch (`[3,8,8,16]` strides `(1024,8,1,64)` → `[192,16]`) during forward_prefill decomposition; pre-existing converter gap, under diagnosis

Unsupported, with the exact failure surfaced:

Family	Class	Reason
RWKV7	`RWKV7CausalLM`	rejected by tokenizer-class validation, not an architecture check — `RWKV7CausalLM` does implement `call_with_cache`, but its `RWKVTokenizer` is neither a `SentencePieceTokenizer` nor a `BytePairTokenizer` subclass, so export hits the same "unsupported tokenizer class" `ValueError` any such model would
PARSeq	`PARSeqCausalLM`	rejected by tokenizer-class validation, not an architecture or task check — `PARSeqCausalLM` implements `call_with_cache`, but `PARSeqTokenizer` subclasses the base `Tokenizer` directly (neither `SentencePieceTokenizer` nor `BytePairTokenizer`), so export hits the same "unsupported tokenizer class" `ValueError` any such model would (`parseq_causal_lm_test.py::test_litertlm_export_unsupported_tokenizer`)

Fail-fast errors (raised before heavy tracing where possible):

Condition	Exception
Backend not PyTorch	`ValueError`
Path not `*.litertlm`	`ValueError`
Model has no `call_with_cache`	`ValueError`
Unsupported KV-cache structure (`cache_structure != "single_stacked"`; e.g. Qwen3.5, see Limitations)	`ValueError`
`num_kv_heads`/`head_dim`/`num_layers`/`cache_length` not determinable from backbone attrs	`ValueError`
`litert-torch`/`litert-lm-builder` missing	`ImportError`/`RuntimeError` (tests skip)
`quant_config` not a `QuantConfig`	`ValueError`
Invalid `backend_constraint`	`ValueError` (`cpu`/`gpu`/`npu`/`gpu_artisan`)
`hf_tokenizer_path` not an existing `.json`	`ValueError`
`hf_tokenizer_path` vocabulary looks grossly incompatible with the model (heuristic, §6.5)	`ValueError`
Unsupported tokenizer class	`ValueError`
`separate_vision_encoder=True` without an encoder	`ValueError`
`separate_vision_encoder=True` for Gemma3n	`ValueError`
Vision model with bucketing	`ValueError` (need `prefill_seq_len == cache_length`, §6.6)
`prefill_seq_len > cache_length`	`ValueError`
BF16 compute dtype	`UserWarning` (untested; convert to FP16)

7. Design Decisions & Rationale

Each decision below answers a constraint from §4;
the trade-off is stated so reviewers can weigh it.

Decision	Why	Trade-off
Adapter translates stacked → flat KV cache (R1)	Matches the runtime's per-layer tensor contract without changing KerasHub's cache model.	Many bound tensors per decode step, plus an estimated ~4x logical-cache-size copy cost via `torch.stack` (static graph-node count, not profiled — §6.1); AUX split-cache is future work.
CPU-only tracing (R3)	Keeps device constants out of the graph → one portable bundle for CPU/GPU/NPU.	Export ignores the local accelerator; a process-global default-device swap needs a lock (§6.1).
Two signatures: prefill + decode (R2)	Mirrors how the runtime actually generates; lets prefill skip logits and decode skip prompt work.	Two traces to maintain; bucketing multiplies prefill traces.
Bucketing is text-only (R4)	The vision-attention-mask computation requires `cache_length == input length`, enforced for every vision-capable family, not just Gemma3 (§6.6).	Every prefill on a vision-capable bundle runs at full `cache_length` — even text-only use of a vision-capable model loses the bucketing TTFT win (§6.6).
Separate vision encoder is opt-in	Baked-in is simplest and correct for one-shot use; separation helps multi-turn/image reuse.	Opt-in surface; unsupported for Gemma3n.
Bundle HF `tokenizer.json`, never lossy-convert (R5)	Preserves exact token identity (~1% mismatch otherwise).	HF runtime features (constrained decoding, streaming) are narrower today.
PyTorch-backend-only export (R6)	`litert_torch` consumes PyTorch modules.	JAX/TF users must switch backend to export; trade-off with a TF-converter path discussed in §8.2.
Traceability patches scoped to export (R4)	Makes non-traceable Keras ops exportable without changing eager/training behavior.	Export-time monkeypatching; restored in `finally`; why not upstreamed: §8.3.
In-conversion + documented post-export quant	Covers both one-call quant and external `ai-edge-quantizer` flows.	Two quant stories to document.

8. Alternatives Considered

The rows below are grouped in two ways: alternatives with a short, settled
rejection reason, and three alternatives substantial enough to warrant a real
cost/benefit discussion.

Alternative	Why not sufficient
Status quo — `format="litert"` only	Produces a bare `.tflite`; every team re-writes tokenizer, KV-cache, prefill/decode, sampling, and chat glue per model (see §2.1).
Export raw TFLite + a separate how-to	Still leaves cache management, tokenization, and chat templating unsolved on device.
Lossy BPE → SentencePiece conversion	~1% token-ID mismatch breaks output fidelity; rejected in favor of bundling HF JSON (R5).
Build a Keras-specific mobile SDK	Fragments the Google AI Edge stack; LiteRT-LM already runs on-device LLMs.
Single max-length prefill (no buckets)	Pads every prompt to the maximum → wasted TTFT; bucketing recovers it for text models (§6.2).
Dynamic-shape export (no static signatures)	Not a design trade-off but a hard technical wall: `torch.export`/`litert_torch` require static shapes; dynamic shapes produce unbacked-symbol failures.

8.1 Weight-mapping onto litert_torch's own model implementations

KerasHub already does the mirror-image of this for loading pretrained
checkpoints: keras_hub/src/utils/transformers/convert_gemma.py-style modules
map named HF safetensors weight tensors onto KerasHub variables (per-layer
loader.port_weight(keras_variable=..., hf_weight_key=..., hook_fn=...)
calls, with small reshape/transpose hooks reconciling layout differences) —
roughly 100-200 lines per family. The alternative here would be the same kind
of mapping in the opposite direction: instead of adapting KerasHub's own
model objects via a translation adapter, map their trained weights onto
litert_torch's own authored implementation of each architecture (if one
exists), then export that.

This repo's only integration points with litert_torch are its tracing API
(litert_torch.signature(...), litert_torch.convert()) and its
quantization API (litert_torch.generative.quantize.*). Upstream
ai-edge-torch's generative package does ship hand-authored,
export-optimized implementations for several of the same families (e.g.
Gemma and Llama variants), so such a target exists for a subset of KerasHub's
family list. Either
way, the directional trade-off still holds: a weight-mapping approach requires
maintaining a second per-architecture mapping (comparable in shape to the
convert_gemma.py pattern above) that has to track both sides' internal
weight naming, for every family, versus this PR's approach of adapting the
same model objects users already load and train with — so a new family gets
export "for free" once it has a LiteRTLMExportSpec entry
(§6.8), rather than a second hand-authored
implementation to keep in sync.

8.2 The TF-backend path (`tf.function` static signatures → TFLiteConverter)

This exporter targets the PyTorch backend because the chosen upstream
converter, litert_torch, consumes PyTorch modules — that is the whole of
the hard constraint (the backend guard in export.py states: "LiteRT-LM
export is only supported with the PyTorch backend.").

A tf.function → TFLiteConverter path would have one real advantage: it
might have avoided the eight traceable-ops patches this PR requires for
torch.export (§6.3). Against that, it would still
need its own implementation of the KV-cache flattening/adapter logic
(§6.1), and converter investment
for new model work — including LiteRT-LM's own tooling — is concentrated on
the MLIR-based JAX/PyTorch paths industry-wide, with the legacy TF-based
TFLiteConverter in maintenance mode. On balance the PyTorch path tracks
upstream's own direction.

8.3 Upstreaming the traceable-ops fixes into Keras's torch backend

traceable_ops.py documents why each op needs a replacement (e.g.
torch.nn.functional.scaled_dot_product_attention lowers to fused ATen ops
litert_torch cannot translate; torch.nn.functional.one_hot inserts
runtime assertions that become ops litert_torch cannot lower). The fixes
ship as export-scoped monkeypatches rather than permanent changes to Keras's
torch backend for one decisive reason: at least one patch is a deliberate
behavior narrowing, not a pure bugfix — dot_product_attention's
replacement does not perform the implicit grouped-query-attention broadcast
the original might, a real behavior change for any caller relying on that
broadcast (Limitations). Making that a shared Keras
default would hit every Keras torch-backend user (training and inference),
not just this export path, and would need cross-repo release coordination
between keras-hub and keras-team/keras. An export-scoped patch, restored in
finally, confines the change to the tracing window.

10. Limitations

9. Testing Strategy

Test coverage for this feature is uneven across families and coverage types.
This table is the one authoritative summary; other sections link here
instead of restating test facts.

Coverage	Scope	Detail
Per-family export smoke test	All 25 `CausalLM` families (bloom, falcon, gemma, gemma3, gemma3n, gemma4, gemma4_assistant, gpt2, gpt_neo_x, gpt_oss, llama, llama3, mistral, mixtral, opt, pali_gemma, parseq, phi3, qwen, qwen3, qwen3_5, qwen3_moe, qwen_moe, rwkv7, smollm3)	Export succeeds (or fails with the documented error for unsupported or export-blocked families); output file exists with correct signature/metadata shape. `verify_numerics=False` for every multimodal family except Gemma4 (which runs `verify_multimodal_numerics=True` but is currently a strict `xfail` — see the next row and §10.3).
Numeric parity — prefill KV, decode logits, decode KV	Plain text Gemma; Gemma3 vision (baked-in)	`atol=rtol=1e-4` (`export_test.py`); measured diff for the Gemma3-vision case is closer to `1e-6`. Gemma4's end-to-end multimodal parity test (`gemma4_causal_lm_test.py::test_litertlm_export`) exists and is non-vacuous, but is currently a strict `xfail` — root-caused to the harness's no-delegate builtin interpreter, not the export (§10.3).
Numeric parity — separate vision encoder	Gemma3 (`separate_vision_encoder=True`)	Encoder-feature parity at `1e-4`; end-to-end (encoder → adapter → prefill) KV parity at `1e-3`.
Op-level parity (patched op vs. original, on eager tensors — not a `torch.export` test)	All 8 traceable-ops patches: `slice`, `dot_product_attention`, `one_hot`, `repeat`, `arange`, `take`, `scatter_update`, `amax`	`traceable_ops_test.py`; verifies the patch itself in isolation, not the full export pipeline.
Runtime-generation smoke test (bundle loads in the LiteRT-LM engine and produces non-empty, multi-token output)	4 hand-written tests: tiny random-weight Gemma, GPT2 (auto HF tokenizer), Llama3 (auto HF tokenizer), Qwen3 (auto HF tokenizer) — all text-only	Each calls `_verify_litertlm_generation` directly; no test anywhere passes `verify_generation=True` to the standard `run_litertlm_export_test` harness. Correctness of the generated text is not checked (weights are random) — only that the engine produces some non-empty output.
Registry/spec correctness	Model-type mapping, chat-stop-token lookups, cache-structure descriptions	`model_specs_test.py`; unit-level, no export or tracing involved.

Not covered today:

Passing numeric parity for Gemma4, Gemma3n, or PaliGemma (the remaining
multimodal families) — Gemma4's end-to-end parity test is a strict xfail
on the merged-vision/audio-position divergence (§10.3);
PaliGemma has only smoke/shape tests (its dedicated multimodal-numerics
test is a strict xfail on a cache-sizing defect,
§10.3); Gemma3n's export test is itself a strict xfail
(§6.8).
Systematic multi-step runtime generation for any family through the
standard test harness — verify_generation=True has zero callers; the 4
direct-call smoke tests above are the only exercise of the real runtime
engine, and none of them are multimodal or part of the per-family test
matrix.
GPU/NPU delegate execution — only backend_constraint string
validation/normalization is tested, not actual GPU/NPU inference.
Quantized-bundle numeric parity — quant recipes are exercised for
shape/success, not for accuracy versus the unquantized model.

10.1 Export-time

Limitation	Detail
PyTorch backend only	`keras.config.backend() == "torch"`; JAX/TF unsupported. See §8.2 for the trade-off with a TF-converter path.
Non-SentencePiece tokenizers	Auto-serialization (§6.5) covers all `BytePairTokenizer` subclasses; other classes need `hf_tokenizer_path`, which is now sanity-checked against the model's vocabulary size (heuristic, §6.5) rather than only checked for existing as a `.json` file.
Fixed prefill length per signature	Mitigated by bucketing for text; every vision-capable family requires `prefill_seq_len == cache_length` for every bucket, including text-only use of a vision-capable model — see §6.6 for the mechanism and consequence.
Batch size fixed at 1	Every traced signature (`prefill`, `decode`, and any vision/audio encoder) is hardcoded to batch size 1; `CausalLM.export()` exposes no `batch_size` kwarg for `format="litertlm"`. Batched on-device inference is not supported.
Large unquantized bundles	Gemma4 2B FP32 with baked-in encoders is observed > 20 GB. Use `separate_vision_encoder=True` + quantization.
Architecture coverage	Standard KV-cache transformers only; RWKV incompatible.
Qwen3.5 hybrid-cache export	Its hybrid full-attention/linear-attention layers use a dual cache (attention KV + linear-attention conv/recurrent state) the adapter's single stacked-KV-tensor `cache_structure` can't represent. `export_to_litertlm` raises a documented `ValueError` before any tracing — via `LiteRTLMExportSpec.describe_unsupported_cache_structure()`, generalized so any future non-`"single_stacked"` family gets a correct message without another `export.py` branch. Tested directly (`test_litertlm_export` in `qwen3_5_causal_lm_test.py`), not `xfail`. Blocked on adapter support for hybrid attention+conv cache layouts; see Future Work.
`dot_product_attention` patch drops implicit GQA broadcast	The traceability patch (§6.3) expects query/key/value head counts to already match (or be broadcast-compatible) before the call, unlike an implicit grouped-query-attention broadcast the original op might perform. Every current KerasHub attention layer already repeats KV heads up to the query head count before calling this op, so there is no observed behavior change today — but this is a permanent, intentional divergence, not a planned fix; a future attention layer relying on an implicit in-op broadcast would silently get wrong (shape-broadcast) results.
KV-cache stack/unstack copy cost	`stack_kv_cache` has non-zero copy cost from its `torch.stack` calls, plus further copying inside Keras's own `call_with_cache` cache-update mechanism, independent of this exporter. See §6.1 for the estimate (a static graph-node count, not a profiled measurement); fixing it needs a different signature contract or Keras cache-internals changes, both larger than this PR's scope.
Chat-turn stop tokens: residual gap	Family-aware for Gemma (+Gemma3/3n/4/PaliGemma), Llama3, Qwen/Qwen3 (+MoE/3.5), and Phi3 (§6.4) — but base Llama (v1/v2) and every other registered or unregistered family (Mistral, Mixtral, GPT2, Bloom, GPT-NeoX, GPT-OSS, OPT, Falcon, SmolLM3) still get only the tokenizer's primary EOS as a stop token. GPT-OSS (its chat turn boundary is a judgment call inside OpenAI's harmony format) and SmolLM3 (its keras-hub tokenizer never registers `<
BF16 compute dtype	Not validated; the exporter emits a `UserWarning` when detected. Convert to FP16 first.
Export time	~linear in signature count (buckets + decode) and model depth; large vocab adds constant serialization cost; sequence length has little effect.
No multi-step-decode/runtime-generation CI coverage	See Testing Strategy — no per-family or multimodal test exercises multi-step runtime generation through the standard harness.

10.2 Runtime (LiteRT-LM)

Limitation	Detail
Emulator unsupported	Fails on x86_64 emulators; physical ARM64 required.
CPU-only validated	`Backend.CPU()` works; GPU/NPU delegates exist but aren't validated across models.
No APK asset bundling	275 MB–1 GB+ exceeds base-module limits; push/download at runtime.
Chat template baked	App controls content (roles, system prompt, history); runtime owns the template string (by `LlmModelType`).
First-load latency	Engine init compiles the TFLite graph on first run.
Over-length prompts	Aborts if no bucket fits; no graceful truncation.
HF tokenizer gaps	Constrained decoding is SentencePiece-only; `HF_Tokenizer_Zlib` streaming load not yet implemented.
`GREEDY` sampler type unimplemented	litertlm-android 0.13.1 and host `litert_lm` 0.13.1 do not implement proto sampler type 3 (`GREEDY`). Bundles exported with `GREEDY_SAMPLER_CONFIG` therefore emit `TOP_K` with `k=1` instead. Pass an explicit `SamplerConfig` when sampling behavior must differ from the runtime default.

10.3 Multimodal

Default is baked-in encoders → a single PREFILL_DECODE model.
separate_vision_encoder=True (Gemma3/Gemma4/PaliGemma) emits
VISION_ENCODER/VISION_ADAPTER and has PREFILL_DECODE consume embeddings.
Gemma3n separate vision export is unsupported (MobileNetV5 projection is inside
the backbone, see §6.6).
Audio stays baked-in — no separate-audio path exists; the rationale is stated
once in §6.6 (upstream-gated).
On-device consumption contracts (measured on litertlm-android 0.13.1, Pixel
9). Device verification of dummy bundles surfaced these hard runtime
constraints:
- The runtime drives image input only through a separate
  VISION_ENCODER section. A baked-in bundle (vision inside
  PREFILL_DECODE) cannot be fed an image on-device at all
  (NOT_FOUND: TF_LITE_VISION_ENCODER). Export with
  separate_vision_encoder=True for any image use case today. The Gemma3
  and PaliGemma separate-vision dummy bundles run image input end-to-end
  on-device (stable across repeated runs).
- The runtime calls the vision encoder once per image and requires its
  input tensor to be 3- or 4-D raw pixels; the exported signatures are
  single-image [B, H, W, 3] (encoder) and [B, tokens, dim] (adapter)
  to match.
- Gemma4's patch_values style is not consumable on-device on 0.13.1:
  the runtime feeds raw images only — it cannot produce pixel_values /
  pixel_position_ids — and its channel check (3 or 4) rejects the
  patch-shaped input. Gemma4 vision is baked-in-or-nothing until the
  runtime grows a patch-input contract.
- Audio input requires a separate TF_LITE_AUDIO_ENCODER_HW section
  (NOT_FOUND otherwise), which this exporter deliberately does not
  produce — so audio cannot be exercised on-device today regardless of
  bundle shape (§6.6 for why).
Numeric parity, per family (full matrix: Testing Strategy).
Gemma3 is the multimodal family with passing numeric-parity coverage,
alongside plain text Gemma: Gemma3 at 1e-4 (baked-in and
separate-encoder). Gemma4's end-to-end multimodal parity test is a strict
xfail — and the failure is not in the export. The exported graph is
semantically correct (torch.export matches Keras eager bit-for-bit at
every position); the no-delegate builtin TFLite interpreter this harness
deliberately uses mis-executes the vision encoder's unflatten RESHAPE
region under its default memory planner (a buffer is recycled while still
live). The same bundle and inputs match eager to 9.1e-06 under XNNPACK
and to 2.4e-07 under the builtin interpreter with
experimental_preserve_all_tensors=True — an upstream ai_edge_litert
memory-planning bug, not a keras-hub export defect.
PaliGemma's export test is smoke/shape-level only
(verify_numerics=False). Gemma3n's export test is a strict xfail on an
upstream converter bug (§6.8) — its
earlier int64/int32 MLIR failure is fixed (Appendix),
but conversion still fails downstream of tracing.
PaliGemma numeric parity is blocked by a real cache-sizing defect,
isolated by a dedicated strict-xfail test
(test_litertlm_export_multimodal_numerics): call_with_cache concatenates
image + text embeddings (e.g. 16 text + 16 image = 32 positions), but the
bundle exports cache_length = prefill_seq_len = text length only, so a
real multimodal prompt overflows the cache. Basic PaliGemma export still
works and stays green — this is a disclosed limitation, not an export
blocker.
PaliGemma vision inputs are wired into the graph (images →
vit_encoder → embeddings), but PaliGemmaSpec populates no LlmMetadata
vision fields — there is no dedicated LlmModelType vision subtype for
PaliGemma yet, so populate_vision_metadata is a documented no-op
(§6.8). The runtime may not have enough
metadata to actually drive images through a PaliGemma bundle end-to-end.

10.4 On-device performance and memory (Pixel 9)

Measured on a Pixel 9 (Tensor G4, 11.5 GB RAM) with
litertlm-android 0.13.1, CPU backend, battery ≥ 50 %, thermal status
NONE/LIGHT. Preset: gemma3_instruct_270m, exported from this PR's branch at
87785afd. All runs use an explicit greedy sampler (top_k=1, top_p=1.0,
temperature=0.0, seed=0) because the embedded GREEDY_SAMPLER_CONFIG
sampler type is not implemented by the 0.13.1 runtime. maxTokens=384.
Cold runs use pm clear to force a fresh TFLite graph compile; warm runs
reuse the process. Numbers are medians of 5 warm runs; TTFT is engine-reported
time-to-first-token, decode tok/s is engine-reported sustained decode
throughput. RSS is resident set size; PSS is proportional set size. Raw
per-run JSON is retained and can be attached on request.

Bundle	Prompt tokens	Run	Init ms	TTFT ms	Decode tok/s	Peak RSS MB	PSS end MB
fp32_default	32	cold	4105	3778	11.2	3264	3085
fp32_default	32	warm	2275	2765	12.2	3302	3178
fp32_default	64	cold	2049	2709	11.6	3266	3146
fp32_default	64	warm	162	2647	12.4	1649	1524
fp32_default	256	cold	2057	2883	12.6	3267	3147
fp32_default	256	warm	149	2696	12.2	1654	1528
int8_drq	32	cold	1500	2915	14.7	1267	1149
int8_drq	32	warm	1009	2278	16.4	1307	1181
int8_drq	64	cold	839	2593	18.4	1273	1148
int8_drq	64	warm	203	2644	16.4	898	768
int8_drq	256	cold	1736	3349	13.8	1274	1147
int8_drq	256	warm	1085	2753	15.2	1273	1181
bucketed	32	cold	3933	2004	8.8	2986	2794
bucketed	32	warm	396	1066	9.4	1669	1545
bucketed	64	cold	2944	1362	9.2	3284	3164
bucketed	64	warm	307	990	9.9	1668	1542
bucketed	256	cold	2351	4200	9.9	3291	3163
bucketed	256	warm	259	3491	9.6	1672	1544
prefill256	32	cold	4304	6243	6.6	2859	2762
prefill256	32	warm	3354	4781	6.9	3272	2897
prefill256	64	cold	3113	5767	6.4	3270	3143
prefill256	64	warm	225	2952	10.3	1645	1525
prefill256	256	cold	2008	4822	9.9	3268	3144
prefill256	256	warm	4470	4519	8.5	3270	3163

Observations:

INT8-DRQ (dynamic-range quantization) is the best production trade-off: ~15–18 decode tok/s and ~1.2
GB peak RSS, roughly half the memory of FP32.
FP32 reaches ~11–12 decode tok/s at ~3.2 GB peak RSS.
Bucketing wins TTFT at short prompts (32/64 tokens: ~1.0–1.4 s vs
~2.6–3.7 s fixed), at the cost of slightly lower decode throughput.
Fixed prefill_seq_len=256 is pathological for short prompts: TTFT at
32/64 tokens is ~4.8–6.2 s because it still runs a full 256-token prefill.
Cold-start graph compilation adds ~1–4 s to init; warm init is ~150–1100 ms
when memory pressure allows reuse.
Some warm-init outliers (e.g., fp32 pt32 warm 2275 ms, prefill256 pt256 warm
4470 ms) indicate memory pressure or thermal throttling; treat the upper end
of warm-init variance as environment-dependent.

11. Future Work

Explicitly out of scope now, natural next steps:

Arbitrary non-BytePair tokenizers. Write a WordPiece/Unigram → HF
tokenizer.json converter. Unlike the existing BytePair path
(§6.5), which just serializes the HF
tokenizers.Tokenizer object BytePairTokenizer already wraps, KerasHub's
WordPieceTokenizer (built on tensorflow_text) and SentencePieceTokenizer
(wraps the separate sentencepiece library's own .model proto format) do
not wrap an HF tokenizers.Tokenizer object at all. This would be a
from-scratch converter (vocab + normalization/pre-tokenizer rules → the
tokenizer.json schema), not an extension of hf_tokenizer_converter.py's
existing mechanism.
Close HF-tokenizer runtime gaps — constrained decoding + streaming load
for HF_Tokenizer_Zlib.
Split cache + externalized embedder — adopt upstream AUX/EMBEDDER
models when APIs stabilize (cuts per-step tensor binding and the
KV-cache stack/copy cost estimated in
§6.1).
Separate audio encoder export — blocked only on a missing upstream
reference contract for the audio bundle slots, not on tracing
(§6.6). Building one without that contract risks a
bundle that traces but the runtime can't correctly consume.
Gemma3n separate vision encoder — once MobileNetV5 projection can export
independently.
QAT checkpoint handling — fail fast when quant_method=gemma is detected.
GPU/NPU delegate validation — only Backend.CPU() is validated
end-to-end today.
Externalized weight packaging — for very large unquantized models.
INT4 on-device compatibility — investigate CPU-delegate GATHER_ND
failures for INT4 weight-only bundles.
Upstream litert_torch speedups — stable JAX-bridge cache key for
per-op lowerings, guarded FX recompiles, cross-signature lowering cache (cuts
the multiplicative export cost).
Hybrid attention+conv cache support — extend the adapter beyond the
single stacked-KV-tensor cache_structure so families like Qwen3.5
(Limitations) can export. Upstream litert_torch's HF
path (core/cache.py) demonstrates a flat-tensor naming convention for
exactly this shape (k_i/v_i for full-attention layers, c_i/r_i
for linear-attention layers) — a useful reference design, though not
reusable code, since it lives in litert_torch's HF-transformers pipeline
rather than the PyTorch signature/convert API this exporter calls.
Extend chat-turn-stop-token overrides to the remaining registered/
unregistered families (base Llama, Mistral, Mixtral, GPT2, Bloom,
GPT-NeoX, GPT-OSS, OPT, Falcon, SmolLM3) if any of them use a distinct
chat-turn-boundary token in practice (§10.1). GPT-OSS
and SmolLM3 are the concrete known candidates (deferred on, respectively,
OpenAI's harmony chat-format semantics and a missing tokenizer
registration).
prefer_activation_type bundle hint — LitertLmFileBuilder.add_tflite_model
accepts a per-section activation-type hint (e.g. to prefer FP16
activations on a delegate) that the exporter does not currently pass.
Wire it through export_to_litertlm if a delegate-tuning use case needs
it; no correctness stake today (§6.4).

12. Usage & Deployment

This exporter produces the .litertlm file; consuming it on-device is owned
by the LiteRT-LM runtime and its own documentation. This section covers only
what's needed to sanity-check a bundle end-to-end, not a full Android
integration guide — see the LiteRT-LM project's own docs for Gradle
configuration, the full Kotlin API, and Play Asset Delivery setup.

12.1 Python export

Requires the litertlm extra for the export-time dependencies:
pip install keras-hub[litertlm] (or litert-torch/litert-lm-builder
directly). These are not part of keras-hub's unconditional
requirements.txt — export is PyTorch-only, so CI installs the extra only for
the torch-backend matrix leg.

import keras_hub

model = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_1b")
model.export(
    "model.litertlm",
    format="litertlm",
    prefill_seq_len=[32, 64, 128],   # buckets: runtime picks smallest fit
)

12.2 Runtime integration (summary)

Dependency. One runtime library on-device
(com.google.ai.edge.litertlm:litertlm-android); no Keras/PyTorch/HF needed
at inference time.
Hardware. ARM64 only (arm64-v8a); LiteRT-LM ships no x86_64 native
libs, so emulators are unsupported.
Packaging. .litertlm files are memory-mapped and must not be
compressed by the app's asset pipeline (Gradle's noCompress setting or
equivalent).
Delivery. Bundles are 275 MB–1 GB+ (or > 20 GB unquantized for large
multimodal models), far above the Play base-module limit (~150 MB) — they
cannot ship as ordinary in-APK assets. Apps download to app storage on
first launch or use Play Asset Delivery/a CDN; for dev/CI, adb push to
/data/local/tmp and load by absolute path works directly.

Minimal consumption example (illustrative, not authoritative — see
LiteRT-LM's own API docs for the full surface):

val engine = Engine(EngineConfig(modelPath, backend = Backend.CPU()))
engine.initialize()
val conversation = engine.createConversation()
conversation.sendMessageAsync("Hello").collect { print(it) }

Backend status. Backend.CPU() is the only backend validated against
exported bundles; Backend.GPU()/Backend.NPU() exist in the runtime API
but are not yet validated across model families (see
Limitations).
Chat templating. The runtime applies the template by LlmModelType; the
app supplies content (roles, system prompt, history), not template syntax.
Bucketing payoff. See §6.2 for the
measured TTFT improvement from prefill bucketing on a real device.

13. Appendix: API, File Layout & Glossary

13.1 `CausalLM.export()` reference for `format="litertlm"`

model.export(
    "model.litertlm",
    format="litertlm",
    prefill_seq_len=128,
    cache_length=None,
    backend_constraint=None,
    quant_config=quant_config,
    separate_vision_encoder=False,
    hf_tokenizer_path="/path/to/tokenizer.json",
    sampler_config=None,
    llm_model_type=None,
)

Kwarg	Type	Description
`cache_length`	optional `int`	The model's maximum context window (KV-cache length) to export with. Defaults to `backbone.max_sequence_length` if the backbone defines it, else `preprocessor.sequence_length` — the latter emits a `UserWarning`, since most backbones besides the GPT2 family don't define `max_sequence_length`, so that fallback is a tokenization default, not necessarily the real context window.
`prefill_seq_len`	`int` or `list[int]`	Prefill length(s) to trace. Defaults to `cache_length`. Each must be a positive int ≤ `cache_length`. A list traces one prefill signature per bucket; vision-capable models require every value to equal `cache_length` (§6.6).
`backend_constraint`	optional `str`	`"cpu"`, `"gpu"`, `"npu"`, or `"gpu_artisan"` (case-insensitive; normalized before use, §6.3). Applied to every TFLite model in the bundle. Default `None` — no constraint is passed to the builder.
`quant_config`	optional `QuantConfig`	`litert_torch.quantize.quant_config.QuantConfig` for in-conversion quantization.
`separate_vision_encoder`	`bool`	If `True` and the model has a vision encoder, export separate `VISION_ENCODER`/`VISION_ADAPTER` models; `PREFILL_DECODE` consumes embeddings. Default `False`. Either way the bundle is a single, complete multimodal `.litertlm` — `PREFILL_DECODE` always consumes text tokens too; this flag only controls whether vision encoding is baked into that same trace or factored into separate, reusable models. It never produces a vision-only export.
`hf_tokenizer_path`	optional `str`	User HF `tokenizer.json`; skips auto-detection, bundled via `add_hf_tokenizer()`. Sanity-checked against the model's vocabulary size (§6.5).
`sampler_config`	optional `SamplerConfig`	When given, writes `LlmMetadata.sampler_params` (§6.4); omitted by default so the runtime picks its own sampling policy. `GREEDY_SAMPLER_CONFIG` (`top_k=1`, in `model_specs.py`) is the one named preset.
`llm_model_type`	optional `str`	Explicit model-type override for presets indistinguishable by class — currently only `"function_gemma"` (the `function_gemma_instruct_270m` preset, which loads as a plain `Gemma3CausalLM`, §6.4). Unknown values raise `ValueError`. Default: auto-detect — by class for every other family, and by tokenizer function-calling tokens for `function_gemma` itself. Note: the explicit kwarg sets the bare `llm_model_type` oneof only and skips model-specific metadata population (mirroring litert-torch's override semantics); auto-detection is what fills the function-calling block.

Batch size is fixed at 1 for every signature; there is no batch_size kwarg
(Limitations).

Extra **kwargs forward to litert_torch.signature(...). Unknown kwargs may
fail inside litert_torch with less explicit messages — prefer the documented
arguments.

13.2 Core files

File	Role
`causal_lm.py`	`CausalLM.export()` override; routes to `export_to_litertlm()` when `format="litertlm"`.
`export.py`	Main pipeline: validate, extract config, build sample inputs, trace, convert, materialize tokenizer, build metadata, package.
`model_specs.py`	`LiteRTLMExportSpec` registry: one class per model family (model type, cache layout, `vision_input_style`, KV-cache stack/unstack, chat-stop-token ids, and the multimodal adapter hooks), resolved once per export via `resolve_export_spec()`.
`adapter.py`	PyTorch `nn.Module` adapter (stacked ⇄ flat KV cache, delegating to the resolved spec's `stack_kv_cache`/`unstack_kv_cache`): `KerasHubLiteRTAdapter` plus the vision-only adapters.
`traceable_ops.py`	Traceable replacements for the eight non-`torch.export`-friendly Keras torch ops, plus the combined `traceable_ops_scope()`.
`traceable_ops_test.py`	Unit parity tests: each patched op vs. the original Keras torch-backend implementation on eager tensors (no `torch.export` involved).
`hf_tokenizer_converter.py`	Serializes the HF `tokenizers.Tokenizer` object `BytePairTokenizer` already wraps into `tokenizer.json` — not a format conversion, see §6.5.

Supporting upstream model changes. A couple of upstream model files (not
part of the litertlm package, so not in the table above) needed small
changes to make export possible at all:

File	Change
`keras_hub/src/models/llama/llama_attention.py`	Restrict the fused `dot_product_attention` op to GPU/TPU via a `_use_fused_attention_op()` gate, matching the pattern `GemmaAttention`/`MixtralAttention` already used. This is a consistency fix, not a traceability one: `traceable_ops.py`'s `dot_product_attention` patch (§6.3) intercepts the op itself regardless of which branch calls it, so the fused branch would have traced fine too — Llama was simply the one family whose attention layer had not yet adopted the gate the other families already had.
`keras_hub/src/models/gemma3n/gemma3n_backbone.py`	Cast the audio padding-token constant to int32 (was int64), fixing an int64/int32 MLIR `func.call` operand-type mismatch during Gemma3n prefill lowering. This removed one of two Gemma3n blockers; a distinct, pre-existing upstream `litert-torch` aten.view defect still blocks Gemma3n conversion end-to-end (§6.8).

13.3 Adapter classes

Class	Role
`KerasHubLiteRTAdapter`	Root adapter: builds prefill/decode signatures and flat KV-cache I/O, delegating the stacked ⇄ flat translation to `self.export_spec.stack_kv_cache`/`unstack_kv_cache`.
`_PrefillAdapter` / `_DecodeAdapter`	Inner adapters (in `export.py`) giving `litert_torch` clean module boundaries per signature.
`KerasHubVisionEncoderAdapter`	Optional — traces the vision encoder as a separate `VISION_ENCODER` model.
`KerasHubVisionAdapter`	No-op rename of encoder `features` → `mm_embedding` so the `VISION_ADAPTER` signature matches the runtime contract.

Both vision adapters are single generic classes (adapter.py), instantiated
fresh per export() call around whatever model instance is passed in
(export.py: KerasHubVisionEncoderAdapter(model).eval()) — not one
subclass per model family. The per-family piece is vision_input_style on
the resolved LiteRTLMExportSpec, which tells this one generic adapter
which of its optional kwargs (images/pixel_values/pixel_position_ids)
to expect populated.

VISION_ENCODER and VISION_ADAPTER are two separate .tflite models
because the LiteRT-LM bundle format itself defines them as two distinct
model slots (TfLiteModelType.VISION_ENCODER/VISION_ADAPTER in
litert_lm_builder) — Google's own reference PyTorch exporter
(litert_torch/generative/export_hf/core/litert_lm_builder.py) already
populates the bundle the same way. This PR conforms to an existing bundle
contract; it does not introduce the split.

13.4 Glossary

Term	Meaning
GQA	Grouped-query attention.
ATen	PyTorch's low-level operator library.
TTFT	Time-to-first-token.
QAT	Quantization-aware training.
GATHER_ND	A TFLite operator.
HF_Tokenizer_Zlib	`.litertlm` archive section holding a compressed HuggingFace tokenizer.
Unbacked symbols	Dynamic shapes `torch.export` cannot prove bounded; a common export-failure cause.
`LlmModelType`	The `litert_lm_builder` protobuf enum selecting the runtime's chat template and multimodal subtype (`generic_model`, `gemma3`, `gemma3n`, `gemma4`, `qwen3`, `qwen2p5`, `function_gemma`, `fast_vlm`).
`cache_structure`	`LiteRTLMExportSpec` attribute describing the KV-cache tensor shape a family's `call_with_cache` expects; the adapter only supports `"single_stacked"` today (§10.1).
`vision_input_style`	`LiteRTLMExportSpec` capability field describing how a family expects vision input (`"raw_images"`, `"patch_values"`, or `"embedded_pixel_values"`); see §6.6.

google-cla · 2026-04-25T12:35:25Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist

Code Review

This pull request introduces the export_to_litertlm functionality for CausalLM models, enabling the creation of .litertlm bundles that include the model graph, SentencePiece tokenizer, and LLM metadata. The implementation includes a PyTorch adapter to handle KV-cache stacking and a traceable slice_update scope for torch.export compatibility. Additionally, the PR updates the test suite to include LiteRT-LM export verification and marks several model tests as xfail where upstream limitations currently exist. Review feedback highlighted the need to replace fragile string-based model type mapping with robust isinstance checks, avoid hardcoded defaults for cache length and data types, and document dependencies on internal LiteRT APIs.

pctablet505

Issue 1: Hardcoded Float32 Key-Value Caches
The current implementation initializes all internal KV cache tensors using torch.float32. In modern production environments, large language models run inference using float16 or bfloat16. Forcing a float32 cache causes an immediate, unnecessary 2x inflation in memory consumption for on-device context storage. It also introduces computational casting overhead and potential compilation friction between layers.
Recommendation: Update the adapter to dynamically match the KV cache data type to the underlying model precision. Extract the compute data type directly from the model attributes and fall back to float32 only if it is undefined.

Issue 2: Fragile Reliance on Internal LiteRT APIs
The pipeline imports directly from ai_edge_litert.internal.litertlm_builder. Referencing an explicit internal namespace introduces severe maintainability risks, as upstream updates to the Google AI Edge codebase can silently break this integration without warning.
Recommendation: Insert a clear code comment above this import acknowledging the dependency on internal APIs. Additionally, wrap the import in a robust try-except block to provide a clean, descriptive installation or compatibility error message if the internal module signature changes.

Issue 3: Missing Frontend Guardrails for Alternative Backends
Because this implementation leverages litert_torch, it is inherently bound to the PyTorch backend introduced in Keras 3.15. If an engineer attempts to execute this export process while running Keras on top of JAX or TensorFlow, the system will fail deep inside the adapter code with cryptic tracing errors.
Recommendation: Introduce an explicit backend check at the very beginning of the intercepted export function. If the current active backend is not PyTorch, raise an immediate ValueError stating that the LiteRT-LM format is currently restricted to PyTorch execution environments.

Issue 4: Global Side Effects from Backend Patching
The pull request applies a monkey-patch to the core Keras torch backend slice_update function to force the use of torch.index_copy_ during tracer execution. This solves the GuardOnDataDependentSymNode exceptions, but runtime monkey-patching of a core framework package introduces a risk of unpredictable side effects in separate components of a user's program.
Recommendation: If this slice_update adjustment is globally beneficial for PyTorch graph compilation, it should be isolated and submitted directly as an upstream PR to the core Keras library. If it must remain here, ensure it is safely localized or explicitly restored to its original state immediately after the model tracing pipeline completes.

pctablet505 · 2026-05-20T21:08:16Z

/gemini review

gemini-code-assist

Code Review

This pull request implements support for exporting KerasHub CausalLM models to LiteRT-LM bundles, including a PyTorch adapter and a context manager to patch Keras torch-backend operations for torch.export compatibility. The review feedback correctly identifies the need to remove the input_signature parameter for Keras 3 consistency, improve the robustness of patched slicing operations when handling scalar tensors, and extend model type detection to include Llama models for accurate metadata generation.

Adds to CausalLM, enabling export to the LiteRT-LM bundle format with prefill/decode signatures. Key components: - : PyTorch adapter wrapping CausalLM for litert_torch.signature() export. Handles KV cache stacking/unstacking and traceable forward_prefill/forward_decode. - : Full export pipeline using litert_torch.signature('prefill', ...).signature('decode', ...). Builds LlmMetadata protobuf with start/stop tokens and model type. - : Integration tests verifying end-to-end export + numerical correctness. - : Unit tests mocking litert_torch and builder dependencies. - : Torch backend LiteRT test support. Requirements: - NAME litert-torch SYNOPSIS litert-torch COMMAND COMMANDS COMMAND is one of the following: export_hf Exports HuggingFace Transformers model to tflite. added to requirements.txt. - PyTorch backend required (Keras >= 3.15). - Prefill returns only KV caches (no logits), per LiteRT-LM spec. Tested with Gemma tiny models on torch backend.

- causal_lm.py: override export() to intercept format='litertlm' instead of a separate export_to_litertlm() method. - export_test.py: use model.export(path, format='litertlm') in all tests. Remove mock-based backend and validation tests; only real integration tests remain (tiny gemma export + output matching). - task_test.py: no litertlm changes (mock tests removed entirely). - __init__.py: empty (export_to_litertlm is internal-only now). - requirements.txt: no litert-torch (handled as optional with clear error).

litert-torch is the only package users need to install; ai-edge-litert is pulled in automatically as a dependency. ai-edge-torch is not required.

- Replace ai_edge_litert.internal.litertlm_builder with litert_lm_builder - Replace ai_edge_litert.internal.llm_metadata_pb2 with litert_lm_builder.litertlm_builder.llm_metadata_pb2 - Replace ai_edge_litert.internal.litertlm_core with litert_lm_builder.litertlm_core in tests - Clean up model type mapping with _set_llm_model_type() helper - Add backend_constraint validation - Add comments clarifying add_tflite_model as the official public API

…, model dtype caches - _set_llm_model_type: use isinstance checks with lazy imports instead of fragile startswith on class names - _get_cache_config: raise ValueError instead of hardcoding 2048 fallback - _build_sample_inputs: accept dtype param, use model.compute_dtype for KV cache and mask tensors instead of hardcoded float32

For basic text-only CausalLM export, generic_model is sufficient. Specific LlmModelType values (gemma3, qwen3, etc.) are only needed for multimodal/function-calling features. If a subclass needs a specific type, it can override _build_llm_metadata.

- forward_prefill: use input_pos[0] as cache_update_index instead of hardcoded 0, so prefill appends to existing cache on subsequent turns. - _patched_slice_update: replace torch.scatter_ with loop-based index_copy_ to avoid XNNPack runtime crashes on Android. - _build_llm_metadata: detect gemma3 model type and add <end_of_turn> as a stop token so generation stops after the assistant's turn. - Export script: increase cache_length to 1024 for multi-turn support while keeping prefill_seq_len at 128 for fast prefill.

- prefill_seq_len now accepts int | list[int] for bucketing - quant_config is forwarded to litert_torch.convert() for in-conversion quantization - Add comprehensive docstrings for bucketing and quantization recipes - Add test_export_with_bucketing() to verify multi-signature export - Update CausalLM.export() docstring to document litertlm kwargs

- Detect vision_encoder / audio_encoder on model backbone and build corresponding sample inputs for prefill signature tracing. - Populate LlmMetadata vision fields for gemma3, gemma3n, and gemma4 (image tokens, patch size, pool size, tensor height/width). - Populate LlmMetadata audio fields for gemma3n and gemma4. - KerasHubLiteRTAdapter.forward_prefill now runs vision/audio encoders and passes embeddings to call_with_cache via dynamic kwargs. - _detect_llm_model_type: add lazy isinstance checks for Gemma4, Gemma3, Qwen3, Qwen2, Llama; fix Llama to return generic_model (protobuf has no llama field). - Add 6 export tests: backend constraints, multimodal bucketing, model type metadata, text-only signatures, tiny Gemma3 multimodal export, and Keras vs TFLite output parity. - Update CausalLM docstring to note multimodal bucketing restriction.

- adapter.py: Accept pixel_values/pixel_position_ids for Gemma4 vision encoder - export.py: Add _build_gemma4_vision_sample_inputs for patch-based inputs - export.py: Use audio_input_feat_size from config instead of hardcoded 128 These changes allow Gemma4 (vision+audio) models to export to .litertlm format when using KERAS_BACKEND=torch.

…ring Gemma3n's LiteRT-LM export fails in litert-torch conversion before any numeric verification can run: an aten.view shape/stride mismatch ([3,64,8] strides (512,1,64) -> [192,8]) during forward_prefill decomposition. This is a pre-existing export bug independent of this change, not a numeric-parity gap, so it blocks wiring verify_multimodal_numerics for this family until the export itself is fixed. Marked xfail(strict=True) so a future fix self-corrects the label instead of silently masking a regression.

… defect Adds test_litertlm_export_multimodal_numerics as a new, separate xfail(strict=True) test rather than touching the existing (green) test_litertlm_export, since the underlying issue is a latent export defect, not a harness gap: call_with_cache concatenates image+text embeddings to length 32 (text 16 + image_sequence_length 16), but the bundle is exported with cache_length = prefill_seq_len = 16, so the eager Keras reference overflows the undersized cache with an IndexError in CachedGemmaAttention before parity can even be measured. This means the exported bundle's vision-token handling is currently unverified and very likely wrong; tracked here rather than fixed, since resizing the export cache for concatenation-style vision families is an export-architecture change outside this task's scope.

…inimal-litertlm

…nimal-litertlm

The torch CI leg installs from requirements.txt with torch>=2.6.0 and torchvision>=0.16.0. pip resolves torch 2.13.0 from PyPI (CUDA build) but torchvision 0.28.0 only ships +cpu on download.pytorch.org, so the CPU-built torchvision::nms fake kernel does not match the installed torch ABI: RuntimeError: operator torchvision::nms does not exist which also breaks transformers' lazy Qwen2ForCausalLM import at collection time. Pin both to the official matched pair (torchvision 0.28.0 <-> torch 2.13.0).

…d branches

… citation Gemma4's exported audio path feeds a pre-extracted log-mel spectrogram straight into backbone.audio_encoder (Gemma4AudioConverter runs mel extraction on the host side), so the bundle must tell the LiteRT-LM runtime to perform mel extraction rather than skip it. Direction confirmed against the flag's runtime consumer, ai-edge-litert's AudioPreprocessorMiniAudio::Preprocess (support/preprocessor/audio_preprocessor_miniaudio.cc:351): the `!SkipMelSpectrogramExtraction()` branch runs STFT + log-mel extraction; the else branch emits framed raw PCM instead (confirmed by audio_preprocessor_miniaudio_test.cc's SkipMelSpectrogramExtraction test). Field docstring: audio_preprocessor.h:177. The two sources the task named first (compiled proto3 descriptor comments, litert-torch export_hf's Gemma4 metadata_builder.py) do not carry this information -- recorded honestly in the code comment. Also replaces the previously bare `del audio_cfg` in Gemma4Spec.populate_audio_metadata with a comment explaining the computed audio config is used only for trace-input sizing, since the installed litert-lm-builder 0.13.0 Gemma4 subtype has no proto field to carry those derived values. Gemma3n is out of scope: its subtype has no skip_mel field in the installed proto (independently re-verified this pass; token/ image-geometry fields only), matching the plan's rescope decision. Adds a metadata round-trip test asserting the field's value as a contract/regression guard. skip_mel_spectrogram_extraction is a proto3 scalar bool with no field presence, so False is also the wire default -- the test's purpose is catching a future flip to True, not proving serialization occurred; this is stated in both the code comment and the test docstring. Full litertlm regression gate re-run before and after on this exact tree: 74 passed -> 75 passed (exactly the one new test), 16 skipped, 3 xfailed, 72 subtests passed unchanged in both runs, xfail reason strings byte-identical. Multimodal numeric-parity magnitudes (gemma3/gemma4 prefill-KV and decode-logits max-abs-err) stayed in the same ~1e-6 noise band before and after, well under the 1e-4 threshold -- consistent with a metadata-only change touching no traced graph.

Mirrors litert-torch's export_hf module, which packs an optional eoi.tflite model producing get_input_embeddings()(eoi_token_ids) and adds it as an END_OF_VISION bundle section whenever present. KerasHub's own vision adapter is a plain rename and does not fold an end-of-image embedding into mm_embedding the way upstream's gemma3/gemma3n adapters do, so this adds a new KerasHubEndOfImageAdapter (adapter.py) that supplies the same embedding as its own input-less signature, gated by a new LiteRTLMExportSpec.end_of_vision_token field (model_specs.py) set for Gemma3/Gemma4 (reusing their existing end-of-image token constants; PaliGemma has no such token and stays opted out) and resolved to a real token id via the existing _lookup_token_id helper before export (get_end_of_vision_token_ids). export.py wires the new edge model through _trace_and_convert/_assemble_bundle only on the separate-vision-encoder path; every other export path is unaffected. Adds a bundle-section assertion test confirming the section, and its tflite embedding output, are now present where they were not before. This is an upstream-parity addition; whether the runtime consumes the section is a separate, later on-device question.

…sion_bucketing flag

…le_bundle

…ma3n rejection from it Adds LiteRTLMExportSpec.supports_separate_vision (default True, overridden to False on Gemma3nSpec) and switches the separate-vision-encoder rejection in export_to_litertlm from the implicit vision_input_style == "embedded_pixel_values" proxy to an explicit spec.supports_separate_vision consult. Makes the 4-family x {baked, separate} support matrix readable directly from each family's spec declaration in model_specs.py instead of a rejection branch buried in export.py. Behavior is unchanged: verified the old proxy condition and the new flag agree for every registered spec class. Adds registry-integrity tests locking the matrix (model_specs_test.py) and a new Gemma3n separate-vision rejection test (gemma3n_causal_lm_test.py) -- no such test previously existed in the tree.

…NFIG preset

…ext families Phi-3's chat template ends every turn with `<|end|>` (HF microsoft/Phi-3-mini-4k-instruct tokenizer_config.json), distinct from its primary EOS `<|endoftext|>` (Phi3Tokenizer's registered `end_token`). Add `Phi3Spec` + a registry entry so on-device chat generation can stop at a turn boundary; verified against the real phi3_test_vocab.spm fixture (`<|end|>`=18, `<|endoftext|>`=15). Mistral, Mixtral, base Llama, GPT2, Bloom, GPT-NeoX, and OPT get no spec override: their tokenizers register no distinct chat-turn-stop token (instruct templates for Mistral/Mixtral terminate on the primary EOS itself; the rest are pure base LMs with no chat template), so the existing EOS-only default is already correct. Documented in a comment instead of adding no-op subclasses. SmolLM3 and GPT-OSS are documented as deferred: SmolLM3's keras-hub tokenizer does not register HF's `<|im_end|>` in any vocabulary (including the test fixture), so an override would be an unverifiable no-op; GPT-OSS's harmony-format stop-token choice among `<|return|>`/`<|end|>`/`<|call|>` remains a genuine judgment trap independent of its now-unblocked export status. Falcon is untouched (export-blocked, tracked separately).

litertlm-android 0.13.1 and host litert_lm 0.13.1 do not implement sampler type 3 (GREEDY). Bundles exported with GREEDY_SAMPLER_CONFIG crashed at load time with 'Sampler type: 3 not implemented yet'. Emit TOP_K with k=1, which is functionally equivalent and is implemented by both runtimes. Update the SamplerConfig docstring and roundtrip test accordingly.

The function_gemma_instruct_270m preset loads as a plain Gemma3CausalLM but its Gemma3Tokenizer vocabulary contains the function-calling special tokens <start_function_call> and <end_function_call>. Detect these tokens in resolve_export_spec and auto-select FunctionGemmaSpec so users no longer need to pass llm_model_type='function_gemma' explicitly. Also populate the open_quote/close_quote fields with the preset's <escape> token, mirroring litert-torch's Gemma4 metadata builder.

The multimodal parity helper previously zero-filled vision_indices, vision_mask, audio_indices, audio_mask, and audio_mel_mask. Because Gemma3 and Gemma4 restore slot 0 to the text embedding, zero indices made the entire vision/audio merge path an identity, so the tests compared only the text path. Synthesize real placement indices (capped to the sequence length) and masks, add a sensitivity assertion proving that changing indices changes the KV cache, and compare decode-step KV as the text helper does. Gemma4's parity now fails with measured prefill-KV error ~2.76 when the vision/audio towers are actually exercised; mark it strict-xfail with the real cause so the bug is tracked rather than hidden.

…gation Use actual encoder output token counts instead of config theoretical maxima for placement indices, and mirror the test preprocessor's all-1s pixel_position_ids. The Gemma4 parity failure is now precisely localized to vision/audio merged token positions (max abs diff ~1.33), with text positions matching; update the xfail reason accordingly.

- Skip populate_function_gemma_metadata when llm_model_type is an explicit override, mirroring litert-torch's override behavior. - Handle NumPy integer axes (np.int64/32) in _patched_amax. - Update Falcon xfail reason to current torch.fx _ModuleNotInstalledAsSubmoduleError. - Update Gemma3n xfail reason's shape numbers to current [3,8,8,16] strides (1024,8,1,64) -> [192,16].

…t-minimal-litertlm Conflict resolutions: - requirements.txt: keep master's litert-torch + ai-edge-litert entries; replace the torch==2.13.0/torchvision==0.28.0 hard pins with torch>=2.12 + unpinned torchvision (litert-torch 0.9.1 on PyPI requires torch<2.13.0, so the 2.13 pin was unresolvable; resolves to the matched 2.12.1/0.27.1 CPU pair). Comment now covers only LiteRT-LM-specific deps living in the litertlm extra. - actions.yml: keep master's keras-3.15 pin and SHA-pinned actions; re-apply the torch-leg 'Install LiteRT-LM export dependencies' step. - qwen3_moe_causal_lm_test.py: drop the module-level KERAS_BACKEND=jax override (backend selection belongs to the environment); keep master's non-strict torch xfail on test_litert_export (aten._assert_async).

…tching the runtime contract On-device verification (Pixel 9, litertlm-android 0.13.1) showed every separate-vision bundle was unconsumable by the runtime: - The runtime rejects vision encoder inputs that are not 3- or 4-D (vision_litert_compiled_model_executor.cc); our encoder signature was traced 5-D [B, N, H, W, 3]. It is now traced single-image [B, H, W, 3] and KerasHubVisionEncoderAdapter reintroduces the N=1 axis via reshape before calling the KerasHub encoder. - The runtime chains encoder -> adapter per image; the adapter signature was traced [B * max_images, tokens, dim], which mismatches the single-image encoder output the runtime feeds it. It is now traced [B, tokens, dim]. With both fixes, gemma3 and pali_gemma separate-vision dummy bundles run image input end-to-end on device (3/3 stable runs each). Host-side parity harness and export_test.py now drive the encoder/adapter per image, mirroring the runtime.

…mantics An explicit llm_model_type override intentionally skips model-specific metadata population (mirroring litert-torch, which skips its metadata builder whenever litert_lm_model_type_override is set). The test still expected the populated function-calling block from the explicit kwarg; assert the designed bare-oneof behavior instead. The populated path is covered by tokenizer auto-detection tests in model_specs_test.py.

Diagnosis (repro evidence in the session report): the divergence was not in the export. torch.export matches Keras eager bit-exactly; the no-delegate builtin TFLite interpreter mis-executes the vision encoder's unflatten reshape under its default memory planner (XNNPACK 9.1e-06 vs builtin ~2.75 vs builtin+preserve_all_tensors 2.4e-07). Upstream ai_edge_litert memory-planning bug, not a keras-hub export defect.

…l, pin jax_enable_x64 for litert conversion - gemma4 test_litertlm_export_skip_mel_metadata lacked the torch-backend guard; failed on jax/tf CI legs - gemma3n litertlm xfail is litert-torch version-dependent (passes on PyPI 0.9.1, fails on 0.10.0); make it non-strict - litert_torch conversion via its JAX bridge requires jax_enable_x64=True for consistent int64 dtypes; with x64=False the bridge silently downcasts some constants to int32 and MLIR verification fails with i32/i64 'func.call' operand mismatches (6 test_litert_export failures on CI). Master passes because test_case.py's unrestored import leaks x64=True; the branch's litertlm preservation restores False, exposing the ordering dependence. Pin x64=True during conversion in both the test harness and the litertlm export path (restored afterwards). - ruff format fixes (test_case.py, adapter.py, traceable_ops_test.py)

…keras-hub#2861) Root cause is not an fx submodule-registration issue: Falcon builds ALiBi over the query length while attention scores span the KV-cache length, so call_with_cache fails whenever cache length > query length. Reproducible on master in eager mode with no export tooling.

github-actions Bot added the Gemma Gemma model specific issues label Apr 25, 2026

github-actions Bot requested review from divyashreepathihalli and sachinprasadhs April 25, 2026 12:35

gemini-code-assist Bot reviewed Apr 25, 2026

View reviewed changes

Comment thread keras_hub/src/utils/litertlm/export.py Outdated

Comment thread keras_hub/src/utils/litertlm/export.py Outdated

Comment thread keras_hub/src/utils/litertlm/export.py Outdated

Comment thread keras_hub/src/utils/litertlm/export.py Outdated

pctablet505 marked this pull request as draft April 25, 2026 12:40

pctablet505 removed request for divyashreepathihalli and sachinprasadhs April 25, 2026 12:41

pctablet505 force-pushed the torch-backend-litert-minimal-litertlm branch from e7a3523 to 0e54c94 Compare May 8, 2026 03:20

pctablet505 force-pushed the torch-backend-litert-minimal-litertlm branch 2 times, most recently from 0ffb29a to 2630a5a Compare May 20, 2026 16:16

pctablet505 commented May 20, 2026

View reviewed changes

gemini-code-assist Bot reviewed May 20, 2026

View reviewed changes

pctablet505 force-pushed the torch-backend-litert-minimal-litertlm branch from 5dcd75e to b320cd6 Compare May 27, 2026 15:55

pctablet505 force-pushed the torch-backend-litert-minimal-litertlm branch 3 times, most recently from 3769558 to 739af85 Compare June 23, 2026 11:58

pctablet505 added 12 commits June 25, 2026 17:25

Clarify optional dependency error messages

5fee0b2

litert-torch is the only package users need to install; ai-edge-litert is pulled in automatically as a dependency. ai-edge-torch is not required.

Skip LiteRT-LM export tests when optional deps are missing

df9e919

Add litert-torch and litert-lm-builder to test requirements

aa68716

pctablet505 added 26 commits July 19, 2026 12:38

Merge branch 'ws2-7b-litertlm-2026-07-19' into torch-backend-litert-m…

596e048

…inimal-litertlm

Merge branch 'ws3-3-litertlm-2026-07-19' into torch-backend-litert-mi…

f35d944

…nimal-litertlm

feat(litertlm): add audio_input_style spec field, drop signature sniff

7aa1328

refactor(litertlm): dispatch vision on spec style, drop presence-base…

92d7f8a

…d branches

docs(litertlm): scope vision bucketing ban family-wide, add allows_vi…

5803ffa

…sion_bucketing flag

refactor(litertlm): drop dead cache_layout attr, keyword-only _assemb…

fdbf76f

…le_bundle

style(litertlm): fix test import order and format drift for ruff CI

a36e09b

feat(litertlm): add conditional sampler_params with GREEDY_SAMPLER_CO…

8d27713

…NFIG preset

feat(litertlm): map function_gemma preset to function_gemma model type

7675522

feat(litertlm): reject Gemma4AssistantCausalLM MTP draft model at export

98fd87f

pctablet505 mentioned this pull request Jul 24, 2026

Builtin no-delegate interpreter corrupts reused buffer under default memory planner (RESHAPE→RMSNorm region) google-ai-edge/LiteRT#8889

Open

pctablet505 added 3 commits July 24, 2026 16:20

chore: drop accidentally committed AGENTS.md and gitignore it

f5f7a26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Implement KerasHub → LiteRT-LM export with prefill/decode signatures#2705

Implement KerasHub → LiteRT-LM export with prefill/decode signatures#2705
pctablet505 wants to merge 148 commits into
keras-team:masterfrom
pctablet505:torch-backend-litert-minimal-litertlm

pctablet505 commented Apr 25, 2026 •

edited

Loading

Uh oh!

google-cla Bot commented Apr 25, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pctablet505 left a comment

Uh oh!

pctablet505 commented May 20, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

pctablet505 commented Apr 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Design Doc: KerasHub → LiteRT-LM Export

Contents

1. Objective

2. Background & Motivation

2.1 The gap today

2.2 What .litertlm provides

2.3 Terminology

3. Goals & Non-Goals

Goals

Non-Goals

4. Requirements & Constraints

5. Solution Overview

5.1 How this PR would split (proposed, pending maintainer approval)

The one public entry point

6. Detailed Design

6.1 The adapter and KV-cache translation

6.2 Prefill and decode signatures

6.3 Export pipeline

6.4 Bundle packaging and metadata

6.5 Tokenizer handling

6.6 Multimodal design

The mental model

Terms, defined once (used everywhere below)

Per-family mode matrix

Dataflow: baked-in mode (default)

Dataflow: separate mode (separate_vision_encoder=True)

Restrictions

6.7 Quantization

6.8 Model support and detection

7. Design Decisions & Rationale

8. Alternatives Considered

8.1 Weight-mapping onto litert_torch's own model implementations

8.2 The TF-backend path (tf.function static signatures → TFLiteConverter)

8.3 Upstreaming the traceable-ops fixes into Keras's torch backend

10. Limitations

9. Testing Strategy

10.1 Export-time

10.2 Runtime (LiteRT-LM)

10.3 Multimodal

10.4 On-device performance and memory (Pixel 9)

11. Future Work

12. Usage & Deployment

12.1 Python export

12.2 Runtime integration (summary)

13. Appendix: API, File Layout & Glossary

13.1 CausalLM.export() reference for format="litertlm"

13.2 Core files

13.3 Adapter classes

13.4 Glossary

Uh oh!

google-cla Bot commented Apr 25, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pctablet505 left a comment

Choose a reason for hiding this comment

Uh oh!

pctablet505 commented May 20, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

pctablet505 commented Apr 25, 2026 •

edited

Loading

2.2 What `.litertlm` provides

Dataflow: separate mode (`separate_vision_encoder=True`)

8.2 The TF-backend path (`tf.function` static signatures → TFLiteConverter)

13.1 `CausalLM.export()` reference for `format="litertlm"`