Gemma-3 2B on Android: KV-cache, prefill, and quantization
Meta description: KV-cache math, Snapdragon 7 vs 8 prefill benchmarks, and INT4 vs INT8 tradeoffs for running Gemma-3 2B on mid-range Android.
TL;DR
Running Gemma-3 2B on a 4 GB Android device is feasible, but only with INT4 quantization and a carefully budgeted context window. The KV-cache grows at roughly 104 KB per token on the 2B architecture. Prefill latency on Snapdragon 7-series degrades sharply beyond 512 tokens. Keep your combined token budget under 1,500 and you can ship a production-grade on-device LLM experience on mid-range hardware.
Why on-device LLM architecture decisions matter now
Google’s MediaPipe LLM Inference API gives Android engineers a clean task-based interface for local language model inference — no REST endpoint, no cloud round-trip, no data leaving the device. For privacy-sensitive use cases, this matters.
But the documentation glosses over the operational reality: memory is tight, context windows are expensive, and prefill latency on mid-range silicon can make or break perceived UX.
KV-cache: where your memory budget actually goes
Most teams benchmark the model weights and forget the KV-cache. That’s where the memory surprises come from.
The KV (Key-Value) cache stores the intermediate attention state for every token in context. It’s what allows autoregressive decoding to avoid recomputing the full attention matrix on every step — a worthwhile tradeoff, but one with a linear memory cost.
The KV-cache formula
KV cache (bytes) = 2 × num_kv_heads × head_dim × num_layers × context_len × dtype_bytesFor Gemma-3 2B, which uses grouped query attention with 4 KV heads, 256-dimensional heads, and 26 transformer layers, the per-token cost in FP16 (2 bytes per value) is:
2 × 4 × 256 × 26 × 1 token × 2 bytes = 106,496 bytes
106,496 ÷ 1,024 = 104 KB per token (using 1 KB = 1,024 bytes)| Context Length | KV-Cache Size |
|---|---|
| 256 tokens | ~26 MB |
| 512 tokens | ~52 MB |
| 1,024 tokens | ~104 MB |
| 2,048 tokens | ~208 MB |
| 4,096 tokens | ~416 MB |
At 2,048 tokens, the cache alone consumes 208 MB. Add model weights and activations and you’re already threatening the foreground memory limit on a 4 GB device.
Quantization tradeoffs: INT4 vs INT8
Gemma-3 2B has approximately 2 billion parameters. Quantization affects the memory picture significantly:
| Quantization | Weight Size | 4 GB Device Feasible | Quality Delta |
|---|---|---|---|
| FP16 (baseline) | ~4 GB | No | — |
| INT8 | ~2 GB | Marginal (6 GB+ recommended) | <1% perplexity |
| INT4 | ~1 GB | Yes (with context budget) | 2–4% perplexity |
INT4 is the only practical path on 4 GB devices. The quality regression is real but tolerable for most inference workloads. If your application requires high-fidelity multi-step reasoning, target INT8 and 6 GB+ devices. MediaPipe bundles quantization into the .task model format — you select the correct file at integration time and the runtime handles the rest.
NPU vs GPU: why MediaPipe uses the GPU delegate
A common question from engineers targeting mid-range devices: why not route inference through the NPU for lower power draw and faster decode?
MediaPipe LLM Inference API dispatches through the GPU delegate (OpenGL ES / Vulkan compute) by default. Direct NPU/HTP (Hexagon Tensor Processor) access for LLM workloads isn’t exposed through the public MediaPipe LLM API — NNAPI delegation works for smaller vision tasks but transformer-based inference still routes through the GPU. Qualcomm’s AI Hub SDK is the alternative if you need HTP acceleration on Snapdragon silicon, at the cost of a heavier integration and different model formats.
The GPU path in MediaPipe is well-optimized on 8-series silicon, so the NPU gap is narrowest there. On 7-series devices, watch for MediaPipe updates — NPU support for LLM tasks is on the roadmap and will meaningfully improve both latency and battery budget when it ships.
Prefill latency: Snapdragon 7 vs 8 series
Prefill is the phase where the model processes your entire prompt in parallel before generating the first output token. The gap between tiers is severe.
| SoC | GPU | 128-tok Prefill | 512-tok Prefill | 1024-tok Prefill |
|---|---|---|---|---|
| Snapdragon 8 Gen 3 | Adreno 750 | ~0.6s | ~2.0s | ~4.5s |
| Snapdragon 8 Gen 2 | Adreno 740 | ~0.8s | ~2.6s | ~5.5s |
| Snapdragon 7 Gen 2 | Adreno 644 | ~1.8s | ~6.5s | ~14s |
| Snapdragon 7s Gen 3 | Adreno 720 | ~2.0s | ~7.2s | ~16s |
Figures are estimates derived from relative GPU compute throughput ratios across SoC tiers. Verify GPU model numbers against Qualcomm’s official SoC specification pages before citing. Capture device-specific baselines with Perfetto or Android GPU Inspector — real-world results vary with thermal state, MediaPipe version, and concurrent workloads.
A 14-second blank screen before the first token is a UX failure. On 7-series devices, keep your system prompt under 300 tokens.
Wiring the MediaPipe LLM Inference API
val options = LlmInference.LlmInferenceOptions.builder()
.setModelPath("/data/local/tmp/gemma3-2b-it-int4.task")
.setMaxTokens(1500) // pre-allocated KV-cache budget
.setMaxTopK(40)
.setTemperature(0.8f)
.build()
val llmInference = LlmInference.createFromOptions(context, options)
llmInference.generateResponseAsync(prompt) { partial, done ->
// stream tokens to UI
}setMaxTokens is your single most important configuration knob — it determines how much KV-cache MediaPipe pre-allocates at initialization. Set it too high on a 4 GB device and the OOM killer will terminate your process before the first inference. I recommend 1,200–1,500 for mid-range INT4 deployments.
Before committing to a device tier, measure the RSS delta before and after LlmInference.createFromOptions() using adb shell dumpsys meminfo. Android Studio’s Memory Profiler visualizes the allocation spike in real time; Perfetto gives you a system-level view of memory pressure across the full stack. The .task file size is not your runtime footprint — don’t use it as a proxy.
The practical token budget on 4 GB devices
Working backwards from a ~1.6 GB ceiling for a foreground app on a 4 GB device:
- INT4 weights: ~1,000 MB
- Activations + MediaPipe runtime overhead: ~200 MB
- Remaining for KV-cache: ~400 MB
- Theoretical maximum context: ~3,900 tokens
In practice, Android memory pressure pushes the safe operational limit to 1,200–1,500 tokens. Beyond that, you’re gambling on the low-memory killer.
3 actionable takeaways
-
Set
setMaxTokensconservatively and load-test at initialization. On 4 GB INT4 deployments, start at 1,200 tokens. Memory is allocated upfront — there’s no recovery path from a mid-session OOM. -
Optimize for prefill, not just decode, on 7-series SoCs. Users tolerate slower token generation better than a long blank wait before the first token appears. Keep system prompts under 300 tokens and design your loading state accordingly.
-
Profile with Android Studio Memory Profiler and Perfetto, not file size. The
.taskfile size is not your runtime footprint. The RSS delta before and afterLlmInference.createFromOptions()is — measure that number before you commit to a device tier.
Tags: android · kotlin · mobile · architecture · multiplatform