On-Device AI

Streaming tokens into Compose without recomposing the world

Streaming output updates state far faster than a normal screen. Narrow what reads the changing value, and let the smallest possible scope recompose.

MobileByteSensei 8 min read
One value changing fast, and everything that reads it
on-devicecomposeperformance

Most screens change state a few times per interaction. A streaming response changes it continuously, for several seconds, while the user watches.

Nothing about that is conceptually hard. It is just that every structural shortcut which was free at one update per tap becomes a per-frame cost at thirty, and they are all in code that looked fine.

The case this post is about

One 400-token response at 30 tokens a second, read at two different depths:

read at the screen (text passed down as a parameter):
  recomposing: ChatScreen, Header, MessageList, Message[0], Message[1], StreamingMessage, TokenText, InputBar, SendButton
  9 nodes x 30 tok/s x 13s = 3,510 recompositions
read at the leaf (lambda deferred to TokenText):
  recomposing: TokenText
  1 nodes x 30 tok/s x 13s = 390 recompositions

Nine times the work for one changing string, and eight of those nine nodes recompose to produce byte-identical output. The input bar re-renders while the user is typing in it.

Compose is doing exactly what it was told — the value is read at the top, so the top is what invalidates. This post follows those two numbers, because the throttle everybody reaches for first takes 3,510 to 1,170 and leaves the scope exactly as wide as it was.

The symptom

The response arrives as a stream. You accumulate it into state and render it, which is correct and obvious, and the screen gets slow in a way that does not look like a rendering problem — scrolling stutters, the input lags, an unrelated animation drops frames.

The cause is almost always scope. The changing value is read somewhere higher in the tree than the place that displays it, so a large subtree recomposes on every token: the surrounding card, the scroll container, the sibling elements that did not change and cannot know that.

Compose is doing exactly what it was told. The instruction was wrong.

Where the value is read, and what that costs

Why the obvious fix failed

The obvious fix is to update less often — buffer the stream and emit every 100ms instead of every token.

It works, and it is a throttle rather than a fix: the recomposition is still too wide, you have just made it happen less. The scroll still stutters, less often. And you have traded away the one thing streaming output is for, which is the sense that something is happening immediately.

The second attempt is to hoist the state higher so "everything can see it". That is the opposite of the correct direction and it is a very natural instinct, because hoisting is usually the right answer for shared state. Here the value is not shared — one element displays it.

The actual mechanism

Make the smallest possible region the only thing that reads the changing value.

Read the value where it is displayed, not where it is owned. State can live in the ViewModel; what matters is where it is read during composition, because that read is what registers the dependency. Passing it down through three layers as a parameter means all three recompose.

Pass a lambda rather than a value when you must cross a boundary. A deferred read moves the dependency to the place that invokes it, which is the leaf you actually wanted.

Separate the stable parts from the streaming part in the tree. The header, the controls and the surrounding chrome should not be siblings-in-the-same-recomposition-scope as the text. Hoisting the stable things out of the changing scope is the move, which feels backwards and is the point.

Keep the accumulation off the composition. Building the growing string is work; doing it inside composition means doing it again on every recomposition, including ones triggered by something else.

Recomposition scope is "which nodes read this value", and that is countable — which is the only way to tell a narrowed read from a read that merely looks narrower:

class Scope(val name: String, val children: List<Scope> = emptyList()) {
    var reads = false
    fun all(): List<Scope> = listOf(this) + children.flatMap { it.all() }
}

// A node recomposes if it reads the value, or if any ANCESTOR of it does — a parent
// recomposing re-invokes its children whether or not they care.
fun recomposed(root: Scope, ancestorRead: Boolean = false): List<String> {
    val hit = ancestorRead || root.reads
    return (if (hit) listOf(root.name) else emptyList()) +
        root.children.flatMap { recomposed(it, hit) }
}

fun frames(root: Scope, tokensPerSecond: Int, seconds: Int) {
    val hit = recomposed(root)
    println("  recomposing: ${hit.joinToString()}")
    println("  %d nodes x %d tok/s x %ds = %,d recompositions"
        .format(hit.size, tokensPerSecond, seconds, hit.size * tokensPerSecond * seconds))
}

The same screen with the read at two depths — a 400-token response at 30 tokens/second:

fun screen() = Scope("ChatScreen", listOf(
    Scope("Header"),
    Scope("MessageList", listOf(Scope("Message[0]"), Scope("Message[1]"),
        Scope("StreamingMessage", listOf(Scope("TokenText"))))),
    Scope("InputBar", listOf(Scope("SendButton"))),
))

