Repository navigation
moonshine v2: export CLI from a Hugging Face snapshot (no Python, runs from the published module) - #446
Merged
Conversation
MoonshineV2Config gains slidingWindows -- explicit (left, right) context per encoder layer in the HF encoder_config.sliding_windows convention. The German tiny streaming checkpoint uses [[17,5],[17,5],[17,1],[17,1],[17,5],[17,5]]: its middle layers look one frame ahead, which the edge-layer rule (lookahead on the first/last N layers, causal in between) cannot express. Null keeps the English tiny behaviour unchanged. decoderDim / decoderFfnDim let the decoder run at a different width than the encoder (the small checkpoints: 620 / 512); the adapter then carries the bias-free proj bridging the memory into decoder space, the encoder attention takes an explicit head dim when nHeads * headDim != dim, and the frontend takes the stream width as a parameter. Defaults are the English tiny values, so existing callers and the exported graphs are unchanged. Tests: MoonshineV2ConfigTest (common; JVM + linuxX64).
…nable from the published module MoonshineV2ExportCli (:llm-inference:moonshine:exportMoonshineV2) turns a moonshine_streaming snapshot (config.json, model.safetensors, tokenizer.json) into the five StableHLO graphs of the streaming contract -- frontend, encoder, adapter, masked fixed-pad prefill, dynamic-cache with_past -- plus dec_embed.bin, vocab.bin and manifest.json. Geometry, per-layer attention bands, vocabulary size and the positional-table length come from the snapshot. With MOONSHINE_GRAPH=all the export fails if a checkpoint tensor was not used. Until now the export lived in jvmTest and read per-tensor .bin dumps produced by Python scripts, so it could not be run from the published artifact. The checkpoint is now read directly (skainet-io-safetensors) through MoonshineV2HfWeightMap (common): DSL parameter -> checkpoint tensor, with the three non-trivial cases (transposed filterbank, zero-centred encoder norm scales, absent biases). jvmMain carries the host tooling; commonMain keeps its lang-core + transformer-core footprint. Parity with the previous flow, checked locally against real checkpoints: German tiny -- all five graphs, dec_embed.bin and vocab.bin byte-identical. English tiny -- encoder, adapter, prefill, with_past, both tables byte-identical; the frontend has the same structure and differs only in conv-weight constants by <= 7.2e-7 (the old flow took them from a weight-norm-folded ONNX). Tests: MoonshineV2HfWeightMapTest (common: every parameter of every module maps, tiny and split-width); MoonshineV2ExportTest (JVM: full export of a synthetic snapshot generated from the model parameter list, tiny and split-width, entry signatures, dynamic vs static cache, tables round-trip, unused-tensor report).
Merged
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
1 participant
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
MoonshineV2ExportCli— a Hugging Facemoonshine_streamingsnapshot in, the complete streaming export out, runnable from the published module and without Python.frontend.mlir@moonshine_v2_frontend[1, samples]→ featuresencoder.mlir@moonshine_v2_encoderadapter.mlir@moonshine_v2_adapterprefill.mlir@moonshine_v2_decoder_prefillwith_past.mlir@moonshine_v2_decoder_with_pastdec_embed.bin,vocab.binmanifest.jsonGeometry, per-layer attention bands, vocabulary size and the positional-table length are read from the snapshot — nothing to configure per language or size. With
MOONSHINE_GRAPH=allthe export fails if any checkpoint tensor was not used. Env contract in the KDoc ofMoonshineV2ExportCli.kt; shapes default to a 1.28 s encoder window and a 5.12 s utterance.Why
The export lived in
jvmTest, gated on env variables, reading per-tensor.bindumps produced by Python scripts that are not in this repository. Nobody could run it from the published artifact. Same move asFunctionGemmaExportCli.How
MoonshineV2HfWeightMap(common): DSL parameter → checkpoint tensor for frontend, encoder, adapter, decoder. The DSL's[out, in]linears match the checkpoint, so almost everything is a plain copy; the three exceptions are explicit: transposed filterbank, zero-centred encoder norm scales (gamma + 1), biases the model does not have (zeros).MoonshineV2Checkpoint(jvm):config.json→MoonshineV2Config, tensors throughskainet-io-safetensors(F32 / F16 / BF16), tracks what was read.MoonshineV2ExportHarness(jvm): the five traces + the two tables + manifest.jvmMaincarries the host tooling;commonMainkeeps its lang-core + transformer-core footprint. The existingjvmTestbake tests are untouched.Verification
Parity with the previous flow on real checkpoints (local, not in CI):
dec_embed.bin,vocab.binFull export of one checkpoint: ~22 s.
CI tests:
MoonshineV2HfWeightMapTest(common, JVM + linuxX64: every parameter of every module maps, tiny and split-width);MoonshineV2ExportTest(JVM: a synthetic snapshot is generated from the model's own parameter list and exported end to end — tiny and split-width, entry signatures, dynamic vs static cache, tables round-trip, unused-tensor report). Module totals locally: jvmTest 27/27, linuxX64Test 7/7;:llm-inference:moonshine:buildgreen for all targets.