On-Device AI Series (Part 5): LiteRT-LM

Put your phone in airplane mode. Open the app, type a question, and watch the answer arrive one token at a time — no spinner waiting on a network round-trip, no API key, no per-token bill, and nothing you typed ever leaving the device. LiteRT-LM removes the genuinely hard parts of running an LLM on-device — KV-cache management, token streaming, backend selection — but it doesn’t remove your job so much as relocate it.

What’s left on your plate is a short, specific list: sizing a combined input+output token budget, owning your own sampling defaults, hand-building system prompts and tool calling out of raw text, and one native-library collision that presents as a SIGSEGV rather than a build error. Know those going in and the API itself is a clean three-step pattern.

We’ll get there in that order:

Why you’d choose this runtime and what it costs you versus the cloud.
The Engine/Session model you need to read the code at all.
Real implementation samples — streaming, system prompts and tool calling, multimodal inputs, thinking mode, and CPU-vs-GPU benchmarking.
The anti-patterns to avoid.
A developer-friendliness rating on the same rubric as Parts 1–4.

Why Use LiteRT-LM?

You reach for LiteRT-LM instead of hand-rolling generation on top of raw LiteRT when:

You need multi-turn conversation, not single-shot inference — session state and KV-cache bookkeeping are handled for you, and resetting a conversation is a session swap, not a model reload.
You need streaming output — token-by-token delivery for a responsive chat UI, instead of a blocking call that returns everything at once.
You’re choosing between CPU and GPU per device — the explicit backend parameter turns that into a runtime decision instead of a build-time guess.
You want a pre-converted model without doing your own PyTorch-to-LiteRT conversion work — the Model Zoo covers Gemma, Qwen, Llama, and more out of the box.
You’re willing to own sampling — the engine won’t pick sane decoding defaults for you; that’s on the dialect config you write once per model family.

The tradeoff mirrors LiteRT’s from Part 4: full on-device privacy and zero network dependency, in exchange for owning the latency, memory, and sampling decisions a cloud LLM API would otherwise absorb for you. And there’s a quality tradeoff underneath that one: an on-device model is nowhere near as capable as the giant models running behind a cloud API, but it’s also free to run at inference time — no per-token bill, no network round-trip.

None of this locks you into one solution for the whole app, either. Nothing stops you from running LiteRT-LM and a cloud LLM side by side — routing simple, latency-sensitive, or privacy-sensitive requests on-device, and falling back to the cloud when a task needs more capability than a local Gemma variant can deliver.Under the Hood: Engine, Session, and a Build-Time Trap

Under the Hood: Engine, Session, and a Build-Time Trap

Before writing code, understand the two-object mental model:

Engine owns the loaded model and backend selection — expensive to create, meant to be reused.
Session owns only the conversation state riding on top of it — cheap to create, and disposable per conversation or per benchmark run.

LiteRT-LM doesn’t replace the LiteRT engine from Part 4; it wraps it with exactly the orchestration a chat needs and nothing else.

⚠️ GOTCHA: Native Library Collision

litert AAR ───────┐
                  ├── Same .so name → SIGSEGV
litertlm AAR ────┘

That layering has one sharp edge: both litert and litertlm-android bundle libLiteRtClGlAccelerator.so, but the versions aren’t interchangeable. Using the wrong one with liblitertlm_jni.so can cause a SIGSEGV in nativeCreateEngine.

Fix it in Gradle, not application code: extract the matching .so from the litertlm-android AAR into generated jniLibs. Source-set libraries take precedence during merging, and pickFirsts provides a fallback. In short: it’s a native library collision, not a runtime bug.

The Toolset: What Can LiteRT-LM Actually Do?

Runs LLMs on-device — The Model Zoo ships pre-converted Gemma, Qwen, Llama, Phi, and more as .litertlm files; a Gemma 4 (E2B/E4B) pair backs the examples below. Check local paths (app storage, Downloads, dedicated models folder) before triggering a fetch for manual side-loading.
Manages KV-cache and session state — A Session holds conversation context so you send only the new turn instead of replaying the transcript every time.
Streams tokens as they’re generated — A callback fires with partial output as soon as the model produces it.
Targets CPU or GPU explicitly — Backend is a config value on Engine, not something baked in at build time.
Resets a conversation without reloading the model — Closing and recreating a Session on the same Engine wipes the KV cache in milliseconds.

Hands-on Implementation Samples

Both examples below are trimmed from a real chat feature and a real CPU/GPU benchmarking screen.

Loading a Model and Streaming a Response

The shape is always: configure an Engine, open a Session, send a prompt, collect tokens.

