rlx-moshi

Native Rust inference for Kyutai Moshi — a real-time speech-to-speech foundation model. Uses rlx-mimi for the Mimi codec and either eager CPU or Candle GPU (Metal/CUDA) for the 7B temporal + DepFormer stack.

Architecture

Variants

Moshiko and Moshika share the same 7B architecture. MoshiVariant picks the runtime codebook layout; MoshiCheckpoint picks the quant preset; together they route to the correct kyutai/moshiko-* or kyutai/moshika-* HuggingFace repo.

Quick start

# Mimi codec (~385 MB) + Moshi LM (~14 GB bf16)
just fetch-mimi
just fetch-moshi

# Q8 GGUF (~8 GB) — smaller on disk, CPU eager or GPU Candle
just fetch-moshi-q8
export RLX_MOSHI_CHECKPOINT=q8
export RLX_MOSHI_DIR=.cache/moshiko-q8

# One-way generation (CPU eager — slow; use Metal for production)
just moshi -- --prompt "Hello, I'm Moshi." --out-wav /tmp/moshi.wav --device cpu

# GPU path (~70 ms/frame on Apple Silicon with Q4 MLX or bf16)
just fetch-moshi-q4
just moshi -- --checkpoint q4 --device metal --prompt "Hello." --out-wav /tmp/moshi-metal.wav --max-steps 8

# Full-duplex: user WAV in → Moshi WAV out
just moshi -- --duplex --in-wav speech.wav --out-wav /tmp/moshi-reply.wav --max-steps 50 --device metal

# Moshika (female voice)
just fetch-moshika
just moshi -- --variant moshika-one-way --prompt "Hi." --out-wav /tmp/moshika.wav --device metal

# WebSocket duplex server (tokio + axum)
just moshi-ws -- --host 127.0.0.1 --port 8998 --device metal

Checkpoints

There is no smaller Moshi architecture — only quantization variants:

Pass --variant moshika or --variant moshika-one-way to use Moshika repos (kyutai/moshika-*) with the same presets.

Default cache directories

Unless RLX_MOSHI_DIR is set, the CLI picks a voice + preset scoped path:

Fetch recipes:

just fetch-moshi              # Moshiko bf16
just fetch-moshi-q8           # Moshiko Q8 GGUF
just fetch-moshi-q4           # Moshiko MLX Q4
just fetch-moshi-mlx-bf16     # Moshiko MLX bf16
just fetch-moshika            # Moshika bf16
just fetch-moshika-q4         # Moshika MLX Q4

(--device mlx uses Apple GPU via Candle Metal; native rlx-mlx quantized matmul is future work.)

Environment

Backends

cargo build -p rlx-moshi --features "gpu-lm,compiled-lm,mlx-lm,metal"
cargo build -p rlx-mimi --features "gpu-codec,metal"

Mimi codec GPU

Session passes --device to both LM and Mimi. GPU Mimi uses Kyutai Candle weights from tokenizer-e351c8d8-checkpoint125.safetensors (bundled in kyutai/moshiko-* / kyutai/moshika-* repos, or set RLX_MOSHI_DIR). HF kyutai/mimi layout stays on CPU eager.

cargo build -p rlx-mimi --features "gpu-codec,metal"
RLX_MIMI_GPU_SMOKE=1 RLX_MOSHI_DIR=.cache/moshiko cargo test -p rlx-mimi --test gpu_codec_smoke --features "gpu-codec,metal"

Streaming (channels + tokio)

Duplex streaming feeds 1920-sample (80 ms @ 24 kHz) PCM chunks and receives Moshi output incrementally:

use rlx_moshi::{
    GenerationConfig, MoshiSession, MoshiVariant, StreamCommand, StreamEvent,
    spawn_duplex_stream,
};
use rlx_runtime::Device;

let session = MoshiSession::open_on(".cache/moshiko", ".cache/mimi", MoshiVariant::Moshiko, Device::Metal)?;
let handle = spawn_duplex_stream(session, "", GenerationConfig::default())?;

// worker thread emits StreamEvent::OutputPcm { samples, .. }
handle.cmd_tx.send(StreamCommand::Pcm(chunk))?;
handle.cmd_tx.send(StreamCommand::Finish)?;

Library

use rlx_moshi::{
    MoshiCheckpoint, MoshiSession, MoshiVariant, MoshiVoice, GenerationConfig,
    default_moshi_dir_for,
};
use rlx_runtime::Device;

let variant = MoshiVariant::MoshikaOneWay;
let checkpoint = MoshiCheckpoint::Q4MlxSafetensors;
let moshi_dir = default_moshi_dir_for(variant, checkpoint); // .cache/moshika-mlx-q4

let mut session = MoshiSession::open_with_checkpoint(
    &moshi_dir,
    ".cache/mimi",
    variant,
    Device::Metal,
    checkpoint,
)?;
let result = session.generate_one_way("Hello!", &GenerationConfig::default())?;

Tests

export RLX_MOSHI_DIR=.cache/moshiko
export RLX_MIMI_DIR=.cache/mimi
cargo test -p rlx-moshi --release

# Checkpoint routing (no weights)
cargo test -p rlx-moshi --test checkpoint_routing --release

just test-moshi-weights

# Q8 GGUF key coverage (needs model.q8.gguf)
RLX_MOSHI_Q8_GGUF=.cache/moshiko-q8/model.q8.gguf cargo test -p rlx-moshi --test gguf_keys --release

# MLX key mapping (needs model.q4.safetensors)
RLX_MOSHI_MLX_Q4=.cache/moshiko-mlx-q4/model.q4.safetensors cargo test -p rlx-moshi --test mlx_keys --release

# GPU quick check
RLX_MOSHI_GPU_SMOKE=1 cargo test -p rlx-moshi --test gpu_smoke --features "gpu-lm,compiled-lm,metal" --release

# Streaming duplex → WAV → Whisper (slow; needs weights)
RLX_MOSHI_STREAM_E2E=1 just test-moshi-stream-whisper