Generic Android JNI runtime for the DSL → StableHLO → IREE compiled path
(transformers#305) — the compiled-path counterpart to skainet-backend-jni-cpu
(engine repo) for the eager path: both serve any model without knowing
which one is calling them.
This module drives any DSL-compiled vmfb that follows the redecode graph
contract established by :llm-inference:smollm2's SmolLm2ExportHarness
(the first producer of a compatible triple — see its
docs/smollm2-vmfb.md):
- one fixed-
seqexported function,tensor<1xSEQxi32> -> tensor<SEQxi32> - the DSL's in-graph
argMaxalready applied (small per-step output, no host-side argmax over a[SEQ, vocab]logits tensor) - weights external — bound at session-create time from a
.irpaparameter archive under scope"model", not baked into the vmfb
Unlike :llm-runtime:gemma-iree (Kotlin/Native, drives iree-run-module as
a subprocess on a Linux board — structurally unusable on Android, which runs
on ART/JVM, not K/N), this is the first real JNI-against-the-IREE-C-API code
in this repo.
val decoder = IreeRedecodeDecoder.fromAssets(
context,
vmfbAsset = "smollm2/smollm2-gen-arm64.vmfb", // pick per Build.SUPPORTED_ABIS
irpaAsset = "smollm2/smollm2.irpa",
functionName = "module.smollm2",
seq = 24,
cacheDirName = "skainet_smollm2",
)
val generated = decoder.generate(promptTokenIds, eosTokenId = tokenizer.eosTokenId)
decoder.close()No tokenizer dependency — pair this with whatever BPE tokenizer your app already has.
| Class | Role |
|---|---|
IreeRedecodeDecoder |
Facade: fromAssets(context, vmfbAsset, irpaAsset, functionName, seq, cacheDirName, device), generate(promptIds, eosTokenId, maxNewTokens) |
IreeRedecodeSession |
Raw JNI wrapper: padded [seq] token buffer in → [seq] predicted-next-token ids out. Package/class name is the JNI symbol contract with the .so — do not move/rename. |
device is just an IREE HAL driver string:
IreeRedecodeSession.DEFAULT_DEVICE("local-task", CPU) — the primary, numerically-verified path (see the smollm2 docs linked above).IreeRedecodeSession.VULKAN_DEVICE("vulkan", GPU — Mali, Adreno, etc., portable SPIR-V) — the.soin this module is built with--vulkanand the Vulkan HAL driver compiled in, so it's ready whenever a caller has avulkan-spirv-compiled vmfb. SmolLM2's own export currently does not compile for Vulkan (a real IREE 3.11.0 SPIR-V codegen gap on its token-embedding gather — see the smollm2 docs) — that's a limitation of that specific export, not of this runtime.
native/iree_redecode_jni.c — see its header doc for the full IREE C API
call sequence (session create → parse .irpa → io_parameters VM module →
append before the compiled bytecode module, since its util.global
initializers resolve against it at link time → invoke by name → transfer
result → rank/length-guard against the redecode contract).
Cross-built via SKaiNET-iree-toolchain's
skainet/iree-android:3.11.0 image:
native/build-iree-redecode.sh arm64-v8a --vulkan
native/build-iree-redecode.sh armeabi-v7a --vulkan
# copy native/out/libskainet_iree_redecode.so to src/main/jniLibs/<abi>/Because weights are external (unlike skainet-embedder-android's
baked-in-weights precedent), the .so links three IREE targets beyond the
default iree_runtime_unified: iree_modules_io_parameters_parameters,
iree_io_parameter_index, iree_io_parameter_index_provider,
iree_io_formats_irpa_irpa — any external-weights vmfb needs these, so
they're unconditional in the build script, not per-model.
The .so is a checked-in artifact, not built by CI — regenerate it by
re-running the script when iree_redecode_jni.c changes. Migrating to
AGP's own externalNativeBuild/CMake (matching skainet-backend-jni-cpu's
in-repo build, against IREE's export-android-sdk output) is real future
work: it means vendoring/consuming the whole IREE runtime build, not a
handful of hand-written kernel .c files, and deserves its own
investigation.
- Generic, model-agnostic native shim + Kotlin API
- Both ABIs (
arm64-v8a,armeabi-v7a) cross-built, JNI symbols verified - Both CPU (
local-task) and GPU (vulkan) HAL drivers compiled into the.so - Standalone module build verified (
./gradlew :llm-runtime:iree-android:assembleRelease) - Verified running on a physical Android device or emulator (none available in the environment this was built in)
- Wired into
kllama'sregisterPlatformBackendsfacade for zero-config app consumption (matches how NEON kernels are auto-discovered on the eager path) — natural follow-up, not done here
The stateful counterpart of the redecode session, for the FunctionGemma contract addendum
(gemma_prefill_at, gemma_prefill_with_past, gemma_with_past; host-gather variants with an
emb input right after the tokens — the embedding rows are read from the with-past archive by the
native side, so callers pass token ids only):
val spec = IreeKvSpec.functionGemma270m(chunk = 32) // or IreeKvSpec.fromManifest(json)
val session = IreeKvSession(spec, IreeKvSession.VULKAN_DEVICE,
"$dir/gemma-with-past-hostgather-valhall4.vmfb", "$dir/gemma-with-past.irpa",
"$dir/gemma-prefill-with-past-hostgather-valhall4.vmfb", "$dir/gemma-prefill-with-past.irpa",
"$dir/gemma-prefill-at-hostgather-valhall4.vmfb", "$dir/gemma-prefill-at.irpa")
val decoder = IreeKvDecoder(session, prefillSeq = 1024)
val catalog = decoder.prefillPrefix(catalogPromptIds) // once per process (25 s on a MagentaTV One, 843 tokens)
session.releasePrefill() // drops the prefill archive mapping
val ids = decoder.generate(catalog, utteranceIds, eosTokenId = 106, maxNewTokens = 32) // per turnWhat the native side does per call: sliding layers (l % globalLayerPeriod != period-1) only ever
see their last slidingWindow cache positions through zero-copy tail views; RoPE cos/sin tables
(split-half, sign folded into the first half) and the chunk graph's per-head additive masks
(causal band, padding, window) are built for the absolute positions; K/V outputs replace the
retained views; snapshot()/restore() retain/release views without copying. Every failure is
thrown as IllegalStateException with the formatted IREE status (also logged under skainet_iree_kv).
Measured (MagentaTV One, Mali via Vulkan, bf16 archives, chunk 32, 843-token catalog prefix,
16 decode tokens): open 12.4 s, prefix 25.2 s once, then p50 5.9 s per utterance (one chunk call
≈ 2.0 s + 16 × 0.245 s), restore 0 ms, RSS ≈ 1.5 GB in the 32-bit process. Rebuild the library
with native/build-iree-kv.sh <abi> --vulkan (same image and links as the redecode .so).