Files
rustytorch/crates/models/rtx-csm/docs/moonshine_port_notes.md
T
osobhandClaude Opus 4.7 c848abef22 rtx-csm: Phase 8.4 spike — Moonshine v2 inspector + port notes
First step on the Tier 2.2 (Moonshine v2 candle port) item from the
Phase 8 plan. Full port is honestly multi-session work (~12-15 hours
of focused implementation across encoder, decoder, generation loop,
tokenizer, weight mapping, smoke test). This commit ships the
foundation so future sessions start from concrete data instead of
arxiv reading.

Two ships:

1. examples/moonshine_inspect — downloads UsefulSensors/moonshine-tiny
   from HF, parses safetensors header, dumps all 160 tensors grouped
   by prefix with shapes + dtypes. Verified output: 27.1 M params,
   108.4 MB safetensors (F32), encoder + decoder layers laid out as
   expected.

2. docs/moonshine_port_notes.md — captures every architectural fact
   established by the inspector + HF config.json:
   - Hyperparameter table (hidden=288, 6+6 layers, vocab=32768,
     partial_rotary=0.9, etc.)
   - Tensor layout per layer (encoder, decoder)
   - Architecture summary (raw waveform input, 3-layer Conv1d stem,
     SwiGLU decoder MLP via fused fc1, tied LM head)
   - Ordered porting tasks with hour estimates totaling ~12-15 h
   - Risks / unknowns (conv strides not in config, tied output head
     question, quality-vs-Kyutai concern)
   - Recommended order of attack for the next session

The full port itself is deferred. Ship the foundation now so the
remaining work has a clean handoff.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
2026-04-27 10:21:30 -07:00

8.6 KiB
Raw Blame History

Moonshine v2 candle port notes

Working notes for the planned (Tier 2.2) candle port of Moonshine v2. Captures architectural facts established by examples/moonshine_inspect so future sessions can start from concrete data, not arxiv reading.

Why this exists

  • Standalone Whisper-tiny runs at 0.020× realtime (50× faster than Kyutai STT 1B), but linking whisper-rs into the same binary as candle/CSM caused a 2-3× CSM regression (Phase 7.6).
  • Silero V5 via voice_activity_detector / ort blocked by protobuf 3.21 vs sentencepiece-sys 3.14 conflict (Phase 8.1.3).
  • A pure-candle Moonshine port avoids both classes of conflict and delivers comparable Whisper-class latency (50 ms TTFT per arXiv 2602.12241).

Source of truth for hyperparameters

UsefulSensors/moonshine-tiny/config.json:

{
  "architectures": ["MoonshineForConditionalGeneration"],
  "is_encoder_decoder": true,
  "hidden_size": 288,
  "intermediate_size": 1152,
  "vocab_size": 32768,
  "max_position_embeddings": 194,
  "encoder_num_hidden_layers": 6,
  "encoder_num_attention_heads": 8,
  "encoder_num_key_value_heads": 8,
  "encoder_hidden_act": "gelu",
  "decoder_num_hidden_layers": 6,
  "decoder_num_attention_heads": 8,
  "decoder_num_key_value_heads": 8,
  "decoder_hidden_act": "silu",
  "rope_theta": 10000.0,
  "partial_rotary_factor": 0.9,
  "pad_head_dim_to_multiple_of": 8,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "pad_token_id": 2,
  "decoder_start_token_id": 1
}

Tensor layout (from moonshine_inspect)

Total: 160 tensors, 27.1 M parameters, 108.4 MB safetensors (F32).

Encoder (68 tensors, 6 layers)

model.encoder.conv1.weight                                       [?]
model.encoder.conv1.bias
model.encoder.conv2.weight                                       [?]
model.encoder.conv2.bias                                         [576]
model.encoder.conv3.weight                                       [?]
model.encoder.conv3.bias

(per layer X = 0..6:)
model.encoder.layers.X.self_attn.q_proj.weight                   [288, 288]
model.encoder.layers.X.self_attn.k_proj.weight                   [288, 288]
model.encoder.layers.X.self_attn.v_proj.weight                   [288, 288]
model.encoder.layers.X.self_attn.o_proj.weight                   [288, 288]
model.encoder.layers.X.input_layernorm.weight                    [288]
model.encoder.layers.X.post_attention_layernorm.weight           [288]
model.encoder.layers.X.mlp.fc1.weight                            [1152, 288]
model.encoder.layers.X.mlp.fc1.bias                              [1152]
model.encoder.layers.X.mlp.fc2.weight                            [288, 1152]
model.encoder.layers.X.mlp.fc2.bias                              [288]

model.encoder.layer_norm.weight                                  [288]

Decoder (92 tensors, 6 layers, ~15.3 tensors/layer)

model.decoder.embed_tokens.weight                                [32768, 288]

