feat(nemotron_voicechat): load the VoiceChat checkpoint and run offline duplex inference (FastConformer+RNNT, Nemotron-H dual-head, EAR-TTS, codec) #1374

Summary

Port the four networks of NemotronLabs VoiceChat and the offline timeline loop that ties them together: a 16 kHz WAV goes in, the assistant's text, the function-channel text, the user transcript, and a 22.05 kHz response WAV come out. This sub-issue owns every module, config struct, weight-key rule, and the non-streaming (full-history) loop; the cache-aware online session and the WebSocket endpoint are the next two sub-issues and only add state management on top of the forward functions defined here.

Current behavior

• src/models/detection.rs has no nemotron_voicechat arm ("nemotron_h" at line 300 and "nemotron_h_nano_omni" at 301 are the neighbours).
• src/models/nemotron_h.rs:1491-1500 (NemotronHModel { embeddings, layers, norm_f, lm_head, block_types, sequence_state }) loads backbone.embeddings, backbone.layers.N.*, backbone.norm_f.weight, lm_head (:2176-2500) and has forward_stage(inputs, layer_range, has_embedding, has_lm_head, caches) (:1583-1662) which accepts pre-computed embeddings when has_embedding == false but returns the hidden state BEFORE norm_f when has_lm_head == false (:1656-1661). There is no second head and no public final-norm entry point.
• src/models/gemma3.rs:499-560 (TransformerBlock::from_weights(weights, args, layer_idx) with the model.layers.{idx} prefix hardcoded at line 552) and Gemma3Model::forward_with_caches_and_embeddings (:735-760) always multiply injected embeddings by sqrt(hidden_size) (line 755-757); the TTS backbone must not scale and has no embed_tokens / lm_head.
• src/audio/whisper_mel.rs:200 computes a Whisper-specific 80/128-mel log spectrogram (Slaney filters, log10, Whisper normalization); the VoiceChat frontend needs preemphasis, a 25 ms / 10 ms Hann STFT of size 512, natural log(x + 2^-24), and no normalization.
• src/models/kokoro/lstm.rs:43-58 (load(w, prefix, hidden), forward(x, t)) and src/models/kokoro/stft.rs:70-179 (stft, istft, irdft_basis) are the only LSTM and STFT/iSTFT implementations; both are pub(crate) and tied to Kokoro's key names and window sizes.
• The tiktoken/HF tokenizer loaders in src/tokenizer/mod.rs handle the LLM's tokenizer.json; nothing reads rnnt_tokenizer/vocab.json or the config's rnnt_vocabulary list.

Expected behavior

./target/release/mlxcel generate -m models/NemotronLabs-VoiceChat-11B-4bit \
  --audio question.wav --output-audio response.wav -p "Be concise and answer in one sentence." -n 0

prints [user] <transcript> and the assistant text, and writes response.wav (22050 Hz mono PCM16) containing the spoken answer in the built-in Aria voice. -p is the system prompt (empty means none); -n is ignored because the output length equals the input timeline length (plus --extra-decoding-seconds, default 3, of appended silence so the model can finish answering).

config.json

Top level: model_type "nemotron_voicechat", bos_token_id 1, eos_token_id 2, pad_token_id 12, silence_token_id 11, rnnt_blank_id 1024, input_sample_rate 16000, output_sample_rate 22050, frame_duration 0.08, function_channel_weight 2.0, speaker "Aria", rnnt_vocabulary (list of 1024 SentencePiece pieces), quantization / quantization_config (per-module map: group_size 64, bits 4 with one entry per quantized linear; only stt_model.embed_tokens, stt_model.lm_head, stt_model.function_head and the Nemotron-H mixer projections are quantized; the speech encoder, RNNT, TTS and codec stay bf16).

text_config: model_type "nemotron_h", hidden_size 4480, intermediate_size 15680, num_hidden_layers 56, num_attention_heads 40, num_key_value_heads 8, head_dim 128, mamba_num_heads 128, mamba_head_dim 80, ssm_state_size 128, conv_kernel 4, n_groups 8, use_conv_bias true, mamba_proj_bias false, mlp_bias false, attention_bias false, layer_norm_epsilon 1e-5, vocab_size 131072, max_position_embeddings 131072, hybrid_override_pattern "M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-" (attention at layers 14, 21, 30, 39; - layers are MLPs).

