Blazing Fast On-Device GenAI with LiteRT-LM

Speed is the difference between an on-device LLM that feels like a gimmick and one that feels like a product. Gemma 4’s release proved small models can reason, call functions, and understand images on a phone—but only if the runtime around them is ruthlessly optimized. That runtime is LiteRT-LM, Google AI Edge’s orchestration layer for on-device LLMs.

LiteRT-LM powers Google’s own demanding production use cases: Chrome, ChromeOS, the Pixel Watch, and the viral Google AI Edge Gallery app. It delivers state-of-the-art inference speed across Android, iOS, and the open web. In this tutorial, you’ll learn what makes it fast, and how to apply those techniques in your own apps.

Prerequisites

• A device or emulator with GPU support (Android, iOS, or a WebGPU-capable browser)
• A Gemma 4 model file (.litertlm format) from the LiteRT Hugging Face community
• The LiteRT-LM CLI for quick experiments, or the flutter_gemma package for Flutter apps

The Performance Stack

LiteRT-LM achieves its speed by attacking latency at every layer of the stack:

┌─────────────────────────────────────────────┐
│        LiteRT-LM orchestration layer         │
│  Multi-Token Prediction │ Session Mgmt       │
├─────────────────────────────────────────────┤
│   Accelerated kernels: XNNPACK + ML Drift    │
├─────────────────────────────────────────────┤
│  CPU │ GPU (OpenCL/Metal) │ NPU │ WebGPU     │
└─────────────────────────────────────────────┘

LiteRT (formerly TensorFlow Lite) provides the accelerated kernel foundation through XNNPACK and ML Drift. LiteRT-LM adds the LLM-specific layers: quantization-friendly runtimes, memory-mapped per-layer embeddings, multi-token prediction, and advanced session management. Together they eliminate the two things that kill mobile AI: latency and memory pressure.

1. Multi-Token Prediction (MTP): 2.2x Speedup

Token-by-token decoding is the fundamental bottleneck of LLM inference. Each step requires a full forward pass, and mobile models can only generate tens of tokens per second. LiteRT-LM’s answer is speculative decoding built on Gemma 4’s native Multi-Token Prediction drafters.

Here’s the trick: a lightweight MTP drafter guesses the next several tokens in one pass. The primary model then verifies all the guesses in parallel—a much cheaper operation than generating them one at a time. LiteRT-LM enforces memory locality by keeping the drafter and the primary model on the same hardware IP, so the shared KV cache never crosses processor boundaries:

Standard:  prompt → t1 → t2 → t3 → t4      (4 sequential passes)
MTP:       prompt → [t1 t2 t3 t4] → verify  (1 draft + 1 verify pass)

The result is up to a 2.2x speedup in end-to-end generation without any loss of output quality. When you hear about Gemma 4 E2B hitting 52 tokens/sec on Android GPUs (OpenCL) and 56 tokens/sec on iOS (Metal), MTP is doing much of that work.

2. A Lean Memory Footprint

Phones don’t have 80GB of VRAM. LiteRT-LM keeps memory usage shockingly low through several techniques:

• 2-bit and 4-bit weight quantization with per-layer control.
• Memory-mapped per-layer embeddings (PLEs) kept out of active memory, loaded only when needed.
• On-demand encoders: image and audio encoders load only when a task requires them, keeping text-only workloads feather-light.
• Weight caching: the ~2.58GB Gemma 4 E2B model runs with a physical memory footprint of just 607MB on Apple mobile CPUs via XNNPACK’s weight caching.

# LiteRT-LM Python binding — quantization is handled for you
import litert_llm

engine = litert_llm.create_engine(
    model_path="gemma-4-E2B-it.litertlm",
    enable_mtp=True,          # multi-token prediction
    backend="gpu",            # or "cpu" / "npu"
)

The practical upshot: Gemma 4 E2B runs on devices you’d never associate with AI, and memory is left over for your app’s real UI.

3. Session Management: Resume Without Re-Prefill

Long conversations are expensive. Every time a user relaunches your app, re-prefilling the entire context wastes compute. LiteRT-LM’s native session save/restore serializes the large KV cache state so conversations resume instantly:

// Flutter: persist an agent session
final bytes = await conversation.saveSession();
await File('agent_session.bin').writeAsBytes(bytes);

