Add Qwen3-TTS streaming endpoint for voice synthesis (#40)

* Task completion checkpoint

* Task completion checkpoint

* Task completion checkpoint

* Add Qwen3-TTS research document for live TTS replacement

Research findings for replacing Chatterbox TTS with Qwen3-TTS-12Hz-0.6B-Base:

- Current TTS: Chatterbox-Turbo-ONNX with batch-only generation, no streaming
- Qwen3-TTS: 97ms end-to-end latency, streaming support, 3-second voice cloning
- Voice cloning: Requires 3s reference audio + transcript (Makima voice planned)
- Integration: Python service with WebSocket bridge (no ONNX export available)
- Languages: 10 supported including English and Japanese

Document includes:
- Current architecture analysis (makima/src/tts.rs)
- Qwen3-TTS capabilities and requirements
- Feasibility assessment for live/streaming TTS
- Audio clip requirements for voice cloning
- Preliminary technical approach with architecture diagrams

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* [WIP] Heartbeat checkpoint - 2026-01-27 03:11:15 UTC

* Add Qwen3-TTS research documentation

Comprehensive research on replacing Chatterbox TTS with Qwen3-TTS-12Hz-0.6B-Base:

- Current TTS implementation analysis (Chatterbox-Turbo-ONNX in makima/src/tts.rs)
- Qwen3-TTS capabilities: 97ms streaming latency, voice cloning with 3s reference
- Cross-lingual support: Japanese voice (Makima/Tomori Kusunoki) speaking English
- Python microservice architecture recommendation (FastAPI + WebSocket)
- Implementation phases and technical approach
- Hardware requirements and dependencies

Key findings:
- Live/streaming TTS is highly feasible with 97ms latency
- Voice cloning fully supported with 0.95 speaker similarity
- Recommended: Python microservice with WebSocket streaming

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add comprehensive Qwen3-TTS integration specification

This specification document defines the complete integration of
Qwen3-TTS-12Hz-0.6B-Base as a replacement for the existing Chatterbox-Turbo
TTS implementation. The document covers:

## Functional Requirements
- WebSocket endpoint /api/v1/speak for streaming TTS
- Voice cloning with default Makima voice (Japanese VA speaking English)
- Support for custom voice references
- Detailed client-to-server and server-to-client message protocols
- Integration with Listen page for bidirectional speech

## Non-Functional Requirements
- Latency targets: < 200ms first audio byte
- Audio quality: 24kHz, mono, PCM16/PCM32f
- Hardware requirements: CUDA GPU with 4-8GB VRAM
- Scalability: 10 concurrent sessions per GPU

## Architecture Specification
- Python TTS microservice with FastAPI/WebSocket
- Rust proxy endpoint in makima server
- Voice prompt caching mechanism (LRU cache)
- Error handling and recovery strategies

## API Contract
- Complete WebSocket message format definitions (TypeScript)
- Error codes and responses (TTS_UNAVAILABLE, SYNTHESIS_ERROR, etc.)
- Session state machine and lifecycle management

## Voice Asset Requirements
- Makima voice clip specifications (5-10s WAV, transcript required)
- Storage location: models/voices/makima/
- Metadata format for voice management

## Testing Strategy
- Unit tests for Python TTS service and Rust proxy
- Integration tests for WebSocket flow
- Latency benchmarks with performance targets
- Test data fixtures for various text lengths

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add Qwen3-TTS implementation plan

Comprehensive implementation plan for replacing Chatterbox-TTS with
Qwen3-TTS streaming TTS service, including:

- Task breakdown with estimated hours for each phase
- Phase 1: Python TTS microservice (FastAPI, WebSocket)
- Phase 2: Rust proxy integration (speak.rs, tts_client.rs)
- Detailed file changes and new module structure
- Testing plan with unit, integration, and latency benchmarks
- Risk assessment with mitigation strategies
- Success criteria for each phase

Based on specification in docs/specs/qwen3-tts-spec.md

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add author and research references to TTS implementation plan

Add links to research documentation and author attribution.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* [WIP] Heartbeat checkpoint - 2026-01-27 03:25:06 UTC

* Add Python TTS service project structure (Phase 1.1-1.3)

Create the initial makima-tts Python service directory structure with:
- pyproject.toml with FastAPI, Qwen-TTS, and torch dependencies
- config.py with pydantic-settings TTSConfig class
- models.py with Pydantic message models (Start, Speak, Stop, Ready, etc.)

This implements tasks P1.1, P1.2, and P1.3 from the Qwen3-TTS implementation plan.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add TTS engine and voice manager for Qwen3-TTS (Phase 1.4-1.5)

Implement core TTS functionality:
- tts_engine.py: Qwen3-TTS wrapper with streaming audio chunk generation
- voice_manager.py: Voice prompt caching with LRU eviction and TTL support

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* [WIP] Heartbeat checkpoint - 2026-01-27 03:30:06 UTC

* Add TTS proxy client and message types (Phase 2.1, 2.2, 2.4)

- Add tts_client.rs with TtsConfig, TtsCircuitBreaker, TtsError,
  TtsProxyClient, and TtsConnection structs for WebSocket proxying
- Add TTS message types to messages.rs (TtsAudioEncoding, TtsPriority,
  TtsStartMessage, TtsSpeakMessage, TtsStopMessage, TtsClientMessage,
  TtsReadyMessage, TtsAudioChunkMessage, TtsCompleteMessage,
  TtsErrorMessage, TtsStoppedMessage, TtsServerMessage)
- Export tts_client module from server mod.rs
- tokio-tungstenite already present in Cargo.toml

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add TTS WebSocket handler and route (Phase 2.3, 2.5, 2.6)

- Create speak.rs WebSocket handler that proxies to Python TTS service
- Add TtsState fields (tts_client, tts_config) to AppState
- Add with_tts() builder and is_tts_healthy() methods to AppState
- Register /api/v1/speak route in the router
- Add speak module export in handlers/mod.rs

The handler forwards WebSocket messages bidirectionally between
the client and the Python TTS microservice with proper error handling.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add Makima voice profile assets for TTS voice cloning

Creates the voice assets directory structure with:
- manifest.json containing voice configuration (voice_id, speaker,
  language, reference audio path, and Japanese transcript placeholder)
- README.md with instructions for obtaining voice reference audio

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add Rust-native Qwen3-TTS integration research document

Research findings for integrating Qwen3-TTS-12Hz-0.6B-Base directly into
the makima Rust codebase without Python. Key conclusions:
- ONNX export is not viable (unsupported architecture)
- Candle (HF Rust ML framework) is the recommended approach
- Model weights available in safetensors format (2.52GB total)
- Three components needed: LM backbone, code predictor, speech tokenizer
- Crane project has Qwen3-TTS as highest priority (potential upstream)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* [WIP] Heartbeat checkpoint - 2026-01-27 11:21:43 UTC

* [WIP] Heartbeat checkpoint - 2026-01-27 11:24:19 UTC

* [WIP] Heartbeat checkpoint - 2026-01-27 11:26:43 UTC

* feat: implement Rust-native Qwen3-TTS using candle framework

Replace monolithic tts.rs with modular tts/ directory structure:

- tts/mod.rs: TtsEngine trait, TtsEngineFactory, shared types (AudioChunk,
  TtsError), and utility functions (save_wav, resample, argmax)
- tts/chatterbox.rs: existing ONNX-based ChatterboxTTS adapted to implement
  TtsEngine trait with Mutex-wrapped sessions for Send+Sync
- tts/qwen3/mod.rs: Qwen3Tts entry point with HuggingFace model loading
- tts/qwen3/config.rs: Qwen3TtsConfig parsing from HF config.json
- tts/qwen3/model.rs: 28-layer Qwen3 transformer with RoPE, GQA (16 heads,
  8 KV heads), SiLU MLP, RMS norm, and KV cache
- tts/qwen3/code_predictor.rs: 5-layer MTP module predicting 16 codebooks
- tts/qwen3/speech_tokenizer.rs: ConvNet encoder/decoder with 16-layer RVQ
- tts/qwen3/generate.rs: autoregressive generation loop with streaming support

Add candle-core, candle-nn, candle-transformers, safetensors to Cargo.toml.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* feat: integrate TTS engine into speak WebSocket handler

- Update speak.rs handler to use TTS engine directly from SharedState
  instead of returning a stub "not implemented" error
- Add TtsEngine (OnceCell lazy-loaded) to AppState in state.rs with
  get_tts_engine() method for lazy initialization on first connection
- Implement full WebSocket protocol: client sends JSON speak/cancel/stop
  messages, server streams binary PCM audio chunks and audio_end signals
- Create voices/makima/manifest.json for Makima voice profile configuration
- All files compile successfully with zero errors

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* feat: add /speak TTS page with WebSocket audio playback

Add a new /speak frontend page for text-to-speech via WebSocket.
The page accepts text input and streams synthesized PCM audio through
the Web Audio API. Includes model loading indicator, cancel support,
and connection status. Also adds a loading bar to the listen page
ControlPanel during WebSocket connection.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

1 files changed, 581 insertions, 0 deletions

diff --git a/makima/src/tts/qwen3/model.rs b/makima/src/tts/qwen3/model.rs
new file mode 100644
index 0000000..551893b
--- /dev/null
+++ b/makima/src/tts/qwen3/model.rs
@@ -0,0 +1,581 @@
+//! Qwen3 Language Model transformer backbone.
+//!
+//! Implements the 28-layer transformer with:
+//! - Rotary Position Embeddings (RoPE)
+//! - Grouped Query Attention (GQA) — 16 heads, 8 KV heads
+//! - SiLU-gated MLP
+//! - RMS normalization
+//! - KV cache for autoregressive generation
+//!
+//! Based on the candle-transformers Qwen2 model architecture,
+//! extended for Qwen3-TTS.
+
+use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D};
+use candle_nn::{embedding, linear_no_bias, rms_norm, Embedding, Linear, RmsNorm, VarBuilder};
+
+use super::config::Qwen3LmConfig;
+
+// ---------------------------------------------------------------------------
+// Rotary Position Embeddings
+// ---------------------------------------------------------------------------
+
+/// Precomputed RoPE sin/cos tables.
+#[derive(Debug, Clone)]
+pub struct RotaryEmbedding {
+    cos: Tensor,
+    sin: Tensor,
+}
+
+impl RotaryEmbedding {
+    pub fn new(config: &Qwen3LmConfig, dtype: DType, device: &Device) -> Result<Self> {
+        let head_dim = config.head_dim;
+        let max_seq = config.max_position_embeddings;
+        let theta = config.rope_theta;
+
+        let inv_freq: Vec<f32> = (0..head_dim)
+            .step_by(2)
+            .map(|i| 1.0 / (theta as f32).powf(i as f32 / head_dim as f32))
+            .collect();
+
+        let inv_freq_tensor =
+            Tensor::from_vec(inv_freq, (head_dim / 2,), device)?.to_dtype(DType::F32)?;
+
+        let positions: Vec<f32> = (0..max_seq).map(|p| p as f32).collect();
+        let positions_tensor = Tensor::from_vec(positions, (max_seq, 1), device)?;
+
+        // [max_seq, head_dim/2]
+        let freqs = positions_tensor.matmul(&inv_freq_tensor.unsqueeze(0)?)?;
+        // [max_seq, head_dim] by repeating
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
+
+        let cos = emb.cos()?.to_dtype(dtype)?;
+        let sin = emb.sin()?.to_dtype(dtype)?;
+
+        Ok(Self { cos, sin })
+    }
+
+    /// Apply RoPE to query and key tensors.
+    /// Input shape: [batch, heads, seq_len, head_dim]
+    pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
+        let seq_len = q.dim(2)?;
+        let cos = self.cos.narrow(0, offset, seq_len)?;
+        let sin = self.sin.narrow(0, offset, seq_len)?;
+
+        let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // [1, 1, seq, dim]
+        let sin = sin.unsqueeze(0)?.unsqueeze(0)?;
+
+        let q_rotated = Self::rotate_half(q, &cos, &sin)?;
+        let k_rotated = Self::rotate_half(k, &cos, &sin)?;
+
+        Ok((q_rotated, k_rotated))
+    }
+
+    fn rotate_half(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
+        let half_dim = x.dim(D::Minus1)? / 2;
+        let x1 = x.narrow(D::Minus1, 0, half_dim)?;
+        let x2 = x.narrow(D::Minus1, half_dim, half_dim)?;
+
+        // [-x2, x1] concatenated
+        let neg_x2 = x2.neg()?;
+        let rotated = Tensor::cat(&[&neg_x2, &x1], D::Minus1)?;
+
+        // x * cos + rotated * sin
+        let result = x.broadcast_mul(cos)?.broadcast_add(&rotated.broadcast_mul(sin)?)?;
+        Ok(result)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// KV Cache
+// ---------------------------------------------------------------------------
+
+/// Per-layer key-value cache for autoregressive generation.
+#[derive(Debug, Clone)]
+pub struct KvCache {
+    key: Option<Tensor>,
+    value: Option<Tensor>,
+}
+
+impl KvCache {
+    pub fn new() -> Self {
+        Self {
+            key: None,
+            value: None,
+        }
+    }
+
+    /// Append new key/value tensors and return the full cached sequence.
+    /// Input shapes: [batch, num_kv_heads, new_seq_len, head_dim]
+    pub fn append(&mut self, key: &Tensor, value: &Tensor) -> Result<(Tensor, Tensor)> {
+        let (full_key, full_value) = match (&self.key, &self.value) {
+            (Some(prev_k), Some(prev_v)) => {
+                let k = Tensor::cat(&[prev_k, key], 2)?;
+                let v = Tensor::cat(&[prev_v, value], 2)?;
+                (k, v)
+            }
+            _ => (key.clone(), value.clone()),
+        };
+
+        self.key = Some(full_key.clone());
+        self.value = Some(full_value.clone());
+
+        Ok((full_key, full_value))
+    }
+
+    /// Current cached sequence length.
+    pub fn seq_len(&self) -> usize {
+        self.key
+            .as_ref()
+            .map(|k| k.dim(2).unwrap_or(0))
+            .unwrap_or(0)
+    }
+
+    /// Reset the cache.
+    pub fn reset(&mut self) {
+        self.key = None;
+        self.value = None;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Attention
+// ---------------------------------------------------------------------------
+
+/// Multi-head attention with GQA and RoPE.
+pub struct Qwen3Attention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_heads: usize,
+    num_kv_heads: usize,
+    head_dim: usize,
+    num_kv_groups: usize,
+}
+
+impl Qwen3Attention {
+    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
+        let hidden = config.hidden_size;
+        let num_heads = config.num_attention_heads;
+        let num_kv_heads = config.num_key_value_heads;
+        let head_dim = config.head_dim;
+
+        let q_proj = linear_no_bias(hidden, num_heads * head_dim, vb.pp("q_proj"))?;
+        let k_proj = linear_no_bias(hidden, num_kv_heads * head_dim, vb.pp("k_proj"))?;
+        let v_proj = linear_no_bias(hidden, num_kv_heads * head_dim, vb.pp("v_proj"))?;
+        let o_proj = linear_no_bias(num_heads * head_dim, hidden, vb.pp("o_proj"))?;
+
+        let q_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?;
+        let k_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?;
+
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_heads,
+            num_kv_heads,
+            head_dim,
+            num_kv_groups: config.num_kv_groups(),
+        })
+    }
+
+    /// Forward pass with KV cache and RoPE.
+    /// Input: [batch, seq_len, hidden_size]
+    /// Returns: [batch, seq_len, hidden_size]
+    pub fn forward(
+        &self,
+        hidden_states: &Tensor,
+        rope: &RotaryEmbedding,
+        kv_cache: &mut KvCache,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (batch, seq_len, _) = hidden_states.dims3()?;
+        let offset = kv_cache.seq_len();
+
+        // Project Q, K, V
+        let q = self.q_proj.forward(hidden_states)?;
+        let k = self.k_proj.forward(hidden_states)?;
+        let v = self.v_proj.forward(hidden_states)?;
+
+        // Reshape: [batch, seq, heads*dim] -> [batch, heads, seq, dim]
+        let q = q
+            .reshape((batch, seq_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((batch, seq_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((batch, seq_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // Apply QK normalization (Qwen3 specific)
+        let q = self.apply_head_norm(&q, &self.q_norm)?;
+        let k = self.apply_head_norm(&k, &self.k_norm)?;
+
+        // Apply RoPE
+        let (q, k) = rope.apply(&q, &k, offset)?;
+
+        // Update KV cache
+        let (k, v) = kv_cache.append(&k, &v)?;
+
+        // Expand KV heads for GQA: [batch, kv_heads, seq, dim] -> [batch, heads, seq, dim]
+        let k = self.repeat_kv(&k)?;
+        let v = self.repeat_kv(&v)?;
+
+        // Scaled dot-product attention
+        let scale = (self.head_dim as f64).sqrt();
+        let attn_weights = (q.matmul(&k.transpose(D::Minus2, D::Minus1)?)? / scale)?;
+
+        let attn_weights = match attention_mask {
+            Some(mask) => attn_weights.broadcast_add(mask)?,
+            None => attn_weights,
+        };
+
+        let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
+
+        // Attention output
+        let attn_output = attn_weights.matmul(&v)?;
+
+        // [batch, heads, seq, dim] -> [batch, seq, heads*dim]
+        let attn_output = attn_output
+            .transpose(1, 2)?
+            .reshape((batch, seq_len, self.num_heads * self.head_dim))?;
+
+        self.o_proj.forward(&attn_output)
+    }
+
+    /// Apply RMS norm per-head.
+    fn apply_head_norm(&self, x: &Tensor, norm: &RmsNorm) -> Result<Tensor> {
+        let (b, h, s, d) = x.dims4()?;
+        // Reshape to [b*h*s, d] for norm, then back
+        let flat = x.reshape((b * h * s, d))?;
+        let normed = norm.forward(&flat)?;
+        normed.reshape((b, h, s, d))
+    }
+
+    /// Repeat KV heads for GQA.
+    fn repeat_kv(&self, x: &Tensor) -> Result<Tensor> {
+        if self.num_kv_groups == 1 {
+            return Ok(x.clone());
+        }
+        let (batch, num_kv_heads, seq_len, head_dim) = x.dims4()?;
+        let x = x
+            .unsqueeze(2)?
+            .expand((batch, num_kv_heads, self.num_kv_groups, seq_len, head_dim))?
+            .reshape((batch, self.num_heads, seq_len, head_dim))?;
+        Ok(x)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// MLP
+// ---------------------------------------------------------------------------
+
+/// SiLU-gated feed-forward network.
+pub struct Qwen3Mlp {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+}
+
+impl Qwen3Mlp {
+    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
+        let hidden = config.hidden_size;
+        let intermediate = config.intermediate_size;
+
+        let gate_proj = linear_no_bias(hidden, intermediate, vb.pp("gate_proj"))?;
+        let up_proj = linear_no_bias(hidden, intermediate, vb.pp("up_proj"))?;
+        let down_proj = linear_no_bias(intermediate, hidden, vb.pp("down_proj"))?;
+
+        Ok(Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+        })
+    }
+
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let gate = self.gate_proj.forward(x)?;
+        let gate = candle_nn::Activation::Silu.forward(&gate)?;
+        let up = self.up_proj.forward(x)?;
+        let hidden = (gate * up)?;
+        self.down_proj.forward(&hidden)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Transformer Layer
+// ---------------------------------------------------------------------------
+
+/// A single Qwen3 transformer decoder layer.
+pub struct Qwen3DecoderLayer {
+    self_attn: Qwen3Attention,
+    mlp: Qwen3Mlp,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+}
+
+impl Qwen3DecoderLayer {
+    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
+        let self_attn = Qwen3Attention::new(config, vb.pp("self_attn"))?;
+        let mlp = Qwen3Mlp::new(config, vb.pp("mlp"))?;
+        let input_layernorm =
+            rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let post_attention_layernorm = rms_norm(
+            config.hidden_size,
+            config.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        hidden_states: &Tensor,
+        rope: &RotaryEmbedding,
+        kv_cache: &mut KvCache,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        // Pre-norm attention
+        let residual = hidden_states;
+        let hidden_states = self.input_layernorm.forward(hidden_states)?;
+        let hidden_states =
+            self.self_attn
+                .forward(&hidden_states, rope, kv_cache, attention_mask)?;
+        let hidden_states = (residual + hidden_states)?;
+
+        // Pre-norm MLP
+        let residual = &hidden_states;
+        let hidden_states = self.post_attention_layernorm.forward(&hidden_states)?;
+        let hidden_states = self.mlp.forward(&hidden_states)?;
+        let output = (residual + hidden_states)?;
+
+        Ok(output)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Full Model
+// ---------------------------------------------------------------------------
+
+/// The complete Qwen3 language model for TTS.
+///
+/// Architecture:
+/// - Token embedding layer
+/// - 28 transformer decoder layers
+/// - Final RMS normalization
+/// - LM head (projects to vocab)
+pub struct Qwen3Model {
+    embed_tokens: Embedding,
+    layers: Vec<Qwen3DecoderLayer>,
+    norm: RmsNorm,
+    lm_head: Linear,
+    rope: RotaryEmbedding,
+    config: Qwen3LmConfig,
+    /// Last hidden states (before lm_head), used by code predictor.
+    last_hidden: std::cell::RefCell<Option<Tensor>>,
+}
+
+impl Qwen3Model {
+    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
+        let model_vb = vb.pp("model");
+
+        let embed_tokens = embedding(config.vocab_size, config.hidden_size, model_vb.pp("embed_tokens"))?;
+
+        let mut layers = Vec::with_capacity(config.num_hidden_layers);
+        for i in 0..config.num_hidden_layers {
+            let layer = Qwen3DecoderLayer::new(config, model_vb.pp(format!("layers.{i}")))?;
+            layers.push(layer);
+        }
+
+        let norm = rms_norm(config.hidden_size, config.rms_norm_eps, model_vb.pp("norm"))?;
+
+        // LM head — may or may not share weights with embed_tokens
+        let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, vb.pp("lm_head"))?;
+
+        let dtype = vb.dtype();
+        let device = vb.device().clone();
+        let rope = RotaryEmbedding::new(config, dtype, &device)?;
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            lm_head,
+            rope,
+            config: config.clone(),
+            last_hidden: std::cell::RefCell::new(None),
+        })
+    }
+
+    /// Forward pass through the full model.
+    ///
+    /// `input_ids`: [batch, seq_len] — token IDs
+    /// `kv_caches`: per-layer KV caches
+    /// `attention_mask`: optional causal mask [batch, 1, seq_len, total_seq_len]
+    ///
+    /// Returns logits: [batch, seq_len, vocab_size]
+    pub fn forward(
+        &self,
+        input_ids: &Tensor,
+        kv_caches: &mut [KvCache],
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let mut hidden_states = self.embed_tokens.forward(input_ids)?;
+
+        for (i, layer) in self.layers.iter().enumerate() {
+            hidden_states =
+                layer.forward(&hidden_states, &self.rope, &mut kv_caches[i], attention_mask)?;
+        }
+
+        hidden_states = self.norm.forward(&hidden_states)?;
+
+        // Store last hidden state for code predictor
+        *self.last_hidden.borrow_mut() = Some(hidden_states.clone());
+
+        let logits = self.lm_head.forward(&hidden_states)?;
+        Ok(logits)
+    }
+
+    /// Forward pass with pre-computed embeddings (for first iteration where
+    /// text embeddings are concatenated with audio features).
+    ///
+    /// `inputs_embeds`: [batch, seq_len, hidden_size]
+    pub fn forward_embeds(
+        &self,
+        inputs_embeds: &Tensor,
+        kv_caches: &mut [KvCache],
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let mut hidden_states = inputs_embeds.clone();
+
+        for (i, layer) in self.layers.iter().enumerate() {
+            hidden_states =
+                layer.forward(&hidden_states, &self.rope, &mut kv_caches[i], attention_mask)?;
+        }
+
+        hidden_states = self.norm.forward(&hidden_states)?;
+
+        *self.last_hidden.borrow_mut() = Some(hidden_states.clone());
+
+        let logits = self.lm_head.forward(&hidden_states)?;
+        Ok(logits)
+    }
+
+    /// Get the last hidden states (for the code predictor).
+    pub fn last_hidden_state(&self) -> Option<Tensor> {
+        self.last_hidden.borrow().clone()
+    }
+
+    /// Number of transformer layers.
+    pub fn num_layers(&self) -> usize {
+        self.config.num_hidden_layers
+    }
+
+    /// Hidden size.
+    pub fn hidden_size(&self) -> usize {
+        self.config.hidden_size
+    }
+
+    /// Get token embedding layer (for input preparation).
+    pub fn embed_tokens(&self) -> &Embedding {
+        &self.embed_tokens
+    }
+
+    /// Create a causal attention mask.
+    pub fn make_causal_mask(
+        seq_len: usize,
+        past_len: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Tensor> {
+        let total_len = past_len + seq_len;
+
+        if seq_len == 1 {
+            // Single token: no masking needed (can attend to everything)
+            return Tensor::zeros((1, 1, 1, total_len), dtype, device);
+        }
+
+        // Full causal mask: lower triangular
+        let mask: Vec<f32> = (0..seq_len)
+            .flat_map(|i| {
+                (0..total_len).map(move |j| {
+                    if j <= past_len + i {
+                        0.0
+                    } else {
+                        f32::NEG_INFINITY
+                    }
+                })
+            })
+            .collect();
+
+        Tensor::from_vec(mask, (1, 1, seq_len, total_len), device)?.to_dtype(dtype)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_kv_cache() {
+        let device = Device::Cpu;
+        let mut cache = KvCache::new();
+        assert_eq!(cache.seq_len(), 0);
+
+        let k = Tensor::zeros((1, 8, 5, 128), DType::F32, &device).unwrap();
+        let v = Tensor::zeros((1, 8, 5, 128), DType::F32, &device).unwrap();
+        let (fk, _fv) = cache.append(&k, &v).unwrap();
+        assert_eq!(cache.seq_len(), 5);
+        assert_eq!(fk.dim(2).unwrap(), 5);
+
+        let k2 = Tensor::zeros((1, 8, 1, 128), DType::F32, &device).unwrap();
+        let v2 = Tensor::zeros((1, 8, 1, 128), DType::F32, &device).unwrap();
+        let (fk2, _fv2) = cache.append(&k2, &v2).unwrap();
+        assert_eq!(cache.seq_len(), 6);
+        assert_eq!(fk2.dim(2).unwrap(), 6);
+
+        cache.reset();
+        assert_eq!(cache.seq_len(), 0);
+    }
+
+    #[test]
+    fn test_causal_mask_single_token() {
+        let mask = Qwen3Model::make_causal_mask(1, 10, DType::F32, &Device::Cpu).unwrap();
+        assert_eq!(mask.dims(), &[1, 1, 1, 11]);
+        // All zeros — single token can attend to everything
+        let sum: f32 = mask.sum_all().unwrap().to_scalar().unwrap();
+        assert_eq!(sum, 0.0);
+    }
+
+    #[test]
+    fn test_causal_mask_multi_token() {
+        let mask = Qwen3Model::make_causal_mask(3, 0, DType::F32, &Device::Cpu).unwrap();
+        assert_eq!(mask.dims(), &[1, 1, 3, 3]);
+        // Upper triangle should be -inf
+        let data: Vec<f32> = mask.flatten_all().unwrap().to_vec1().unwrap();
+        // Row 0: [0, -inf, -inf]
+        assert_eq!(data[0], 0.0);
+        assert!(data[1].is_infinite() && data[1] < 0.0);
+        assert!(data[2].is_infinite() && data[2] < 0.0);
+        // Row 1: [0, 0, -inf]
+        assert_eq!(data[3], 0.0);
+        assert_eq!(data[4], 0.0);
+        assert!(data[5].is_infinite() && data[5] < 0.0);
+        // Row 2: [0, 0, 0]
+        assert_eq!(data[6], 0.0);
+        assert_eq!(data[7], 0.0);
+        assert_eq!(data[8], 0.0);
+    }
+}