audio_config.preprocessor: sample_rate 16000, features 128, n_fft 512, window_size 0.025 (400 samples), window_stride 0.01 (160), window "hann", preemph 0.97, log true, normalize "NA", dither 1e-5 (training only; not applied), log_zero_guard_value 2^-24, pad_to 0.

audio_config.encoder: feat_in 128, n_layers 24, d_model 1024, n_heads 8, ff_expansion_factor 4, subsampling_factor 8, subsampling_conv_channels 256, conv_kernel_size 9, causal_downsampling true, conv_context_size "causal", conv_norm_type "layer_norm", self_attention_model "rel_pos", att_context_style "chunked_limited", att_context_size [[70, 0]], pos_emb_max_len 5000, use_bias false, xscaling false.

audio_config.decoder (RNNT prediction net): pred_hidden 640, pred_rnn_layers 2, vocab_size 1024, blank_as_pad true. audio_config.joint: joint_hidden 640, activation "relu", encoder_hidden 1024, pred_hidden 640, num_classes 1024. audio_config.output_dim 4480, audio_config.max_symbols 10.

tts_config: hidden_size 1152, intermediate_size 4608, num_hidden_layers 28, num_attention_heads 16, num_key_value_heads 16, head_dim 72, sliding_window 7500, sliding_window_pattern 6, rms_norm_eps 1e-6, query_pre_attn_scalar 256, rope_global_base_freq 1e6, rope_local_base_freq 1e4, latent_size 512, num_quantizers 31, codebook_size 1024, num_delay_speech_tokens 2, num_iterations 8, guidance_scale 0.2, top_p 0.95, noise_scale 0.001, exponent 3.0, use_gated_fusion_for_text_audio true, use_subword_flag_emb true, use_bos_eos_emb true, use_audio_prompt_frozen_projection true, audio_prompt_duration 3.0; character_encoder {hidden_size 1152, intermediate_size 4608, num_hidden_layers 1, num_attention_heads 16, num_key_value_heads 16, head_dim 72, rms_norm_eps 1e-6, query_pre_attn_scalar 256, attn_logit_softcapping 50, rope_base 1e4, char_vocab_size 257}; mog_head {intermediate_size 4608, low_rank 64, min_log_std -4, num_layers 3, num_predictions 1024, eps 1e-6}.

codec_config: sample_rate 22050, base_channels 384, channel_multipliers [1, 2, 4], downsample_rates [7, 7, 9], blocks_per_stage 3, block_kernel_size 7, latent_dim 512, n_fft 16, hop_length 4, num_quantizers 31, codebook_size 1024. Derived: waveform_to_token_ratio = 4 * 7 * 7 * 9 = 1764 samples per code frame (80 ms at 22050 Hz), stft_channels = 2 * (16/2 + 1) = 18.

Network 1: log-mel frontend and FastConformer (stt_model.perception.*)

Log-mel for a waveform x (f32, 16 kHz): x' = [x0, x1 - 0.97 x0, x2 - 0.97 x1, ...]; STFT with n_fft 512, hop 160, a symmetric 400-point Hann window zero-padded to 512 (56 zeros left, 56 right), reflect padding of 256 samples each side (center = true); power spectrum |X|^2 over 257 bins; mel = filters[128, 257] @ power with Slaney-normalized Slaney-scale filters (hz_to_mel_slaney exists in src/audio/whisper_mel.rs:60); log(mel + 2^-24); no normalization. Output [1, T, 128] with T = 1 + floor(len / 160).

Subsampling (pre_encode): input [B, T, 128, 1] channels-last. Three stride-2 stages; every 3x3 stride-2 conv is preceded by asymmetric padding (left 2, right 1) on both the time and frequency axes: conv.0 = Conv2d(1 -> 256, k3, s2) + ReLU; conv.2 = depthwise Conv2d(256, k3, s2, groups 256), conv.3 = Conv2d(256 -> 256, k1) + ReLU; conv.5 / conv.6 repeat. Output lengths per stage floor((n + 3 - 3) / 2) + 1; frequency 128 -> 64 -> 32 -> 16. Flatten [B, T', 16, 256] channel-major as transpose(0, 1, 3, 2).reshape(B, T', 4096) then out: Linear(4096 -> 1024, bias). Ten ms mel frames become 80 ms encoder frames.