(per layer X = 0..6:)
model.decoder.layers.X.self_attn.q_proj.weight                   [288, 288]
model.decoder.layers.X.self_attn.k_proj.weight                   [288, 288]
model.decoder.layers.X.self_attn.v_proj.weight                   [288, 288]
model.decoder.layers.X.self_attn.o_proj.weight                   [288, 288]
model.decoder.layers.X.encoder_attn.q_proj.weight                [288, 288]
model.decoder.layers.X.encoder_attn.k_proj.weight                [288, 288]
model.decoder.layers.X.encoder_attn.v_proj.weight                [288, 288]
model.decoder.layers.X.encoder_attn.o_proj.weight                [288, 288]
model.decoder.layers.X.input_layernorm.weight                    [288]
model.decoder.layers.X.post_attention_layernorm.weight           [288]
model.decoder.layers.X.final_layernorm.weight                    [288]
model.decoder.layers.X.mlp.fc1.weight                            [2304, 288]   ← SwiGLU gate+up fused
model.decoder.layers.X.mlp.fc2.weight                            [288, 1152]

model.decoder.layer_norm.weight                                  [288]

Architecture summary

  • Encoder: 3-layer Conv1d audio stem (raw 16 kHz waveform → 288-d hidden), 6 transformer layers (Pre-LN, GELU MLP, partial RoPE on 90 % of head_dim), final layer norm. No mel-spec or fbank — input is the raw PCM array.
  • Decoder: token embedding (32 768 vocab), 6 transformer layers each with self-attn + cross-attn-to-encoder + SwiGLU MLP (fc1 outputs 2 × intermediate=2304, split into gate+up; fc2 maps 1152 → 288), final layer norm. Tied output head likely shares embed_tokens.weight.
  • RoPE: partial_rotary_factor=0.9 means rotary embedding applied to first 90 % of head_dim (head_dim = 288/8 = 36; rotary on 32 dims, last 4 plain). pad_head_dim_to_multiple_of=8 rounds up to 40.
  • Generation: encoder-decoder seq2seq. Encode the audio once, then greedy/beam decode. Start token id = 1 (bos), stop on id = 2 (eos).

Porting tasks (~12-15 h, multi-session)

In order of dependency:

  1. Audio preprocessor (~1 h)Conv1d × 3 stem; input shape (B, 1, T_audio), output (B, T_seq, hidden=288) where T_seq = T_audio / stride_total. Stride values not in config; need to read from the conv weight shapes once we have them. Also need the receptive-field math for max_position_embeddings = 194.

  2. Encoder layer (~2 h) — vanilla post-LN Pre-LN transformer with GELU FFN (288 → 1152 → 288 with bias). Reuse candle_nn::Linear + standard attention. Implement RoPE with partial_rotary_factor=0.9 (use candle_transformers::models::rope helpers if compatible with partial application).

  3. Decoder layer (~3 h) — self-attn + cross-attn + SwiGLU MLP. fc1 weight shape [2304, 288] is gate+up fused: split into halves on dim 0 to get gate / up. SiLU activation, multiplied gate * up, then fc2 [288, 1152] projects back. Cross-attn keys/values come from encoder output (cached after the one-time encoder forward).

  4. Tokenizer (~1-2 h) — Moonshine ships a tokenizer.json on HF. Should load directly via the tokenizers crate (already a workspace dep). Vocab 32 768 BPE. Verify special tokens (<|startoftranscript|> etc.) match config.

  5. Generation loop (~2 h) — encoder forward once, decoder greedy loop with KV cache (use candle_nn::Cache or hand-rolled), stop on eos_token_id=2 or max_position_embeddings=194.

  6. Weight mapping + smoke test (~2-3 h)VarBuilder paths matching the model.encoder.layers.X.* / model.decoder.layers.X.* prefixes from inspector dump. First run: encode 10 s audio, decode, compare against the HF reference output (run the Python model on the same WAV, save the transcript).

  7. Profile binary (~30 min) — copy examples/whisper_profile.rs structure: load model, time N transcriptions of /tmp/asr_test.flac, report mean/p50 + realtime factor.

  8. Server integration (~1-2 h, OPTIONAL for first ship) — extend AsrEngine enum in examples/converse_server.rs with a third Moonshine variant. English-only.

Risks / unknowns

  • Conv stem strides not in config. Need to read from weight shapes or the HF source code (modeling_moonshine.py in transformers).
  • Tied output projection? Decoder probably shares embed_tokens with the LM head. Need to verify against transformers code.
  • Cross-attn caching shape. With encoder_num_key_value_heads=8 matching encoder_num_attention_heads, no GQA on encoder side, so cross-attn K/V shape is straightforward: (B, T_enc, 8, head_dim).
  • Quality vs Kyutai 1B. Moonshine-tiny is 27 M; Kyutai 1B is 1 B. Even if 50× faster, the WER may regress. Bench WER on the same audio + reference transcript before committing to swap.
  1. Run moonshine_inspect again, this time with shape printing for the 3 conv layers (currently truncated by the take(8) cap).
  2. Read the HF modeling_moonshine.py to confirm conv strides, output head sharing, and the partial-RoPE application detail.
  3. Start with encoder forward only — drop a minimal MoonshineEncoder that takes raw PCM and returns (B, T_seq, 288). Verify it loads weights cleanly. Don't worry about decoding yet.
  4. Sanity-check the encoder hidden-state norms match HF reference at layer 0 (catch any dtype / RoPE direction bugs early).
  5. Iterate decoder + generation only after encoder parity holds.

Cited sources

  • Paper: arXiv 2410.15608 (Moonshine v1) and 2602.12241 (Moonshine v2)
  • Weights: UsefulSensors/moonshine-tiny on HF
  • Reference impl: transformers >= 4.48.0, MoonshineForConditionalGeneration