// ... later, or after app restart:
final restored = await Conversation.restoreSession(bytes);
final response = await restored.sendMessage('Continue where we left off');

Restoring a session bypasses the heavy prefill phase entirely. That’s both a UX win (seamless continuity) and a cost win (no redundant computation). This is what powers the extended Agent Skills in the AI Edge Gallery app.

4. Constrained Decoding for Reliable Tool Calls

On-device agents depend on structured output—function calls, JSON, schemas. Small models left to free generation produce malformed JSON often enough to break pipelines. LiteRT-LM supports constrained decoding: the runtime restricts sampling to the set of tokens that keep the output valid.

final response = await conversation.sendMessage(
  'Add a meeting at 9am tomorrow',
  responseFormat: JsonSchema(
    schema: {
      'type': 'object',
      'properties': {
        'tool': {'type': 'string'},
        'args': {'type': 'object'},
      },
      'required': ['tool'],
    },
  ),
);

Every token emitted is guaranteed to move the output toward a valid JSON document, so your parser can trust the structure. Combined with Gemma 4’s native function calling, this makes the tool-call loop dramatically more reliable.

5. Dynamic Context Across CPUs and GPUs

On-device workloads are heterogeneous: a phone might switch from NPU to GPU to CPU as battery and thermals dictate. LiteRT-LM supports dynamic context lengths and single-model execution across CPUs and GPUs, so the same model file runs everywhere without conversion. Build once, deploy everywhere, and let the runtime pick the fastest available backend.

Measuring It Yourself

The litert-lm CLI gives you instant visibility into your device’s real performance:

litert lm run \
  --from-huggingface-repo=litert-community/gemma-4-E2B-it-litert-lm \
  gemma-4-E2B-it.litertlm \
  --prompt="Explain speculative decoding in one sentence."

# Prints prefill and decode token rates plus peak memory.

On a Raspberry Pi 5, Gemma 4 E2B hits 99 tokens/sec prefill and 9 tokens/sec decode with a 1432MB peak footprint. On Android GPUs expect 52 tokens/sec decode; on iOS Metal, 56. On a MacBook Pro via WebGPU, up to 76 tokens/sec. These are the numbers that make real-time, on-device voice assistants and translators feasible.

Thinking Mode: Reasoning Without the Bloat

Gemma 4 models support a Thinking Mode scratchpad. LiteRT-LM exposes it so you can stream the raw reasoning to the UI or strip it entirely to save precious KV cache space in multi-turn sessions:

await for (final chunk in conversation.streamMessage(prompt)) {
  if (chunk.reasoning != null) {
    // stream to a collapsible "Thinking" panel, or drop it
  } else {
    // render the final answer
  }
}

In long agentic sessions, dropping reasoning tokens can be the difference between staying within the KV cache budget and running out of memory mid-task.

Putting It All Together

Here’s a complete optimization checklist for your on-device app:

class OnDeviceConfig {
  static const enableMtp = true;        // speculative decoding
  static const useSessionPersistence = true;
  static const stripThinkingInTurns = true;
  static const maxContextTokens = 4096;

  static Future<Conversation> createOptimizedConversation() async {
    final model = await FlutterGemma.getActiveModel(
      maxTokens: maxContextTokens,
      enableMtp: enableMtp,
    );
    return model.createConversation();
  }
}

Combined, these techniques are what let a fully offline app feel instant: MTP hides latency, session persistence removes re-prefill costs, constrained decoding keeps tool calls trustworthy, and the lean footprint leaves the phone responsive.

Conclusion & Next Steps

LiteRT-LM isn’t just a way to run a model on a phone—it’s a complete system for making on-device GenAI production-fast. You now understand the core techniques: speculative decoding, memory-mapped embeddings, KV-cache session management, constrained decoding, and dynamic context.

Next steps:

• Benchmark E2B vs E4B on your exact hardware with the CLI.
• Profile prefill vs decode separately—they have very different optimization levers.
• Experiment with backend selection (CPU/GPU/NPU) to find the thermal sweet spot.
• Build a real-time translator or voice assistant using streaming inference and session persistence.

The combination of Gemma 4 and LiteRT-LM has made on-device AI fast enough that the next generation of mobile apps won’t just call AI—they’ll be it.

References

• Blazing fast on-device GenAI with LiteRT-LM
• LiteRT-LM on GitHub
• LiteRT documentation
• Gemma 4 model family