Speculative decoding on Android: 2–3x token throughput

Meta description: Learn how to wire a Qwen 0.5B draft model to a 7B target in llama.cpp on Android, achieving 40+ tokens/sec on Pixel 9 using the draft/verify loop.

---

TL;DR

Speculative decoding pairs a small, fast draft model with a larger target model to generate multiple candidate tokens per forward pass. On Android with llama.cpp, this yields 2–3x token throughput, hitting 40+ tokens/sec on a Pixel 9 (Snapdragon X Elite, Q4_K_M quantization, GPU backend via NNAPI) with no measurable accuracy loss. Getting there requires careful attention to KV cache coordination, NNAPI limitations, and thread-safe batching to keep the UI responsive.

---

Why speculative decoding belongs on device

On-device inference is bottlenecked by memory bandwidth, not compute. A large model like a 7B parameter network spends most of its time loading weights from RAM into registers — the actual matrix multiplications are comparatively cheap. Speculative decoding exploits this asymmetry.

A draft model (Qwen 0.5B in this configuration) runs roughly 10–15x faster than a 7B target. If the draft model correctly predicts even 3 out of 5 tokens, you’ve effectively tripled throughput while the target model is a verifier, not the primary generator.

---

The draft/verify loop

The architecture:

┌─────────────────────────────────────────┐
│         Draft Model (Qwen 0.5B)         │
│  Generates k candidate tokens in batch  │
└────────────────┬────────────────────────┘
                 │ k tokens
                 ▼
┌─────────────────────────────────────────┐
│        Target Model (7B+)               │
│  Verifies all k+1 positions in one pass │
└────────────────┬────────────────────────┘
                 │ accept/reject per token
                 ▼
         Accepted tokens emitted;
         first rejection resets draft

In practice, you generate a batch of k=4 or k=5 tokens from the draft model, then pass the original context plus all draft tokens into the target model in a single forward pass. The target produces logits for every position simultaneously. You walk left to right, accepting tokens where the draft and target distributions agree (within a threshold), and truncating at the first disagreement.

The key insight: the target model’s single batched forward pass costs roughly the same as generating one token autoregressively. Accepting even two draft tokens doubles effective throughput.

---

KV cache sharing: where most implementations break

Most teams get this wrong. Naively running two independent llama.cpp contexts burns memory and destroys the latency gains. The correct approach is to maintain a shared prefix in the KV cache.

Both models process the confirmed prompt context identically. Once verification accepts tokens up to position n, both KV caches are trimmed to n in lockstep. The draft model then extends from n, and the target re-evaluates only the new draft span — not the entire sequence.

This requires explicit KV cache management calls in llama.cpp:

// After acceptance, sync both caches to confirmed length
// -1 means remove all tokens after confirmed_len through the end of the sequence
llama_kv_cache_seq_rm(ctx_draft, 0, confirmed_len, -1);
llama_kv_cache_seq_rm(ctx_target, 0, confirmed_len, -1);

Skipping this step causes cache divergence, which either corrupts output or forces full re-evaluation, eliminating the speedup entirely.

---

NNAPI constraints that break naive implementations

NNAPI acceleration on Android introduces hard constraints that aren’t obvious from the documentation.

Constraint	Impact	Mitigation
Static input shapes	Batched verification fails if shape changes per step	Pre-allocate fixed k+1 batch size, pad as needed
No dynamic graph recompilation	Model swap mid-session crashes	Initialize both models at startup
INT8 quantization mismatch	Draft/target quantization must be compatible	Use same quantization scheme across both
Memory mapping limits	Two large models may exceed NNAPI buffer limits	Quantize to Q4_K_M or smaller

The most painful failure mode is shape dynamism. NNAPI compiles the model graph at load time against fixed dimensions. If your verification batch varies from 3 to 6 tokens, you’ll hit shape errors at runtime. The fix: always pass a fixed-size batch (e.g., always 5 tokens), padding with a sentinel token and masking the attention accordingly.

---

Batching strategy for a responsive UI thread

Blocking the main thread on Android will trigger ANR errors — inference must run entirely on a background dispatcher. The batching strategy that works in production:

// Coroutine-based producer on inference dispatcher
launch(Dispatchers.Default) {
    val draftTokens = draftModel.generateBatch(context, k = 5)
    val accepted = targetModel.verify(context, draftTokens)
    
    accepted.forEach { token ->
        _tokenFlow.emit(token) // StateFlow consumed by UI
    }
}

Run both models on a dedicated Dispatchers.Default coroutine. Emit accepted tokens to a StateFlow or Channel that the UI collects. This keeps the main thread free while tokens stream in at 40+ per second.

The batching cadence matters too. Emitting individual tokens causes excessive UI recomposition in Jetpack Compose. Buffering accepted tokens per verification round (typically 2–4 at a time) and emitting them as a group reduces recomposition overhead measurably.

---

Three things to get right first

1. Sync KV caches after every verification round. Cache divergence is the silent killer here — trim both caches to confirmed length before the next draft batch, every time.

2. Pre-allocate fixed NNAPI batch shapes at startup. Dynamic input shapes cause runtime failures on NNAPI backends. Commit to a fixed speculation width k and pad every batch to that size.

3. Stream tokens via StateFlow from a background coroutine. The UI thread must never touch inference. Buffering accepted token batches before emitting reduces Compose recomposition and produces smoother perceived streaming at high throughput.

---

#android #mobile #architecture #kotlin #llm