Relative positional encoding: the sinusoid table pe[2*max_len - 1, 1024] over positions max_len-1 .. -(max_len-1); for a length-L input take the centered slice of 2L - 1 rows; xscaling false so the input is not scaled.

Each of 24 blocks (layers.N):

x = x + 0.5 * ff1(LN(norm_feed_forward1)(x))        # ff: linear1 1024->4096 (no bias), SiLU, linear2
x = x + attn(LN(norm_self_att)(x), pos_emb, mask)
x = x + conv(LN(norm_conv)(x))
x = x + 0.5 * ff2(LN(norm_feed_forward2)(x))
x = LN(norm_out)(x)

Attention (self_attn, no bias on linear_q/k/v/out, linear_pos without bias, per-layer pos_bias_u / pos_bias_v [8, 128]): q_u = (q + pos_bias_u), q_v = (q + pos_bias_v), p = linear_pos(pos_emb); bd = rel_shift(q_v @ p^T)[..., :L] * scale where rel_shift left-pads the last axis by one, views as [B, H, pos_len + 1, L], drops the first row and views back; scores = (q_u @ k^T) * scale + bd + mask; softmax; @ v; linear_out. Mask for att_context_size [70, 0] (chunked_limited, chunk size right + 1 = 1): frame i may attend to frames i - 70 ..= i, everything else is -1e30.

Convolution module: pointwise_conv1 (1024 -> 2048, k1, no bias) -> GLU on the channel axis -> left-pad 8 frames -> depthwise conv k9 (groups 1024, no bias) -> LayerNorm (weight key conv.batch_norm, it is a LayerNorm despite the name) -> SiLU -> pointwise_conv2 (k1, no bias).

perception.proj: Linear(1024 -> 4480, bias). The perception returns both the projected [B, T', 4480] (LLM audio channel) and the raw [B, T', 1024] encoder output (RNNT input).

Network 1b: RNNT transcript branch (stt_model.rnnt_decoder.*, stt_model.rnnt_joint.*)

