Skip to content

feat(smollm2): compiled-export module — redecode graph with argMax tail (#305) - #308

Merged
michalharakal merged 1 commit into
developfrom
feat/smollm2-export-harness-305
Aug 12, 2026
Merged

michalharakal merged 1 commit into
developfrom
feat/smollm2-export-harness-305

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Summary

New :llm-inference:smollm2 module — the whisper/functiongemma "one export
surface" shape. SmolLm2ExportHarness.export() traces llamaNetwork() from
the real SmolLM2 GGUF, strips per-layer KV caches, and appends the DSL argMax
tail so the emitted func @smollm2 returns tensor<seqxi32> directly
(small per-step output — no host-side argmax over [seq, vocab] logits).
This is the redecode-graph counterpart to FunctionGemmaExportHarness's
exportRedecode and matches GemmaDecoder's consumption pattern: one
fixed-seq vmfb invoked in a loop over a growing, causally-masked-safe padded
buffer (the two-graph KV-cache decode is deliberately out of scope — see #305).

Depends on :llm-inference:llama for the architecture; owns the export
product, mirroring functiongemma's relationship to gemma.

Verification — the full pipeline, numerically, not just compiled

  1. SmolLm2ExportHarness.export() against the real
    SmolLM2-135M-Instruct-Q8_0.gguf: 393 params, 310 MiB bf16 archive,
    func.func @smollm2(%arg0: tensor<1x24xi32>) -> (tensor<24xi32>).
  2. iree-convert-parameterssmollm2.irpa, iree-compile --iree-hal-target-backends=llvm-cpu → a host vmfb (via
    SKaiNET-iree-toolchain's skainet/iree-compiler image).
  3. Drove the GemmaDecoder-style redecode loop by hand via
    iree-run-module for the prompt "The capital of France is" (tokens
    1,504,3575,282,4649,314): 8 greedy steps decode to "the city of
    Paris, a city of"
    — correct and coherent. This confirms the DSL trace,
    argMax tail, bf16 externalization, and IREE execution are numerically
    sound end to end, not just that the graph compiles.

What's not in this PR

Per #305's split: compiling this graph for Android arm64 and the JNI decode
module that drives it on-device. I have a working (compiled, symbol-verified)
spike of that JNI shim over the IREE C API — cross-built via
SKaiNET-iree-toolchain's skainet/iree-android image — but it isn't
landed anywhere yet pending a decision on where it should live (new repo
vs. existing).

Test plan

  • ./gradlew :llm-inference:smollm2:compileTestKotlinJvm — clean, no warnings
  • SMOLLM2_GGUF=... ./gradlew :llm-inference:smollm2:jvmTest --tests '*SmolLm2ExportHarnessTest*' -PincludeIntegration — passes (393 params, 310 MiB, correct function signature)
  • Manual end-to-end numeric verification via iree-compile + iree-run-module (see above)
  • CI (no model file present, so the integration test self-skips — expected)

Advances #305.

…rgMax tail (#305)

New :llm-inference:smollm2 module (whisper/functiongemma "one export
surface" shape): SmolLm2ExportHarness.export() traces llamaNetwork() from
the real SmolLM2 GGUF, strips per-layer KV caches (a fixed-seq prefill pass
needs none), and appends the DSL argMax tail (ectx.ops.argMax + squeeze) so
the emitted `func @smollm2` returns `tensor<seqxi32>` directly — small
per-step output, matching FunctionGemmaExportHarness's `exportRedecode`
graph and GemmaDecoder's re-decode consumption pattern (one fixed-seq vmfb,
invoked in a loop over a growing, causally-masked-safe padded buffer).
Weights externalize as bf16 (bit-exact truncation, halves the archive).
Depends on :llm-inference:llama for the architecture; owns the export
product, same relationship functiongemma has to gemma.

Verified beyond compilation — the FULL pipeline, numerically:
  1. SmolLm2ExportHarness.export() against the real Q8_0 checkpoint: 393
     params, 310 MiB bf16 archive, func.func @smollm2(%arg0: tensor<1x24xi32>)
     -> tensor<24xi32>.
  2. iree-convert-parameters -> smollm2.irpa (SKaiNET-iree-toolchain compiler
     image), iree-compile --iree-hal-target-backends=llvm-cpu -> a host vmfb.
  3. Drove the GemmaDecoder-style redecode loop by hand via iree-run-module
     for "The capital of France is" (tokenized 1,504,3575,282,4649,314):
     8 greedy steps decode to "the city of Paris, a city of" — correct and
     coherent, confirming the DSL trace, argMax tail, bf16 externalization,
     and IREE execution are all numerically sound end to end.

Follow-up (not in this module): compiling for Android arm64 and the JNI
decode module that drives it on-device (transformers#305 part 2).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@michalharakal
michalharakal merged commit 2df051c into develop Aug 12, 2026
2 checks passed
@michalharakal
michalharakal deleted the feat/smollm2-export-harness-305 branch August 12, 2026 07:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant