SmolLM2-135M compiled leg: DSL → StableHLO → IREE showcase (gemma-iree precedent)

[michalharakal]

Split out from #272 — the eager-path scope of that issue is done (the
SmolLm2InferenceSpike.kt commonTest reproducer, AndroidSmolLm2E2eTest.kt,
and the feature/smoll disposition are all resolved). This issue tracks the
one item that's still open: the compiled path.

Ask

Add the SmolLM2-135M-Instruct compiled leg — DSL → nn-DAG → StableHLO → IREE
— following the :llm-runtime:gemma-iree precedent:

1. A gemma3 (or smollm2) → StableHLO export on the host side, matching
   :llm-inference:gemma's existing gemma3 → StableHLO export that produces
   the .vmfb + .irpa pair.
2. A :llm-runtime:smollm2-iree-shaped module (or extend gemma-iree if the
   architecture is close enough) with the on-device decode loop —
   IreeRuntime + a *Decoder, mirroring GemmaDecoder/GemmaKvDecoder in
   llm-runtime/gemma-iree/src/nativeMain.
3. Wire it into the same cross-target reproducer harness from SmolLM2-135M as the standard cross-target reproducer: commonTest inference spike + showcase for eager and DSL → nn-DAG → StableHLO → IREE #272 so eager
   vs. compiled numbers are directly comparable per target (JVM, Android,
   iOS simulator/device, linuxX64/Arm64, macosArm64).

Why this model

SmolLM2-135M-Instruct (138 MiB, GGUF Q8_0) is already the standard
cross-target reproducer per #272 — small enough for CI, phones and
simulators, llama-architecture. Using the same checkpoint for the compiled
leg gives an apples-to-apples eager-vs-compiled comparison per target instead
of introducing a second model just for this path.

Context