Prediction net: prediction.embed Embedding(1025, 640) (index 1024 is blank, never embedded), prediction.dec_rnn.lstm two stacked LSTM(640) layers. Joint: enc Linear(1024 -> 640, bias), pred Linear(640 -> 640, bias), ReLU, joint_net.2 Linear(640 -> 1025, bias). Greedy decode per encoder frame e[1, 1, 1024]: loop up to max_symbols = 10 times: pred_in = embed(last_token) if last_token != blank else zeros [1, 1, 640]; (out, h') = lstm(pred_in, h); logits = joint_net(relu(enc(e) + pred(out))); tok = argmax; if tok == 1024 break; else last_token = tok, h = h', and append tok unless it is a special piece (<unk>, <pad>, <s>, </s>, or a <xx-XX> language tag). Text = join of pieces with ▁ -> space, stripped. Pieces come from rnnt_vocabulary (config) or rnnt_tokenizer/vocab.json.

LSTM weight layout: weight_ih_l{n} [2560, in], weight_hh_l{n} [2560, 640], gate order i, f, g, o; bias = bias_ih_l{n} + bias_hh_l{n}. src/models/kokoro/lstm.rs implements the same cell; generalize its loader to take explicit keys or add a small Lstm in src/audio/ (shared by both).

Network 2: Nemotron-H with a fused input and two heads (stt_model.embed_tokens, stt_model.llm.*, stt_model.lm_head, stt_model.function_head)

Per timeline position t the LLM input is one embedding row:

fused[t] = E(prev_text[t]) + audio[t] + 2.0 * E(prev_function[t])

where E = stt_model.embed_tokens (quantized), audio[t] is the projected perception frame (or a system-prompt token embedding during the prefix), and prev_*[0] = pad (12). Run the 56 Nemotron-H layers on fused (no embedding lookup inside the LLM), apply norm_f, then text_logits = lm_head(h), function_logits = function_head(h); both heads are [4480 -> 131072] quantized, no bias. Text and function ids are the argmax of the last position (greedy; no sampling).

Key mapping for the existing loader: stt_model.llm.layers.N.X -> backbone.layers.N.X, stt_model.llm.norm_f.weight -> backbone.norm_f.weight, stt_model.embed_tokens.* -> backbone.embeddings.*, stt_model.lm_head.* -> lm_head.*; stt_model.function_head.* is loaded separately as a UnifiedLinear. Mamba conv1d.weight is [channels, 4, 1] (MLX layout) in the converted checkpoints; keep the [0, 2, 1] transpose gate for torch [channels, 1, 4] exports.

Network 3: EAR-TTS (tts_model.tts_model.*, tts_model.audio_prompt_latents.Aria, tts_model.codec_silence_tokens, tts_model._control_codes)

Backbone: 28 Gemma-3-style blocks (backbone.layers.N.{input_layernorm, self_attn.{q,k,v,o}_proj, self_attn.{q,k}_norm, post_attention_layernorm, pre_feedforward_layernorm, mlp.{gate,up,down}_proj, post_feedforward_layernorm}, backbone.norm), hidden 1152, 16 heads of 72 (no GQA), attention scale 256^-0.5, gelu-tanh gated MLP 4608, offset RMSNorm (weight + 1), layer i global when (i + 1) % 6 == 0 (RoPE base 1e6, full KVCache) else sliding (base 1e4, rotating cache of 7500). Inputs are injected embeddings with NO sqrt(hidden) scaling and there is no token embedding or LM head.

Code embedding: depthsum(code[.., 31]) = sum_q rvq_embs[q][code_q] with rvq_embs [31, 1024, 512] extended by one zero row per codebook so index 1024 (mask) contributes nothing; embed_code: Linear(512 -> 1152, no bias).

Text conditioning (embed_subword.*): for each conditioned subword id build its character sequence through the dense character vocabulary (the tokenizer's single-character tokens sorted by id, 256 entries + 1 padding = 257; the constructor must fail if the count differs), embed with embed_tokens Embedding(257, 1152), scale by sqrt(1152), run one T5Gemma encoder layer (pre/post self-attn offset norms, q/k/v/o without bias, RoPE base 1e4, scores tanh(s / 50) * 50 softcap, mask over padded chars, pre/post feed-forward offset norms, gated gelu-tanh MLP 4608), final offset norm, masked mean-pool over characters, proj_embedding: Linear(1152 -> 1152, no bias); then add subword_flag_emb.cont_emb[is_continuation[id]] (Embedding(2, 1152) indexed by the int buffer is_continuation[131073], OOV -> pad_tensor) and bos_eos_emb.special_emb[special_flags[id]] (Embedding(3, 1152) indexed by special_flags[131072], OOV -> pad_tensor).

Classifier-free guidance: every TTS forward runs batch 2: row 0 conditioned on the text embedding, row 1 on null_emb broadcast; code embeddings are duplicated.

Gated fusion (gated_fusion_audio_text): final_norm(sigmoid(residual_scale) * (sigmoid(gate) * audio_proj(code_embed / 31) + (1 - sigmoid(gate)) * text_proj(text_embed))), audio_proj / text_proj Linear(1152, 1152) with bias, gate [1152], residual_scale scalar, final_norm offset RMSNorm.

Warmup (once per session, the Aria prompt): let frames = 37 (audio_prompt_latents.Aria is [1, 37, 1152]). Encode (37 + 1) * 1764 zero samples with the codec to get 38 code frames; set frame 0 and frame 36 to all-mask (1024); keep frames 0..37 (drop the last) as code [1, 37, 31]; subwords = [12] * 37, subword_mask = false except the last two, audio_mask = false except the last one. Then shifted = concat(zeros, code[:, :-1]), code_embed = embed_code(depthsum(shifted)), bos_mask = audio_mask & !prev(audio_mask), pre_bos = cumsum(bos_mask) == 0, code_embed = where(pre_bos, Aria_latent, code_embed) + bos_mask * bos_emb, duplicate for CFG, cond = condition(subwords, subword_mask), inputs = fusion(code_embed, cond), run the backbone to fill the caches. previous_code = code[:, -1:].

Per step (timeline position t >= 1, current text id cur): if cur == 2 (EOS) set previous_code = codec_silence_tokens; code_embed = embed_code(depthsum(previous_code)) duplicated, cond = condition([[cur]], mask true), hidden = backbone(fusion(code_embed, cond), caches) [2, 1, 1152]; then generate_codes(hidden):

masked_i = ceil((1 - (i/8)^3)^(1/3) * 31) for i in 0..8  -> counts_i = masked_i - masked_{i+1} (last: masked_7)
code = [1024] * 31; done = 0
for count in counts (skip zeros):
    emb = embed_code(depthsum(code))
    x = concat(emb + hidden[0], emb + hidden[1])            # cond / uncond rows
    mu, logs = mog_head.infer(x, guidance 0.2, top_p 0.95)
    residual = mu + exp(logs) * N(0, 1) * 0.001
    for q in done..done+count: code[q] = argmin_k ||residual - rvq_embs[q][k]||^2 ; residual -= rvq_embs[q][code[q]]
    done += count

mog_head.infer: 3 x (pre_norm, gated gelu-tanh MLP 1152 -> 4608 -> 1152, post_norm, residual) then the final offset norm (mlp_stack.3.weight); CFG x = cond + 0.2 * (cond - uncond); logits = proj_logits(x) [.., 1024] with top-p 0.95 nucleus masking; pick the component by Gumbel-max on log softmax(logits); mu = low_mat[c] @ (proj_mus.weight.reshape(1024, 64, 1152)[c] @ x) [512]; logs = max(proj_logs(x), -4); return mu * exp(logs) + proj_else(x) and logs. Randomness: seed the generator from the session seed (default 0) so runs are reproducible.

Before codec decoding, replace any of the three _control_codes values in the generated codes by codec_silence_tokens (the per-codebook silence code).

Network 4: codec (tts_model.audio_codec.*)

Encoder (encoder.layers.{0..15}): 0 Conv1d(18 -> 384, k1, no bias); stages (384, 7), (768, 7), (1536, 9): each stage has 3 ConvNeXt blocks then a strided Conv1d to the next width (384 -> 768, k7 s7; 768 -> 1536, k7 s7; 1536 -> 512, k9 s9, no bias). ConvNeXt block (dwconv depthwise k7 with bias, causal left pad 6; norm channel LayerNorm eps 1e-6; pwconv1 k1 C -> 4C; GELU; pwconv2 4C -> C; residual). Spectrogram input: pad the waveform by 6 samples each side (zeros), periodic Hann(16), center = false, hop 4 -> [B, 9 complex bins, T], features = concat(real, imag) -> 18 channels, transpose to [B, T, 18].

PRVQ: 31 codebooks prvq.mus_list.{q} [1024, 512]; encode = iterative nearest-mean residual quantization; decode = sum of selected means. prvq.variance_list.{q}.variance scalars are loaded but unused.

Decoder (decoder.layers.{0..15}): three stages in reverse, each ConvTranspose1d(src -> C, k = rate, stride = rate, no bias) followed by 3 ConvNeXt blocks, then Conv1d(384 -> 18, k1). Output [B, T, 18] -> transpose -> magnitude = 100 * exp(-softplus(-m + ln 100)) over the first 9 channels, phase the last 9, real = mag cos(phase), imag = mag sin(phase) with imag of bin 0 and bin 8 forced to zero; iSTFT (n_fft 16, hop 4, periodic Hann, center = false) then drop 6 samples at each end. One code frame becomes 1764 samples.

Weight layout in the converted checkpoints: Conv1d kernels are [out, k, in] (MLX) and ConvTranspose1d kernels are [out, k, in] after the (1, 2, 0) transpose of torch's [in, out, k]; prvq._variance_list is renamed to prvq.variance_list. Gate the transposes on shape[-1] == in_channels so a re-sanitize is a no-op.

Offline timeline loop

1. Load WAV, resample to 16 kHz mono (src/audio/whisper_mel.rs:176 resample_to_16k or the rubato path in src/audio/preprocessing_resample.rs), append extra_decoding_seconds * 16000 zeros.
2. mel = log_mel(x); (audio_embeds [1, F, 4480], lengths, asr_embeds [1, F, 1024]) = perception(mel); F = lengths[0].
3. System prompt (when non-empty): ids = [1] + encode(prompt) + [2]; audio_embeds = concat(E(ids), audio_embeds) (prompt tokens enter on the audio channel); prompt_frames = len(ids).
4. text[t] = function[t] = 12 for all t < prompt_frames + F; TTS warmup; previous_code = prompt_code[:, -1:].
5. For t in 0..timeline: fused = E(text[t-1]) + audio_embeds[t] + 2 E(function[t-1]) (pad at t = 0); h = llm(fused, caches); if t >= prompt_frames: text[t] = argmax(text_logits), function[t] = argmax(function_logits); if t == 0 continue (no TTS at position 0); TTS step with text[t] -> codes[t].
6. Drop the prompt prefix from text, function, codes; replace control codes; audio = codec.decode(codes) (full-history decode, no cache); transcript = RNNT greedy over asr_embeds[:, :F].
7. Assistant text = decode of text ids with {12, 11, 1, 2} removed (skip_special_tokens = false); function text likewise.

Implementation plan

1. Detection and registry. src/models/detection.rs: "nemotron_voicechat" => Ok(ModelType::NemotronVoiceChat) placed before the nemotron_h arms. src/models/mod.rs: enum, supported list, description ("Nemotron VoiceChat (FastConformer + Nemotron-H + EAR-TTS + codec, full duplex)", "Speech"). LoadedModel::NemotronVoiceChat(models::NemotronVoiceChatModel) in src/loaded_model.rs with a LanguageModel impl that refuses decoder-only calls (pattern: Florence2VlmModel, src/models/florence2/runtime.rs:228). Detection test.

2. Config (src/models/nemotron_voicechat/config.rs). VoiceChatConfig { text: NemotronHConfig, audio: AudioConfig { preprocessor: MelArgs, encoder: ConformerArgs, decoder: PredictArgs, joint: JointArgs, output_dim, max_symbols }, tts: TtsConfig { .., character_encoder: CharEncoderConfig, mog_head: MogConfig }, codec: CodecConfig, bos/eos/pad/silence ids, rnnt_blank_id, input/output sample rates, frame_duration, function_channel_weight, speaker, rnnt_vocabulary: Vec<String>, quantization } with serde defaults equal to the values above; att_context_size accepts [70, 0] or [[70, 0]]. Per-module quantization: read the quantization map and hand (group_size, bits) to UnifiedLinear::from_weights only for keys present in it; everything else loads dense (UnifiedLinear already resolves dense vs quantized from the presence of .scales, so the map is a cross-check and the dense loaders must not require it).

3. Frontend (src/audio/nemotron_mel.rs). fn log_mel_spectrogram(x: &[f32], args: &MelArgs) -> (Vec<f32>, usize /*frames*/) (preemphasis, centered 400-in-512 Hann, reflect pad 256, power, Slaney mel 128, ln(x + 2^-24)). Reuse periodic_hann / hz_to_mel_slaney / reflect_pad from src/audio/whisper_mel.rs by making them pub(crate). Tests: frame_count_is_one_plus_floor_len_over_hop, silence_gives_log_guard_floor (every value equals ln(2^-24)), preemphasis_first_sample_unchanged.

4. Speech encoder (src/audio/fastconformer.rs). FastConformerEncoder::from_weights(weights, "stt_model.perception.encoder", &ConformerArgs) with CausalDwStridingSubsampling, RelPositionalEncoding (table built once for max_len = 5000, extended on demand), RelPosMultiHeadAttention (rel_shift as above), ConformerConvolution, ConformerBlock; forward(mel [1, T, 128], lengths) -> ([1, T', 1024], lengths') and fn chunked_limited_mask(len, left, right). Conv2d kernels are [out, kh, kw, in] (converted) with the shape[-1] == in_channels gate for torch [out, in, kh, kw]; depthwise kernels carry groups = channels (mlxcel_core::conv2d takes a groups argument; copy the existing depthwise call in deepseekocr_sam.rs). Perception { encoder, proj: Linear } returns (projected, lengths, encoded). Tests on random weights: subsampling_lengths_128_mel_frames_give_16, chunked_limited_mask_70_0_allows_self_and_70_left, rel_shift_matches_hand_indexing (a 1x1x3x5 hand-checked case), block_output_shape.

5. RNNT (src/audio/rnnt.rs). PredictNetwork (Embedding + 2-layer LSTM; reuse/generalize src/models/kokoro/lstm.rs), JointNetwork, fn greedy_decode_frame(enc [1,1,1024], state) -> Vec<i32>, fn decode_pieces(ids, vocabulary) -> String with the special-piece filter. Tests: blank_only_frame_emits_nothing, max_symbols_bounds_loop, piece_decode_maps_underscore_to_space_and_drops_lang_tags.

6. LLM glue (src/models/nemotron_h.rs). Add pub fn forward_embeds_to_hidden(&self, inputs_embeds, caches) -> UniquePtr<MlxArray> = forward_stage(.., has_embedding false, has_lm_head false) followed by norm_f; keep lm_head inside the model and expose pub fn lm_head(&self, h); VoiceChat holds function_head: UnifiedLinear beside it. Loader renames keys as listed in the Network 2 section before NemotronHModel::from_weights(config, weights, block_types). Test: forward_embeds_to_hidden_matches_forward_on_embedded_ids on a tiny random config.

7. TTS (src/models/nemotron_voicechat/tts.rs). OffsetRmsNorm, GatedMlp, MogHead, CharAwareSubwordEncoder (with set_vocabulary(&HashMap<String, u32>) building the 257-entry char table and the subword -> chars map), GatedFusion, RvqEarTtsModel { backbone: Gemma3Backbone, bos_emb, null_emb, embed_code, embed_subword, fusion, audio_prompt_projection_W, mog_head, rvq_embs }, SpeechDecoder { codec, tts, control_codes: [i32; 3], aria_latent [1, 37, 1152], codec_silence_tokens [31] }. Add Gemma3Backbone to src/models/gemma3.rs: from_weights(weights, prefix, &ModelArgs) building TransformerBlocks from {prefix}.layers.{i} (parameterize the prefix at line 552) and {prefix}.norm, plus forward_embeds(inputs, caches, scale_inputs: bool); make_caches() returning KVCache for (i+1) % 6 == 0 and RotatingKVCache::new(7500, keep 0) otherwise (src/lib/mlxcel-core/src/cache.rs:4094). Tests: offset_rms_norm_equals_rms_norm_with_weight_plus_one, top_p_keeps_at_least_the_largest_logit, mask_schedule_counts_sum_to_31_for_exponent_3, depthsum_ignores_mask_index, char_encoder_scatters_only_masked_positions, mog_infer_shapes_and_finite (random weights, batch 2).

8. Codec (src/audio/nemotron_codec.rs). ConvNeXtBlock1d, AudioEncoder, AudioDecoder, Prvq, NemotronCodec { encode(wave [B, N]) -> codes [B, 31, T], decode(codes) -> wave [B, 1, N] }, CausalConv1dCache (per-layer overlap buffers, used by the streaming sub-issue), STFT/iSTFT at n_fft 16 hop 4 (generalize src/models/kokoro/stft.rs to take n_fft, hop, window, center, or add a small src/audio/stft.rs). mlxcel_core::conv1d (bridge line src/lib/mlxcel-core/src/lib.rs:2246) with groups = channels for the depthwise conv and mlxcel_core::conv_transpose1d (:2402) for the decoder stage starts. Tests: waveform_to_token_ratio_is_1764, stft_channels_is_18, encode_decode_round_trip_on_silence_is_near_zero (random-init codebooks make this a shape test; the numeric check uses the real checkpoint), decoder_transpose_gate_is_idempotent.

9. Model and offline runtime (src/models/nemotron_voicechat/{mod.rs, model.rs, session.rs}). NemotronVoiceChatModel { config, embed_tokens: UnifiedEmbedding, llm: NemotronHModel, function_head, perception, rnnt: (PredictNetwork, JointNetwork), speech_decoder: SpeechDecoder, tokenizer: MlxcelTokenizer } with load(model_path); VoiceChatResult { text, function_text, user_transcript, audio: Vec<f32>, sample_rate: 22050, text_tokens, function_tokens, audio_codes }; fn generate_offline(&self, wav: &[f32], system_prompt: Option<&str>, extra_decoding_seconds: f32, seed: u64) -> Result<VoiceChatResult> implementing the timeline loop. Model-owned caches: Nemotron-H NemotronLayerCache vector from make_caches, TTS caches from Gemma3Backbone::make_caches.

10. CLI. In src/commands/generate.rs, branch on LoadedModel::NemotronVoiceChat before the text path (pattern: the Florence-2 branch at line 2431): require --audio, treat -p as the system prompt, add --extra-decoding-seconds (f32, default 3.0) and --seed reuse, write --output-audio with src/audio/wav_writer.rs::encode_wav_pcm16(samples, 22050, 1), print [user] ..., the assistant text, and [function] ... when non-empty. Reject --image / --video.

11. Docs. docs/supported-models.md speech section; docs/nemotron-voicechat.md with the architecture summary above and the CLI usage.

Validation

(a) Unit tests listed per step, in src/audio/nemotron_mel_tests.rs, src/audio/fastconformer_tests.rs, src/audio/rnnt_tests.rs, src/audio/nemotron_codec_tests.rs, src/models/nemotron_voicechat/tts_tests.rs, src/models/nemotron_h_tests.rs, src/models/detection_tests.rs.

(b) Real checkpoint: mlx-community/NemotronLabs-VoiceChat-11B-4bit (9.2 GB; rnnt_tokenizer/, tokenizer.json, two safetensors shards) and mlx-community/NemotronLabs-VoiceChat-11B-8bit for the quality check.

# Make a 16 kHz mono question WAV: macOS `say -o /tmp/q.aiff "What is the capital of France?"` then
# `ffmpeg -i /tmp/q.aiff -ar 16000 -ac 1 /tmp/question.wav`, or Kokoro through the server's
# POST /v1/audio/speech (24 kHz output; resample to 16 kHz the same way).
./target/release/mlxcel generate -m models/NemotronLabs-VoiceChat-11B-4bit --audio /tmp/question.wav \
  --output-audio /tmp/response.wav -p "Be concise and answer in one sentence." -n 0

Acceptance: the printed [user] transcript contains "capital of France"; the assistant text mentions Paris; /tmp/response.wav is 22050 Hz mono, non-silent (RMS > 0.01 over the answer region), and its length equals (timeline_frames) * 1764 samples; every intermediate (mel, encoder, logits, codes) is finite; with --seed 0 two runs are byte-identical. Parity: the codec alone round-trips a 1 s 440 Hz tone with SNR > 20 dB through encode then decode on the real weights; the speech encoder output for a 1 s WAV matches the model publisher's own inference code run once on CPU within RMS 5e-3 at bf16; the first 20 assistant text ids match exactly for the same WAV and system prompt.

Acceptance criteria

• nemotron_voicechat resolves to ModelType::NemotronVoiceChat and the 4-bit checkpoint loads with per-module quantization (only the listed modules quantized)
• Log-mel + FastConformer + RNNT transcribe a spoken sentence correctly
• Nemotron-H runs on fused embeddings with norm_f and both heads; text/function ids are greedy argmax
• EAR-TTS warmup with the Aria prompt and per-step code generation produce 31 codes per frame; control codes are replaced
• Codec round-trip SNR > 20 dB on a tone; decoded answer audio is non-silent and 1764 samples per frame
• mlxcel generate --audio --output-audio prints transcript + answer and writes a playable WAV; seeded runs are deterministic
• docs/supported-models.md updated
• detection table in src/models/detection.rs updated with a test
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

• Cache-aware online inference, event streaming, WebSocket (next sub-issues).
• Other speakers; the speaker field is validated to equal "Aria".

Dependencies

• Part of epic feat(nemotron_voicechat): full-duplex speech model (FastConformer + Nemotron-H + EAR-TTS + codec) with a /v1/realtime session #1372