iree-android: stateful KV session (IreeKvSession / IreeKvDecoder, libskainet_iree_kv.so) — prefill once, snapshot, chunk per utterance, step per token (#410)

[michalharakal]

Implements the stateful KV contract of #410 for Android. Stacked on #416 (same module, same error-reporting style) — merge that first. Uses the graphs from #415 and #417.

What

IreeKvSession / IreeKvSpec / IreeKvDecoder over a new libskainet_iree_kv.so (native/iree_kv_jni.c, native/build-iree-kv.sh; both ABIs checked in, built with skainet/iree-android:3.11.0 --vulkan):

• Three IREE sessions on one device (gemma_prefill_at, gemma_prefill_with_past, gemma_with_past — each archive has its own key numbering, so they cannot share a parameter scope). The prefill session can be released once the catalog snapshot exists.
• Device-resident per-layer K/V. The 15 sliding layers only ever see their last slidingWindow (512) positions through zero-copy tail views (iree_hal_buffer_subspan + a new buffer view) — the with-past graphs carry no window mask (found in iree-android: stateful prefill/step contract with a reusable prompt-prefix KV snapshot (and an embeddings-input variant for Vulkan) #410's measurements: without this the token after a 852-token prompt is wrong).
• RoPE tables (split-half, sign folded into the first half, as GemmaKvDecoder.splitHalfCosSin) and the chunk graph's per-head additive masks (causal band, padding, window) are built natively for the absolute positions, so the Kotlin side deals in token ids only.
• Host-gather variants of all three graphs: the embedding rows are read from the with-past archive (the vocab × hidden bf16 entry, located through the parameter index and mmapped) and passed as the emb input — IREE 3.11's SPIR-V backend cannot lower the in-graph gather.
• snapshot() / restore() retain and release the current views — no copy; restore measured at 0 ms.
• Every failure throws IllegalStateException with the formatted iree_status_t (and logs under skainet_iree_kv).
• IreeKvDecoder: prefillPrefix(ids) → Snapshot once, generate(snapshot, utteranceIds, eos, maxNew) per turn (chunk calls of spec.chunk, then greedy steps), with per-stage timing.

Measured in-process (arm32 Android device, Android 14, armeabi-v7a, Mali via Vulkan, bf16 archives, chunk 32)

Golden-8 German voice-command utterances with the 843-token shared tool-catalog prefix (ids tokenised on the host with the official template) and 16 decode tokens:

stage	measured
session open (3 archives)	12.4 s; RSS 1,854 MB → 1,474 MB after releasePrefill()
catalog prefix prefill, 843 tokens, once	25.2 s
per utterance: restore	0 ms
per utterance: one chunk call (7–12 real tokens)	1.9–2.3 s
per utterance: 16 decode tokens	3.9 s (≈ 245 ms/token)
per utterance, total	p50 5.87 s, max 6.25 s

8/8 utterances return a complete <start_function_call>call:… (or the model's refusal prose) within budget; 2/8 token streams are identical to the eager JVM run for all 16 tokens, the rest diverge after the function name (bf16 weights vs the FP32 eager reference).

Notes

• nKvHeads must be 1 for the zero-copy row views (FunctionGemma); other layouts would need a copy.
• Not touched: IreeRedecodeSession (kept for the redecode contract), CMake/Gradle (the .so stays a checked-in artifact as before).