Streaming LLM tokens in Android with Compose

Meta description: Wire MediaPipe’s LLM Inference API to Jetpack Compose with Kotlin Flow back-pressure, correct buffer sizing, and a 16ms recomposition budget that keeps mid-range devices jank-free at 20+ tokens/sec.

---

TL;DR

MediaPipe’s LLM Inference API emits tokens via a callback, not a suspend function. Naively forwarding each token as a StateFlow update triggers one recomposition per token — at 25 tokens/sec on a mid-range device, that blows the 16ms frame budget. The fix: buffer tokens into timed chunks (≤80ms windows), bridge the callback to a Channel, and let Compose diff a single string per frame rather than per token.

---

The streaming contract MediaPipe gives you

MediaPipe’s on-device LLM inference exposes a setListener callback:

inferenceSession.addQueryChunk(prompt)
inferenceSession.generateAsync(
    { partialResult, done ->
        onToken(partialResult, done)
    }
)

Two problems. First, the callback fires on an internal MediaPipe thread — not the main thread, not a coroutine dispatcher you control. Second, there’s no back-pressure signal; the model pushes tokens as fast as it can decode them, regardless of whether your UI has caught up.

I’ve seen both constraints ignored in production systems, and the result is always the same: frame drops that only show up on mid-range devices under thermal throttling, which is exactly the hardware your users actually own.

---

Bridging callbacks to Kotlin Flow

The cleanest bridge is a Channel with a bounded buffer:

fun LlmInferenceSession.tokenFlow(prompt: String): Flow<String> =
    channelFlow {
        val buffer = Channel<String>(capacity = 64, onBufferOverflow = BufferOverflow.SUSPEND)

        addQueryChunk(prompt)
        generateAsync { partial, done ->
            runBlocking { buffer.send(partial) }
            if (done) buffer.close()
        }

        for (token in buffer) send(token)
    }.flowOn(Dispatchers.Default)

A capacity of 64 covers roughly 2–3 seconds of burst output at 25 tok/sec before SUSPEND kicks in and applies natural back-pressure to the MediaPipe thread. runBlocking inside the callback is intentional — it’s the right primitive to block the producer when the consumer is behind, without dropping tokens.

---

The recomposition budget problem

Here’s what most teams get wrong: Compose doesn’t batch state writes coming from collectAsState the same way it batches layout passes. Each emission from your flow can schedule a recomposition independently.

Device tier	Vsync budget	Max safe recompositions/sec	Safe token batch window
Flagship (Snapdragon 8 Gen 3)	8.3ms (120Hz)	~60	16ms
Mid-range (Snapdragon 6 Gen 1)	16ms (60Hz)	~30	33ms
Budget (Helio G88)	16ms (60Hz)	~15	66ms

At 25 tokens/sec, a budget device recomposing on every token will miss frames. The solution is chunked emission:

fun Flow<String>.chunkByTime(windowMs: Long = 80L): Flow<String> =
    flow {
        val accumulator = StringBuilder()
        var lastEmit = System.currentTimeMillis()

        collect { token ->
            accumulator.append(token)
            val now = System.currentTimeMillis()
            if (now - lastEmit >= windowMs) {
                emit(accumulator.toString())
                accumulator.clear()
                lastEmit = now
            }
        }
        if (accumulator.isNotEmpty()) emit(accumulator.toString())
    }

An 80ms window caps recompositions at ~12/sec — safely below the budget-device ceiling — while keeping perceived latency well under the 100ms threshold where users notice UI lag.

---

Wiring it into Compose

@Composable
fun StreamingResponse(session: LlmInferenceSession, prompt: String) {
    val text by remember(prompt) {
        session.tokenFlow(prompt)
            .chunkByTime(windowMs = 80L)
            .runningReduce { acc, chunk -> acc + chunk }
    }.collectAsState(initial = "")

    Text(text = text, modifier = Modifier.fillMaxWidth())
}

runningReduce accumulates the full response string without holding a mutable StringBuilder in composition scope. The remember(prompt) key ensures the flow restarts cleanly on new prompts, avoiding stale emissions.

---

Profiling checklist

Before shipping, run these checks in Android Studio’s Compose Layout Inspector:

• Skipped frames should be zero during active streaming on your lowest-supported device
• Recomposition count on Text should be ≤15/sec
• Log buffer.isEmpty periodically on channelFlow; sustained false means the UI is falling behind

---

Before you ship

Don’t bridge MediaPipe callbacks through MutableStateFlow. A Channel with BufferOverflow.SUSPEND is what prevents dropped tokens when the producer is a native thread — this part isn’t optional.

Tune your chunk window to your minimum supported device, not the Pixel on your desk. 80ms is a safe default; you can tighten it to 33ms if your floor is mid-range 60Hz hardware, but I’d start conservative and profile from there.

Prefer runningReduce over a stateful StringBuilder in composition. It keeps accumulated state in the flow pipeline where it belongs, makes the composable itself stateless, and makes testing significantly less painful.

---

kotlin android jetpackcompose mobile architecture