feat(smollm2): compiled-export module — redecode graph with argMax tail (#305) - #308
Merged
Merged
Conversation
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
New
:llm-inference:smollm2module — the whisper/functiongemma "one exportsurface" shape.
SmolLm2ExportHarness.export()tracesllamaNetwork()fromthe real SmolLM2 GGUF, strips per-layer KV caches, and appends the DSL argMax
tail so the emitted
func @smollm2returnstensor<seqxi32>directly(small per-step output — no host-side argmax over
[seq, vocab]logits).This is the redecode-graph counterpart to
FunctionGemmaExportHarness'sexportRedecodeand matchesGemmaDecoder's consumption pattern: onefixed-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:llamafor the architecture; owns the exportproduct, mirroring
functiongemma's relationship togemma.Verification — the full pipeline, numerically, not just compiled
SmolLm2ExportHarness.export()against the realSmolLM2-135M-Instruct-Q8_0.gguf: 393 params, 310 MiB bf16 archive,func.func @smollm2(%arg0: tensor<1x24xi32>) -> (tensor<24xi32>).iree-convert-parameters→smollm2.irpa,iree-compile --iree-hal-target-backends=llvm-cpu→ a host vmfb (viaSKaiNET-iree-toolchain'sskainet/iree-compilerimage).GemmaDecoder-style redecode loop by hand viairee-run-modulefor the prompt "The capital of France is" (tokens1,504,3575,282,4649,314): 8 greedy steps decode to "the city ofParis, 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'sskainet/iree-androidimage — but it isn'tlanded anywhere yet pending a decision on where it should live (new repo
vs. existing).
Test plan
./gradlew :llm-inference:smollm2:compileTestKotlinJvm— clean, no warningsSMOLLM2_GGUF=... ./gradlew :llm-inference:smollm2:jvmTest --tests '*SmolLm2ExportHarnessTest*' -PincludeIntegration— passes (393 params, 310 MiB, correct function signature)Advances #305.