From 8002b36285e8edc11bf4779f50f5f755f9898943 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Tue, 21 Jul 2026 19:58:47 +0000 Subject: [PATCH 1/8] Moonshine streaming --- .../cpu/.gitignore | 14 + .../cpu/README.md | 153 +++++++ .../cpu/__init__.py | 0 .../cpu/export_moonshine_streaming.py | 199 +++++++++ .../cpu/info.yaml | 28 ++ .../cpu/moonshine_adapter_fp32_cpu.json | 41 ++ .../cpu/moonshine_cross_kv_fp32_cpu.json | 40 ++ .../cpu/moonshine_decoder_kv_fp32_cpu.json | 44 ++ .../cpu/moonshine_decoder_kv_int4_cpu.json | 52 +++ .../cpu/moonshine_decoder_kv_int8_cpu.json | 53 +++ .../cpu/moonshine_decoder_kv_kquant8_cpu.json | 52 +++ .../cpu/moonshine_encoder_fp32_cpu.json | 40 ++ .../cpu/moonshine_encoder_int4_cpu.json | 48 ++ .../cpu/moonshine_encoder_int8_cpu.json | 49 +++ .../cpu/moonshine_encoder_kquant8_cpu.json | 48 ++ .../cpu/moonshine_frontend_fp32_cpu.json | 51 +++ .../cpu/moonshine_model_load.py | 414 ++++++++++++++++++ .../cpu/optimize.py | 372 ++++++++++++++++ .../cpu/requirements.txt | 9 + .../cpu/validate_export.py | 117 +++++ 20 files changed, 1824 insertions(+) create mode 100644 usefulSensors-moonshine-streaming/cpu/.gitignore create mode 100644 usefulSensors-moonshine-streaming/cpu/README.md create mode 100644 usefulSensors-moonshine-streaming/cpu/__init__.py create mode 100644 usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py create mode 100644 usefulSensors-moonshine-streaming/cpu/info.yaml create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_adapter_fp32_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_cross_kv_fp32_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_fp32_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_kquant8_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_fp32_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_kquant8_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_frontend_fp32_cpu.json create mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py create mode 100644 usefulSensors-moonshine-streaming/cpu/optimize.py create mode 100644 usefulSensors-moonshine-streaming/cpu/requirements.txt create mode 100644 usefulSensors-moonshine-streaming/cpu/validate_export.py diff --git a/usefulSensors-moonshine-streaming/cpu/.gitignore b/usefulSensors-moonshine-streaming/cpu/.gitignore new file mode 100644 index 000000000..291419770 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/.gitignore @@ -0,0 +1,14 @@ +# Generated model artifacts +build/ + +# Python bytecode +__pycache__/ +*.pyc + +# Olive cache +.olive-cache/ + +# Temp and log files +*.temp +*.bak +*.log diff --git a/usefulSensors-moonshine-streaming/cpu/README.md b/usefulSensors-moonshine-streaming/cpu/README.md new file mode 100644 index 000000000..2497a9c5e --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/README.md @@ -0,0 +1,153 @@ +# Moonshine Streaming (CPU EP, FP32) + +This recipe exports **usefulsensors/moonshine-streaming-small** to ONNX and +produces CPU-ready ONNX Runtime GenAI artifacts for streaming ASR. + +The streaming model is exported as **five FP32 ONNX components**, each handled +through Olive's declarative `OnnxConversion` pass (dynamo exporter, opset 20) +with the exact input/output names the streaming runner expects: + +- **Frontend** – stateful convolutional feature extractor. Consumes an audio + chunk plus rolling sample/conv buffers and emits log-mel-like features while + updating its buffers. +- **Encoder** – sliding-window transformer encoder (features → hidden states). +- **Adapter** – positional projection (`encoded + pos_emb(arange(T) + pos_offset)`) + producing the decoder memory. +- **Cross-KV** – precomputes per-layer cross-attention K/V from the memory once + per segment. +- **Decoder-KV** – autoregressive decoder with cached self-attention KV and the + precomputed cross-attention KV → logits. + +## Files +- `cpu/moonshine_frontend_fp32_cpu.json` – Olive frontend config (convert only) +- `cpu/moonshine_encoder_fp32_cpu.json` – Olive encoder config (convert only) +- `cpu/moonshine_adapter_fp32_cpu.json` – Olive adapter config (convert only) +- `cpu/moonshine_cross_kv_fp32_cpu.json` – Olive cross-KV config (convert only) +- `cpu/moonshine_decoder_kv_fp32_cpu.json` – Olive decoder-KV config (convert only) +- `cpu/moonshine_encoder_int8_cpu.json` – Olive encoder config (convert → INT8 dynamic quant) +- `cpu/moonshine_decoder_kv_int8_cpu.json` – Olive decoder-KV config (convert → INT8 dynamic quant) +- `cpu/moonshine_encoder_int4_cpu.json` – Olive encoder config (convert → INT4 k-quant) +- `cpu/moonshine_decoder_kv_int4_cpu.json` – Olive decoder-KV config (convert → INT4 k-quant) +- `cpu/moonshine_model_load.py` – model loaders + wrapper modules + dummy inputs +- `cpu/optimize.py` – full pipeline script (Olive × 5 + tokenizer + configs + VAD) +- `cpu/export_moonshine_streaming.py` – standalone exporter (no Olive; for debugging) +- `cpu/validate_export.py` – per-component torch-vs-ONNX numeric check + +## Setup +From repo root: + +```bash +python -m venv .venv +source .venv/bin/activate +pip install -r usefulSensors-moonshine-streaming/cpu/requirements.txt +``` + +## Run + +From the `usefulSensors-moonshine-streaming` directory: + +```bash +cd usefulSensors-moonshine-streaming + +python cpu/optimize.py --output-dir build/moonshine-small +``` + +This runs the full pipeline: +1. **Frontend / Encoder / Adapter / Cross-KV / Decoder-KV** — Olive: `OnnxConversion` (FP32) for each of the five components +2. **Configs** — generates `genai_config.json` + `streaming_config.json` +3. **Tokenizer** — exports `tokenizer.json` + `tokenizer_config.json` +4. **VAD** — downloads Silero VAD ONNX model + +`--output-dir` is resolved relative to the `cpu/` directory unless an absolute +path is given. + +Export the **tiny** variant instead: + +```bash +python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-tiny \ + --output-dir build/moonshine-tiny +``` + +### Quantization (`--quantize`) + +Add `--quantize` to quantize the **encoder** and **decoder_kv** MatMuls +(frontend / adapter / cross_kv stay FP32). IO names and the runtime configs +are unchanged. `--quant-method` picks the algorithm: + +- `--quant-method dynamic` (default) — INT8 RTN **dynamic** quant. Swaps in + `moonshine_encoder_int8_cpu.json` / `moonshine_decoder_kv_int8_cpu.json`, + which chain `OnnxConversion → OnnxDynamicQuantization` (weight matmuls become + `MatMulInteger` + `DynamicQuantizeLinear`), matching the shipped official + `.ort`. Fastest on CPU. +- `--quant-method kquant` — INT4 **weight-only k-quant** (like the nemotron + recipe). Swaps in `moonshine_encoder_int4_cpu.json` / + `moonshine_decoder_kv_int4_cpu.json`, which chain + `OnnxConversion → OnnxKQuantQuantization` (`bits=4`, `block_size=32`, + `accuracy_level=4`; weight matmuls become `MatMulNBits`). Smallest on disk; + weight-only, so activation×activation attention matmuls stay FP32 and the + token-embedding `Gather` is left FP32. + +```bash +# INT8 RTN dynamic (default) +python cpu/optimize.py --quantize --output-dir build/moonshine-small-int8 + +# INT4 k-quant +python cpu/optimize.py --quantize --quant-method kquant \ + --output-dir build/moonshine-small-int4 +``` + +On a 40s clip (CPU EP), both preserve transcription quality; INT8 dynamic is +fastest, INT4 k-quant is smallest: + +| build | encoder | decoder_kv | total | RTF | +|---|---|---|---|---| +| FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× | +| INT8 dynamic (`--quantize`) | 42 MB | 125 MB | ~233 MB | ~9.0× | +| INT4 k-quant (`--quant-method kquant`) | 27 MB | 103 MB | ~196 MB | ~7.7× | +| official `.ort` | 42 MB | 174 MB | — | ~9.4× | + +Transcription is essentially identical to FP32 for both (only quant-noise +wording drift). INT4 k-quant is the smallest model but on CPU it is *not* +faster than INT8 dynamic: `MatMulNBits` is weight-only, so the int4→compute +de-quant overhead offsets the memory-bandwidth win for this small, +compute-bound model. Prefer `dynamic` for speed, `kquant` for the smallest +artifact. + +Or run individual components directly with the Olive CLI: + +```bash +python -m olive run --config cpu/moonshine_frontend_fp32_cpu.json +python -m olive run --config cpu/moonshine_encoder_fp32_cpu.json +python -m olive run --config cpu/moonshine_adapter_fp32_cpu.json +python -m olive run --config cpu/moonshine_cross_kv_fp32_cpu.json +python -m olive run --config cpu/moonshine_decoder_kv_fp32_cpu.json +``` + +## Output +Expected artifacts in `cpu/build/moonshine-small/`: +- `frontend.onnx` (+ `frontend.onnx.data`) +- `encoder.onnx` (+ `encoder.onnx.data`) +- `adapter.onnx` (+ `adapter.onnx.data`) +- `cross_kv.onnx` (+ `cross_kv.onnx.data`) +- `decoder_kv.onnx` (+ `decoder_kv.onnx.data`) +- `genai_config.json` +- `streaming_config.json` +- `tokenizer.json` +- `tokenizer_config.json` +- `silero_vad.onnx` + +## Validation +`validate_export.py` checks each component two ways: exported ONNX vs the +PyTorch reference outputs (tight tolerance), and — optionally — exported ONNX vs +the shipped official graph on identical inputs (loose tolerance for the +int8-quantized cross-KV / decoder-KV graphs). + +It expects the `--mine` directory to contain `.onnx` plus +`refs/.npz` (torch reference in/out dumped by the standalone +`export_moonshine_streaming.py`). `--official` is optional: + +```bash +python cpu/validate_export.py \ + --mine /path/to/moonshine-streaming-small-mine \ + --official /path/to/moonshine-streaming-small-official +``` diff --git a/usefulSensors-moonshine-streaming/cpu/__init__.py b/usefulSensors-moonshine-streaming/cpu/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py new file mode 100644 index 000000000..884c74e62 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py @@ -0,0 +1,199 @@ +"""Standalone exporter: HuggingFace MoonshineStreaming -> five stateful ONNX +graphs for onnxruntime-genai (frontend / encoder / adapter / cross_kv / +decoder_kv). + +Run inside the ``moonshine`` conda env: + + source /home/nebanfic/miniconda3/bin/activate moonshine + python export_moonshine_streaming.py \ + --model usefulsensors/moonshine-streaming-small \ + --output-dir /datadisks/disk3/nebanfic/moonshine-streaming-small-mine + +The exporter uses the TorchDynamo ONNX path (``dynamo=True``) so data-dependent +shapes in the stateful frontend export cleanly, then rewrites every graph's +input/output names to the exact contract genai expects. Reference inputs and +torch outputs for each graph are dumped to ``/refs/`` so the +companion validation script can check numerical parity against the official +``.ort`` graphs without needing transformers. +""" + +from __future__ import annotations + +import argparse +import os + +import numpy as np +import onnx +import torch +from torch.export import Dim + +import moonshine_model_load as mml + +AUTO = Dim.AUTO + + +# --------------------------------------------------------------------------- # +# Per-component export specification # +# --------------------------------------------------------------------------- # +def build_specs(): + """Return the ordered list of component export specs. ``dynamic`` is a + tuple aligned with the positional dummy inputs: each entry is either a + ``{axis: Dim.AUTO}`` dict or ``None`` for a fully static input.""" + return [ + { + "name": "frontend", + "loader": mml.frontend_model_loader, + "dummy": mml.frontend_dummy_inputs, + "input_names": [ + "audio_chunk", "sample_buffer", "sample_len", + "conv1_buffer", "conv2_buffer", "frame_count", + ], + "output_names": [ + "features", "sample_buffer_out", "sample_len_out", + "conv1_buffer_out", "conv2_buffer_out", "frame_count_out", + ], + "dynamic": ({1: AUTO}, None, None, None, None, None), + }, + { + "name": "encoder", + "loader": mml.encoder_model_loader, + "dummy": mml.encoder_dummy_inputs, + "input_names": ["features"], + "output_names": ["encoded"], + "dynamic": ({1: AUTO},), + }, + { + "name": "adapter", + "loader": mml.adapter_model_loader, + "dummy": mml.adapter_dummy_inputs, + "input_names": ["encoded", "pos_offset"], + "output_names": ["memory"], + "dynamic": ({1: AUTO}, None), + }, + { + "name": "cross_kv", + "loader": mml.cross_kv_model_loader, + "dummy": mml.cross_kv_dummy_inputs, + "input_names": ["memory"], + "output_names": ["k_cross", "v_cross"], + "dynamic": ({1: AUTO},), + }, + { + "name": "decoder_kv", + "loader": mml.decoder_kv_model_loader, + "dummy": mml.decoder_kv_dummy_inputs, + "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], + "output_names": [ + "logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross", + ], + "dynamic": ({1: AUTO}, {3: AUTO}, {3: AUTO}, {3: AUTO}, {3: AUTO}), + }, + ] + + +# --------------------------------------------------------------------------- # +# ONNX I/O renaming (make names match the genai contract exactly) # +# --------------------------------------------------------------------------- # +def rename_io(path, input_names, output_names): + model = onnx.load(path) + graph = model.graph + + def remap(value_infos, desired): + mapping = {} + for vi, new in zip(value_infos, desired): + if vi.name != new: + mapping[vi.name] = new + vi.name = new + return mapping + + in_map = remap(graph.input, input_names) + out_map = remap(graph.output, output_names) + rename = {**in_map, **out_map} + if rename: + for node in graph.node: + node.input[:] = [rename.get(n, n) for n in node.input] + node.output[:] = [rename.get(n, n) for n in node.output] + onnx.save(model, path) + return [i.name for i in graph.input], [o.name for o in graph.output] + + +# --------------------------------------------------------------------------- # +# Reference dump for validation # +# --------------------------------------------------------------------------- # +def dump_refs(refs_dir, name, module, dummy_args, input_names, output_names): + os.makedirs(refs_dir, exist_ok=True) + with torch.no_grad(): + outputs = module(*dummy_args) + if not isinstance(outputs, (tuple, list)): + outputs = (outputs,) + arrays = {} + for n, t in zip(input_names, dummy_args): + arrays[f"in__{n}"] = t.detach().cpu().numpy() + for n, t in zip(output_names, outputs): + arrays[f"out__{n}"] = t.detach().cpu().numpy() + np.savez(os.path.join(refs_dir, f"{name}.npz"), **arrays) + + +# --------------------------------------------------------------------------- # +# Main # +# --------------------------------------------------------------------------- # +def export_component(spec, model_name, output_dir, opset): + name = spec["name"] + print(f"\n=== exporting {name} ===") + module = spec["loader"](model_name).eval() + dummy = spec["dummy"](module) + # dummy_inputs_func returns a dict keyed by forward-parameter name (Olive's + # kwargs export path); the standalone exporter uses positional args, so + # flatten to a value tuple in the declared input order. + dummy_args = tuple(dummy.values()) if isinstance(dummy, dict) else tuple(dummy) + onnx_path = os.path.join(output_dir, f"{name}.onnx") + + torch.onnx.export( + module, + dummy_args, + onnx_path, + input_names=spec["input_names"], + output_names=spec["output_names"], + dynamic_shapes=spec["dynamic"], + opset_version=opset, + dynamo=True, + ) + ins, outs = rename_io(onnx_path, spec["input_names"], spec["output_names"]) + onnx.checker.check_model(onnx_path) + print(f" inputs : {ins}") + print(f" outputs: {outs}") + + dump_refs( + os.path.join(output_dir, "refs"), name, module, dummy_args, + spec["input_names"], spec["output_names"], + ) + return onnx_path + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default="usefulsensors/moonshine-streaming-small") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--opset", type=int, default=20) + parser.add_argument("--chunk-samples", type=int, default=8000) + parser.add_argument( + "--only", nargs="*", default=None, + help="Optional subset of component names to export.", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + mml.set_chunk_samples(args.chunk_samples) + + specs = build_specs() + if args.only: + specs = [s for s in specs if s["name"] in args.only] + + for spec in specs: + export_component(spec, args.model, args.output_dir, args.opset) + + print(f"\nDone. Graphs written to {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/usefulSensors-moonshine-streaming/cpu/info.yaml b/usefulSensors-moonshine-streaming/cpu/info.yaml new file mode 100644 index 000000000..e7108185d --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/info.yaml @@ -0,0 +1,28 @@ +name: usefulSensors-moonshine-streaming +provider: usefulsensors +model_id: usefulsensors/moonshine-streaming-small +task: automatic-speech-recognition +framework: ONNX Runtime +execution_provider: CPUExecutionProvider +summary: > + CPU recipe for exporting the Moonshine streaming ASR model to ONNX (FP32) as + five ONNX Runtime GenAI components (frontend, encoder, adapter, cross-KV, + decoder-KV) with the exact IO names the streaming runner expects, plus the + generated genai/streaming configs, tokenizer, and Silero VAD. + +artifacts: + - cpu/moonshine_frontend_fp32_cpu.json + - cpu/moonshine_encoder_fp32_cpu.json + - cpu/moonshine_adapter_fp32_cpu.json + - cpu/moonshine_cross_kv_fp32_cpu.json + - cpu/moonshine_decoder_kv_fp32_cpu.json + - cpu/moonshine_encoder_int8_cpu.json + - cpu/moonshine_decoder_kv_int8_cpu.json + - cpu/moonshine_encoder_int4_cpu.json + - cpu/moonshine_decoder_kv_int4_cpu.json + +scripts: + - cpu/optimize.py + - cpu/moonshine_model_load.py + - cpu/export_moonshine_streaming.py + - cpu/validate_export.py diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_adapter_fp32_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_adapter_fp32_cpu.json new file mode 100644 index 000000000..9b5929d2d --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_adapter_fp32_cpu.json @@ -0,0 +1,41 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "adapter_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["encoded", "pos_offset"], + "output_names": ["memory"], + "dynamic_shapes": { + "encoded": {"1": "adp_seq"}, + "pos_offset": {} + } + }, + "dummy_inputs_func": "adapter_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "adapter.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/adapter.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_cross_kv_fp32_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_cross_kv_fp32_cpu.json new file mode 100644 index 000000000..962ebc76e --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_cross_kv_fp32_cpu.json @@ -0,0 +1,40 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "cross_kv_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["memory"], + "output_names": ["k_cross", "v_cross"], + "dynamic_shapes": { + "memory": {"1": "mem_seq"} + } + }, + "dummy_inputs_func": "cross_kv_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "cross_kv.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/cross_kv.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_fp32_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_fp32_cpu.json new file mode 100644 index 000000000..08ddff498 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_fp32_cpu.json @@ -0,0 +1,44 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "decoder_kv_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], + "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], + "dynamic_shapes": { + "token": {"1": "q_seq"}, + "k_self": {"3": "past_seq"}, + "v_self": {"3": "past_seq"}, + "out_k_cross": {"3": "cross_seq"}, + "out_v_cross": {"3": "cross_seq"} + } + }, + "dummy_inputs_func": "decoder_kv_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "decoder_kv.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/decoder_kv.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json new file mode 100644 index 000000000..d156e610e --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json @@ -0,0 +1,52 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "decoder_kv_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], + "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], + "dynamic_shapes": { + "token": {"1": "q_seq"}, + "k_self": {"3": "past_seq"}, + "v_self": {"3": "past_seq"}, + "out_k_cross": {"3": "cross_seq"}, + "out_v_cross": {"3": "cross_seq"} + } + }, + "dummy_inputs_func": "decoder_kv_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "decoder_kv_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "decoder_kv.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/decoder_kv.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json new file mode 100644 index 000000000..538dc3cbe --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json @@ -0,0 +1,53 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "decoder_kv_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], + "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], + "dynamic_shapes": { + "token": {"1": "q_seq"}, + "k_self": {"3": "past_seq"}, + "v_self": {"3": "past_seq"}, + "out_k_cross": {"3": "cross_seq"}, + "out_v_cross": {"3": "cross_seq"} + } + }, + "dummy_inputs_func": "decoder_kv_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "decoder_kv_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxDynamicQuantization", + "precision": "int8", + "op_types_to_quantize": ["MatMul"], + "per_channel": true, + "quant_preprocess": false, + "save_as_external_data": true, + "external_data_name": "decoder_kv.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/decoder_kv.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_kquant8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_kquant8_cpu.json new file mode 100644 index 000000000..ffe8d8199 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_kquant8_cpu.json @@ -0,0 +1,52 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "decoder_kv_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], + "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], + "dynamic_shapes": { + "token": {"1": "q_seq"}, + "k_self": {"3": "past_seq"}, + "v_self": {"3": "past_seq"}, + "out_k_cross": {"3": "cross_seq"}, + "out_v_cross": {"3": "cross_seq"} + } + }, + "dummy_inputs_func": "decoder_kv_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "decoder_kv_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxKQuantQuantization", + "bits": 8, + "block_size": 32, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "decoder_kv.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/decoder_kv.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_fp32_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_fp32_cpu.json new file mode 100644 index 000000000..cce4bb68a --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_fp32_cpu.json @@ -0,0 +1,40 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "encoder_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["features"], + "output_names": ["encoded"], + "dynamic_shapes": { + "features": {"1": "enc_seq"} + } + }, + "dummy_inputs_func": "encoder_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "encoder.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/encoder.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json new file mode 100644 index 000000000..3ade9b78d --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json @@ -0,0 +1,48 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "encoder_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["features"], + "output_names": ["encoded"], + "dynamic_shapes": { + "features": {"1": "enc_seq"} + } + }, + "dummy_inputs_func": "encoder_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "encoder_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "encoder.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/encoder.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json new file mode 100644 index 000000000..e604ac06e --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json @@ -0,0 +1,49 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "encoder_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["features"], + "output_names": ["encoded"], + "dynamic_shapes": { + "features": {"1": "enc_seq"} + } + }, + "dummy_inputs_func": "encoder_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "encoder_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxDynamicQuantization", + "precision": "int8", + "op_types_to_quantize": ["MatMul"], + "per_channel": true, + "quant_preprocess": false, + "save_as_external_data": true, + "external_data_name": "encoder.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/encoder.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_kquant8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_kquant8_cpu.json new file mode 100644 index 000000000..abd600391 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_kquant8_cpu.json @@ -0,0 +1,48 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "encoder_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": ["features"], + "output_names": ["encoded"], + "dynamic_shapes": { + "features": {"1": "enc_seq"} + } + }, + "dummy_inputs_func": "encoder_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "encoder_fp32.onnx.data" + }, + "quantize": { + "type": "OnnxKQuantQuantization", + "bits": 8, + "block_size": 32, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "encoder.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/encoder.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_frontend_fp32_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_frontend_fp32_cpu.json new file mode 100644 index 000000000..31ca93a1c --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_frontend_fp32_cpu.json @@ -0,0 +1,51 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "usefulsensors/moonshine-streaming-small", + "model_loader": "frontend_model_loader", + "model_script": "cpu/moonshine_model_load.py", + "io_config": { + "input_names": [ + "audio_chunk", "sample_buffer", "sample_len", + "conv1_buffer", "conv2_buffer", "frame_count" + ], + "output_names": [ + "features", "sample_buffer_out", "sample_len_out", + "conv1_buffer_out", "conv2_buffer_out", "frame_count_out" + ], + "dynamic_shapes": { + "audio_chunk": {"1": "audio_len"}, + "sample_buffer": {}, + "sample_len": {}, + "conv1_buffer": {}, + "conv2_buffer": {}, + "frame_count": {} + } + }, + "dummy_inputs_func": "frontend_dummy_inputs" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"] + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "target_opset": 20, + "dynamic": true, + "use_dynamo_exporter": true, + "save_as_external_data": true, + "external_data_name": "frontend.onnx.data" + } + }, + "target": "local_system", + "output_dir": "build/onnx/frontend.onnx", + "no_artifacts": true +} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py b/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py new file mode 100644 index 000000000..078147311 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py @@ -0,0 +1,414 @@ +"""Model loaders, wrapper modules, dummy inputs and I/O specs for exporting the +HuggingFace ``MoonshineStreamingForConditionalGeneration`` model into the five +stateful ONNX graphs consumed by onnxruntime-genai's ``streaming_enc_dec_asr`` +model type (frontend / encoder / adapter / cross_kv / decoder_kv). + +Design principle: reuse the HF submodules verbatim inside thin ``nn.Module`` +wrappers so the exported math is numerically identical to the reference model. +The only component that is *reimplemented* is the frontend, whose two causal +convolutions must be made stateful (left padding replaced by carried buffers) +so the model can run chunk-by-chunk. + +Runs under the ``moonshine`` conda env (transformers>=5.2 + torch + onnxscript). + +Both the standalone exporter (``export_moonshine_streaming.py``) and the Olive +recipe JSONs import from this module. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import MoonshineStreamingForConditionalGeneration +from transformers.models.moonshine_streaming.modeling_moonshine_streaming import ( + apply_rotary_pos_emb, +) + +NEG_INF = float("-inf") + +# Cache full models by name so the five per-component loaders don't each reload +# the ~200MB checkpoint from disk. +_MODEL_CACHE: dict[str, MoonshineStreamingForConditionalGeneration] = {} + + +def load_full_model(model_name: str) -> MoonshineStreamingForConditionalGeneration: + """Load (and cache) the fp32, eval, eager-attention reference model.""" + if model_name not in _MODEL_CACHE: + model = MoonshineStreamingForConditionalGeneration.from_pretrained( + model_name, attn_implementation="eager" + ) + model = model.to(torch.float32).eval() + for p in model.parameters(): + p.requires_grad_(False) + _MODEL_CACHE[model_name] = model + return _MODEL_CACHE[model_name] + + +def model_dims(model: MoonshineStreamingForConditionalGeneration) -> dict: + """Extract the dimensions needed to build dummy inputs / configs, so the + same code works for both the small and tiny checkpoints.""" + enc = model.model.encoder + dec = model.model.decoder + sa0 = dec.layers[0].self_attn + return { + "encoder_dim": enc.embedder.conv1.in_channels, # 620 (small) + "conv1_channels": enc.embedder.conv1.in_channels, # 620 (conv1 input buffer) + "conv2_channels": enc.embedder.conv2.in_channels, # 1240 (conv2 input buffer) + "frame_len": int(enc.embedder.frame_len), # 80 + "left_pad1": int(enc.embedder.conv1.left_pad), # 4 + "left_pad2": int(enc.embedder.conv2.left_pad), # 4 + "num_encoder_layers": len(enc.layers), # 10 + "num_decoder_layers": len(dec.layers), # 10 + "decoder_dim": dec.embed_tokens.embedding_dim, # 512 + "num_kv_heads": sa0.config.num_key_value_heads, # 8 + "head_dim": sa0.head_dim, # 64 + "vocab_size": model.proj_out.out_features, # 32768 + "sample_buffer_size": int(enc.embedder.frame_len) - 1, # 79 + } + + +# --------------------------------------------------------------------------- # +# 1. Frontend (stateful streaming feature extractor) # +# --------------------------------------------------------------------------- # +class FrontendModule(nn.Module): + """Streaming re-implementation of ``MoonshineStreamingEncoderEmbedder``. + + Non-streaming HF applies ``F.pad(x, (left_pad, 0))`` before each causal + conv. For chunked streaming we instead carry the last ``left_pad`` input + columns of every conv in a buffer and run a *valid* (unpadded) convolution + on ``cat(buffer, x)``. Because a normal chunk is 8000 samples = 100 frames + (even), the stride-2 phase stays aligned across chunks and the concatenated + output is bit-for-bit identical to running the embedder on the full signal. + + Sub-frame audio (``total_samples % frame_len``) is carried in + ``sample_buffer`` / ``sample_len`` and prepended to the next chunk so the + framing is contiguous. + """ + + def __init__(self, full: MoonshineStreamingForConditionalGeneration): + super().__init__() + emb = full.model.encoder.embedder + self.cmvn = emb.cmvn + self.comp = emb.comp + self.linear = emb.linear + self.conv1 = emb.conv1 + self.conv2 = emb.conv2 + self.frame_len = int(emb.frame_len) + self.pad1 = int(emb.conv1.left_pad) + self.pad2 = int(emb.conv2.left_pad) + self.buf_size = self.frame_len - 1 # sample_buffer width (79) + + def forward( + self, + audio_chunk, # [1, L] float32 + sample_buffer, # [1, 79] float32 + sample_len, # [1] int64 + conv1_buffer, # [1, C1, 4] float32 (last inputs of conv1) + conv2_buffer, # [1, C2, 4] float32 (last inputs of conv2) + frame_count, # [1] int64 + ): + # In the genai streaming pipeline chunk_samples is a multiple of + # frame_len, so the carried sample_buffer is always empty on input + # (sample_len == 0). We therefore frame the chunk directly. The + # correct leftover state is still emitted so a final (flush) chunk of + # non-multiple length threads/records its remainder before reset. + total = audio_chunk.shape[1] + n_frames = total // self.frame_len + used = n_frames * self.frame_len + + frames = audio_chunk.narrow(1, 0, used).reshape(1, -1, self.frame_len) # [1, nf, 80] + leftover = audio_chunk.narrow(1, used, total - used) # [1, rem] + rem = leftover.shape[1] + sample_buffer_out = F.pad(leftover, (0, self.buf_size - rem)) # [1, 79] + zero_i64 = frame_count * 0 + sample_len_out = zero_i64 + rem # [1] + + # ---- per-frame feature transform (framing-invariant) ---- + h = self.cmvn(frames) # [1, nf, 80] + h = self.comp(h) + h = F.silu(self.linear(h)) # [1, nf, encoder_dim] + h = h.transpose(1, 2) # [1, encoder_dim, nf] + + # ---- stateful causal conv 1 ---- + c1_in = torch.cat([conv1_buffer, h], dim=2) # [1, C, pad1 + nf] + conv1_buffer_out = c1_in[:, :, c1_in.shape[2] - self.pad1:] + h = F.conv1d( + c1_in, self.conv1.weight, self.conv1.bias, + stride=self.conv1.stride, dilation=self.conv1.dilation, + ) + h = F.silu(h) + + # ---- stateful causal conv 2 ---- + c2_in = torch.cat([conv2_buffer, h], dim=2) # [1, C2, pad2 + o1] + conv2_buffer_out = c2_in[:, :, c2_in.shape[2] - self.pad2:] + h = F.conv1d( + c2_in, self.conv2.weight, self.conv2.bias, + stride=self.conv2.stride, dilation=self.conv2.dilation, + ) + features = h.transpose(1, 2) # [1, feat_len, encoder_dim] + + frame_count_out = frame_count + n_frames + return ( + features, + sample_buffer_out, + sample_len_out, + conv1_buffer_out, + conv2_buffer_out, + frame_count_out, + ) + + +# --------------------------------------------------------------------------- # +# 2. Encoder (sliding-window bidirectional transformer, mask built from T) # +# --------------------------------------------------------------------------- # +class EncoderModule(nn.Module): + """Reuses the HF encoder layers + final norm. The per-layer sliding-window + attention mask is rebuilt from the dynamic sequence length because the + genai encoder graph takes no external mask input.""" + + def __init__(self, full: MoonshineStreamingForConditionalGeneration): + super().__init__() + enc = full.model.encoder + self.layers = enc.layers + self.final_norm = enc.final_norm + self.windows = [tuple(int(v) for v in w) for w in enc.config.sliding_windows] + + @staticmethod + def _sliding_mask(seq_len, left, right, device, dtype): + idx = torch.arange(seq_len, device=device) + dist = idx.unsqueeze(1) - idx.unsqueeze(0) # q - k, [T, T] + allowed = ((dist >= 0) & (dist < left)) | ((dist < 0) & (-dist < right)) + mask = torch.zeros(seq_len, seq_len, dtype=dtype, device=device) + mask = mask.masked_fill(~allowed, NEG_INF) + return mask.unsqueeze(0).unsqueeze(0) # [1, 1, T, T] + + def forward(self, features): # [1, T, encoder_dim] + hidden = features + seq_len = hidden.shape[1] + for layer, (left, right) in zip(self.layers, self.windows): + mask = self._sliding_mask(seq_len, left, right, hidden.device, hidden.dtype) + hidden = layer(hidden, attention_mask=mask) + return self.final_norm(hidden) # [1, T, encoder_dim] + + +# --------------------------------------------------------------------------- # +# 3. Adapter (positional embedding + projection to decoder dim) # +# --------------------------------------------------------------------------- # +class AdapterModule(nn.Module): + """memory = proj(encoded + pos_emb(arange(T) + pos_offset)). Mirrors the + top of ``MoonshineStreamingDecoder.forward``.""" + + def __init__(self, full: MoonshineStreamingForConditionalGeneration): + super().__init__() + dec = full.model.decoder + self.pos_emb = dec.pos_emb + self.proj = dec.proj + + def forward(self, encoded, pos_offset): # [1,T,enc], [1] + seq_len = encoded.shape[1] + offset = pos_offset.reshape(()).to(torch.long) + positions = torch.arange(seq_len, device=encoded.device) + offset + hidden = encoded + self.pos_emb(positions) + return self.proj(hidden) # [1, T, decoder_dim] + + +# --------------------------------------------------------------------------- # +# 4. Cross-KV (project memory into per-layer cross-attention key/value) # +# --------------------------------------------------------------------------- # +class CrossKvModule(nn.Module): + """Stack every decoder layer's ``encoder_attn`` k/v projection of memory.""" + + def __init__(self, full: MoonshineStreamingForConditionalGeneration): + super().__init__() + self.layers = full.model.decoder.layers + a0 = self.layers[0].encoder_attn + self.num_heads = a0.config.num_key_value_heads + self.head_dim = a0.head_dim + + def forward(self, memory): # [1, T, decoder_dim] + seq_len = memory.shape[1] + keys, values = [], [] + for layer in self.layers: + attn = layer.encoder_attn + k = attn.k_proj(memory).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + v = attn.v_proj(memory).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + keys.append(k) + values.append(v) + k_cross = torch.stack(keys, dim=0) # [L, 1, H, T, D] + v_cross = torch.stack(values, dim=0) + return k_cross, v_cross + + +# --------------------------------------------------------------------------- # +# 5. Decoder-KV (autoregressive step with self-KV cache + precomputed cross) # +# --------------------------------------------------------------------------- # +class DecoderKvModule(nn.Module): + """One decoder step over ``S`` query tokens. Self-attention appends to the + incoming self-KV cache (position derived from its length); cross-attention + reuses the precomputed per-layer cross K/V. Reuses every HF submodule + (projections, layernorms, MLP, rotary embedding, lm head).""" + + def __init__(self, full: MoonshineStreamingForConditionalGeneration): + super().__init__() + dec = full.model.decoder + self.embed_tokens = dec.embed_tokens + self.layers = dec.layers + self.norm = dec.norm + self.rotary_emb = dec.rotary_emb + self.proj_out = full.proj_out + a0 = dec.layers[0].self_attn + self.num_heads = a0.config.num_key_value_heads + self.head_dim = a0.head_dim + + def forward(self, token, k_self, v_self, out_k_cross, out_v_cross): + # token [1,S] int64 ; *_self [L,1,H,Ts,D] ; out_*_cross [L,1,H,Tc,D] + seq_len = token.shape[1] + past_len = k_self.shape[3] + hidden = self.embed_tokens(token) # [1, S, dec_dim] + + position_ids = torch.arange( + past_len, past_len + seq_len, device=token.device + ).unsqueeze(0) # [1, S] + cos, sin = self.rotary_emb(hidden, position_ids) + + # causal mask over [S query positions, past_len + S key positions] + total = past_len + seq_len + q_abs = position_ids.reshape(seq_len, 1) # [S, 1] + k_abs = torch.arange(total, device=token.device).reshape(1, total) + causal = torch.zeros(seq_len, total, dtype=hidden.dtype, device=token.device) + causal = causal.masked_fill(k_abs > q_abs, NEG_INF).unsqueeze(0).unsqueeze(0) + + new_k, new_v = [], [] + for i, layer in enumerate(self.layers): + # ---- self attention (causal, rotary, cached) ---- + residual = hidden + hs = layer.input_layernorm(hidden) + sa = layer.self_attn + q = sa.q_proj(hs).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = sa.k_proj(hs).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + v = sa.v_proj(hs).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + k_full = torch.cat([k_self[i], k], dim=2) # [1, H, Ts+S, D] + v_full = torch.cat([v_self[i], v], dim=2) + new_k.append(k_full) + new_v.append(v_full) + scores = torch.matmul(q, k_full.transpose(2, 3)) * sa.scaling + causal + ctx = torch.matmul(torch.softmax(scores, dim=-1), v_full) + ctx = ctx.transpose(1, 2).reshape(1, seq_len, -1) + hidden = residual + sa.o_proj(ctx) + + # ---- cross attention (precomputed k/v, full) ---- + residual = hidden + hs = layer.post_attention_layernorm(hidden) + ca = layer.encoder_attn + qc = ca.q_proj(hs).view(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + scores = torch.matmul(qc, out_k_cross[i].transpose(2, 3)) * ca.scaling + ctx = torch.matmul(torch.softmax(scores, dim=-1), out_v_cross[i]) + ctx = ctx.transpose(1, 2).reshape(1, seq_len, -1) + hidden = residual + ca.o_proj(ctx) + + # ---- feed forward ---- + residual = hidden + hs = layer.final_layernorm(hidden) + hidden = residual + layer.mlp(hs) + + hidden = self.norm(hidden) + logits = self.proj_out(hidden) # [1, S, vocab] + out_k_self = torch.stack(new_k, dim=0) # [L, 1, H, Ts+S, D] + out_v_self = torch.stack(new_v, dim=0) + return logits, out_k_self, out_v_self, out_k_cross, out_v_cross + + +# --------------------------------------------------------------------------- # +# Olive-style per-component loaders # +# --------------------------------------------------------------------------- # +def frontend_model_loader(model_name): + return FrontendModule(load_full_model(model_name)) + + +def encoder_model_loader(model_name): + return EncoderModule(load_full_model(model_name)) + + +def adapter_model_loader(model_name): + return AdapterModule(load_full_model(model_name)) + + +def cross_kv_model_loader(model_name): + return CrossKvModule(load_full_model(model_name)) + + +def decoder_kv_model_loader(model_name): + return DecoderKvModule(load_full_model(model_name)) + + +# --------------------------------------------------------------------------- # +# Dummy inputs (used both for tracing/export and Olive dummy_inputs_func) # +# --------------------------------------------------------------------------- # +def _dims_from_module(module): + """Recover the dims needed for dummy inputs from a wrapper instance.""" + full = _MODEL_CACHE[next(iter(_MODEL_CACHE))] + return model_dims(full) + + +def frontend_dummy_inputs(model): + d = model_dims(_any_full()) + return { + "audio_chunk": torch.randn(1, model_chunk_samples(), dtype=torch.float32), + "sample_buffer": torch.zeros(1, d["sample_buffer_size"], dtype=torch.float32), + "sample_len": torch.zeros(1, dtype=torch.int64), + "conv1_buffer": torch.zeros(1, d["conv1_channels"], d["left_pad1"], dtype=torch.float32), + "conv2_buffer": torch.zeros(1, d["conv2_channels"], d["left_pad2"], dtype=torch.float32), + "frame_count": torch.zeros(1, dtype=torch.int64), + } + + +def encoder_dummy_inputs(model): + d = model_dims(_any_full()) + return {"features": torch.randn(1, 48, d["encoder_dim"], dtype=torch.float32)} + + +def adapter_dummy_inputs(model): + d = model_dims(_any_full()) + return { + "encoded": torch.randn(1, 24, d["encoder_dim"], dtype=torch.float32), + "pos_offset": torch.zeros(1, dtype=torch.int64), + } + + +def cross_kv_dummy_inputs(model): + d = model_dims(_any_full()) + return {"memory": torch.randn(1, 24, d["decoder_dim"], dtype=torch.float32)} + + +def decoder_kv_dummy_inputs(model): + d = model_dims(_any_full()) + L, H, D = d["num_decoder_layers"], d["num_kv_heads"], d["head_dim"] + past, cross = 6, 24 + return { + "token": torch.ones(1, 4, dtype=torch.int64), + "k_self": torch.randn(L, 1, H, past, D, dtype=torch.float32), + "v_self": torch.randn(L, 1, H, past, D, dtype=torch.float32), + "out_k_cross": torch.randn(L, 1, H, cross, D, dtype=torch.float32), + "out_v_cross": torch.randn(L, 1, H, cross, D, dtype=torch.float32), + } + + +# Helpers so dummy funcs work without receiving the model name -------------- # +_CHUNK_SAMPLES = 8000 + + +def model_chunk_samples(): + return _CHUNK_SAMPLES + + +def set_chunk_samples(value): + global _CHUNK_SAMPLES + _CHUNK_SAMPLES = int(value) + + +def _any_full(): + if not _MODEL_CACHE: + raise RuntimeError("No model loaded yet; call a *_model_loader first.") + return next(iter(_MODEL_CACHE.values())) diff --git a/usefulSensors-moonshine-streaming/cpu/optimize.py b/usefulSensors-moonshine-streaming/cpu/optimize.py new file mode 100644 index 000000000..53c473201 --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/optimize.py @@ -0,0 +1,372 @@ +"""End-to-end Olive optimization pipeline for Moonshine Streaming ASR. + +Exports the HuggingFace ``MoonshineStreamingForConditionalGeneration`` model +into the five stateful ONNX graphs consumed by onnxruntime-genai's +``streaming_enc_dec_asr`` model type, then generates the runtime config files +and fetches the tokenizer + Silero VAD. Everything is reproducible from the +Torch checkpoint -- no pre-built ``.ort`` graphs are used. + + frontend / encoder / adapter / cross_kv / decoder_kv + -> OnnxConversion (FP32, dynamo exporter, dynamic sequence axes) + +Usage: + # Full pipeline (small model -> build/onnx) + python cpu/optimize.py + + # Tiny model + python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-tiny \ + --output-dir build/onnx-tiny + + # Or run a single component directly through the Olive CLI: + python -m olive run --config cpu/moonshine_frontend_fp32_cpu.json +""" + +import argparse +import json +import shutil +import sys +import tempfile +from pathlib import Path + +# Ensure the recipe root is importable so the Olive model_script resolves. +_SCRIPT_DIR = Path(__file__).resolve().parent +_RECIPE_ROOT = _SCRIPT_DIR.parent +if str(_RECIPE_ROOT) not in sys.path: + sys.path.insert(0, str(_RECIPE_ROOT)) + +MODEL_NAME = "usefulsensors/moonshine-streaming-small" +DEFAULT_OUTPUT_DIR = "build/onnx" + +# chunk_samples is the streaming window fed to the frontend each step. It must +# be a multiple of frame_len (80) so the stateful conv phase stays aligned and +# the carried sample_buffer is empty between chunks. 8000 = 0.5s @ 16kHz. +CHUNK_SAMPLES = 8000 + +# Ordered (config file, output basename) for the five component graphs. +COMPONENTS = [ + ("moonshine_frontend_fp32_cpu.json", "frontend.onnx"), + ("moonshine_encoder_fp32_cpu.json", "encoder.onnx"), + ("moonshine_adapter_fp32_cpu.json", "adapter.onnx"), + ("moonshine_cross_kv_fp32_cpu.json", "cross_kv.onnx"), + ("moonshine_decoder_kv_fp32_cpu.json", "decoder_kv.onnx"), +] + +# With --quantize, swap the encoder + decoder_kv to a quantized config +# (frontend/adapter/cross_kv always stay FP32). --quant-method picks the +# algorithm: +# "dynamic" -> INT8 RTN dynamic quant (MatMulInteger + DynamicQuantizeLinear), +# matching the official .ort. +# "kquant" -> INT4 weight-only k-quant (MatMulNBits), like the nemotron +# recipe; smaller and uses least-squares refinement. +# "kquant8" -> INT8 weight-only k-quant (MatMulNBits, bits=8); same k-quant +# least-squares refinement as "kquant" but 8-bit weights for +# higher accuracy at a larger artifact than INT4. +# "kquant8-enc" -> same INT8 weight-only k-quant as "kquant8" but applied to +# the ENCODER ONLY; decoder_kv stays FP32 (use when decoder +# quantization degrades transcription quality). +QUANTIZED_CONFIGS = { + "dynamic": { + "encoder.onnx": "moonshine_encoder_int8_cpu.json", + "decoder_kv.onnx": "moonshine_decoder_kv_int8_cpu.json", + }, + "kquant": { + "encoder.onnx": "moonshine_encoder_int4_cpu.json", + "decoder_kv.onnx": "moonshine_decoder_kv_int4_cpu.json", + }, + "kquant8": { + "encoder.onnx": "moonshine_encoder_kquant8_cpu.json", + "decoder_kv.onnx": "moonshine_decoder_kv_kquant8_cpu.json", + }, + "kquant8-enc": { + # Encoder only -> INT8 k-quant; decoder_kv is omitted so it falls + # through to the FP32 config in COMPONENTS. + "encoder.onnx": "moonshine_encoder_kquant8_cpu.json", + }, +} +_QUANT_LABELS = { + "dynamic": "OnnxConversion -> INT8 DynamicQuant (RTN)", + "kquant": "OnnxConversion -> INT4 k-quant (MatMulNBits)", + "kquant8": "OnnxConversion -> INT8 k-quant (MatMulNBits)", + "kquant8-enc": "OnnxConversion -> INT8 k-quant, encoder only (MatMulNBits)", +} + +# Streaming pipeline hyper-parameters that are NOT stored in the model weights. +# They describe the onnxruntime-genai runtime contract and match the published +# usefulsensors/moonshine-streaming model card. +PIPELINE = { + "max_seq_len": 448, + "tokens_per_second": 6.5, + "max_segment_memory_frames": 500, + "min_segment_memory_frames": 250, + "left_context_frames": 160, +} +VAD = { + "filename": "silero_vad.onnx", + "threshold": 0.5, + "silence_duration_ms": 500, + "prefix_padding_ms": 200, +} + +# genai tokenizer_config.json for the TokenizersBackend (token strings only; +# no model weights). Mirrors the published streaming model's runtime tokenizer. +TOKENIZER_CONFIG = { + "backend": "tokenizers", + "bos_token": "", + "eos_token": "", + "is_local": True, + "model_max_length": 4096, + "pad_token": "", + "processor_class": "MoonshineStreamingProcessor", + "tokenizer_class": "TokenizersBackend", + "unk_token": "", +} + + +def _resolve(path: str) -> Path: + p = Path(path) + return p if p.is_absolute() else _SCRIPT_DIR / p + + +def _run_olive_pipeline(config_name, model_name, output_dir, output_subdir): + """Run one Olive pipeline from a JSON config, overriding the model path and + output directory so a single set of configs works for any checkpoint.""" + from olive import run as olive_run + + with open(_SCRIPT_DIR / config_name) as f: + config = json.load(f) + + config["input_model"]["model_path"] = model_name + config["output_dir"] = str(_resolve(output_dir) / output_subdir) + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", dir=str(_SCRIPT_DIR), delete=False + ) as tmp: + json.dump(config, tmp, indent=4) + tmp_path = tmp.name + try: + olive_run(tmp_path) + finally: + Path(tmp_path).unlink(missing_ok=True) + + +def run_olive_pipelines(model_name, output_dir, quantize=False, quant_method="dynamic"): + configs = QUANTIZED_CONFIGS.get(quant_method, {}) if quantize else {} + for i, (config_name, subdir) in enumerate(COMPONENTS, 1): + if subdir in configs: + config_name = configs[subdir] + label = _QUANT_LABELS[quant_method] + else: + label = "OnnxConversion, FP32" + print(f"=== Stage 1.{i}: Olive {subdir} ({label}) ===") + _run_olive_pipeline(config_name, model_name, output_dir, subdir) + print() + + +def _derive_params(model_name): + """Read every architecture value the config files need from the Torch model.""" + from cpu.moonshine_model_load import load_full_model, model_dims + + model = load_full_model(model_name) + d = model_dims(model) + enc = model.model.encoder + ec = model.config.encoder_config + emb = enc.embedder + + stride = int(emb.conv1.stride[0]) * int(emb.conv2.stride[0]) # subsampling = 4 + sample_rate = int(ec.sample_rate) + frame_len = d["frame_len"] + sliding_windows = [list(w) for w in ec.sliding_windows] + + return { + **d, + "sample_rate": sample_rate, + "conv1_out": int(emb.conv1.out_channels), # 1240 (small) + "conv2_out": int(emb.conv2.out_channels), # 620 (== encoder_dim) + "num_heads": int(ec.num_attention_heads), + "bos_token_id": int(model.config.bos_token_id), + "eos_token_id": int(model.config.eos_token_id), + "pad_token_id": int(model.config.pad_token_id), + "decoder_start_token_id": int(model.config.decoder_start_token_id), + # subsampling maps one memory frame back to `stride` input frames + "seconds_per_memory_frame": frame_len / sample_rate * stride, # 0.02 + # total future frames the encoder depends on = sum of per-layer right ctx + "total_lookahead": int(sum(int(w[1]) for w in sliding_windows)), # 16 + } + + +def generate_configs(model_name, output_dir): + """Write genai_config.json and streaming_config.json from derived params.""" + print("=== Stage 2: Generating config files ===") + p = _derive_params(model_name) + dst = _resolve(output_dir) + dst.mkdir(parents=True, exist_ok=True) + + genai_config = { + "model": { + "type": "streaming_enc_dec_asr", + "bos_token_id": p["bos_token_id"], + "eos_token_id": p["eos_token_id"], + "pad_token_id": p["pad_token_id"], + "decoder_start_token_id": p["decoder_start_token_id"], + "vocab_size": p["vocab_size"], + "sample_rate": p["sample_rate"], + "chunk_samples": CHUNK_SAMPLES, + "encoder": { + "hidden_size": p["encoder_dim"], + "num_attention_heads": p["num_heads"], + "num_hidden_layers": p["num_encoder_layers"], + "head_size": p["head_dim"], + }, + "decoder": { + "hidden_size": p["decoder_dim"], + "num_attention_heads": p["num_heads"], + "num_hidden_layers": p["num_decoder_layers"], + "head_size": p["head_dim"], + }, + "vad": dict(VAD), + "moonshine": { + "frontend_filename": "frontend.onnx", + "encoder_filename": "encoder.onnx", + "adapter_filename": "adapter.onnx", + "cross_kv_filename": "cross_kv.onnx", + "decoder_kv_filename": "decoder_kv.onnx", + "sample_buffer_size": p["sample_buffer_size"], + "conv1_buffer_size": p["left_pad1"], + "conv2_buffer_size": p["left_pad2"], + "total_lookahead": p["total_lookahead"], + "max_seq_len": PIPELINE["max_seq_len"], + "tokens_per_second": PIPELINE["tokens_per_second"], + "seconds_per_memory_frame": p["seconds_per_memory_frame"], + "max_segment_memory_frames": PIPELINE["max_segment_memory_frames"], + "min_segment_memory_frames": PIPELINE["min_segment_memory_frames"], + "left_context_frames": PIPELINE["left_context_frames"], + }, + }, + "search": { + "diversity_penalty": 0.0, + "do_sample": False, + "max_length": PIPELINE["max_seq_len"], + "min_length": 0, + "num_beams": 1, + "num_return_sequences": 1, + "past_present_share_buffer": False, + "repetition_penalty": 1.0, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0, + }, + } + with open(dst / "genai_config.json", "w") as f: + json.dump(genai_config, f, indent=2) + print(" [OK] genai_config.json") + + streaming_config = { + "encoder_dim": p["encoder_dim"], + "decoder_dim": p["decoder_dim"], + "depth": p["num_encoder_layers"], + "nheads": p["num_heads"], + "head_dim": p["head_dim"], + "vocab_size": p["vocab_size"], + "bos_id": p["bos_token_id"], + "eos_id": p["eos_token_id"], + "frame_len": p["frame_len"], + "total_lookahead": p["total_lookahead"], + "d_model_frontend": p["encoder_dim"], + "c1": p["conv1_out"], + "c2": p["conv2_out"], + "frontend_state_shapes": { + "sample_buffer": [1, p["sample_buffer_size"]], + "sample_len": [1], + "conv1_buffer": [1, p["conv1_channels"], p["left_pad1"]], + "conv2_buffer": [1, p["conv2_channels"], p["left_pad2"]], + "frame_count": [1], + }, + } + with open(dst / "streaming_config.json", "w") as f: + json.dump(streaming_config, f, indent=2) + print(" [OK] streaming_config.json") + print() + + +def export_tokenizer(model_name, output_dir): + """Fetch tokenizer.json from the HF repo and write the genai tokenizer_config.""" + from huggingface_hub import hf_hub_download + + print("=== Stage 3: Exporting tokenizer ===") + dst = _resolve(output_dir) + dst.mkdir(parents=True, exist_ok=True) + cached = hf_hub_download(model_name, "tokenizer.json") + shutil.copy2(cached, dst / "tokenizer.json") + with open(dst / "tokenizer_config.json", "w") as f: + json.dump(TOKENIZER_CONFIG, f, indent=2) + print(f" [OK] tokenizer.json + tokenizer_config.json") + print() + + +def download_silero_vad(output_dir): + """Download the Silero VAD ONNX model from onnx-community/silero-vad.""" + from huggingface_hub import hf_hub_download + + print("=== Stage 4: Downloading Silero VAD ===") + dst = _resolve(output_dir) + dst.mkdir(parents=True, exist_ok=True) + cached = hf_hub_download(repo_id="onnx-community/silero-vad", filename="onnx/model.onnx") + shutil.copy2(cached, dst / "silero_vad.onnx") + size_mb = (dst / "silero_vad.onnx").stat().st_size / (1024 * 1024) + print(f" [OK] silero_vad.onnx ({size_mb:.1f} MB)") + print() + + +def cleanup(output_dir): + """Remove Olive's per-run model_config.json (not used by genai).""" + stray = _resolve(output_dir) / "model_config.json" + if stray.exists(): + stray.unlink() + + +def main(): + parser = argparse.ArgumentParser( + description="Export & optimize Moonshine Streaming ASR for CPU (onnxruntime-genai)." + ) + parser.add_argument("--model-name", default=MODEL_NAME, + help="HuggingFace model id (small or tiny).") + parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR, + help=f"Output model directory (default: {DEFAULT_OUTPUT_DIR}).") + parser.add_argument("--skip-vad", action="store_true", + help="Skip downloading Silero VAD (VAD is off by default at runtime).") + parser.add_argument("--quantize", action="store_true", + help="Quantize the encoder + decoder_kv MatMuls " + "(frontend/adapter/cross_kv stay FP32).") + parser.add_argument("--quant-method", choices=["dynamic", "kquant", "kquant8", "kquant8-enc"], default="dynamic", + help="Algorithm used when --quantize is set: 'dynamic' = INT8 RTN " + "dynamic quant (MatMulInteger, matches the official .ort); " + "'kquant' = INT4 weight-only k-quant (MatMulNBits, like nemotron); " + "'kquant8' = INT8 weight-only k-quant (MatMulNBits, bits=8); " + "'kquant8-enc' = INT8 k-quant on the encoder only (decoder_kv stays FP32).") + args = parser.parse_args() + + run_olive_pipelines(args.model_name, args.output_dir, + quantize=args.quantize, quant_method=args.quant_method) + generate_configs(args.model_name, args.output_dir) + export_tokenizer(args.model_name, args.output_dir) + if not args.skip_vad: + vad_dest = _resolve(args.output_dir) / "silero_vad.onnx" + try: + download_silero_vad(args.output_dir) + except Exception as exc: + print(f" Warning: Silero VAD download failed ({exc}).\n" + f" Download manually from https://huggingface.co/onnx-community/silero-vad\n" + f" and place silero_vad.onnx at: {vad_dest}") + cleanup(args.output_dir) + + out = _resolve(args.output_dir) + files = sorted(f for f in out.iterdir() if f.is_file()) + total_mb = sum(f.stat().st_size for f in files) / (1024 * 1024) + print(f"=== Done! Model -> {out} ===") + print(f" Total size: {total_mb:.1f} MB") + for f in files: + print(f" {f.name} ({f.stat().st_size / (1024 * 1024):.1f} MB)") + + +if __name__ == "__main__": + main() diff --git a/usefulSensors-moonshine-streaming/cpu/requirements.txt b/usefulSensors-moonshine-streaming/cpu/requirements.txt new file mode 100644 index 000000000..61f31c9ae --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/requirements.txt @@ -0,0 +1,9 @@ +huggingface_hub>=0.23.0 +numpy>=1.26.0 +olive-ai @ git+https://github.com/microsoft/Olive.git@main +onnx>=1.19.1 +onnxruntime>=1.24.0 +onnxruntime-genai>=0.13.0 +onnxscript>=0.5.0 +torch>=2.9.0 +transformers>=5.2.0 diff --git a/usefulSensors-moonshine-streaming/cpu/validate_export.py b/usefulSensors-moonshine-streaming/cpu/validate_export.py new file mode 100644 index 000000000..b6913ffbf --- /dev/null +++ b/usefulSensors-moonshine-streaming/cpu/validate_export.py @@ -0,0 +1,117 @@ +"""Validate the exported MoonshineStreaming ONNX graphs. + +Two checks per component: + 1. exported .onnx vs torch reference outputs (from refs/*.npz) -> proves + the ONNX export is faithful to the torch wrappers (tight tolerance). + 2. exported .onnx vs official .ort graph on identical inputs -> proves + the wrappers are semantically correct vs the shipped model. The official + cross_kv / decoder graphs are int8-quantized, so these diffs are compared + with a loose tolerance. + +Runs in the ``moonshine`` env (or any env with onnxruntime + numpy): + + python validate_export.py \ + --mine /datadisks/disk3/nebanfic/moonshine-streaming-small-mine \ + --official /datadisks/disk3/nebanfic/moonshine-streaming-small-official +""" + +from __future__ import annotations + +import argparse +import glob +import os + +import numpy as np +import onnxruntime as ort + +COMPONENTS = ["frontend", "encoder", "adapter", "cross_kv", "decoder_kv"] +# Components whose official graph is quantized -> only expect approximate parity. +QUANTIZED_OFFICIAL = {"cross_kv", "decoder_kv"} + + +def make_session(path): + so = ort.SessionOptions() + so.log_severity_level = 3 + return ort.InferenceSession(path, so, providers=["CPUExecutionProvider"]) + + +def find_official(official_dir, name): + for ext in (".onnx", ".ort"): + cand = os.path.join(official_dir, name + ext) + if os.path.exists(cand): + return cand + return None + + +def run_session(sess, feeds): + names = [o.name for o in sess.get_outputs()] + outs = sess.run(names, feeds) + return dict(zip(names, outs)) + + +def summarize(tag, a, b): + a = np.asarray(a, dtype=np.float64) + b = np.asarray(b, dtype=np.float64) + if a.shape != b.shape: + return f" {tag:24s} SHAPE MISMATCH {a.shape} vs {b.shape}" + diff = np.abs(a - b) + denom = np.maximum(np.abs(b), 1e-6) + return (f" {tag:24s} max_abs={diff.max():.3e} mean_abs={diff.mean():.3e}" + f" max_rel={np.max(diff / denom):.3e}") + + +def validate_component(name, mine_dir, official_dir): + print(f"\n=== {name} ===") + refs_path = os.path.join(mine_dir, "refs", f"{name}.npz") + mine_path = os.path.join(mine_dir, f"{name}.onnx") + if not (os.path.exists(refs_path) and os.path.exists(mine_path)): + print(" (missing exported graph or refs; skipped)") + return + + refs = np.load(refs_path) + feeds = {k[len("in__"):]: refs[k] for k in refs.files if k.startswith("in__")} + torch_out = {k[len("out__"):]: refs[k] for k in refs.files if k.startswith("out__")} + + mine_sess = make_session(mine_path) + mine_out = run_session(mine_sess, feeds) + + print(" [onnx vs torch reference]") + for oname, tval in torch_out.items(): + if oname in mine_out: + print(summarize(oname, mine_out[oname], tval)) + + if official_dir: + off_path = find_official(official_dir, name) + if off_path is None: + print(f" [official] no graph found for {name}") + return + off_sess = make_session(off_path) + # Match the official graph's expected input names/order. + off_in_names = {i.name for i in off_sess.get_inputs()} + off_feeds = {k: v for k, v in feeds.items() if k in off_in_names} + missing = off_in_names - set(off_feeds) + if missing: + print(f" [official] missing inputs {missing}; skipped") + return + off_out = run_session(off_sess, off_feeds) + tol_note = " (quantized: loose)" if name in QUANTIZED_OFFICIAL else "" + print(f" [onnx vs official{tol_note}]") + for oname in mine_out: + if oname in off_out: + print(summarize(oname, mine_out[oname], off_out[oname])) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mine", required=True) + p.add_argument("--official", default=None) + p.add_argument("--only", nargs="*", default=None) + args = p.parse_args() + + comps = args.only or COMPONENTS + for name in comps: + validate_component(name, args.mine, args.official) + + +if __name__ == "__main__": + main() From 4cce501ed6bfed0cce819dd868a37eae9264005c Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Tue, 21 Jul 2026 20:01:50 +0000 Subject: [PATCH 2/8] Remove dump_refs / validate_export (dead after refs removal) --- .../cpu/export_moonshine_streaming.py | 28 +---- .../cpu/validate_export.py | 117 ------------------ 2 files changed, 1 insertion(+), 144 deletions(-) delete mode 100644 usefulSensors-moonshine-streaming/cpu/validate_export.py diff --git a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py index 884c74e62..318b903cc 100644 --- a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py +++ b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py @@ -11,10 +11,7 @@ The exporter uses the TorchDynamo ONNX path (``dynamo=True``) so data-dependent shapes in the stateful frontend export cleanly, then rewrites every graph's -input/output names to the exact contract genai expects. Reference inputs and -torch outputs for each graph are dumped to ``/refs/`` so the -companion validation script can check numerical parity against the official -``.ort`` graphs without needing transformers. +input/output names to the exact contract genai expects. """ from __future__ import annotations @@ -22,7 +19,6 @@ import argparse import os -import numpy as np import onnx import torch from torch.export import Dim @@ -117,23 +113,6 @@ def remap(value_infos, desired): return [i.name for i in graph.input], [o.name for o in graph.output] -# --------------------------------------------------------------------------- # -# Reference dump for validation # -# --------------------------------------------------------------------------- # -def dump_refs(refs_dir, name, module, dummy_args, input_names, output_names): - os.makedirs(refs_dir, exist_ok=True) - with torch.no_grad(): - outputs = module(*dummy_args) - if not isinstance(outputs, (tuple, list)): - outputs = (outputs,) - arrays = {} - for n, t in zip(input_names, dummy_args): - arrays[f"in__{n}"] = t.detach().cpu().numpy() - for n, t in zip(output_names, outputs): - arrays[f"out__{n}"] = t.detach().cpu().numpy() - np.savez(os.path.join(refs_dir, f"{name}.npz"), **arrays) - - # --------------------------------------------------------------------------- # # Main # # --------------------------------------------------------------------------- # @@ -162,11 +141,6 @@ def export_component(spec, model_name, output_dir, opset): onnx.checker.check_model(onnx_path) print(f" inputs : {ins}") print(f" outputs: {outs}") - - dump_refs( - os.path.join(output_dir, "refs"), name, module, dummy_args, - spec["input_names"], spec["output_names"], - ) return onnx_path diff --git a/usefulSensors-moonshine-streaming/cpu/validate_export.py b/usefulSensors-moonshine-streaming/cpu/validate_export.py deleted file mode 100644 index b6913ffbf..000000000 --- a/usefulSensors-moonshine-streaming/cpu/validate_export.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Validate the exported MoonshineStreaming ONNX graphs. - -Two checks per component: - 1. exported .onnx vs torch reference outputs (from refs/*.npz) -> proves - the ONNX export is faithful to the torch wrappers (tight tolerance). - 2. exported .onnx vs official .ort graph on identical inputs -> proves - the wrappers are semantically correct vs the shipped model. The official - cross_kv / decoder graphs are int8-quantized, so these diffs are compared - with a loose tolerance. - -Runs in the ``moonshine`` env (or any env with onnxruntime + numpy): - - python validate_export.py \ - --mine /datadisks/disk3/nebanfic/moonshine-streaming-small-mine \ - --official /datadisks/disk3/nebanfic/moonshine-streaming-small-official -""" - -from __future__ import annotations - -import argparse -import glob -import os - -import numpy as np -import onnxruntime as ort - -COMPONENTS = ["frontend", "encoder", "adapter", "cross_kv", "decoder_kv"] -# Components whose official graph is quantized -> only expect approximate parity. -QUANTIZED_OFFICIAL = {"cross_kv", "decoder_kv"} - - -def make_session(path): - so = ort.SessionOptions() - so.log_severity_level = 3 - return ort.InferenceSession(path, so, providers=["CPUExecutionProvider"]) - - -def find_official(official_dir, name): - for ext in (".onnx", ".ort"): - cand = os.path.join(official_dir, name + ext) - if os.path.exists(cand): - return cand - return None - - -def run_session(sess, feeds): - names = [o.name for o in sess.get_outputs()] - outs = sess.run(names, feeds) - return dict(zip(names, outs)) - - -def summarize(tag, a, b): - a = np.asarray(a, dtype=np.float64) - b = np.asarray(b, dtype=np.float64) - if a.shape != b.shape: - return f" {tag:24s} SHAPE MISMATCH {a.shape} vs {b.shape}" - diff = np.abs(a - b) - denom = np.maximum(np.abs(b), 1e-6) - return (f" {tag:24s} max_abs={diff.max():.3e} mean_abs={diff.mean():.3e}" - f" max_rel={np.max(diff / denom):.3e}") - - -def validate_component(name, mine_dir, official_dir): - print(f"\n=== {name} ===") - refs_path = os.path.join(mine_dir, "refs", f"{name}.npz") - mine_path = os.path.join(mine_dir, f"{name}.onnx") - if not (os.path.exists(refs_path) and os.path.exists(mine_path)): - print(" (missing exported graph or refs; skipped)") - return - - refs = np.load(refs_path) - feeds = {k[len("in__"):]: refs[k] for k in refs.files if k.startswith("in__")} - torch_out = {k[len("out__"):]: refs[k] for k in refs.files if k.startswith("out__")} - - mine_sess = make_session(mine_path) - mine_out = run_session(mine_sess, feeds) - - print(" [onnx vs torch reference]") - for oname, tval in torch_out.items(): - if oname in mine_out: - print(summarize(oname, mine_out[oname], tval)) - - if official_dir: - off_path = find_official(official_dir, name) - if off_path is None: - print(f" [official] no graph found for {name}") - return - off_sess = make_session(off_path) - # Match the official graph's expected input names/order. - off_in_names = {i.name for i in off_sess.get_inputs()} - off_feeds = {k: v for k, v in feeds.items() if k in off_in_names} - missing = off_in_names - set(off_feeds) - if missing: - print(f" [official] missing inputs {missing}; skipped") - return - off_out = run_session(off_sess, off_feeds) - tol_note = " (quantized: loose)" if name in QUANTIZED_OFFICIAL else "" - print(f" [onnx vs official{tol_note}]") - for oname in mine_out: - if oname in off_out: - print(summarize(oname, mine_out[oname], off_out[oname])) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--mine", required=True) - p.add_argument("--official", default=None) - p.add_argument("--only", nargs="*", default=None) - args = p.parse_args() - - comps = args.only or COMPONENTS - for name in comps: - validate_component(name, args.mine, args.official) - - -if __name__ == "__main__": - main() From b778e5fbee991e77e7a709de4b19064f2ab24627 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Tue, 21 Jul 2026 20:12:46 +0000 Subject: [PATCH 3/8] Drop moonshine INT4 k-quant variant (accuracy too low) --- .../cpu/README.md | 25 ++------- .../cpu/moonshine_decoder_kv_int4_cpu.json | 52 ------------------- .../cpu/moonshine_encoder_int4_cpu.json | 48 ----------------- .../cpu/optimize.py | 28 +++++----- 4 files changed, 19 insertions(+), 134 deletions(-) delete mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json delete mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json diff --git a/usefulSensors-moonshine-streaming/cpu/README.md b/usefulSensors-moonshine-streaming/cpu/README.md index 2497a9c5e..48d9a61aa 100644 --- a/usefulSensors-moonshine-streaming/cpu/README.md +++ b/usefulSensors-moonshine-streaming/cpu/README.md @@ -26,8 +26,6 @@ with the exact input/output names the streaming runner expects: - `cpu/moonshine_decoder_kv_fp32_cpu.json` – Olive decoder-KV config (convert only) - `cpu/moonshine_encoder_int8_cpu.json` – Olive encoder config (convert → INT8 dynamic quant) - `cpu/moonshine_decoder_kv_int8_cpu.json` – Olive decoder-KV config (convert → INT8 dynamic quant) -- `cpu/moonshine_encoder_int4_cpu.json` – Olive encoder config (convert → INT4 k-quant) -- `cpu/moonshine_decoder_kv_int4_cpu.json` – Olive decoder-KV config (convert → INT4 k-quant) - `cpu/moonshine_model_load.py` – model loaders + wrapper modules + dummy inputs - `cpu/optimize.py` – full pipeline script (Olive × 5 + tokenizer + configs + VAD) - `cpu/export_moonshine_streaming.py` – standalone exporter (no Olive; for debugging) @@ -79,39 +77,22 @@ are unchanged. `--quant-method` picks the algorithm: which chain `OnnxConversion → OnnxDynamicQuantization` (weight matmuls become `MatMulInteger` + `DynamicQuantizeLinear`), matching the shipped official `.ort`. Fastest on CPU. -- `--quant-method kquant` — INT4 **weight-only k-quant** (like the nemotron - recipe). Swaps in `moonshine_encoder_int4_cpu.json` / - `moonshine_decoder_kv_int4_cpu.json`, which chain - `OnnxConversion → OnnxKQuantQuantization` (`bits=4`, `block_size=32`, - `accuracy_level=4`; weight matmuls become `MatMulNBits`). Smallest on disk; - weight-only, so activation×activation attention matmuls stay FP32 and the - token-embedding `Gather` is left FP32. ```bash # INT8 RTN dynamic (default) python cpu/optimize.py --quantize --output-dir build/moonshine-small-int8 - -# INT4 k-quant -python cpu/optimize.py --quantize --quant-method kquant \ - --output-dir build/moonshine-small-int4 ``` -On a 40s clip (CPU EP), both preserve transcription quality; INT8 dynamic is -fastest, INT4 k-quant is smallest: +On a 40s clip (CPU EP), INT8 dynamic preserves transcription quality: | build | encoder | decoder_kv | total | RTF | |---|---|---|---|---| | FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× | | INT8 dynamic (`--quantize`) | 42 MB | 125 MB | ~233 MB | ~9.0× | -| INT4 k-quant (`--quant-method kquant`) | 27 MB | 103 MB | ~196 MB | ~7.7× | | official `.ort` | 42 MB | 174 MB | — | ~9.4× | -Transcription is essentially identical to FP32 for both (only quant-noise -wording drift). INT4 k-quant is the smallest model but on CPU it is *not* -faster than INT8 dynamic: `MatMulNBits` is weight-only, so the int4→compute -de-quant overhead offsets the memory-bandwidth win for this small, -compute-bound model. Prefer `dynamic` for speed, `kquant` for the smallest -artifact. +Transcription is essentially identical to FP32 for INT8 dynamic (only +quant-noise wording drift). Or run individual components directly with the Olive CLI: diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json deleted file mode 100644 index d156e610e..000000000 --- a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int4_cpu.json +++ /dev/null @@ -1,52 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "usefulsensors/moonshine-streaming-small", - "model_loader": "decoder_kv_model_loader", - "model_script": "cpu/moonshine_model_load.py", - "io_config": { - "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], - "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], - "dynamic_shapes": { - "token": {"1": "q_seq"}, - "k_self": {"3": "past_seq"}, - "v_self": {"3": "past_seq"}, - "out_k_cross": {"3": "cross_seq"}, - "out_v_cross": {"3": "cross_seq"} - } - }, - "dummy_inputs_func": "decoder_kv_dummy_inputs" - }, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "cpu", - "execution_providers": ["CPUExecutionProvider"] - } - ] - } - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "target_opset": 20, - "dynamic": true, - "use_dynamo_exporter": true, - "save_as_external_data": true, - "external_data_name": "decoder_kv_fp32.onnx.data" - }, - "quantize": { - "type": "OnnxKQuantQuantization", - "bits": 4, - "block_size": 32, - "accuracy_level": 4, - "save_as_external_data": true, - "external_data_name": "decoder_kv.onnx.data" - } - }, - "target": "local_system", - "output_dir": "build/onnx/decoder_kv.onnx", - "no_artifacts": true -} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json deleted file mode 100644 index 3ade9b78d..000000000 --- a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int4_cpu.json +++ /dev/null @@ -1,48 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "usefulsensors/moonshine-streaming-small", - "model_loader": "encoder_model_loader", - "model_script": "cpu/moonshine_model_load.py", - "io_config": { - "input_names": ["features"], - "output_names": ["encoded"], - "dynamic_shapes": { - "features": {"1": "enc_seq"} - } - }, - "dummy_inputs_func": "encoder_dummy_inputs" - }, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "cpu", - "execution_providers": ["CPUExecutionProvider"] - } - ] - } - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "target_opset": 20, - "dynamic": true, - "use_dynamo_exporter": true, - "save_as_external_data": true, - "external_data_name": "encoder_fp32.onnx.data" - }, - "quantize": { - "type": "OnnxKQuantQuantization", - "bits": 4, - "block_size": 32, - "accuracy_level": 4, - "save_as_external_data": true, - "external_data_name": "encoder.onnx.data" - } - }, - "target": "local_system", - "output_dir": "build/onnx/encoder.onnx", - "no_artifacts": true -} diff --git a/usefulSensors-moonshine-streaming/cpu/optimize.py b/usefulSensors-moonshine-streaming/cpu/optimize.py index 53c473201..adcc583d5 100644 --- a/usefulSensors-moonshine-streaming/cpu/optimize.py +++ b/usefulSensors-moonshine-streaming/cpu/optimize.py @@ -56,11 +56,9 @@ # algorithm: # "dynamic" -> INT8 RTN dynamic quant (MatMulInteger + DynamicQuantizeLinear), # matching the official .ort. -# "kquant" -> INT4 weight-only k-quant (MatMulNBits), like the nemotron -# recipe; smaller and uses least-squares refinement. -# "kquant8" -> INT8 weight-only k-quant (MatMulNBits, bits=8); same k-quant -# least-squares refinement as "kquant" but 8-bit weights for -# higher accuracy at a larger artifact than INT4. +# "kquant8" -> INT8 weight-only k-quant (MatMulNBits, bits=8); uses +# least-squares refinement for higher accuracy than dynamic +# quant at a similar disk size. # "kquant8-enc" -> same INT8 weight-only k-quant as "kquant8" but applied to # the ENCODER ONLY; decoder_kv stays FP32 (use when decoder # quantization degrades transcription quality). @@ -69,10 +67,6 @@ "encoder.onnx": "moonshine_encoder_int8_cpu.json", "decoder_kv.onnx": "moonshine_decoder_kv_int8_cpu.json", }, - "kquant": { - "encoder.onnx": "moonshine_encoder_int4_cpu.json", - "decoder_kv.onnx": "moonshine_decoder_kv_int4_cpu.json", - }, "kquant8": { "encoder.onnx": "moonshine_encoder_kquant8_cpu.json", "decoder_kv.onnx": "moonshine_decoder_kv_kquant8_cpu.json", @@ -85,7 +79,6 @@ } _QUANT_LABELS = { "dynamic": "OnnxConversion -> INT8 DynamicQuant (RTN)", - "kquant": "OnnxConversion -> INT4 k-quant (MatMulNBits)", "kquant8": "OnnxConversion -> INT8 k-quant (MatMulNBits)", "kquant8-enc": "OnnxConversion -> INT8 k-quant, encoder only (MatMulNBits)", } @@ -337,14 +330,25 @@ def main(): parser.add_argument("--quantize", action="store_true", help="Quantize the encoder + decoder_kv MatMuls " "(frontend/adapter/cross_kv stay FP32).") - parser.add_argument("--quant-method", choices=["dynamic", "kquant", "kquant8", "kquant8-enc"], default="dynamic", + parser.add_argument("--quant-method", choices=["dynamic", "kquant8", "kquant8-enc"], default="dynamic", help="Algorithm used when --quantize is set: 'dynamic' = INT8 RTN " "dynamic quant (MatMulInteger, matches the official .ort); " - "'kquant' = INT4 weight-only k-quant (MatMulNBits, like nemotron); " "'kquant8' = INT8 weight-only k-quant (MatMulNBits, bits=8); " "'kquant8-enc' = INT8 k-quant on the encoder only (decoder_kv stays FP32).") args = parser.parse_args() + # This recipe is architecture-locked to the usefulsensors/moonshine-streaming + # family (tiny / small). Other checkpoints may have different frontend + # buffer sizes, encoder/decoder dims, or tokenizer configs and won't + # produce a runnable genai model. Warn but don't hard-fail so local + # forks can still opt in. + if not args.model_name.startswith("usefulsensors/moonshine-streaming-"): + print( + f" Warning: --model-name '{args.model_name}' is outside the " + f"'usefulsensors/moonshine-streaming-*' family this recipe was " + f"built for; export may fail or produce an unusable model." + ) + run_olive_pipelines(args.model_name, args.output_dir, quantize=args.quantize, quant_method=args.quant_method) generate_configs(args.model_name, args.output_dir) From 964020865680af9ead40356433085f73516c4d52 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Fri, 24 Jul 2026 22:23:30 +0000 Subject: [PATCH 4/8] More changes --- .../cpu/README.md | 33 +++++++----- .../cpu/moonshine_decoder_kv_int8_cpu.json | 53 ------------------- .../cpu/moonshine_encoder_int8_cpu.json | 49 ----------------- .../cpu/optimize.py | 18 ++----- 4 files changed, 25 insertions(+), 128 deletions(-) delete mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json delete mode 100644 usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json diff --git a/usefulSensors-moonshine-streaming/cpu/README.md b/usefulSensors-moonshine-streaming/cpu/README.md index 48d9a61aa..39ef1c402 100644 --- a/usefulSensors-moonshine-streaming/cpu/README.md +++ b/usefulSensors-moonshine-streaming/cpu/README.md @@ -24,8 +24,8 @@ with the exact input/output names the streaming runner expects: - `cpu/moonshine_adapter_fp32_cpu.json` – Olive adapter config (convert only) - `cpu/moonshine_cross_kv_fp32_cpu.json` – Olive cross-KV config (convert only) - `cpu/moonshine_decoder_kv_fp32_cpu.json` – Olive decoder-KV config (convert only) -- `cpu/moonshine_encoder_int8_cpu.json` – Olive encoder config (convert → INT8 dynamic quant) -- `cpu/moonshine_decoder_kv_int8_cpu.json` – Olive decoder-KV config (convert → INT8 dynamic quant) +- `cpu/moonshine_encoder_kquant8_cpu.json` – Olive encoder config (convert → INT8 k-quant) +- `cpu/moonshine_decoder_kv_kquant8_cpu.json` – Olive decoder-KV config (convert → INT8 k-quant) - `cpu/moonshine_model_load.py` – model loaders + wrapper modules + dummy inputs - `cpu/optimize.py` – full pipeline script (Olive × 5 + tokenizer + configs + VAD) - `cpu/export_moonshine_streaming.py` – standalone exporter (no Olive; for debugging) @@ -72,27 +72,34 @@ Add `--quantize` to quantize the **encoder** and **decoder_kv** MatMuls (frontend / adapter / cross_kv stay FP32). IO names and the runtime configs are unchanged. `--quant-method` picks the algorithm: -- `--quant-method dynamic` (default) — INT8 RTN **dynamic** quant. Swaps in - `moonshine_encoder_int8_cpu.json` / `moonshine_decoder_kv_int8_cpu.json`, - which chain `OnnxConversion → OnnxDynamicQuantization` (weight matmuls become - `MatMulInteger` + `DynamicQuantizeLinear`), matching the shipped official - `.ort`. Fastest on CPU. +- `--quant-method kquant8` (default) — INT8 **weight-only k-quant**. Swaps in + `moonshine_encoder_kquant8_cpu.json` / `moonshine_decoder_kv_kquant8_cpu.json`, + which chain `OnnxConversion → OnnxKQuantQuantization` (`bits=8`, `block_size=32`; + weight matmuls become `MatMulNBits`). Weight-only, so activation×activation + attention matmuls stay FP32. +- `--quant-method kquant8-enc` — same INT8 k-quant, **encoder only**; + `decoder_kv` stays FP32. Use when decoder quantization degrades transcription + quality and you can afford the larger decoder. ```bash -# INT8 RTN dynamic (default) -python cpu/optimize.py --quantize --output-dir build/moonshine-small-int8 +# INT8 k-quant on encoder + decoder_kv (default) +python cpu/optimize.py --quantize --output-dir build/moonshine-small-kquant8 + +# INT8 k-quant on encoder only, FP32 decoder +python cpu/optimize.py --quantize --quant-method kquant8-enc \ + --output-dir build/moonshine-small-kquant8-enc ``` -On a 40s clip (CPU EP), INT8 dynamic preserves transcription quality: +On a 40s clip (CPU EP), INT8 k-quant preserves transcription quality: | build | encoder | decoder_kv | total | RTF | |---|---|---|---|---| | FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× | -| INT8 dynamic (`--quantize`) | 42 MB | 125 MB | ~233 MB | ~9.0× | +| INT8 k-quant (`--quantize`) | ~40 MB | ~120 MB | ~230 MB | ~8– 9× | | official `.ort` | 42 MB | 174 MB | — | ~9.4× | -Transcription is essentially identical to FP32 for INT8 dynamic (only -quant-noise wording drift). +Transcription is essentially identical to FP32 (only quant-noise wording +drift). Or run individual components directly with the Olive CLI: diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json deleted file mode 100644 index 538dc3cbe..000000000 --- a/usefulSensors-moonshine-streaming/cpu/moonshine_decoder_kv_int8_cpu.json +++ /dev/null @@ -1,53 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "usefulsensors/moonshine-streaming-small", - "model_loader": "decoder_kv_model_loader", - "model_script": "cpu/moonshine_model_load.py", - "io_config": { - "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"], - "output_names": ["logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross"], - "dynamic_shapes": { - "token": {"1": "q_seq"}, - "k_self": {"3": "past_seq"}, - "v_self": {"3": "past_seq"}, - "out_k_cross": {"3": "cross_seq"}, - "out_v_cross": {"3": "cross_seq"} - } - }, - "dummy_inputs_func": "decoder_kv_dummy_inputs" - }, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "cpu", - "execution_providers": ["CPUExecutionProvider"] - } - ] - } - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "target_opset": 20, - "dynamic": true, - "use_dynamo_exporter": true, - "save_as_external_data": true, - "external_data_name": "decoder_kv_fp32.onnx.data" - }, - "quantize": { - "type": "OnnxDynamicQuantization", - "precision": "int8", - "op_types_to_quantize": ["MatMul"], - "per_channel": true, - "quant_preprocess": false, - "save_as_external_data": true, - "external_data_name": "decoder_kv.onnx.data" - } - }, - "target": "local_system", - "output_dir": "build/onnx/decoder_kv.onnx", - "no_artifacts": true -} diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json b/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json deleted file mode 100644 index e604ac06e..000000000 --- a/usefulSensors-moonshine-streaming/cpu/moonshine_encoder_int8_cpu.json +++ /dev/null @@ -1,49 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "usefulsensors/moonshine-streaming-small", - "model_loader": "encoder_model_loader", - "model_script": "cpu/moonshine_model_load.py", - "io_config": { - "input_names": ["features"], - "output_names": ["encoded"], - "dynamic_shapes": { - "features": {"1": "enc_seq"} - } - }, - "dummy_inputs_func": "encoder_dummy_inputs" - }, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "cpu", - "execution_providers": ["CPUExecutionProvider"] - } - ] - } - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "target_opset": 20, - "dynamic": true, - "use_dynamo_exporter": true, - "save_as_external_data": true, - "external_data_name": "encoder_fp32.onnx.data" - }, - "quantize": { - "type": "OnnxDynamicQuantization", - "precision": "int8", - "op_types_to_quantize": ["MatMul"], - "per_channel": true, - "quant_preprocess": false, - "save_as_external_data": true, - "external_data_name": "encoder.onnx.data" - } - }, - "target": "local_system", - "output_dir": "build/onnx/encoder.onnx", - "no_artifacts": true -} diff --git a/usefulSensors-moonshine-streaming/cpu/optimize.py b/usefulSensors-moonshine-streaming/cpu/optimize.py index adcc583d5..f9b4073fa 100644 --- a/usefulSensors-moonshine-streaming/cpu/optimize.py +++ b/usefulSensors-moonshine-streaming/cpu/optimize.py @@ -54,19 +54,13 @@ # With --quantize, swap the encoder + decoder_kv to a quantized config # (frontend/adapter/cross_kv always stay FP32). --quant-method picks the # algorithm: -# "dynamic" -> INT8 RTN dynamic quant (MatMulInteger + DynamicQuantizeLinear), -# matching the official .ort. # "kquant8" -> INT8 weight-only k-quant (MatMulNBits, bits=8); uses -# least-squares refinement for higher accuracy than dynamic -# quant at a similar disk size. +# least-squares refinement for higher accuracy than plain RTN +# at a similar disk size. # "kquant8-enc" -> same INT8 weight-only k-quant as "kquant8" but applied to # the ENCODER ONLY; decoder_kv stays FP32 (use when decoder # quantization degrades transcription quality). QUANTIZED_CONFIGS = { - "dynamic": { - "encoder.onnx": "moonshine_encoder_int8_cpu.json", - "decoder_kv.onnx": "moonshine_decoder_kv_int8_cpu.json", - }, "kquant8": { "encoder.onnx": "moonshine_encoder_kquant8_cpu.json", "decoder_kv.onnx": "moonshine_decoder_kv_kquant8_cpu.json", @@ -78,7 +72,6 @@ }, } _QUANT_LABELS = { - "dynamic": "OnnxConversion -> INT8 DynamicQuant (RTN)", "kquant8": "OnnxConversion -> INT8 k-quant (MatMulNBits)", "kquant8-enc": "OnnxConversion -> INT8 k-quant, encoder only (MatMulNBits)", } @@ -142,7 +135,7 @@ def _run_olive_pipeline(config_name, model_name, output_dir, output_subdir): Path(tmp_path).unlink(missing_ok=True) -def run_olive_pipelines(model_name, output_dir, quantize=False, quant_method="dynamic"): +def run_olive_pipelines(model_name, output_dir, quantize=False, quant_method="kquant8"): configs = QUANTIZED_CONFIGS.get(quant_method, {}) if quantize else {} for i, (config_name, subdir) in enumerate(COMPONENTS, 1): if subdir in configs: @@ -330,9 +323,8 @@ def main(): parser.add_argument("--quantize", action="store_true", help="Quantize the encoder + decoder_kv MatMuls " "(frontend/adapter/cross_kv stay FP32).") - parser.add_argument("--quant-method", choices=["dynamic", "kquant8", "kquant8-enc"], default="dynamic", - help="Algorithm used when --quantize is set: 'dynamic' = INT8 RTN " - "dynamic quant (MatMulInteger, matches the official .ort); " + parser.add_argument("--quant-method", choices=["kquant8", "kquant8-enc"], default="kquant8", + help="Algorithm used when --quantize is set: " "'kquant8' = INT8 weight-only k-quant (MatMulNBits, bits=8); " "'kquant8-enc' = INT8 k-quant on the encoder only (decoder_kv stays FP32).") args = parser.parse_args() From e094b2282baffe7a8266d3974004ccab9ec55463 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Fri, 24 Jul 2026 22:32:40 +0000 Subject: [PATCH 5/8] More cjanges --- .../cpu/.gitignore | 14 -------- .../cpu/README.md | 36 ++++++------------- .../cpu/info.yaml | 18 +++++----- 3 files changed, 20 insertions(+), 48 deletions(-) delete mode 100644 usefulSensors-moonshine-streaming/cpu/.gitignore diff --git a/usefulSensors-moonshine-streaming/cpu/.gitignore b/usefulSensors-moonshine-streaming/cpu/.gitignore deleted file mode 100644 index 291419770..000000000 --- a/usefulSensors-moonshine-streaming/cpu/.gitignore +++ /dev/null @@ -1,14 +0,0 @@ -# Generated model artifacts -build/ - -# Python bytecode -__pycache__/ -*.pyc - -# Olive cache -.olive-cache/ - -# Temp and log files -*.temp -*.bak -*.log diff --git a/usefulSensors-moonshine-streaming/cpu/README.md b/usefulSensors-moonshine-streaming/cpu/README.md index 39ef1c402..c67d2daad 100644 --- a/usefulSensors-moonshine-streaming/cpu/README.md +++ b/usefulSensors-moonshine-streaming/cpu/README.md @@ -1,7 +1,9 @@ -# Moonshine Streaming (CPU EP, FP32) +# Moonshine Streaming (CPU EP) -This recipe exports **usefulsensors/moonshine-streaming-small** to ONNX and -produces CPU-ready ONNX Runtime GenAI artifacts for streaming ASR. +This recipe exports **usefulsensors/moonshine-streaming-small** or +**usefulsensors/moonshine-streaming-tiny** to ONNX and produces CPU-ready +ONNX Runtime GenAI artifacts for streaming ASR. Pick the variant with +`--model-name` (see [Run](#run)). The streaming model is exported as **five FP32 ONNX components**, each handled through Olive's declarative `OnnxConversion` pass (dynamo exporter, opset 20) @@ -29,7 +31,6 @@ with the exact input/output names the streaming runner expects: - `cpu/moonshine_model_load.py` – model loaders + wrapper modules + dummy inputs - `cpu/optimize.py` – full pipeline script (Olive × 5 + tokenizer + configs + VAD) - `cpu/export_moonshine_streaming.py` – standalone exporter (no Olive; for debugging) -- `cpu/validate_export.py` – per-component torch-vs-ONNX numeric check ## Setup From repo root: @@ -92,11 +93,12 @@ python cpu/optimize.py --quantize --quant-method kquant8-enc \ On a 40s clip (CPU EP), INT8 k-quant preserves transcription quality: -| build | encoder | decoder_kv | total | RTF | -|---|---|---|---|---| -| FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× | -| INT8 k-quant (`--quantize`) | ~40 MB | ~120 MB | ~230 MB | ~8– 9× | -| official `.ort` | 42 MB | 174 MB | — | ~9.4× | +| variant | build | encoder | decoder_kv | total | RTF | +|---|---|---|---|---|---| +| small | FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× | +| small | INT8 k-quant (`--quantize`) | ~40 MB | ~120 MB | ~230 MB | ~8–9× | +| small | official `.ort` | 42 MB | 174 MB | — | ~9.4× | +| tiny | INT8 k-quant (`--quantize`) | ~8 MB | ~64 MB | ~94 MB | ~17× | Transcription is essentially identical to FP32 (only quant-noise wording drift). @@ -123,19 +125,3 @@ Expected artifacts in `cpu/build/moonshine-small/`: - `tokenizer.json` - `tokenizer_config.json` - `silero_vad.onnx` - -## Validation -`validate_export.py` checks each component two ways: exported ONNX vs the -PyTorch reference outputs (tight tolerance), and — optionally — exported ONNX vs -the shipped official graph on identical inputs (loose tolerance for the -int8-quantized cross-KV / decoder-KV graphs). - -It expects the `--mine` directory to contain `.onnx` plus -`refs/.npz` (torch reference in/out dumped by the standalone -`export_moonshine_streaming.py`). `--official` is optional: - -```bash -python cpu/validate_export.py \ - --mine /path/to/moonshine-streaming-small-mine \ - --official /path/to/moonshine-streaming-small-official -``` diff --git a/usefulSensors-moonshine-streaming/cpu/info.yaml b/usefulSensors-moonshine-streaming/cpu/info.yaml index e7108185d..6dd3ed215 100644 --- a/usefulSensors-moonshine-streaming/cpu/info.yaml +++ b/usefulSensors-moonshine-streaming/cpu/info.yaml @@ -1,14 +1,17 @@ name: usefulSensors-moonshine-streaming provider: usefulsensors model_id: usefulsensors/moonshine-streaming-small +# The recipe also supports usefulsensors/moonshine-streaming-tiny via +# `--model-name usefulsensors/moonshine-streaming-tiny`. task: automatic-speech-recognition framework: ONNX Runtime execution_provider: CPUExecutionProvider summary: > - CPU recipe for exporting the Moonshine streaming ASR model to ONNX (FP32) as - five ONNX Runtime GenAI components (frontend, encoder, adapter, cross-KV, - decoder-KV) with the exact IO names the streaming runner expects, plus the - generated genai/streaming configs, tokenizer, and Silero VAD. + CPU recipe for exporting the Moonshine streaming ASR model to ONNX (FP32 or + INT8 k-quant) as five ONNX Runtime GenAI components (frontend, encoder, + adapter, cross-KV, decoder-KV) with the exact IO names the streaming runner + expects, plus the generated genai/streaming configs, tokenizer, and Silero + VAD. artifacts: - cpu/moonshine_frontend_fp32_cpu.json @@ -16,13 +19,10 @@ artifacts: - cpu/moonshine_adapter_fp32_cpu.json - cpu/moonshine_cross_kv_fp32_cpu.json - cpu/moonshine_decoder_kv_fp32_cpu.json - - cpu/moonshine_encoder_int8_cpu.json - - cpu/moonshine_decoder_kv_int8_cpu.json - - cpu/moonshine_encoder_int4_cpu.json - - cpu/moonshine_decoder_kv_int4_cpu.json + - cpu/moonshine_encoder_kquant8_cpu.json + - cpu/moonshine_decoder_kv_kquant8_cpu.json scripts: - cpu/optimize.py - cpu/moonshine_model_load.py - cpu/export_moonshine_streaming.py - - cpu/validate_export.py From 9f799e6e13df89fe7762998d041b0c9edd47292e Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Fri, 24 Jul 2026 22:37:24 +0000 Subject: [PATCH 6/8] Make tiny default --- .../cpu/README.md | 18 +++++++++--------- .../cpu/export_moonshine_streaming.py | 6 +++--- .../cpu/info.yaml | 6 +++--- .../cpu/optimize.py | 2 +- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/usefulSensors-moonshine-streaming/cpu/README.md b/usefulSensors-moonshine-streaming/cpu/README.md index c67d2daad..7d9b76679 100644 --- a/usefulSensors-moonshine-streaming/cpu/README.md +++ b/usefulSensors-moonshine-streaming/cpu/README.md @@ -1,7 +1,7 @@ # Moonshine Streaming (CPU EP) -This recipe exports **usefulsensors/moonshine-streaming-small** or -**usefulsensors/moonshine-streaming-tiny** to ONNX and produces CPU-ready +This recipe exports **usefulsensors/moonshine-streaming-tiny** (default) or +**usefulsensors/moonshine-streaming-small** to ONNX and produces CPU-ready ONNX Runtime GenAI artifacts for streaming ASR. Pick the variant with `--model-name` (see [Run](#run)). @@ -48,7 +48,7 @@ From the `usefulSensors-moonshine-streaming` directory: ```bash cd usefulSensors-moonshine-streaming -python cpu/optimize.py --output-dir build/moonshine-small +python cpu/optimize.py --output-dir build/moonshine-tiny ``` This runs the full pipeline: @@ -60,11 +60,11 @@ This runs the full pipeline: `--output-dir` is resolved relative to the `cpu/` directory unless an absolute path is given. -Export the **tiny** variant instead: +Export the **small** variant instead: ```bash -python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-tiny \ - --output-dir build/moonshine-tiny +python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-small \ + --output-dir build/moonshine-small ``` ### Quantization (`--quantize`) @@ -84,11 +84,11 @@ are unchanged. `--quant-method` picks the algorithm: ```bash # INT8 k-quant on encoder + decoder_kv (default) -python cpu/optimize.py --quantize --output-dir build/moonshine-small-kquant8 +python cpu/optimize.py --quantize --output-dir build/moonshine-tiny-kquant8 # INT8 k-quant on encoder only, FP32 decoder python cpu/optimize.py --quantize --quant-method kquant8-enc \ - --output-dir build/moonshine-small-kquant8-enc + --output-dir build/moonshine-tiny-kquant8-enc ``` On a 40s clip (CPU EP), INT8 k-quant preserves transcription quality: @@ -114,7 +114,7 @@ python -m olive run --config cpu/moonshine_decoder_kv_fp32_cpu.json ``` ## Output -Expected artifacts in `cpu/build/moonshine-small/`: +Expected artifacts in `cpu/build/moonshine-tiny/`: - `frontend.onnx` (+ `frontend.onnx.data`) - `encoder.onnx` (+ `encoder.onnx.data`) - `adapter.onnx` (+ `adapter.onnx.data`) diff --git a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py index 318b903cc..aa623c23e 100644 --- a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py +++ b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py @@ -6,8 +6,8 @@ source /home/nebanfic/miniconda3/bin/activate moonshine python export_moonshine_streaming.py \ - --model usefulsensors/moonshine-streaming-small \ - --output-dir /datadisks/disk3/nebanfic/moonshine-streaming-small-mine + --model usefulsensors/moonshine-streaming-tiny \ + --output-dir /datadisks/disk3/nebanfic/moonshine-streaming-tiny-mine The exporter uses the TorchDynamo ONNX path (``dynamo=True``) so data-dependent shapes in the stateful frontend export cleanly, then rewrites every graph's @@ -146,7 +146,7 @@ def export_component(spec, model_name, output_dir, opset): def main(): parser = argparse.ArgumentParser() - parser.add_argument("--model", default="usefulsensors/moonshine-streaming-small") + parser.add_argument("--model", default="usefulsensors/moonshine-streaming-tiny") parser.add_argument("--output-dir", required=True) parser.add_argument("--opset", type=int, default=20) parser.add_argument("--chunk-samples", type=int, default=8000) diff --git a/usefulSensors-moonshine-streaming/cpu/info.yaml b/usefulSensors-moonshine-streaming/cpu/info.yaml index 6dd3ed215..086f92106 100644 --- a/usefulSensors-moonshine-streaming/cpu/info.yaml +++ b/usefulSensors-moonshine-streaming/cpu/info.yaml @@ -1,8 +1,8 @@ name: usefulSensors-moonshine-streaming provider: usefulsensors -model_id: usefulsensors/moonshine-streaming-small -# The recipe also supports usefulsensors/moonshine-streaming-tiny via -# `--model-name usefulsensors/moonshine-streaming-tiny`. +model_id: usefulsensors/moonshine-streaming-tiny +# The recipe also supports usefulsensors/moonshine-streaming-small via +# `--model-name usefulsensors/moonshine-streaming-small`. task: automatic-speech-recognition framework: ONNX Runtime execution_provider: CPUExecutionProvider diff --git a/usefulSensors-moonshine-streaming/cpu/optimize.py b/usefulSensors-moonshine-streaming/cpu/optimize.py index f9b4073fa..12c67c3fa 100644 --- a/usefulSensors-moonshine-streaming/cpu/optimize.py +++ b/usefulSensors-moonshine-streaming/cpu/optimize.py @@ -34,7 +34,7 @@ if str(_RECIPE_ROOT) not in sys.path: sys.path.insert(0, str(_RECIPE_ROOT)) -MODEL_NAME = "usefulsensors/moonshine-streaming-small" +MODEL_NAME = "usefulsensors/moonshine-streaming-tiny" DEFAULT_OUTPUT_DIR = "build/onnx" # chunk_samples is the streaming window fed to the frontend each step. It must From 9cfdde57d745f25c8aeebab2e644a31c05eb3975 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Fri, 24 Jul 2026 22:37:50 +0000 Subject: [PATCH 7/8] Fix comm --- usefulSensors-moonshine-streaming/cpu/optimize.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/usefulSensors-moonshine-streaming/cpu/optimize.py b/usefulSensors-moonshine-streaming/cpu/optimize.py index 12c67c3fa..466515911 100644 --- a/usefulSensors-moonshine-streaming/cpu/optimize.py +++ b/usefulSensors-moonshine-streaming/cpu/optimize.py @@ -10,12 +10,12 @@ -> OnnxConversion (FP32, dynamo exporter, dynamic sequence axes) Usage: - # Full pipeline (small model -> build/onnx) + # Full pipeline (tiny model -> build/onnx) python cpu/optimize.py - # Tiny model - python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-tiny \ - --output-dir build/onnx-tiny + # Small model + python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-small \ + --output-dir build/onnx-small # Or run a single component directly through the Olive CLI: python -m olive run --config cpu/moonshine_frontend_fp32_cpu.json From 98e07d80221722eefc8b57d54067c805e1fe5ffb Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Fri, 24 Jul 2026 22:52:10 +0000 Subject: [PATCH 8/8] Fix Copilot comments --- .../cpu/export_moonshine_streaming.py | 6 +++--- .../cpu/moonshine_model_load.py | 14 +++++++++----- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py index aa623c23e..1ccdfa74c 100644 --- a/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py +++ b/usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py @@ -2,12 +2,12 @@ graphs for onnxruntime-genai (frontend / encoder / adapter / cross_kv / decoder_kv). -Run inside the ``moonshine`` conda env: +Run inside a Python env with the moonshine transformers integration installed: - source /home/nebanfic/miniconda3/bin/activate moonshine + conda activate moonshine # or: source .venv/bin/activate python export_moonshine_streaming.py \ --model usefulsensors/moonshine-streaming-tiny \ - --output-dir /datadisks/disk3/nebanfic/moonshine-streaming-tiny-mine + --output-dir build/moonshine-streaming-tiny The exporter uses the TorchDynamo ONNX path (``dynamo=True``) so data-dependent shapes in the stateful frontend export cleanly, then rewrites every graph's diff --git a/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py b/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py index 078147311..815c5f588 100644 --- a/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py +++ b/usefulSensors-moonshine-streaming/cpu/moonshine_model_load.py @@ -81,9 +81,13 @@ class FrontendModule(nn.Module): (even), the stride-2 phase stays aligned across chunks and the concatenated output is bit-for-bit identical to running the embedder on the full signal. - Sub-frame audio (``total_samples % frame_len``) is carried in - ``sample_buffer`` / ``sample_len`` and prepended to the next chunk so the - framing is contiguous. + Sub-frame audio (``total_samples % frame_len``) is emitted in + ``sample_buffer_out`` / ``sample_len_out`` so a final (flush) chunk of + non-multiple length records its remainder before reset. The input + ``sample_buffer`` / ``sample_len`` slots exist for graph shape parity but + are not consumed in the current genai contract: ``chunk_samples`` is a + multiple of ``frame_len``, so ``sample_len`` is always ``0`` on input and + the chunk is framed directly. See ``forward()``. """ def __init__(self, full: MoonshineStreamingForConditionalGeneration): @@ -102,8 +106,8 @@ def __init__(self, full: MoonshineStreamingForConditionalGeneration): def forward( self, audio_chunk, # [1, L] float32 - sample_buffer, # [1, 79] float32 - sample_len, # [1] int64 + sample_buffer, # [1, 79] float32 (unused; kept for graph parity) + sample_len, # [1] int64 (unused; assumed 0 by contract) conv1_buffer, # [1, C1, 4] float32 (last inputs of conv1) conv2_buffer, # [1, C2, 4] float32 (last inputs of conv2) frame_count, # [1] int64