val engineConfig = EngineConfig(
    modelPath = modelPath,
    backend = Backend.GPU(), // or Backend.CPU()
    // maxNumTokens is the SUM of input + output tokens — the KV-cache
    // ceiling, not an output-only limit. The engine's small default
    // crashes once a loaded chat's prefill replay outgrows it.
    // 16_384 worked for Gemma 4 here; the right ceiling is model-
    // dependent, so re-check it against whatever model you load.
    maxNumTokens = 16_384,
)

val engine = Engine(engineConfig)

engine.initialize()

val session = engine.createSession(SessionConfig(samplerConfig = dialect.samplerConfig))

session.generateContentStream(
    listOf(InputData.Text(prompt)),
    object : ResponseCallback {
        override fun onNext(response: String) {
            // Fires on an engine thread - hop to your UI dispatcher before touching state.
            trySend(AgentChunk.Text(response))
        }
        override fun onDone() { /* generation complete */ }
        override fun onError(throwable: Throwable) { /* surface it */ }
    }
)

Unpacking

maxNumTokens is the total context budget (input + output KV-cache), not an output-token limit. Google AI Edge Gallery uses 10,000 tokens as a conservative warning threshold because higher values increase memory usage across Android devices.
A bigger budget only delays context exhaustion; you still need to truncate or summarize old messages in long conversations.
Engine is expensive; Session is cheap. Initialize one Engine and reuse it. Create a new Session for each conversation.
To start a new chat, reset the Session, not the Engine. Keep the model loaded.

💡 Experience note: Set an explicit SamplerConfig to avoid repetitive greedy output. A solid default is topK = 64, topP = 0.95, temperature = 1.0. Keep Gemma’s turn-formatting tokens in mind and clean up malformed special tokens during streaming.

<a href="https://medium.com/media/dc7aa2b4bc666e5f1ed187f3f5faf441/href">https://medium.com/media/dc7aa2b4bc666e5f1ed187f3f5faf441/href</a>

System Prompts and Tool Calling

