When I tried building an on-device AI app with Gemma 4, the pitch was clear: model weights on the device, no server, no API calls, works offline. Getting it to actually run fast was a different problem.

This post covers what I learned working with LiteRT-LM 0.12.0 and Gemma 4 E2B on Android in Kotlin. Some of it is configuration. Some of it is understanding what the bottleneck actually is before reaching for a fix. If you're building with Gemma 4 E2B on Android and inference feels too slow to ship, here are the tricks that actually helped.

1. Basic Setup

Add the dependency:

// build.gradle
implementation("com.google.ai.edge.litertlm:litertlm-android:0.12.0")

The model file itself comes from Hugging Face. The litert-community/gemma-4-E2B-it-litert-lm repository hosts the .litertlm format that LiteRT-LM expects. This is not a GGUF file. Using the wrong format will cause a silent failure on model load, so confirm the file extension before downloading.

The model is gated on Hugging Face, so you'll need an access token. A read token is enough. If your app handles the download directly (via DownloadManager or a similar mechanism), pass the token as an Authorization header in the request rather than entering it interactively. The full LiteRT-LM Android API reference is here.

Initialize the engine:

val options = LlmInferenceOptions.builder()
    .setModelPath(modelPath)
    .setMaxTokens(512)
    .setTopK(40)
    .setTemperature(0.8f)
    .setRandomSeed(101)
    .build()

val llmInference = LlmInference.createFromOptions(context, options)

2. GPU Backend and Why It Silently Falls Back to CPU

LiteRT-LM supports three backends: CPU, GPU (via OpenCL), and NPU. GPU is where you get meaningful speed on Android.

The problem is that OpenCL support is not universal. Mid-range and budget chips from Qualcomm and MediaTek often don't expose OpenCL to the Android application layer. If you initialize with Backend.GPU() on one of these devices, the engine falls back to CPU without throwing an error by default.

If you don't log this, you'll spend time optimizing prompts thinking you're on GPU when you're not.

Check which backend actually initialized:

try {
    val config = EngineConfig.builder()
        .setModelPath(modelPath)
        .setBackend(Backend.GPU())
        .build()
    engine = Engine(config)
    Log.d("Inference", "GPU backend initialized")
} catch (e: Exception) {
    val config = EngineConfig.builder()
        .setModelPath(modelPath)
        .setBackend(Backend.CPU())
        .build()
    engine = Engine(config)
    Log.d("Inference", "CPU fallback: ${e.message}")
}

On CPU with Gemma 4 E2B, expect roughly 2 to 5 tokens per second on mid-range hardware. On GPU-capable devices via OpenCL, LiteRT-LM benchmarks show around 52 tokens per second on a Samsung S26 Ultra. The delta between CPU and GPU is not incremental, it is a different category of usability.

If your target users are running budget Android devices, plan your UX around CPU speeds. Streaming tokens as they arrive, showing a "thinking" indicator early, and capping output length all reduce how slow it feels even when the hardware is constrained.

One more thing on backends: NPU initialization is not just a silent fallback situation. On some devices, attempting Backend.NPU() can cause a native process crash (SIGKILL or SIGSEGV) due to driver fragmentation across Android hardware. If you want to expose NPU as an option, treat it as an experimental toggle rather than a default path, and always have the GPU-to-CPU chain as the safe baseline.

3. Prefill Is the First Bottleneck, Not Decoding

Most discussions about LLM inference speed focus on decode speed (tokens per second). On mobile, the more immediate pain point is often prefill: the time before the model generates the first token.

Prefill is proportional to the size of your input prompt. Every character you inject into the system prompt has to be processed before generation starts. If you're doing context injection (pasting a document or manual into the prompt), this cost hits on every single query.

A rough example. A 50,000 character document injected into a system prompt is approximately 12,000 to 15,000 tokens. On CPU, processing that input alone takes several seconds before the model produces anything. A user taps submit and waits in silence.

Gemma 4 E2B supports a 128K context window, and that number is real. But mobile hardware is bound by prefill latency and KV cache limits long before you hit 128K. The theoretical capacity and the practical ceiling on a 4GB device are very different numbers.

Practical fixes:

Set a hard character budget on injected context and enforce it at the application layer:

val contextBudget = 6000 // characters, not tokens
val injectedContext = sourceDocument.take(contextBudget)

6,000 characters is roughly 1,500 tokens. That's enough context to be useful for most domain-specific queries while keeping prefill manageable on CPU.

If you're building a document Q&A feature, extract only the relevant section rather than injecting the full document. A keyword match or simple sentence scoring function in Kotlin can identify the most relevant passage and inject that instead of the whole file. This is not full RAG. It's a practical middle ground that works without vector databases.

4. Multi-Token Prediction: The Feature That Makes Gemma 4 Worth It on Mobile

Multi-Token Prediction (MTP) is one of the things that genuinely sets Gemma 4 apart from earlier versions for on-device use. It was introduced with the Gemma 4 model family specifically, and it changes what's achievable on mobile hardware in a meaningful way.

