深入 Android 端侧 LLM 推理的流式输出全链路：从 Token 生成到 Compose UI 增量渲染的实时交互架构
去年做端侧智能助手时踩了个坑：模型推理速度远快于 UI 消费速度，界面卡顿、内存暴涨。直觉反应是加 buffer 缓冲，结果越加越糟。后来发现，问题不在于缓冲大小，而在于整条链路缺乏统一的背压（Backpressure）机制。下面是我在这条链路上踩过的坑和最终落地的架构。
链路全景：Token 从生成到上屏经过三关
端侧 LLM 的流式输出拆成三个阶段：
- 推理层：LLM 引擎（MediaPipe LLM Inference / llama.cpp）逐 Token 生成，输出到回调或 Flow
- 传输层：将 Token 流从 Native 层传递到 Kotlin 层，涉及线程切换和缓冲策略
- 渲染层：Compose UI 接收 Token 增量，做增量重组而非全量刷新
每个阶段的时间尺度不同：推理层 10-50ms 生成一个 Token，渲染层在主线程上每 16ms 一帧。不做协调，要么掉帧，要么 OOM。
KV Cache 的内存困境
端侧 LLM 推理最吃内存的是 KV Cache（Key-Value Cache）。每生成一个 Token，增量大致为 2 × 层数 × KV 维度 × 精度字节数。
具体数字因模型架构差异很大。同样 2B 参数的模型，用 Multi-Query Attention（1 个 KV head）和 Grouped-Query Attention（8 个 KV head），KV Cache 能差 8 倍。以 Gemma 2B（INT4 量化）为例：18 层、2048 维、MQA 架构，512 Token 的回复 KV Cache 约 10MB，尚可接受。但如果是 Llama 8B 这类 GQA 模型，同样 512 Token 直奔 70MB。多轮对话场景上下文拉到 4096 时，即便 INT4 量化，KV Cache 也能上 140MB——叠加上模型权重，6GB RAM 的设备就很吃力了。
// 估算 KV Cache 大小
fun estimateKVCache(
numLayers: Int, hiddenDim: Int, kvHeads: Int,
seqLen: Int, bytesPerElem: Int = 2 // FP16
): Long {
val perToken = 2L * numLayers * hiddenDim * kvHeads * bytesPerElem
return perToken * seqLen
}
// Gemma 2B: 2 * 18 * 2048 * 8 * 2 * 512 ≈ 576MB (含 K+V)
// INT4 量化后约 144MB，仍不可忽视实际项目中的优化思路：
- 滑动窗口：只保留最近 N 个 Token 的 KV Cache，旧 Cache 直接丢弃。对话场景下用户很少关心 2000 个 Token 前的上下文。
- 量化 KV Cache：使用 INT8 甚至 INT4 存储 KV Cache，精度损失在可接受范围内。
- 主动回收：推理结束后立即释放 Native 内存，不要等 GC。
Flow 背压：从「能跑」到「不崩」
推理引擎的回调跑在 Native 线程。直接回调到 Kotlin 层，如果用 Channel.UNLIMITED 或没有背压策略，生产者持续发射 Token，消费者来不及处理，内存就会持续堆积。
我最初用的是 MutableSharedFlow，很快就炸了——Flow 默认没有背压，生产者不会暂停。
最后换成 callbackFlow + Channel.RENDEZVOUS（零缓冲通道），推理线程在消费者没准备好时自动挂起：
fun streamTokens(modelPath: String, prompt: String): Flow<String> = callbackFlow {
val inference = MediaPipeLlmInference.create(modelPath)
inference.setTokenCallback(object : TokenCallback {
override fun onToken(token: String) {
// trySendBlocking 会阻塞 Native 线程直到消费者接收
// 天然形成背压：推理速度 = UI 消费速度
trySendBlocking(token)
}
override fun onComplete() {
close()
}
})
inference.generate(prompt)
awaitClose { inference.close() }
}trySendBlocking 是这个方案的关键：Channel 满时阻塞调用线程，恰好让 Native 推理线程停下来，形成天然背压。推理引擎内部生成下一个 Token 前会等待回调返回，整条链路速度由最慢的一环决定。
这套方案在 200+ Token 长回复的场景下，内存峰值从 380MB 降到 120MB，没有再出现过 OOM。
Compose 增量重组：每帧只刷新变化的部分
拿到 Token 流后，最直接的做法是拼成完整文本整体刷新。但 512 个 Token 的文本，每次新增一个 Token 就全量重新排版，主线程直接卡死。
更好的方式是用增量重组（Incremental Recomposition）。Compose 的 Snapshot 系统天然支持这一点，核心是用 mutableStateListOf 管理 Token 列表而非拼接字符串：
@Composable
fun StreamingChatBubble(tokens: SnapshotStateList<String>) {
Column {
// 已确认的 Token：不再重组
Text(
text = buildAnnotatedString {
tokens.dropLast(1).forEach { append(it) }
}
)
// 最新 Token：带打字机效果，独立重组
if (tokens.isNotEmpty()) {
Text(
text = tokens.last(),
modifier = Modifier.alpha(1f)
// 仅这一个 Text 参与重组
)
}
}
}再加上 derivedStateOf 避免不必要的重组传播：
val displayText by remember {
derivedStateOf {
tokens.joinToString("")
}
}
// displayText 只在 tokens 列表结构变化时重新计算
// 列表内元素修改不会触发重组实测效果：500 Token 流式输出，主线程每帧耗时从 80ms 降到 4ms 以内，帧率稳定在 58-60fps。核心原因是只重组最后一个 Token 对应的 Composable，前面的静态文本 Compose 会智能跳过。
三端协同的完整时序
把三个环节串起来：
[推理线程] 生成 Token → trySendBlocking → Channel.RENDEZVOUS
↓ (背压：推理线程暂停，直到...)
[消费协程] 接收 Token → 追加到 mutableStateListOf
↓ (Snapshot 通知)
[主线程] Compose 重组 → 仅刷新最后一个 Text 节点
↓ (16ms 内完成)
[消费协程] 下一次 trySendBlocking 解除阻塞 → 推理继续这个设计里，推理速度由 UI 渲染速度决定，而不是反过来。KV Cache 的释放时机与推理完成同步，不会有悬空内存。
实测数据与调优建议
在 Pixel 6（8GB RAM）上跑 Gemma 2B INT4，200 次对话测试（平均每轮 300 Token）：
| 指标 | 优化前 | 优化后 |
|---|---|---|
| 内存峰值 | 380MB | 120MB |
| 主线程帧耗时 P99 | 82ms | 5ms |
| 端到端延迟（首 Token） | 1.2s | 0.9s |
| OOM 率 | 12% | 0% |
调优中总结的几点心得：
- Buffer 不是解药：
Channel.BUFFERED默认容量 64，UI 持续慢于推理的话，64 个 Token 后照样卡。RENDEZVOUS才是正解。 - KV Cache 窗口设 1024 足够：移动端对话场景，1024 Token 上下文覆盖率在 95% 以上。设太大除了浪费内存没有收益。
- 量化优先于剪枝：INT4 量化的精度损失远小于激进剪枝，端侧部署优先选量化方案。
- 用
derivedStateOf做防抖：Token 生成太快（< 5ms）时，加一个简单防抖——50ms 内累积的 Token 批量提交到 State，减少重组次数。但要权衡交互实时感，延迟超过 100ms 用户就能感知到卡顿。
这套架构已经在三个端侧 AI 功能（智能回复、文档摘要、代码补全）上线，稳定性经住了线上考验。如果你也在做类似的事，从背压机制入手是性价比最高的切入点。