LiteRT-LM has no system role and no function-calling schema in its API — Session.generateContentStream takes a flat list of InputData. Everything you’d recognize from a cloud LLM SDK as “system prompt” or “tools” has to be built by hand, out of plain text, formatted the way your specific model was trained to read it. That’s the job of a small ChatDialect abstraction: one implementation per model family, holding the literal turn-marker tokens and a function that assembles them./<turn|>and<|tool_call>` conventions:

override val toolCallOpen: String = "<|tool_call>"
override val toolCallClose: String = "<tool_call|>"
override val stopSequences: List<String> = listOf("<turn|>", "<|tool_response>")

override fun firstTurn(systemPrompt: String, userMessage: String): String = buildString {
    if (systemPrompt.isNotBlank()) {
        append("<|turn>systemn")
        append(systemPrompt)
        append("<turn|>n")
    }
    append("<|turn>usern$userMessage<turn|>n")
    append("<|turn>modeln")
}

A system prompt is text prepended inside a <|turn>system … <turn|> block.

⚠️ Note: Gemma 1–3 use <start_of_turn>role … <end_of_turn>, whereas Gemma 4 uses <|turn>role … <turn|>. Ensure you use the exact prompt spec for your model checkpoint.

For tool calling, tools declare JSON schema strings flattened into system prompts as <|tool>{…}<tool|> blocks:

// LlmReasonTool.kt — a tool's "schema" is a plain JSON string you write yourself
override val definition = ToolDefinition(
    name = "llm_reason",
    description = "Reasons over text with a cloud LLM: summarize, synthesize, extract...",
    parametersJson = """{"type":"object","properties":{"prompt":{"type":"string"}},"required":["prompt"]}""",
)

Detecting tool calls involves parsing output strings for special tokens:

dialect  = the ChatDialect for the model we configured (e.g. GemmaDialect)
toolCall = dialect.parseToolCall(modelOutputSoFar)
if toolCall == null:
    if dialect.hadToolCallAttempt(modelOutputSoFar):
        # model opened a "<|tool_call>" block but never closed it with valid JSON
        # → treat as malformed and retry the call
        retryWithCorrection()
    else:
        # ordinary text turn - stream it to the UI as-is
        emitText(dialect.stripSpecialTokens(modelOutputSoFar))
else:
    result = toolRegistry.execute(toolCall.name, toolCall.argsJson)
    # feed the result back in as the next turn, then let the model continue
    continueWith(dialect.toolResponseTurn(result))

<a href="https://medium.com/media/190924ba54c374399ec382de00a5f46f/href">https://medium.com/media/190924ba54c374399ec382de00a5f46f/href</a>

Multimodal Input

Multimodal input is not an InputData variant on Session. The important detail is that multimodal input is not simply another InputData variant. Session -> generateContentStream() Multimodal generation: Engine -> Conversation -> sendMessageAsync()

This helps readers grasp why Session cannot be used directly for multimodal inputs before explaining the single-session native constraint.” Passing InputData.Image directly into Session.generateContentStream() throws a runtime exception (“Image must be preprocessed before being used in SessionAdvanced.”).

Instead, LiteRT-LM provides the higher-level Conversation API:

val conversation = engine.createConversation()

val contents = Contents.of(
    Content.ImageBytes(jpegBytes),
    Content.Text(prompt)
)

conversation.sendMessageAsync(contents).collect { message ->
    val text = message.contents.contents
        .filterIsInstance<Content.Text>()
        .joinToString("") { it.text }
    trySend(AgentChunk.Text(text))
}

Session and Conversation share the same underlying native slot (FAILED_PRECONDITION: A session already exists). To mix them, temporarily close the Session, perform the multimodal turn on a Conversation, and reopen the Session:

session.close()

val conversation = engine.createConversation()
val reply = runImageTurn(conversation, jpegBytes, prompt)
conversation.close()

session = engine.createSession(
    SessionConfig(samplerConfig = dialect.samplerConfig)
)

turnHistory.forEach {
    session.runPrefill(listOf(InputData.Text(it)))
}

<a href="https://medium.com/media/245e5c141fde5480e54125c101eaec19/href">https://medium.com/media/245e5c141fde5480e54125c101eaec19/href</a>

Thinking Mode

The last capability worth covering is thinking: a supported model generates a reasoning block before its final answer, the same way it generates a tool call — as plain text, using tokens it was fine-tuned to produce, with the application responsible for telling the two apart.

Like system prompts and tool calling, thinking is not a universal LiteRT-LM API. It’s a property of the model and its prompt format. Gemma 4 turns it on with a one-line system turn containing a dedicated toggle token, then wraps its reasoning in a <|channel>thought … <channel|> block before the final response:

// thinking is opt-in per turn, via the model's own toggle token
override fun firstTurn(
    systemPrompt: String,
    userMessage: String,
    thinkingEnabled: Boolean,
): String = buildString {
    if (thinkingEnabled) {
        // Gemma 4's dedicated toggle — an empty system turn containing only <|think|>
        append("<|turn>systemn<|think|><turn|>n")
    }
    if (systemPrompt.isNotBlank()) {
        append("<|turn>systemn$systemPrompt<turn|>n")
    }
    append("<|turn>usern$userMessage<turn|>n")
    append("<|turn>modeln")
}

Your app splits reasoning from the final answer:

// ThinkingParser.kt — splitting streamed response
private val THOUGHT_BLOCK =
    Regex("<\|channel>thought(.*?)<channel\|>", RegexOption.DOT_MATCHES_ALL)

data class ParsedResponse(val thinking: String?, val content: String)

fun parseResponse(raw: String): ParsedResponse {
    val thinking = THOUGHT_BLOCK.find(raw)?.groupValues?.get(1)?.trim()
    val content = raw.replace(THOUGHT_BLOCK, "").trim()
    return ParsedResponse(thinking = thinking, content = content)
}

// Store ONLY content in conversation history for subsequent turns
conversationHistory.add(Turn.Model(parsed.content)) // never parsed.thinking

That split matters most at the point where you build the next turn’s prompt. Carrying a prior turn’s reasoning block forward is pure wasted context — the model doesn’t need to re-read its own scratch work, and on a 16K-token budget, a long thought block is exactly the kind of thing that pushes a real conversation past maxNumTokens sooner than you’d expect:

// Building history for the next turn — only ParsedResponse.content goes in
conversationHistory.add(Turn.Model(parsed.content)) // never parsed.thinking

Key Takeaways for Thinking Mode

Thinking costs tokens: Thought blocks count toward maxNumTokens.
Strip thoughts from history: Omit prior reasoning from future prompts to preserve context space.
Separate reasoning in UI: Display thought blocks in collapsible or distinct UI components.
Per-request toggle: Enable thinking dynamically based on query complexity.
Buffer streaming tokens: Handle mid-marker splits across stream chunks.

<a href="https://medium.com/media/b14ff21eec2b65a29bfc3b90f8a09863/href">https://medium.com/media/b14ff21eec2b65a29bfc3b90f8a09863/href</a>

CPU vs. GPU Benchmarking

Evaluating model performance across different backends requires isolated execution to accurately capture latency and throughput metrics.

// Setup engine configuration with the target backend (CPU or GPU)
val engineConfig = EngineConfig(
    modelPath = config.modelPath,
    backend = nativeBackend,
    maxNumTokens = 4096
)

val engine = Engine(engineConfig)
engine.initialize()
val sessionConfig = SessionConfig(
    samplerConfig = config.modelConfig.dialect.samplerConfig
)
// Execute iterations using isolated sessions
repeat(config.iterations) {
    val session = engine.createSession(sessionConfig)
    generateOnce(session, config)
    session.close() // Clear state/KV-cache after each run
}
engine.close()

Here is how the Gemma:E2B-it model performed when shifting the workload from CPU to GPU execution:

<a href="https://medium.com/media/b4747841dc40013f7f98257c84469492/href">https://medium.com/media/b4747841dc40013f7f98257c84469492/href</a>

CPU Execution: Completed generation in 14.4 seconds at a throughput of 2.9 tokens/sec.
GPU Execution: Completed generation in 9.5 seconds (~34% faster) at a throughput of 5.0 tokens/sec (~72% speedup).

While exact performance depends on hardware specifications, GPU acceleration consistently yields significant speed gains. Expect even sharper performance improvements when targeting dedicated NPUs.

Anti-Patterns to Avoid

These caused real crashes, hangs, or bad output during development:

❌ Copying a tutorial’s maxNumTokens — it covers both input and output, so the default can overflow during longer chats.
❌ Leaving SamplerConfig unset — greedy decoding can cause repetitive output on small models.
❌ Reloading the model to reset a chat — close the Session; keep the Engine warm.
❌ Reusing a Session in benchmarks — KV-cache state carries over and skews results.
❌ Trusting GPU capability checks — a supported GPU can still fail at runtime; fall back to CPU.
❌ Updating UI state directly in onNext — callbacks run off the UI thread; dispatch to the UI first.
❌ Trusting tool-call JSON — parse defensively and handle incomplete blocks.
❌ Using litert and litertlm-android without the jniLibs fix — conflicting .so files can cause a SIGSEGV in nativeCreateEngine.

Rating LiteRT-LM: How Developer-Friendly Is It?

Using the same rubric from Parts 1–4:

The Final Score: 7 / 10 (B ✅ Good)

Android Integration (7.5/10): A standard Gradle dependency, but you’ll hit the shared native-library conflict with plain LiteRT if both are in the same app — budget for a one-time build-time fix, not a runtime workaround.
API Simplicity (7.5/10): Engine → Session → generateContentStream is a clean three-step pattern. Owning sampler defaults yourself is the one real gap between “it compiles” and “it produces coherent output.”
Kotlin-First Design (6/10): Callback-based (ResponseCallback), not Flow-native — wrapping it in callbackFlow is on you for idiomatic coroutine use.
Model Compatibility (7/10): Solid Model Zoo coverage on paper (Gemma, Qwen, Llama, Phi, and more), but you’re limited to what’s been converted — no arbitrary .tflite LLM the way raw LiteRT accepts any model.
Performance & HW Acceleration (8.5/10): Explicit CPU/GPU backend selection makes it trivial to actually measure the tradeoff instead of guessing. NPU support isn’t yet exposed.
Documentation & Community (6/10): A newer surface than core LiteRT — expect to read source and benchmark yourself more than you’d like for anything beyond the basic loop.
Offline Capability (10/10): Fully on-device inference once the model file is on the device — no network path during generation.
Maintenance & Stability (6/10): Actively evolving. Constrained decoding (JSON-schema/grammar-constrained tool calls) already exists in the C++ core but isn’t yet exposed in the Android Kotlin binding — worth watching if structured tool-calling is on your roadmap.

Conclusion

LiteRT-LM earns its place in this series by solving one specific problem LiteRT’s raw tensors don’t: a conversation isn’t a single inference call, it’s a sequence of them that all have to remember each other. Engine and Session exist to carry that memory — KV-cache, sampler state, streaming — so you’re not hand-rolling the bookkeeping raw LiteRT would leave on your plate.

None of that removes your job, though; it just relocates it — exactly as promised at the top. maxNumTokens is a combined input+output budget you have to size for a real conversation, not the small default a tutorial hands you. Sampling defaults to greedy decoding that collapses into repetition unless you configure topK/topP/temperature yourself. And the one failure that looks nothing like your code — a SIGSEGV in nativeCreateEngine from two colliding copies of libLiteRtClGlAccelerator.so — is a Gradle packaging fix, not a runtime one. Know those three going in and the rest of the API is a clean three-step pattern: configure, open a session, stream tokens.

Placed on the ladder this series has been climbing, LiteRT-LM is the generative-AI specialization of the LiteRT foundation from Part 4 — session state and streaming bolted onto the same low-level engine, so you no longer build a chat runtime up from tensors by hand. If your use case is “run an LLM on this phone,” this is where you start; drop to raw LiteRT only for architectures it doesn’t cover, and go cloud only when the model genuinely doesn’t fit on-device.

Youtube

Love you all.

Stay tune for upcoming blogs.

Take care.

On-Device AI Series (Part 5): LiteRT-LM was originally published in ProAndroidDev on Medium, where people are continuing the conversation by highlighting and responding to this story.