Standard autoregressive inference generates one token per forward pass. The processor moves model parameters from memory to compute units, generates one token, then does it again. On mobile hardware, the data movement cost dominates over the actual computation.

MTP uses speculative decoding to work around this. A lightweight drafter model proposes several tokens ahead of time. The primary model then verifies those proposals in a single parallel forward pass. If the proposed tokens are correct, the model accepts them all plus generates one more. If the drafter was wrong at some position, it rejects from that point and takes over. Output quality doesn't change because the primary model has final say over every token.

LiteRT-LM bundles the MTP drafter inside the same .litertlm model artifact. Both models run on the same hardware backend, sharing KV cache in local memory. This avoids the cross-device data transfer overhead that would otherwise cancel out part of the gain.

Google's benchmarks show up to a 2.2x decode speedup with MTP enabled on the GPU backend. See their full breakdown here. For the dedicated MTP announcement and how the drafter was designed for the Gemma 4 family specifically, see this post. Enabling it is two lines of configuration:

val options = LlmInferenceOptions.builder()
    .setModelPath(modelPath)
    .setMaxTokens(512)
    .setUseMtp(true)       // enable MTP drafter
    .setTopK(40)
    .setTemperature(0.8f)
    .build()

The gains are more pronounced for predictable completions. For creative or open-ended generation where the drafter has low acceptance rates, the speedup is smaller. For structured or domain-constrained outputs, acceptance rates are higher and the gains are closer to the ceiling.

One important caveat: if you're on CPU, disable MTP. The 2.2x gain assumes parallel GPU execution where the drafter and target model run simultaneously. On CPU they run sequentially, and the overhead of running two models back to back outweighs the benefit. Check which backend actually initialized before deciding whether to enable it.

val useMtp = backend == Backend.GPU() // only enable on GPU
val options = LlmInferenceOptions.builder()
    .setModelPath(modelPath)
    .setUseMtp(useMtp)
    .build()

5. Thinking Mode: When to Use It and When Not To

Gemma 4 supports a reasoning mode where the model generates an internal scratchpad before producing its final response. LiteRT-LM exposes this directly. The reasoning output appears between <|think|> and </think> tags in the stream.

Thinking mode improves output quality for multi-step or diagnostic tasks. It costs tokens. On CPU, those extra 200 to 400 reasoning tokens represent meaningful latency before the user sees a final answer.

The practical approach: enable thinking on tasks where accuracy matters, disable it on conversational turns where it doesn't.

fun buildSystemPrompt(requiresReasoning: Boolean): String {
    return if (requiresReasoning) {
        "<|think|> You are a diagnostic expert. Think through this step by step before answering."
    } else {
        "You are a helpful assistant. Answer clearly and concisely."
    }
}

If you're displaying the thinking stream in the UI (as a visible "reasoning" component), the latency becomes part of the experience rather than dead time. The user sees the model working. This matters more on CPU where the stream is slow enough to read.

If you strip thinking tokens before displaying the response, you're paying the token cost with no UX return. In that case, disable it.

6. Constrained Decoding for Structured Output

LiteRT-LM supports constrained decoding, which enforces a JSON schema on the model's output. Instead of parsing free text and hoping the model follows your format instructions, you define the schema and the engine guarantees compliance.

This is useful for any feature that needs to render structured results rather than prose. A diagnosis card, a checklist, a decision tree. The model produces valid JSON every time.

val schema = """
{
  "type": "object",
  "properties": {
    "diagnosis": { "type": "string" },
    "confidence": { "type": "string", "enum": ["high", "medium", "low"] },
    "action": { "type": "string" },
    "escalate": { "type": "boolean" }
  },
  "required": ["diagnosis", "confidence", "action", "escalate"]
}
"""

val options = LlmInferenceOptions.builder()
    .setModelPath(modelPath)
    .setResponseSchema(schema)
    .setMaxTokens(300)
    .build()

The max tokens value matters here. A constrained JSON response is short. Setting a generous 2,000 token budget for a response that will always be under 100 tokens keeps the KV cache allocated longer than necessary. Set it tight.

7. Session Save and Restore

LiteRT-LM supports serializing and restoring the KV cache state across sessions. For applications with persistent context (a long document loaded once, or a multi-turn workflow), this means the prefill phase only happens once. On return sessions, the engine restores the cached state and skips the expensive input processing step.

For document Q&A specifically, this is worth implementing. The user loads a document, the prefill runs once and the state is serialized to disk. Every subsequent question in that session resumes from the cached state rather than reprocessing the document from scratch. The Google AI Edge Gallery app is the most complete open-source example of session management in a real LiteRT-LM application.

Summary

The model download is the upfront cost. After that, LiteRT-LM gives you enough knobs to tune inference for the specific hardware and use case you're targeting. The techniques above don't require changes to model weights or training. They're all configuration and prompt engineering decisions available in 0.12.0.

For the full LiteRT-LM technical deep-dive from the Google AI Edge team, including iOS and WebGPU benchmarks, read the official post.