Add Moonshine Streaming (moonshine_streaming) speech-to-text support - #589
Open
justinchuby wants to merge 3 commits into
Open
Add Moonshine Streaming (moonshine_streaming) speech-to-text support#589justinchuby wants to merge 3 commits into
justinchuby wants to merge 3 commits into
Conversation
Add complete support for the HuggingFace `moonshine_streaming` architecture, replicating `MoonshineStreamingForConditionalGeneration` as separate encoder and cached-decoder ONNX graphs. Validated against moonshine-ai/moonshine-streaming-tiny at pinned revision f8e9dfd8c562c257c151a907b7b7f2fe8ff8511a. Moonshine Streaming is not a variant of offline Moonshine's front end; it replaces it entirely and changes the encoder's positional scheme: * Framing front end: the raw waveform is reshaped into fixed 5 ms frames (80 samples at 16 kHz), per-frame mean/RMS normalised (CMVN), compressed by a learned `asinh(exp(log_k) * x)` gain, and projected by a bias-free linear layer followed by SiLU. The frame validity mask is `sum(mask) // frame_len`. * Causal downsampling: two left-padded stride-2 convolutions replace the centred convolution stem, so no future frame ever leaks backwards. The padding mask is pushed through the same receptive field with an "any valid" reduction, matching upstream's `conv1d(mask, ones, stride) > 0`. * No encoder RoPE: encoder self-attention is purely content based. Ordering comes from the causal stem plus per-layer asymmetric `(left, right)` sliding windows, where `right` is the strict lookahead in encoder frames and `right == 0` makes a layer fully causal. This is what bounds streaming latency, so the schedule is modelled per layer rather than as a single uniform window. * Unit-offset LayerNorm: encoder norms are affine-free LayerNorm scaled by `gamma + 1`, so the checkpoint stores `gamma`, not `weight`. * Context adapter: the decoder adds a learned absolute position table (`pos_emb`) to the encoder output and projects it to the decoder width when the two differ, before cross-attention. The decoder itself (partial interleaved RoPE, cached causal self-attention, cross-attention, fused gate/up SiLU MLP, bias-free LayerNorms) is identical to offline Moonshine, so `MoonshineAttention`, `MoonshineDecoderLayer` and `MoonshineDecoderModel` are reused directly rather than duplicated. ONNX initializer names match the checkpoint 1:1, so no renames beyond the shared `model.` prefix strip are needed. Supporting changes: * `MoonshineAttention` gains keyword-only `hidden_size`/`head_dim` so an encoder of a different width can reuse it, and `qkv_bias`/`o_bias` so the projections follow upstream's gating instead of being hardcoded bias-free. Upstream gates Q/K/V on `attention_bias` and keeps the decoder's output projection bias-free unconditionally; both defaults are unchanged, which is what every published Moonshine checkpoint uses. * `Conv1d.padding` accepts a `(left, right)` pair so a convolution can be made causal without a separate `Pad` node. * `SpeechToTextConfig.encoder_output_size` lets the decoder graph declare the real encoder width instead of assuming it equals the decoder's. * `scripts/generate_golden.py` now forwards the pinned `revision` on the speech-to-text path, which was previously unpinned. * ORT GenAI export rejects `moonshine_streaming` for the same reason it rejects `moonshine`: the runtime cannot host a variable-length raw-waveform encoder. Tests: L1 graph build and weight alignment, L2 YAML/full-config build, L3 synthetic parity (heavy padding, three window schedules, with and without attention bias) and real-weight parity on real nonzero speech, plus L4 and L5 goldens generated independently from the HuggingFace pipeline. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Comment on lines
+110
to
+116
| def forward( # type: ignore[override] | ||
| self, | ||
| op: OpBuilder, | ||
| hidden_states: ir.Value, | ||
| mask: ir.Value, | ||
| dtype: ir.DataType, | ||
| ) -> tuple[ir.Value, ir.Value]: |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds full moonshine_streaming (Moonshine Streaming) speech-to-text architecture support to mobius, including a new streaming encoder + reuse of the existing Moonshine decoder, plus pinned-revision goldens and extensive parity coverage to validate correctness end-to-end.
Changes:
- Introduces
MoonshineStreamingConfigandMoonshineStreamingForConditionalGeneration, implementing the raw-waveform framing frontend, causal conv downsampling with mask propagation, per-layer asymmetric sliding-window attention, unit-offset LayerNorm, and the decoder context adapter. - Updates the speech-to-text task contract to support encoders whose output width differs from the decoder (
encoder_output_size), and extends shared components (Conv1d asymmetric padding; MoonshineAttention bias gating / encoder-width reuse). - Adds L1–L5 tests and pinned golden artifacts for
moonshine-ai/moonshine-streaming-tiny, plus registry + documentation updates and ORT GenAI export rejection parity with offline Moonshine.
Reviewed changes
Copilot reviewed 21 out of 21 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
tests/moonshine_streaming_integration_test.py |
New numerical parity and generation tests (synthetic + real weights) for Moonshine Streaming. |
tests/build_graph_test.py |
Adds L1 graph/initializer/IO contract tests for the new streaming encoder+decoder package. |
tests/_test_configs.py |
Registers a representative tiny config for moonshine_streaming in the shared config matrix. |
testdata/golden/audio/moonshine-streaming-tiny.json |
Adds pinned golden logits/top-k reference for the audio case. |
testdata/golden/audio/moonshine-streaming-tiny_generation.json |
Adds pinned golden generation token sequence + transcript. |
testdata/cases/audio/moonshine-streaming-tiny.yaml |
Adds L4/L5 YAML test case for the streaming tiny checkpoint (pinned revision). |
src/mobius/tasks/_speech_to_text.py |
Uses config.encoder_output_size for decoder encoder_hidden_states input shape. |
src/mobius/models/moonshine.py |
Extends MoonshineAttention to support configurable projection bias and encoder-width reuse. |
src/mobius/models/moonshine_streaming.py |
New model implementation for streaming encoder + context-adapted decoder packaging. |
src/mobius/models/__init__.py |
Exports MoonshineStreamingForConditionalGeneration. |
src/mobius/integrations/transformers/_config_resolver_test.py |
Adds tests ensuring HF config routing/extraction for moonshine_streaming. |
src/mobius/integrations/ort_genai/auto_export.py |
Rejects ORT GenAI export for moonshine_streaming (same constraint as moonshine). |
src/mobius/integrations/ort_genai/auto_export_test.py |
Parametrizes rejection test over moonshine and moonshine_streaming. |
src/mobius/components/_whisper.py |
Extends Conv1d.padding to accept asymmetric (left, right) padding for causal convs. |
src/mobius/components/_whisper_test.py |
Adds unit tests for new Conv1d padding behavior. |
src/mobius/_registry.py |
Registers moonshine_streaming model type and default model id alias. |
src/mobius/_configs/_base.py |
Adds encoder_output_size to SpeechToTextConfig and introduces MoonshineStreamingConfig extraction/validation. |
src/mobius/_configs/__init__.py |
Exports MoonshineStreamingConfig. |
src/mobius/__init__.py |
Re-exports MoonshineStreamingConfig from the package root. |
scripts/generate_golden.py |
Threads pinned revision into speech-to-text golden generation model/processor loads. |
README.md |
Lists “Moonshine Streaming” in supported Speech-to-Text models. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+3220
to
+3221
| #: Initial value of the learned asinh compression gain (stored as ``log_k``). | ||
| encoder_compression_k: float = 0.75 |
The existing parity cases all use batch 1, so nothing exercised the encoder frame mask, the causal-convolution mask propagation or the cross-attention bias across rows with different real audio lengths. Add a three-row case with 150, 90 and 40 valid frames that asserts the downsampled mask matches HuggingFace exactly, that the three rows genuinely keep different frame counts, and that encoder hidden states and decoder logits agree per row. Also make the shared decoder feed helper batch-aware: `position_ids` was always built with a batch dimension of 1, which the graph rejects for larger batches because `RotaryEmbedding` requires the cos/sin cache to carry the same batch dimension as its input. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
…16+CUDA Two correctness gaps found in final review. **Encoder-context accumulation.** HuggingFace's `MoonshineStreamingDecoder` adds its absolute position table to `encoder_hidden_states` with `+=`, which mutates the caller's tensor: a single decoder call moves it by up to 1.73, and a 34-step decode leaves it 58.8 away from the encoder output. Cached cross-attention means the mutated tensor is never read again, so the generated tokens are in fact unaffected — verified by decoding three ways (`generate()`, a manual loop over a deliberately accumulating buffer, and a manual loop with a pristine clone per step), all of which produce identical tokens. But a golden reference must not *depend* on a dead-store argument. `scripts/generate_golden.py` now wraps the speech-to-text reference in `non_mutating_encoder_context`, which clones `encoder_hidden_states` on decoder entry, so the committed golden provably describes "encoder context is added exactly once" — the semantics the exported ONNX decoder implements, since it takes the encoder output as a read-only graph input. Regenerating L4 and L5 under the guard reproduces both files byte for byte, which is the evidence that the previous references were not silently accumulating. Three tests lock this down: * the guard is proven non-vacuous — the unguarded decoder is asserted to mutate its argument, and the guarded one is asserted not to; * the ONNX decoder is asserted not to accumulate: after a full decode its encoder input buffer is unchanged and replaying step 0 reproduces the original step-0 logits exactly; * every decode step is compared against a non-mutating HuggingFace reference on the **full** vocabulary distribution, not just the argmax, and the resulting token sequence is checked against the committed L5 golden. **float16 on CUDA.** ORT's CUDA float16 fused attention kernel returns encoder frame 0 as exactly zero for this model — that frame is the sparsest masked query row, seeing only 4 of 456 keys under the `(16, 4)` window. The same graph is accurate on CPU float16 and on CUDA in float32 and bfloat16, and four controlled experiments (mask fill magnitude, additive-vs-bool mask, an isolated asymmetric-padding convolution, and an all-allow bias) place the fault in the kernel rather than the graph. Previously the combination was still selectable and silently produced an encoder whose first frame is wrong. `ArchitectureConfig` gains an optional `validate_execution_provider` hook that the builder calls with the requested EP once the dtype is known, and `MoonshineStreamingConfig` overrides it to reject float16 on CUDA with an actionable message. The rejection is deliberately not a graph-side workaround: compensating for an execution-provider defect inside the model would hide it and diverge from upstream. float32 (the default) and bfloat16 remain available on CUDA, and float16 remains available on CPU. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.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.
Adds complete support for the HuggingFace architecture
moonshine_streaming, validated end-to-end againstmoonshine-ai/moonshine-streaming-tinyat pinned revisionf8e9dfd8c562c257c151a907b7b7f2fe8ff8511a. The revision is threaded through config extraction, the processor, weight download, golden generation, the L4/L5 test harness and the CLI.What's different from offline Moonshine
Moonshine Streaming is not a variant of Moonshine's front end — it replaces it and changes the encoder's positional scheme. All of the following are reproduced exactly:
asinh(exp(log_k) * x)compression, bias-free framing linear + SiLUpadding_mask.sum(-1) // frame_len, applied multiplicativelypads=[4, 0]), so no future frame leaks backwards. The mask is pushed through the same receptive field with an "any valid" reduction, matching upstream'sconv1d(mask, ones, stride) > 0[(16,4), (16,4), (16,0), (16,0), (16,4), (16,4)].rightis the strict lookahead in encoder frames;right == 0is fully causal. Modelled per layer, not as one uniform window, because this is what bounds streaming latencygamma + 1; the checkpoint storesgamma, notweightpos_embto the encoder output, then projects to the decoder width when they differproj_outThe window schedule is genuinely load-bearing: with identical weights, HuggingFace's own encoder output differs by up to 3.29 between
[[2,1],[3,0]],[[16,4],[16,0]]and a global window — and the ONNX graph matches HuggingFace for all three.ONNX initializer names match the checkpoint 1:1; no renames beyond the shared
model.prefix strip.Reuse
The decoder is the same architecture as offline Moonshine, so
MoonshineAttention,MoonshineDecoderLayerandMoonshineDecoderModelare reused directly instead of duplicated. Everything the streaming variant actually changes is new code.Encoder context is added exactly once
HuggingFace's
MoonshineStreamingDecoderadds its absolute position table toencoder_hidden_stateswith+=, mutating the caller's tensor — a single decoder call moves it by up to 1.73, and a 34-step decode leaves it 58.8 away from the encoder output. Cached cross-attention means the mutated tensor is never read again, so generation is in fact unaffected: decoding three ways (generate(), a manual loop over a deliberately accumulating buffer, and a manual loop with a pristine clone each step) produces identical tokens.A golden reference must not depend on a dead-store argument, so
scripts/generate_golden.pywraps the speech-to-text reference innon_mutating_encoder_context, which clonesencoder_hidden_stateson decoder entry. Regenerating L4 and L5 under the guard reproduces both files byte for byte — that is the evidence the previous references were not silently accumulating. The exported ONNX decoder takes the encoder output as a read-only graph input and adds the position table on every call, which is the same semantics.Three tests lock this down:
Supporting changes
MoonshineAttentiongains keyword-onlyhidden_size/head_dim(so an encoder of a different width can reuse it) andqkv_bias/o_bias. Previously all four projections were hardcoded bias-free whileMoonshineConfigextractedattention_bias, so a bias-enabled checkpoint would have silently mis-loaded. Upstream gates Q/K/V onattention_biasand keeps the decoder'so_projbias-free unconditionally, while the streaming encoder gates all four. Defaults unchanged (False).Conv1d.paddingaccepts a(left, right)pair so a convolution can be made causal without a separatePadnode.SpeechToTextConfig.encoder_output_sizelets the decoder graph declare the real encoder width instead of assuming it equals the decoder's.ArchitectureConfig.validate_execution_provider— a new optional builder hook (see below).scripts/generate_golden.pynow forwards the pinnedrevisionon the speech-to-text path — it was previously unpinned.moonshine_streamingfor the same reason it rejectsmoonshine.float16 on CUDA is explicitly rejected
ORT's CUDA float16 fused attention kernel returns encoder frame 0 as exactly zero (
std 0.00000where fp32 givesstd 0.52855). That frame is the sparsest masked query row — it sees only 4 of 456 keys under the(16, 4)window. The same graph on CPU float16 givesstd 0.52861, and CUDA float32/bfloat16 are correct with identical masks.Root-caused with four controlled experiments:
-65504→-1e4: still zero → not the fill magnitude.Convrepro withpads=[4,0], stride 2, fp16: CUDA == CPU bit-for-bit → not the causal convolution.The combination was previously still selectable and silently produced a wrong first encoder frame.
ArchitectureConfignow exposes an optionalvalidate_execution_provider(ep)hook that the builder calls with the requested EP once the dtype is known, andMoonshineStreamingConfigoverrides it to reject float16 + CUDA with an actionable message. Verified through the publicbuild()API and therefore the CLI; float32/CUDA, bfloat16/CUDA and float16/CPU all still build.This is deliberately not a graph-side workaround. Compensating for an execution-provider defect inside the model would hide it and diverge from upstream.
float32remains the default and is unaffected.Evidence
Environment for every command (this machine has an editable
mobiusinstall pointing at a different worktree):L1 / L2 / L3
build_graph_test.py -k moonshine+weight_alignment_test.py— new cases for causalpads=[4,0],gammanaming,projonly when widths differ, window-length validation, bias gatingarch_validation_test.py [moonshine_streaming]config-downloads / full-graph-builds / shapes;yaml_schema_test.py252 passedattention_bias; batch-3 ragged padding (150 / 90 / 40 valid frames) with per-row parity; decoder prefill + cached decode652-129742-0006.flac, peak 0.59, 456 encoder frames)Real-weight fp32, ONNX encoder → ONNX decoder (not an HF intermediate):
max|diff|8.3e-6, encoder mask exactmax|diff|2.0e-5HF generate()token-for-token (34 tokens), identical transcript:Cauliflower mayonnaise, take cold boiled cauliflower, break into branches, adding salt, pepper and vinegar to season.rtol=atol=1e-3for every steptests/moonshine_streaming_integration_test.py: 23 passed.L4 / L5 goldens
Generated independently from the upstream HuggingFace pipeline under the non-mutating guard and committed:
testdata/golden/audio/moonshine-streaming-tiny.json— top19243testdata/golden/audio/moonshine-streaming-tiny_generation.json— 33 tokens (≥ 20)The YAML case sets
exact_match: true, so L5 asserts full sequence length + exact token IDs + exact decoded transcript. Both pass on CPU and on CUDA (MOBIUS_TEST_DEVICE=cuda MOBIUS_TEST_BUILD_EP=cuda).Multi-dtype and multi-EP (real audio; reference = HF fp32,
|hidden|max6.70,|logits|max24.53)MatMulkernelCPU-vs-CUDA f32: encoder
max|diff|0.041, logitmax|diff|0.011, tokens equal. CUDA CLI build with--execution-provider cudasucceeds.Olive quantization (CPU)
NF4 (
OnnxBnb4Quantization, olive 0.13.0) on the decoder: 138.5 MB → 60.2 MB (0.435), loads, and reproduces the exact golden token sequence and transcript.CLI
The emitted
audio_processor.jsoncarriespad_to_multiple_of: 80, exactly the frame alignment the encoder's reshape requires. A test asserts the processor contract directly (do_normalize False, 16 kHz,return_attention_mask True,pad_to_multiple_of == config.frame_length,len(input_values) % frame_length == 0).Full suite
tests/build_graph_test.py tests/cli_test.py src/→ 4892 passed, 7 failed. All 7 failures were reproduced on a stashed baseline of this same branch and are pre-existing and unrelated (gguf q4_0, deepseek_v4 qmoe ×2, decompose_attention softcap ×2, qwen_image CUDA low-precision ×2).lintrunner init+lintrunner f+lintrunner -aclean with the pinned RUFF 0.16.2 / RUFF-FORMAT actually running.Waivers
write_ort_genai_configrejectsmoonshine_streamingexactly as it already rejectsmoonshine— the runtime cannot host a variable-length raw-waveform encoder. A parametrized test asserts the rejection for both. Theonnx-genaiworkflow runtime is supported and emits a complete package.MatMulCPU kernel (NOT_IMPLEMENTEDatencoder/embedder/linear/MatMul). bf16 is validated on CUDA instead.Intentional mixed precision
Per-frame CMVN and the
asinhgain are computed in float32 and cast to the build dtype before the framing linear;log_kis marked_keep_float32. Squared raw-waveform amplitudes are ~1e-6 and collapse into float16 subnormals, so evaluating that stage in fp16 would be numerically meaningless. Everything downstream runs in the build dtype.