DEV Community

SoftwareDevs mvpfactory.io
SoftwareDevs mvpfactory.io

Posted on Originally published at mvpfactory.io

Wiring Gemma 3n to Android's NNAPI via ExecuTorch: Multimodal On-Device Inference Under 2GB RAM

What We Are Building

By the end of this tutorial you will have a working Android inference pipeline that runs Gemma 3n — Google's multimodal on-device model — using ExecuTorch's XNNPACK and NNAPI delegates. We will cover every layer that teams consistently underestimate: model export with dynamic shapes, PT2E quantization to hit the sub-2GB target, Android's ImageReader preprocessing pipeline, and the threading model that keeps inference off the main thread without starving the UI compositor.

This is not a toy demo. Let me show you a pattern I use in every on-device AI project.


Prerequisites

  • Android device with NNAPI support (API 27+, mid-range or above)
  • Python environment with PyTorch 2.x and the ExecuTorch SDK installed
  • Familiarity with Kotlin coroutines and Android Camera2 basics
  • The Gemma 3n PyTorch checkpoint (available via Google's model hub)

Why ExecuTorch Over TFLite?

Most teams reach for TFLite out of habit, then discover its quantization story for large generative models is painful and its multimodal graph support is immature. ExecuTorch — Meta's production on-device inference runtime — delegates computation at the operator level: XNNPACK handles float/quantized CPU math while NNAPI offloads eligible ops to NPU or GPU accelerators.

Gemma 3n's MatMul-heavy transformer blocks are exactly what NPU accelerators are built for. That gap becomes obvious once you profile on real hardware.


Step 1: Export the Model with Dynamic Shapes

Start with the PyTorch checkpoint. ExecuTorch uses torch.export with a two-phase flow:

import torch
from executorch.exir import to_edge, EdgeCompileConfig

exported = torch.export.export(
    model,
    args=(image_tensor, input_ids, attention_mask),
    dynamic_shapes={
        "image_tensor": {0: torch.export.Dim("batch")},
        "input_ids": {1: torch.export.Dim("seq_len", max=512)},
    }
)

edge_program = to_edge(
    exported,
    compile_config=EdgeCompileConfig(_check_ir_validity=True)
)
Enter fullscreen mode Exit fullscreen mode

The to_edge step lowers to ExecuTorch's portable IR. From there, apply delegate backends before serializing to .pte.


Step 2: Quantize for the Sub-2GB Target

Gemma 3n in bf16 lands well above the 2GB ceiling for mid-range devices. Here is the minimal setup to get this working — PT2E quantization with per-channel int8 weights and dynamic activations:

from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
from executorch.backends.xnnpack.quantizer import XNNPACKQuantizer

quantizer = XNNPACKQuantizer().set_global(
    get_symmetric_quantization_config(is_per_channel=True, is_dynamic=True)
)
prepared = prepare_pt2e(exported_program, quantizer)
# Run calibration with representative image+text pairs
converted = convert_pt2e(prepared)
Enter fullscreen mode Exit fullscreen mode

Int8 weights with dynamic int8 activations on linear layers typically cut memory by ~4x versus fp32. Combined with Gemma 3n's architecture-level efficiency, staying under 2GB on a 4GB device is achievable while leaving headroom for the app runtime.


Step 3: Image Preprocessing via Android's ImageReader

Android does not have CVPixelBuffer. The equivalent pipeline runs through ImageReaderBitmap → normalized FloatBuffer:

val imageReader = ImageReader.newInstance(width, height, ImageFormat.YUV_420_888, 2)

imageReader.setOnImageAvailableListener({ reader ->
    val image = reader.acquireLatestImage() ?: return@setOnImageAvailableListener
    val bitmap = image.toBitmap() // extension via ImageUtils
    val tensor = bitmap.toNormalizedFloatTensor(mean = IMAGENET_MEAN, std = IMAGENET_STD)
    image.close()
    inferenceQueue.offer(tensor)
}, backgroundHandler)
Enter fullscreen mode Exit fullscreen mode

The YUV→RGB conversion is the expensive step. Keep it off the main thread — always.


Step 4: The Threading Model

Here is the architecture that keeps inference from starving your compositor:

Main Thread        ──→  UI updates only
ImageReader Thread ──→  YUV decode + normalization → CoroutineChannel
Inference Thread   ──→  ExecuTorch .forward() (pinned, high priority)
Result Thread      ──→  Dispatches to Main via Dispatchers.Main.immediate
Enter fullscreen mode Exit fullscreen mode

In Kotlin, pinning inference to a single high-priority thread avoids lock contention on the ExecuTorch module:

private val inferenceDispatcher = Executors.newSingleThreadExecutor { thread ->
    thread.apply { priority = Thread.MAX_PRIORITY - 1 }
}.asCoroutineDispatcher()

suspend fun runInference(imageTensor: FloatArray, tokens: IntArray): String =
    withContext(inferenceDispatcher) {
        module.forward(imageTensor, tokens)
    }
Enter fullscreen mode Exit fullscreen mode

Gotchas

Here is what will save you hours.

Quantization scheme locks in at export time. PT2E quantization must be applied pre-delegation. Retrofitting it after the .pte is serialized means re-exporting from scratch — an expensive loop on large multimodal models.

NNAPI delegation coverage is worth profiling explicitly. Use adb shell dumpsys nnapi and ExecuTorch's op partitioning logs to verify what percentage of ops actually land on the NPU. CPU fallback for unsupported ops silently kills your latency targets. There is no warning when it happens — the docs do not mention this, but it is the most common performance regression we see in production.

The ImageReader pipeline is a real bottleneck, not scaffolding. YUV conversion and tensor normalization are consistently underestimated. Benchmark end-to-end latency from camera frame to model output, not just model execution time, to find your actual ceiling.


Conclusion

Gemma 3n on Android is production-viable with ExecuTorch — but only if you treat every layer of the pipeline as a first-class engineering concern. Export with dynamic shapes, quantize before delegation, keep YUV conversion on a background thread, and pin your inference dispatcher. Get those four right and sub-2GB multimodal inference on mid-range hardware is not a stretch goal — it is the baseline.

For further reading: ExecuTorch documentation and the Gemma model card are the authoritative references for anything not covered here.

Top comments (0)