Add Fun-ASR-Nano and SenseVoiceSmall support - #236
Conversation
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
There was a problem hiding this comment.
Pull request overview
This PR adds initial fun_asr support to the mobius speech stack by introducing a new SANM-based audio encoder, a 3-model speech-language task, registry wiring, an example inference script, and accompanying tests. It fits into the existing Components → Models → Tasks → Registry flow similarly to the existing Qwen3-ASR path, but for a Fun-ASR-specific frontend and encoder architecture.
Changes:
- Add Fun-ASR model/task implementation, including SANM attention components and registry/task exports.
- Add build/runtime tests, test configs, and golden-case metadata for the new speech model.
- Add an end-to-end example script and tighten SkipLayerNorm fusion safety checks.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
tests/model_coverage_test.py |
Updates coverage skip list for fun_asr. |
tests/build_graph_test.py |
Adds graph-build and ORT smoke tests for the 3-model Fun-ASR pipeline. |
tests/_test_configs.py |
Adds a tiny synthetic config entry for fun_asr. |
testdata/cases/speech/fun-asr-nano.yaml |
Adds golden-test metadata for Fun-ASR. |
src/mobius/tasks/_fun_asr_speech_language.py |
Introduces the new 3-model speech-language task wiring. |
src/mobius/tasks/__init__.py |
Exports and registers the new task name. |
src/mobius/rewrite_rules/_skip_layer_norm.py |
Adds rank guards to SkipLayerNorm fusion checks. |
src/mobius/models/fun_asr_test.py |
Adds colocated runtime and weight-routing tests for Fun-ASR. |
src/mobius/models/fun_asr.py |
Implements the Fun-ASR model, embedding fusion, decoder, and weight preprocessing. |
src/mobius/models/__init__.py |
Exports the new Fun-ASR model class. |
src/mobius/components/_sanm_attention.py |
Adds SANM/FSMN attention and encoder building blocks. |
src/mobius/components/__init__.py |
Re-exports SANM components. |
src/mobius/_registry.py |
Registers fun_asr, default task metadata, and test model mapping. |
src/mobius/_configs.py |
Extends AudioConfig with Fun-ASR-specific fields. |
examples/fun_asr.py |
Adds a full Fun-ASR preprocessing/inference CLI example. |
69930f2 to
7050b9e
Compare
- Add fun-asr-speech-language and audio-ctc to schema.json task_type enum - Fix empty except blocks: add explanatory comments in fun_asr.py and sensevoice_small.py so intent is clear to maintainers - Fix stale docstring in FunASRAudioEncoder: remove temporal pooling description (T→T//2) that no longer matches the implementation (sequence length is preserved) 221 YAML schema tests pass, 4 fun_asr build tests pass, lint clean. Signed-off-by: Justin Chu <justinchu@microsoft.com>
titaiwangms
left a comment
There was a problem hiding this comment.
Reviewed locally. Solid Fun-ASR scaffolding (SANM components, 3-model split task, golden tests, example script), but several issues and significant scope concerns.
🔴 Critical
-
Stale rebase regresses Gemma4 KV head logic. Main (#234) intentionally made
num_global_key_value_headsapply wheneveris_sliding == False(independent ofattention_k_eq_v) so Gemma4 GGUF — which sets per-layer KV heads but never setsattention_k_eq_v— works. This PR reverts that to gate on_use_alternative_attention, which will silently break Gemma4 GGUF inference (full layers will fall back tonum_key_value_headsand KV cache shapes will mismatch the weights). Same regression intasks/_gemma4.py::_make_gemma4_kv_cache_inputs. The gemma4 hunks shouldn't be part of this PR. -
PR is
mergeable_state: dirty— there are merge conflicts with main. Rebase is required, and that rebase is the same operation that should drop the gemma4 reverts above.
🟡 Major
-
Scope creep — SenseVoiceSmall is a separate model. The PR title/body advertise Fun-ASR only, but it also adds
models/sensevoice_small.py,tasks/_audio_ctc.py,examples/sensevoice_small.py, two golden files, and registry entries. Two unrelated models in one PR makes review and revert harder; recommend splitting SenseVoiceSmall into a follow-up PR. -
Multi-batch audio fusion is silently incorrect.
FunASREmbeddingModel.forwardbuilds gather indices viaCumSum(axis=1)over a per-batch mask, butaudio_featuresis a single flat(num_audio_tokens, hidden)table shared across the batch. Withbatch > 1, each row's cumsum restarts at 1, so all rows index the same prefix ofaudio_features— wrong audio is scattered into rows ≥ 1. Either assertbatch == 1, or document the constraint, or compute per-batch offsets. Today it works only because ASR inference is single-utterance; nothing in the graph enforces it.
🟢 Minor / Nits
-
_skip_layer_norm.pyrank guard skips unknown shapes.if inp.shape is not Nonemeans an Add input with no static shape passes the guard and can still produce a 1-D skip that ORT rejects at runtime. Either fail-closed when the shape is unknown, or document the assumption. -
SANMAttentionandAdaptorAttentionignore attention masks. No padding mask is plumbed through the encoder; for the example pipeline this is OK because audio is right-padded zeros and SANM is local enough, but if anyone later batches variable-length audio the encoder will attend across pad frames silently. -
Hardcoded fallback
audio_token_id = 151676inFunASREmbeddingModel.__init__— fine as a default for this checkpoint, but worth a# TODOor making the absence loud. -
_FSMNBlockuses(kernel - 1) // 2left-pad — asymmetric for even kernels. Default 11 is odd so it's symmetric; just a footgun if kernel becomes configurable.
✅ Good
- SANM components are clean, well-commented, model-agnostic, and follow the codebase's
op-as-first-arg +_test.pycolocation conventions. - Weight-name alignment is mostly structural (
nn.ModuleList, matching attribute names) with only a smallpreprocess_weightsrename map. - FSMN conv-weight transpose
(C, K, 1) → (C, 1, K)is detected by shape rather than name — robust. - Golden L4/L5 tests + 20 unit tests with deterministic seeds is good coverage.
- Architecture-diff bot reports zero changes to existing model graphs ✅.
Recommendation
Request changes. Rebase onto current main, drop the gemma4 + _gemma4 task hunks (they regress GGUF), decide whether SenseVoiceSmall should be a separate PR, and either fix or document the batch>1 audio-fusion limitation.
- Add fun-asr-speech-language and audio-ctc to schema.json task_type enum - Fix empty except blocks: add explanatory comments in fun_asr.py and sensevoice_small.py so intent is clear to maintainers - Fix stale docstring in FunASRAudioEncoder: remove temporal pooling description (T→T//2) that no longer matches the implementation (sequence length is preserved) 221 YAML schema tests pass, 4 fun_asr build tests pass, lint clean. Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Rebase onto latest main, dropping Gemma4 commits (already in PR #239) - Document batch=1 limitation in FunASREmbeddingModel - Document audio_token_id=0 collision workaround - Add explicit shape-unknown comment in SkipLayerNorm rank guard Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
f816967 to
11518e5
Compare
Add SANMAttention, SANMFFN, and SANMEncoderLayer components implementing the Self-Attention with Normalization and Memory (SANM) architecture used by Fun-ASR Nano speech encoder. Key features: - Fused QKV projection with split into Q, K, V - FSMN memory block: depthwise Conv1d on values with residual - Scaled dot-product attention via op.Attention - Simple ReLU FFN (w_1 → ReLU → w_2) - Pre-norm encoder layer with optional residual skip when in_size != out_size Attribute names match Fun-ASR checkpoint naming (self_attn.linear_q_k_v, self_attn.fsmn_block.weight, feed_forward.w_1/w_2, norm1, norm2) to minimize preprocess_weights renames. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Implement FunASRForConditionalGeneration for the Fun-ASR-Nano speech recognition model. The model uses a 3-model ONNX split: 1. Audio encoder (SenseVoiceEncoderSmall): 3 stacks of SANM layers with temporal pooling that halves the sequence length. 2. Embedding: Audio adaptor (MLP + transformer blocks projecting encoder dim to LLM dim) + text/audio embedding fusion. 3. Decoder: Qwen3-based text decoder with KV cache. New files: - src/mobius/models/fun_asr.py: Model classes (encoder, adaptor, embedding, decoder, top-level conditional generation) - src/mobius/tasks/_fun_asr_speech_language.py: Custom task for the 3-model split (audio encoder takes LFR fbank features with different input shape than mel spectrograms) Also fixes two pre-existing issues discovered during development: - SANM _FSMNBlock: Refactor from plain nn.Module() + dynamic attribute to proper subclass with forward() method so parameter names are qualified correctly in the ONNX graph. - SkipLayerNormalization rewrite rule: Add rank check on Add inputs to prevent incorrectly fusing bias-Add + LayerNorm (1D bias) into SkipLayerNormalization which ORT rejects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Create examples/fun_asr.py for FunAudioLLM/Fun-ASR-Nano-2512 speech recognition. Follows the same pattern as examples/qwen3_asr.py but with: - Fbank + LFR frontend instead of WhisperFeatureExtractor - Manual config construction from config.yaml + Qwen3-0.6B/config.json (Fun-ASR lacks standard HF model_type in config.json) - Standard RoPE position_ids (not MRoPE) - Supports zh/en/ja languages Builds 3 ONNX models (audio_encoder, embedding, decoder) via build_from_module with FunASRSpeechLanguageTask. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
initialize_rope() returns None when rope_type is None, causing TypeError in decoder forward(). Set rope_type='default' in both the example config builder and the test config. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
model.pt wraps weights in a 'state_dict' key. Unwrap before passing to preprocess_weights. Also use weights_only=False since the checkpoint contains non-tensor metadata. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Add 20 tests covering the 3-model pipeline (audio_encoder → embedding → decoder) with random weights through ORT: - TestFunASRPipelineShapes: Verify output shapes across all 3 models - TestFunASRTemporalPooling: Edge cases for T//2 pooling (T=2,4,50,100,200) - TestFunASRFullPipeline: End-to-end pipeline with various input configurations (short audio, no prefix/suffix, long prefix, KV cache decode step) - TestFunASRDeterminism: Same input produces same output - TestFunASRWeightNames: preprocess_weights routes all weight patterns correctly (encoder→audio_tower, adaptor→embedding, model.*→decoder) Also adds testdata/cases/speech/fun-asr-nano.yaml with skip_reason noting the model uses non-standard HF packaging (model.pt + config.yaml). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
audio_token_id was None in AudioConfig, causing op.Constant(value_int=None) to create an empty Constant node that ORT rejects. Set audio_token_id=151676 in the example config and add fallback in the model code. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Checkpoint uses 'llm.model.*' and 'llm.lm_head.*' prefixes, not 'model.*' and 'lm_head.*'. Update preprocess_weights to strip the 'llm.' prefix before routing to decoder/embedding models. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Restructure the Fun-ASR 3-model split so the audio adaptor lives inside FunASRAudioEncoder instead of FunASREmbeddingModel: - audio_encoder ONNX: fbank → SANM encoder → adaptor → LLM-dim features - embedding ONNX: input_ids + audio_features (LLM-dim) → inputs_embeds - decoder ONNX: unchanged Changes: - FunASRAudioEncoder: add self.adaptor, call it at the end of forward() - FunASRAudioAdaptor.forward: accept 3D (batch, seq, dim) input directly - FunASREmbeddingModel: remove audio_adaptor, scatter features directly - preprocess_weights: route audio_adaptor.* → audio_tower.adaptor.* - Task: use config.hidden_size for audio encoder output dim - Example: update comments for LLM-dim audio features Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Use Fun-ASR's actual inference format (from funasr source):
- System: 'You are a helpful assistant.'
- User: '语音转写成{language}:' + fake tokens (zeros) for audio
- Position-based overwrite of audio embeddings (not token-ID scatter)
- Language names in Chinese: 中文, 英文, 日文
Remove unused Qwen3-ASR token ID constants.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Keep only zh/en/ja — the three languages the Nano model supports well. Remove dialect entries (Hakka, Cantonese, Wu, etc.) that were added prematurely. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
1. Update stale docstrings: remove temporal pooling references (removed in 3ee45ab), document sequence-length-preserving behavior 2. Remove fun_asr from model_coverage_test.py skip list (has test_model_id) 3. Replace torch.load(weights_only=False) with safetensors loading from justinchuby/Fun-ASR-Nano-2512 4. Update registry test_model_id to justinchuby/Fun-ASR-Nano-2512 5. Reduce default chunk-length from 600s to 30s (safe for 6000-frame PE) 6. Fix token_id=0 decode collision: extract embed_tokens weight table and use numpy lookup during autoregressive decode steps, avoiding the ONNX embedding model's audio placeholder matching Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Cherry-picked from sensevoice-small-support branch. Implements SenseVoiceSmall with SANM encoder + CTC head + language query tokens. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Model is now available at justinchuby/Fun-ASR-Nano-2512 with safetensors format. Updated model_id and notes to reflect current state. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- examples/sensevoice_small.py: CTC-based ASR with LFR frontend, language control, and greedy CTC decode - Fix query token broadcast: use op.Expand instead of int64 Mul - Verified with justinchuby/SenseVoiceSmall-Hakka: correctly transcribes Cantonese audio to '这几个字都表达不到我想讲的意思' Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Generated from justinchuby/Fun-ASR-Nano-2512 with 652-129742-0006.flac test audio. English transcription matches expected: 'Cauliflower mayonnaise take cold boiled cauliflower break into branches...' Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Load and apply CMVN stats from am.mvn (shift + rescale on LFR features) - Add --chunk-length flag for long audio (default 30s) - With CMVN, output now matches FunASR reference (7 chars vs 7 chars) - Chunked mode works on 90s+ recordings Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Config loading: try config.json first, fall back to config.yaml - Token loading: try tokens.json, fall back to sentencepiece BPE model - Verified with mlx-community/SenseVoiceSmall: - Cantonese: '呢几个字都表达唔到我想讲嘅意思' (matches Fun-ASR) - English: 'cauliflower mayonnaise take cold boiled cauliflower...' - Hakka 30s: 46 chars of Cantonese transcription Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- YAML test case at testdata/cases/speech/sensevoice-small.yaml - test_model_id: mlx-community/SenseVoiceSmall (standard safetensors) - Golden data generated from mlx-community/SenseVoiceSmall (L4+L5) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Config loading: try config.json (mlx format) before config.yaml - Tokenizer loading: try Qwen3-0.6B subfolder, fall back to root - Transpose FSMN conv weights if shape is (C, K, 1) → (C, 1, K) - Verified with mlx-community/Fun-ASR-Nano-2512-fp16: - Cantonese: '呢几个字都表达唔到我想讲嘅意思。' (matches) - English: 'Cauliflower mayonnaise...' (matches golden) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Golden data now generated from FunASR PyTorch SDK (iic/SenseVoiceSmall) instead of our ONNX model. This ensures the golden reference represents ground truth for ONNX parity testing. Fun-ASR-Nano golden unchanged — Fun-ASR-Nano is not in the FunASR SDK registry (custom 3-model architecture). Its ONNX pipeline was validated against PyTorch encoder (cosine 0.998, ratio 1.01x). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Add left padding of (lfr_m-1)//2 = 3 frames before LFR stacking (FunASR convention). This centers the first output frame on the first input frame, matching the reference implementation. Impact: - SenseVoice-Hakka output now exactly matches FunASR reference - Fun-ASR English output gains proper punctuation - Golden data regenerated with corrected LFR Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
DRY refactor: - Extract shared audio preprocessing into examples/asr_utils.py: load_audio_file, compute_fbank, apply_lfr, load_cmvn, preprocess_audio - Both fun_asr.py and sensevoice_small.py now import from asr_utils Review comment fixes: - Fix stale docstring: seq_len//2 → seq_len in FunASRAudioEncoder - Fix AudioCTCTask: use config.dtype instead of hardcoded FLOAT - Add explanatory comments to all empty except clauses Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- #8: Add comment in registry noting Fun-ASR requires build_from_module() with manual config (config.yaml, not auto-detected by build()) - #10: Document audio_token_id=0 collision in embedding model docstring, advising callers to bypass for decode steps Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Add fun-asr-speech-language and audio-ctc to schema.json task_type enum - Fix empty except blocks: add explanatory comments in fun_asr.py and sensevoice_small.py so intent is clear to maintainers - Fix stale docstring in FunASRAudioEncoder: remove temporal pooling description (T→T//2) that no longer matches the implementation (sequence length is preserved) 221 YAML schema tests pass, 4 fun_asr build tests pass, lint clean. Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Rebase onto latest main, dropping Gemma4 commits (already in PR #239) - Document batch=1 limitation in FunASREmbeddingModel - Document audio_token_id=0 collision workaround - Add explicit shape-unknown comment in SkipLayerNorm rank guard Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
11518e5 to
2a68c48
Compare
titaiwangms
left a comment
There was a problem hiding this comment.
Update LGTM on the items I raised earlier — thanks for the quick turnaround. A few remaining concerns from this revision:
Confirmed addressed ✅
- KV-head regression in Gemma4 attention/task is gone — diff vs main no longer touches the
_use_alternative_attentiongating logic. 👍 - batch=1 limitation now in
FunASREmbeddingModelclass docstring. audio_token_id=0collision documented in the module docstring.- SkipLayerNorm rank-guard now has a comment explaining unknown-shape Adds intentionally pass through.
Still open
-
Gemma4 MoE refactor still slipping in. The new
_remap_moe_expert_weightshelper (gemma4.py +33 / -38) is unrelated to Fun-ASR — Fun-ASR uses a Qwen3 decoder. It also introduces a behavioral change forGemma4CausalLMModel.preprocess_weights: the old code'sif ".router.scale" in key:now becomes... and ".per_expert_scale" not in key, so any.router.per_expert_scaleweights are no longer multiplied byhidden_size^-0.5. That may well be the right fix (matches the pathGemma4Modelalready had after #239), but it should land in a separate PR with its own test rather than ride along here. If you do keep it, please add a one-line note in the PR body and a test that pins the new behavior forGemma4CausalLMModel. -
Title typo: "Add Fun-ASR-Nano and SneseAudio-ASR support" → presumably "SenseVoice-ASR" (or "SenseVoiceSmall").
-
batch=1 is documented but not enforced. The new docstring is good, but a runtime guard would prevent a future caller from silently producing wrong outputs:
# In FunASREmbeddingModel.forward, after computing batch dim batch = op.Shape(input_ids, start=0, end=1) # ... add a check or rely on an op that fails for batch > 1
At minimum, add an
assertin the example pipeline (examples/fun_asr.py) and a unit test that verifies the model raises/returns wrong shapes for batch>1 so the constraint is mechanically enforced somewhere.
Recommendation
Drop or split out the Gemma4 hunk and fix the title typo; then I'm comfortable approving. The batch>1 enforcement is nice-to-have, not blocking.
|
@titaiwangms Thanks - will handle batch>1 in a follow up. The gemma4 changes should have already been removed. |
titaiwangms
left a comment
There was a problem hiding this comment.
@justinchuby You're right — apologies, my local main was stale at 353dcdfb. After git fetch origin main (now at d39b13b8 from #240), git diff origin/main..pr-236 -- src/mobius/models/gemma4.py src/mobius/tasks/_gemma4.py is empty. No Gemma4 changes in this PR. False alarm on my part.
Remaining items from my last review reduce to:
- Title typo ("SneseAudio" → "SenseVoice").
- batch>1 enforcement — acknowledged as follow-up. 👍
Otherwise LGTM.
This pull request introduces support for the SenseVoiceSmall speech recognition model, including a new ASR example script and shared audio preprocessing utilities. It also expands the model registry and configuration to accommodate Fun-ASR and SenseVoiceSmall models, and exposes new SANM encoder components. The most important changes are grouped below:
New ASR Example and Utilities:
examples/sensevoice_small.py, a complete CLI script for running ONNX-based speech recognition with SenseVoiceSmall, including CTC decoding, language control, and chunked inference.examples/asr_utils.pywith reusable audio preprocessing functions: audio loading, log-mel fbank computation, LFR stacking, CMVN loading, and full frontend pipeline.Model Registry and Configuration Updates:
FunASRForConditionalGenerationandSenseVoiceSmallModelin the model registry, with proper task types and default HuggingFace model IDs. [1] [2] [3] [4]AudioConfigwith additional fields for Fun-ASR/SenseVoice encoder configuration (e.g.,tp_num_blocks,adaptor_proj_dim).Component Exposure:
SANMAttention,SANMEncoderLayer,SANMFFN) inmobius.components.__init__and imported them from the internal module. [1] [2]