iree-android: stateful prefill/step contract with a reusable prompt-prefix KV snapshot (and an embeddings-input variant for Vulkan)

未关闭
  • 哪里不对？
#410 3 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
5/5
预计耗时
一周以上
新手友好度
30/100
Issue 类型
功能
描述清晰度
基本清楚
活跃度
停滞
技术栈
android, kotlin, machine-learning
领域
backend-api-design, mobile, performance

调研方向

Start by reading the Android runtime contract and the FunctionGemma export path mentioned in FunctionGemmaExportHarness.kt:258–383; compare the exported gemma_prefill and gemma_with_past graphs with the runtime’s current stateless step(IntArray) path. The proposal also spans device-side KV lifetime, JNI input and error handling, Vulkan embeddings, and head slicing. Done means implementing and validating the stateful prefill, snapshot, step, and release contract, including the embeddings-input variant, against the supplied parity golden sets on the target device.

由索引模型根据 Issue 内容生成。

描述

Motivation, with numbers

runtime-iree-android only offers the stateless redecode contract (step(IntArray) -> IntArray, whole sequence recomputed). On an arm32 Android device (32-bit process, 4× ARMv8, Mali) with FunctionGemma-270M FP32 and IREE 3.11.0:

SEQ CPU arm32 per step Vulkan per step (head sliced, embeddings gathered on the host)
64 10.4 s (6.6 s with the head sliced) 3.16 s
256 41.9 s 7.76 s
1024 77.3 s (sliced head) 29.9 s

A 15-token tool call against an 853-token catalog prompt therefore costs ≈ 19 min on the CPU and ≈ 7.5 min on the GPU. Prompt shortening does not fix it (short descriptions: −29 % tokens, −2/8 accuracy). With a KV contract that prefills the catalog once per process, snapshots the KV state, and per turn prefills only the utterance (≈ 36 tokens) plus ~15 decode steps, the same call projects to 8–29 s on the GPU. There is no other route to an interactive NLU on this class of device, and no arm64 firmware option (the product stays 32-bit).

Proposal

  1. Contract: prefill(ids) -> (lastId, kvHandle), snapshot(kvHandle) -> kvHandle', step(kvHandle, id) -> id, release(kvHandle); KV kept device-side (HAL buffers), never round-tripped through the JVM. The FunctionGemma export already emits gemma_prefill/gemma_with_past graphs (FunctionGemmaExportHarness.kt:258–383) — this is the runtime side of those.
  2. Head sliced to the target position in every graph (separate issue; 36 % of the step and the 32-bit memory cap).
  3. Embeddings-input variant (prefill(ids, embeddings)): IREE 3.11.0's SPIR-V backend cannot legalize the token-embedding gather (failed to legalize operation 'vector.step' on the gather dispatch). Gathering the 640-float rows on the host and passing tensor<1xSEQx640xf32> makes the same graph compile and run on Mali at 2.4–2.6× the CPU speed. The JNI currently accepts only IntArray.
  4. Error reporting through JNI (see the null-return issue).

I can provide the device, the measured graphs (MLIR with the sliced head and the host-gather input), the iree-benchmark-module scripts, and the golden sets for parity.

Context. Measured on 2026-09-03 while bringing FunctionGemma-270M up as an NLU cartridge on an arm32 Android device (Android 14, armeabi-v7a-only 32-bit process, 4× ARMv8 @ 2.0 GHz, Mali GPU) with SKaiNET 0.53.0 + SKaiNET-transformers 0.53.0 from Maven Central. Harness and raw result files: a local harness (can share on request).

主要语言
Kotlin
星标
8
派生
3
平均合并
2 天 3 小时
30 天内合并 PR
43

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南，环境需要你自己搭建：先看它的 README，通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

SKaiNET-developers/SKaiNET-transformers 的其他 Issue

查看 SKaiNET-developers/SKaiNET-transformers 的全部 Issue

相似的 Issue

更多 Kotlin Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。