• Field report that motivated the reproducer effort: SKaiNET-developers/SKaiNET's mobile
  field report (context in SmolLM2-135M as the standard cross-target reproducer: commonTest inference spike + showcase for eager and DSL → nn-DAG → StableHLO → IREE #272).
• gemma-iree precedent: llm-runtime/gemma-iree (IreeRuntime.kt,
  GemmaDecoder.kt, GemmaKvDecoder.kt in src/nativeMain), paired with the
  host-side gemma3 → StableHLO export in :llm-inference:gemma.

[michalharakal]

Updated scoping after digging into the actual tooling landscape (not just this
repo's gemma-iree):

The JNI-on-Android path is real, not hypothetical

skainet-embedder-android already ships exactly this shape for
MiniLM sentence embeddings: DSL-authored model → StableHLO → single .vmfb
→ a thin JNI C shim (native/skainet_embedder_jni.c, ~140 lines) over the
IREE C API → AAR with jniLibs (arm64-v8a/armeabi-v7a) + vmfb/vocab as
assets → a :demo Activity app. Built via
SKaiNET-iree-toolchain's bin/iree-run android build-so. This directly
answers the open question from before: no new native-build infrastructure
needs inventing, there's a working template to copy.

The one thing embedder's shim doesn't need that SmolLM2 does: autoregressive
KV-cache decode instead of one static-shape forward pass. That algorithm
already exists, just on the wrong transport — gemma-iree's
GemmaKvDecoder.kt (two-graph prefill/with_past loop, dynamic cache dim,
board-verified token-for-token) drives it via subprocess iree-run-module
calls. Porting that loop into a stateful JNI shim — keeping the IREE
session's K/V device buffers alive across calls instead of round-tripping
through files — is the real new engineering here.

Revised split (three parts, each independently landable)

1. Host-side SmolLM2 StableHLO export (SKaiNET-transformers,
   llm-inference/llama): no StableHLO/vmfb harness exists for the llama
   architecture yet (unlike gemma/whisper/bert) — needs the
   trace → StableHloConverterFactory → external-params → .irpa pattern
   from RealGemmaBakeIrpaTest.kt, replicated against LlamaNetworkLoader.
   The HLO converters (attention/RoPE/RMSNorm) are already
   architecture-agnostic, so this is plausible with today's compiler, not
   blocked on new compiler work.
2. Android KV-cache decode module: new JNI shim porting
   GemmaKvDecoder's two-graph loop to the IREE C API (stateful session,
   not subprocess), packaged as an AAR the way skainet-embedder-android
   packages its embedder — likely its own repo/module rather than living
   inside gemma-iree, since it's Android-only (JNI, not Kotlin/Native).
3. Android demo app (SKaiNET-examples): a compiled-path counterpart to
   AndroidNeonLlmDemo, same UI shape, wired to the module from (2) instead
   of the eager JNI NEON runtime — ideally with a NEON-vs-IREE-compiled
   toggle mirroring that demo's NEON-vs-scalar one.

Starting on (1) — it's self-contained, host-only, and doesn't block on any
Android-side build tooling.

[michalharakal]

Part 1 (host-side SmolLM2 StableHLO export) is up: #306. Verified end-to-end against the real Q8_0 checkpoint — see the PR description for the export numbers. Parts 2 (Android JNI KV-cache decode module) and 3 (compiled-path demo app) are still open.

[michalharakal]

Part 1 continued: #308 adds a reusable :llm-inference:smollm2 export module (redecode graph, in-graph argMax tail) and verifies the FULL pipeline numerically — DSL trace → StableHLO → bf16 externalized weights → iree-compile → iree-run-module — with a hand-driven redecode loop producing a correct, coherent completion for a real prompt. Also have a working, symbol-verified JNI shim over the IREE C API (Part 2), cross-built via SKaiNET-iree-toolchain, not yet landed pending a repo-placement decision.

[michalharakal]

Generalized the JNI/IREE Android runtime and moved it into SK-TR: #309.

• No longer SmolLM2-specific — `llm-runtime/iree-android` drives any
  DSL-compiled fixed-seq redecode vmfb with an in-graph argMax tail (the
  one hardcoded string, the function name, is now a runtime parameter).
  Plain `com.android.library`, matching the engine repo's
  `skainet-backend-jni-cpu` precedent — no KMP precedent exists for
  jniLibs bundling in this codebase.
• Both ABIs (`arm64-v8a`, `armeabi-v7a`) and both HAL drivers (CPU
  `local-task`, GPU `vulkan`) built and symbol-verified.
• `llm-inference/smollm2/docs/smollm2-vmfb.md` — a tool-generated
  reference (real `iree-run-module`/`iree-dump-parameters` output, not
  hand-written) for the SmolLM2 export this runtime was verified against.
  Also documents that Vulkan compilation of this export currently fails
  — a real IREE 3.11.0 SPIR-V codegen gap on the token-embedding gather,
  not a runtime limitation.
• The local Android demo consuming it (decoder+demo, now demo-only) was
  updated accordingly and verified to package correctly (both .sos, both
  ABI vmfbs, shared irpa, in the debug APK).

Still not verified anywhere: running on a physical device or emulator — no
device/emulator available in the environment this was built in.

[michalharakal]

#306, #308, #309 all merged into `develop`.

Landed:

• `:llm-inference:smollm2` — the redecode-graph StableHLO export (argMax
  tail), verified numerically via `iree-run-module` (host CPU), plus a
  tool-generated reference doc for the compiled artifact.
• `:llm-runtime:iree-android` — a generic (not SmolLM2-specific) JNI
  runtime over the real IREE C API, both ABIs (`arm64-v8a`,
  `armeabi-v7a`), both HAL drivers (CPU `local-task`, GPU `vulkan`)
  compiled in and symbol-verified.

Still open, tracked outside this repo / left for follow-up:

• On-device verification — nothing has run on a physical device or
  emulator yet.
• Vulkan for SmolLM2 specifically is blocked on a real IREE 3.11.0 SPIR-V
  codegen gap (token-embedding gather), not this runtime.
• The Android demo app consuming this (local, not yet in a public repo)
  still needs a placement decision — SKaiNET-examples vs. elsewhere.
• Wiring `iree-android` into `kllama`'s `registerPlatformBackends` for
  zero-config app consumption, matching how the eager NEON path is
  auto-discovered.