println("read at the screen (text passed down as a parameter):")
screen().also { it.reads = true }.let { frames(it, 30, 13) }

println("\nread at the leaf (lambda deferred to TokenText):")
screen().also { s -> s.all().first { it.name == "TokenText" }.reads = true }.let { frames(it, 30, 13) }
read at the screen (text passed down as a parameter):
  recomposing: ChatScreen, Header, MessageList, Message[0], Message[1], StreamingMessage, TokenText, InputBar, SendButton
  9 nodes x 30 tok/s x 13s = 3,510 recompositions

read at the leaf (lambda deferred to TokenText):
  recomposing: TokenText
  1 nodes x 30 tok/s x 13s = 390 recompositions

Nine times the work for one changing string, and every one of the other eight nodes recomposes to produce byte-identical output. Note what the throttle fix does to these numbers: emitting every 100ms instead of every token takes the first case from 3,510 to 1,170 — a real reduction, still three times the correct version, and the input bar still recomposes while the user is typing in it.

The same tree, with the read at two different depths

The fix

The structural rule I would apply: a value that changes at frame rate should have exactly one reader, and that reader should be a leaf.

Everything else follows from it. If a second thing genuinely needs to react — a character count, a scroll-to-bottom, an enabled state — derive it so that the derivation only republishes when its own result changes, rather than every time the source does. Most such derivations change far less often than the string: a boolean "is non-empty" changes exactly once during an entire stream.

That asymmetry is the whole opportunity. The text changes constantly; almost everything computed from the text changes rarely. Pushing the recomposition boundary to sit between those two facts is what makes a streaming screen cheap.

Scroll position is the case where this gets subtle, because it is the one piece of behaviour that genuinely does need to react at frame rate — a transcript that grows should stay pinned to the bottom. The temptation is to drive that from the text itself, which puts a frame-rate dependency into whatever owns the scroll container, usually the widest scope on the screen. The cheaper shape is to react to the thing you actually care about — the content's measured height — and to let the scroll effect observe that rather than the string.

The second-order version of the same mistake is worth watching for: a key that changes every token. If the growing text is used as an identity anywhere — a list key, a remember key, an animation target — then every token invalidates whatever that identity was protecting, and the symptom looks like state resetting rather than like a performance problem. That one is particularly unpleasant to debug, because the screen behaves correctly in every test where the text is fixed.

What changes constantly, and what changes once

How we proved it

The boundary, stated plainly, because this article is architecture rather than measurement.

Not measured: I have not profiled a streaming on-device model rendering into a Compose screen. I have no frame-time numbers, no comparison of buffering intervals, and I am not quoting any — the figures depend on the device, the surrounding tree and the token rate, and a number from my machine would mislead more than it informed.

What I can point at is the general discipline this codebase enforces around Compose state, where the recurring defect is a value read at the wrong level rather than an expensive leaf. The gates here catch structural problems — a handler that does nothing, a route nothing reaches, a string that bypasses resources — and the pattern across all of them is the same: the defect is in where something is declared, not in how it is computed.

What I would measure before changing anything: turn on recomposition counts and find out which composables are recomposing per token. Almost always the answer is "more of them than you think", and the list tells you exactly where the boundary is wrong. Doing that first is cheaper than every optimisation in this article and will usually make three of them unnecessary.

What I expect to be specific to inference

Token rate is not steady. A stream from a local model arrives in bursts shaped by compute and thermal state, so a buffering interval tuned on one device may be wrong on another. Anything frame-rate-dependent should be measured across device classes rather than tuned on the best one.

The first token is the latency that matters. Users judge time-to-first-token far more than tokens-per-second, which argues for rendering the very first fragment as fast as possible even if the rest is buffered more aggressively.

Cancellation has to be immediate. A user who navigates away or edits the prompt should stop paying for generation instantly — and a stream that keeps producing into a composable that no longer exists is both wasted compute and a leak.

What this generalizes to

When a value changes orders of magnitude more often than the rest of your state, the question is not how to render it faster. It is how few things are allowed to notice.

Narrow the read to one leaf, derive anything else so it republishes only when its own value changes, and keep the accumulation out of the render path. Then measure what actually recomposes rather than reasoning about it — the gap between where you think the dependency is and where it is turns out to be the entire problem, and it is directly observable.

One post a week, no digest

New writing when it lands. No roundups, no course pitch, no "10 tips".

Follow by RSS