On-Device LLM in Compose: MediaPipe, StateFlow, ViewModel

Meta description: Wire MediaPipe LLM Inference to Jetpack Compose — streaming tokens via StateFlow, coroutine scoping, and GPU cleanup patterns for production.

---

TL;DR

Google’s MediaPipe LLM Inference API gives you a production-grade, hardware-accelerated path to on-device LLMs without wrestling raw NNAPI or llama.cpp bindings. The real engineering challenge is plumbing: streaming tokens into Compose without triggering recomposition storms, keeping inference off the main thread, and cleaning up GPU resources when users navigate away mid-generation. This post walks through the architecture that handles all three.

---

Why MediaPipe over raw llama.cpp?

GPU memory leaks from abandoned inference sessions are the top silent crash source in on-device LLM apps. That risk is meaningfully lower with MediaPipe than with the alternatives.

Approach	Setup complexity	GPU/NPU delegation	Streaming API	Memory management
Raw llama.cpp JNI	High	Manual	Manual callbacks	DIY
NNAPI direct	Very high	Built-in	None	DIY
MediaPipe LLM Inference	Low	Automatic	Built-in async	Handled by session lifecycle

MediaPipe abstracts LlmInference.Session with built-in GPU delegation, a documented lifecycle, and a callback-based streaming API that maps cleanly onto Kotlin coroutines. In my experience building production systems, the 2–3 days you spend on integration pay back within the first month of debugging you avoid.

---

The session lifecycle

Most teams get this backwards: they treat LlmInference as a singleton and Session as throwaway. It should be the inverse.

// ViewModel init — create once, reuse across prompts
private val inference = LlmInference.create(context, options)

// Per-conversation — create fresh, close explicitly
private var session: LlmInference.Session? = null

fun startSession() {
    session?.close()
    session = inference.createSession()
}

LlmInference holds the loaded model weights in GPU/NPU memory — expensive to create, must live for the ViewModel’s lifetime. Session carries conversation state (KV cache) and must be closed to release that slice of GPU memory when the conversation ends or the user navigates away.

---

Streaming tokens into Compose without recomposition storms

The naive approach — updating a MutableStateFlow<String> by concatenating each token — works until it doesn’t. At 30+ tokens/second, you get a recomposition on every emission. The fix: buffer at the ViewModel layer, not the UI layer.

private val _tokenBuffer = MutableStateFlow("")
val outputText: StateFlow<String> = _tokenBuffer
    .sample(50) // emit at most every 50ms
    .stateIn(viewModelScope, SharingStarted.Lazily, "")

.sample(50) meaningfully reduces recompositions on mid-range devices in my testing during active generation. Your Compose Text composable reads outputText via collectAsStateWithLifecycle() and recomposes at a human-perceivable rate rather than at inference speed.

On the collection side, MediaPipe’s streaming callback feeds the buffer:

fun generate(prompt: String) {
    inferenceJob = viewModelScope.launch {
        session?.generateResponseAsync(prompt) { partialResult, done ->
            // This callback runs on MediaPipe's internal thread, not the coroutine's thread
            _tokenBuffer.update { it + partialResult }
            if (done) _isGenerating.value = false
        }
    }
}

The viewModelScope.launch wrapper here is not about thread dispatch — generateResponseAsync registers a callback and returns immediately, so the coroutine body completes almost instantly. The value is structured cancellation: it gives you a Job handle that, when cancelled, lets you signal intent and gate any post-callback work. The callback itself runs on MediaPipe’s internal thread regardless of which dispatcher you choose.

---

Cancellation and GPU cleanup

When a user taps back mid-generation, the inference callback keeps firing into a ViewModel that’s been cleared. Two layers of defense:

override fun onCleared() {
    inferenceJob?.cancel()
    session?.close()
    inference.close()
    super.onCleared()
}

And in the composable, cancel on DisposableEffect:

DisposableEffect(Unit) {
    onDispose { viewModel.cancelGeneration() }
}

inferenceJob?.cancel() stops any post-callback coroutine work from proceeding and signals that generation should halt. session?.close() and inference.close() release the actual GPU allocations. Skipping either step produces a memory leak that’s invisible in small tests and catastrophic in production — especially on devices with shared CPU/GPU memory.

---

Takeaways

1. Scope LlmInference to the ViewModel, Session to the conversation. Swapping these kills performance or leaks GPU memory. The distinction is non-negotiable.

2. Buffer token emissions with .sample(50) before exposing to Compose. This is the single highest-leverage change for UI smoothness on mid-range hardware during active generation.

3. Implement two-layer cancellation: Job.cancel() + session.close(). The Job provides structured cancellation of coroutine-scoped work; closing the session releases hardware resources. You need both — they do different things.

---

#android #jetpackcompose #kotlin #mobile #architecture