Android GPU inference for real-time semantic search

---

TL;DR

Quantize a sentence-transformer to INT8, route inference through NNAPI or Vulkan compute delegate based on device capability at runtime, build an in-process HNSW index with hnswlib-kotlin, and profile the full pipeline. On a Pixel 9, this stack achieves sub-40ms end-to-end retrieval — fast enough for keystroke-level semantic search without a cloud round-trip.

---

The problem

200–400ms server round-trips are the default for on-device search. They don’t have to be. Users perceive latency above ~100ms as lag, and above ~300ms as broken. For a local notes app, document retrieval, or an on-device RAG pipeline feeding an edge LLM, that server hop is a tax you don’t need to pay — if you architect the inference and retrieval stack correctly.

---

Step 1: Quantizing the sentence-transformer to INT8

Start with a compact sentence-transformer (e.g., all-MiniLM-L6-v2, 384-dim output). Export to TFLite with full INT8 quantization using a representative dataset for calibration:

def representative_dataset():
    for sample in calibration_samples:
        yield [np.array(sample, dtype=np.float32)]

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()

Full INT8 quantization requires a representative dataset for calibration. With it, model size drops from ~90MB (FP32) to ~23MB, and inference arithmetic shifts to integer ops — exactly what NNAPI and Vulkan compute shaders accelerate.

---

Step 2: NNAPI vs Vulkan delegate — runtime routing

Most teams pick one delegate and ship it. That’s wrong. The right approach is capability detection at runtime.

Delegate	Best on	Latency (Pixel 9)	Notes
NNAPI	Devices with DSP/NPU	~8–14ms per embedding	Falls back to CPU silently
Vulkan compute	GPU-capable devices	~10–18ms per embedding	More predictable; explicit control
CPU (XNNPACK)	Fallback	~35–60ms per embedding	Always available

Define the capability checks explicitly — don’t leave these as magic calls:

fun isNnapiSupported(): Boolean =
    Build.VERSION.SDK_INT >= Build.VERSION_CODES.P

fun isVulkanSupported(context: Context): Boolean =
    CompatibilityList().isDelegateSupportedOnThisDevice(GpuDelegateFactory.Options())

Then route dynamically, using explicit variable assignment to avoid delegate chaining ambiguity:

fun buildInterpreter(context: Context): Interpreter {
    val options = Interpreter.Options()
    return when {
        isNnapiSupported() -> {
            val delegate = NnApiDelegate()
            options.addDelegate(delegate)
            Interpreter(model, options)
        }
        isVulkanSupported(context) -> {
            val gpuOptions = GpuDelegate.Options().apply {
                inferencePreference = GpuDelegate.Options.INFERENCE_PREFERENCE_FAST_SINGLE_ANSWER
            }
            val delegate = GpuDelegate(gpuOptions)
            options.addDelegate(delegate)
            Interpreter(model, options)
        }
        else -> {
            options.setNumThreads(4)
            Interpreter(model, options)
        }
    }
}

On Pixel 9, NNAPI with the Tensor Processing Unit achieves ~8ms median embedding latency for a 128-token input. Vulkan lands at ~12ms but with lower variance — important for P99 budgets.

---

Step 3: HNSW index sizing on memory-constrained devices

The memory math for 100,000 vectors at 384 dimensions in FP32:

• Raw vectors: 100k × 384 × 4 bytes ≈ 147MB
• HNSW graph overhead (M=16, efConstruction=200): ~60–80MB
• Total: ~210–230MB

On a device where the LLM context window may hold 1–2GB of KV cache, this is workable but requires discipline. Keep embeddings in INT8 (147MB → ~37MB) and cap index size at 50k entries for a safe ceiling under 100MB total.

Add hnswlib-kotlin to your project first:

// build.gradle.kts
implementation("io.github.jelmerk:hnswlib-kotlin:1.1.0")

Then build and tune the index:

val index = HnswIndex.newBuilder(384, DistanceFunctions.FLOAT_COSINE_DISTANCE, 100_000)
    .withM(16)
    .withEfConstruction(200)
    .withEf(64) // query-time ef — tune for recall vs latency
    .build()

Tune ef at query time. ef=64 delivers ~97% recall at 384 dims with ~2–4ms query latency. Dropping to ef=32 cuts that to ~1ms at the cost of ~3% recall — acceptable for most retrieval-augmented use cases.

---

Step 4: End-to-end latency budget

Profiling on Pixel 9 (Tensor G4, 8GB RAM, Android 15):

Stage	P50	P99
Tokenization	1ms	2ms
Embedding (NNAPI INT8)	9ms	14ms
HNSW query (ef=64)	3ms	6ms
Result marshal to Kotlin	<1ms	1ms
Total	~14ms	~23ms

Sub-40ms is achievable and comfortable. Even on mid-range devices without a dedicated NPU, the Vulkan path plus XNNPACK fallback keeps P99 under 60ms.

---

Takeaways

Route delegates at runtime. Hard-coding NNAPI ships fine on Pixel and breaks silently on devices without an NPU. Detect, route, and instrument each path separately in production telemetry.

Keep embeddings INT8 end-to-end. FP32 vectors double your index memory budget for negligible recall improvement at 384 dimensions. Quantize at export with a representative calibration dataset and store quantized vectors in the index.

Profile the full pipeline, not just inference. Tokenization and marshaling overhead is real. The latency ceiling is a system property — measure it as one, and set ef based on your end-to-end P99 budget, not a micro-benchmark of the HNSW query alone.