GitHub - FarhanAliRaza/moonshine-mojo

Moonshine ASR - Mojo Implementation

A high-performance implementation of Moonshine speech recognition in Mojo, achieving 5.3x speedup over the Python baseline.

Performance

Benchmarked on RTX 3060 with 10s audio. RTF = Real-Time Factor (lower is better).

Features

• Pure Mojo implementation with GPU acceleration
• KV-cache optimization for decoder (3.5x decoder speedup)
• Server mode for persistent model loading
• Voice typing TUI tool included

Quick Start

Prerequisites

• Pixi package manager
• NVIDIA GPU with CUDA support
• ~500MB disk space for model weights

Installation

# Clone the repository
git clone https://github.com/FarhanAliRaza/moonshine-mojo
cd moonshine-mojo

# Download and convert weights
python scripts/convert_weights.py

# Build the Mojo binary
pixi run build

Usage

Single file transcription:

Server mode (keeps weights loaded):

Voice typing TUI:

uv run python tools/voice_type.py

Architecture

Audio (16kHz) → Conv Stem (384x downsample) → Encoder (8 layers) → Decoder (8 layers) → Tokens → Text

Model: Moonshine Base (62M parameters)

Project Structure

moonshine-mojo/
├── main.mojo           # Entry point, CLI handling
├── encoder.mojo        # Conv stem + transformer encoder
├── decoder.mojo        # Autoregressive decoder with KV-cache
├── ops.mojo            # GPU kernels (attention, RoPE, norms, etc.)
├── config.mojo         # Model configuration
├── weight_loader.mojo  # Binary weight loading
├── tokenizer.mojo      # Token decoding
├── tools/
│   └── voice_type.py   # Voice typing TUI
├── scripts/
│   ├── convert_weights.py      # Convert safetensors to binary
│   ├── benchmark_comparison.py # Run Python vs Mojo benchmark
│   └── baseline_benchmark.py   # Python baseline benchmark
├── weights/            # Model weights (not in git)
├── test_audio/         # Test audio files
└── assets/             # Images for documentation

Benchmarking

Run the full benchmark comparison:

uv run python scripts/benchmark_comparison.py

This generates benchmark_results.html with detailed results.

Key Optimizations

1. KV-Cache (3.5x decoder speedup)

• Self-attention cache: Stores post-RoPE K,V values, grows with each token
• Cross-attention cache: Encoder projections computed once, reused for all tokens
• Reduces complexity from O(n²) to O(n) for K,V computations

2. GPU Kernels

All operations run on GPU using MAX's optimized kernels:

• Flash attention for self and cross attention
• Fused RoPE (partial rotary embeddings)
• Custom kernels for conv1d, RMSNorm, SiLU-GLU

3. Binary Weight Format

Custom format for fast loading:

• Direct memory mapping
• Pre-transposed linear weights
• No framework overhead

Voice Typing Tool

A terminal UI for real-time voice transcription:

uv run python tools/voice_type.py

Controls:

• Space - Start/stop recording
• c - Copy transcript to clipboard
• x - Clear transcript
• q - Quit

Development

# Run with debug output
pixi run mojo main.mojo

# Run benchmarks
pixi run mojo benchmarks.mojo

# Run Python baseline
uv run python scripts/baseline_benchmark.py --audio-files test_audio/beckett.wav

References

• Moonshine Paper - Original research
• Moonshine GitHub - Reference implementation
• Mojo - Programming language
• MAX - ML framework

License

MIT License - See LICENSE file for details.