perf(nemotron_voicechat): bring streaming frame time under the 80 ms budget on M1 Ultra #2045

Problem / Background

Epic #1372 (PRs #2037, #2039, #2041) set out to run Nemotron VoiceChat in real time on one Apple Silicon machine. The streaming session does not meet that goal: on an idle Apple M1 Ultra (128 GB) each 80 ms frame takes slightly longer than 80 ms, so a live /v1/realtime session falls behind by 4% (4-bit) or 12% (8-bit) of the elapsed audio. PR #2044 recorded the numbers in docs/nemotron-voicechat.md (Performance).

Measured on main cdd531d5, "What is the capital of France?" prompt, 3 s extra decoding, 56 frames with the first 5 dropped as cold, 1 warm-up plus 3 measured runs (the runs agreed within 0.01):

Stage columns are p50 in ms. realtime_factor is mean frame time divided by 80 ms (src/models/nemotron_voicechat/streaming/profile.rs:123-127); above 1.0 is slower than real time. Only the LLM is quantized; the encoder, RNNT, EAR-TTS and codec are bf16 in every checkpoint (docs/nemotron-voicechat.md:5), which is why only the language stage moves between 4-bit and 8-bit.

Reproduction (alone on the machine):

mlxcel generate -m <NemotronLabs-VoiceChat-11B-4bit> --audio question.wav --stream --profile --seed 0

Goal: 4-bit realtime_factor below 1.0 on the same machine and methodology, with output parity preserved.

Current Behavior

Paths below are relative to src/models/nemotron_voicechat/ unless they start with src/ or audio/ (src/audio/).

Per frame, VoiceChatStreamingSession::step_audio_frame (streaming/session.rs:326-367) runs the stages strictly in sequence, and every stage ends in a host sync because the profiler times each stage after forcing its outputs (streaming/profile.rs:19-20).

Dependency graph inside frame t:

• perception(t) (streaming/perception.rs:76-100) produces encoded (RNNT input) and projected (LLM audio channel).
• rnnt(t) (streaming/session.rs:338-343) needs only encoded; nothing downstream reads its result.
• language(t) (streaming/session.rs:222-233) needs projected plus the previous frame's text and function ids; it returns host ids through two item_i32 calls (llm.rs:189-190).
• tts(t) (streaming/session.rs:237-252) needs the host text_id(t) (the subword char encoder looks the id up on the host, tts/subword.rs:249-255) and previous_code(t-1); it ends with eval(code) at streaming/session.rs:252.
• codec(t) (streaming/session.rs:292-314) needs code(t), and nothing else reads its output.
• function_id(t) is consumed only by fused_input of frame t+1 (streaming/session.rs:222-225).

So within a frame only rnnt can overlap language (0.4 ms, negligible), and the function-head readback can move off the critical path. Across frames, codec(t) could overlap frame t+1, but perception(t+1) cannot be hoisted in a live session because that audio has not arrived. MLX runs one Metal queue, so overlap only fills CPU graph-build and sync bubbles; it does not run two GPU-bound stages in parallel.

Host syncs per audio frame today: perception materialize eval_all (streaming/perception.rs:99, audio/fastconformer/streaming.rs:155-167), 1 to max_symbols RNNT try_eval plus item_i32 (audio/rnnt/mod.rs:245-246), two LLM item_i32 (llm.rs:189-190), TTS eval(code) (streaming/session.rs:252), codec validate_codes try_eval plus readback of codes that are already evaluated (audio/nemotron_codec/mod.rs:208-211), host_f32 of the waveform (streaming/session.rs:300, session.rs:70-77), and host_i32(&clean) (streaming/session.rs:313). That is 7 or more round trips per frame.

Per-Stage Analysis

"Exact" means the change cannot alter any value (same ops in the same dtype and order), so --stream --seed 0 output stays byte-identical. "Numerics" means results may differ in the last bits and the change needs the parity tests plus a documented decision. Gains are estimates to be confirmed by the profiling step.

TTS (36.0 ms, largest stage)

The EAR-TTS step is launch-bound, not bandwidth-bound. Weights read per step are about 1.2 GB for the 28-layer Gemma 3 backbone at batch 2 plus about 0.5 GB for five MoG head passes, roughly 3 ms of memory traffic on this machine. The op count is what is large. With the defaults (Q 31, 8 iterations, exponent 3) TtsConfig::mask_schedule (tts/config.rs:204-217) yields counts [0, 0, 0, 1, 1, 3, 4, 22]: 5 non-zero passes, 31 sequential RVQ iterations in total. Each pass in generate_codes (tts/model.rs:301-328) runs embed_codes (31 slice, reshape, take, add in tts/rvq.rs:113-117), a 3-layer MoG MLP plus argsort top-p (tts/mog_head.rs:40-56, 150-212), and encode_step (31 column slices and casts at tts/rvq.rs:138-140, then per unmasked codebook a transpose, matmul, scale, subtract, argmin, take, subtract at tts/rvq.rs:142-151). Together with the backbone this is on the order of a few thousand kernels per step.

Perception (22.7 ms for one 80 ms chunk)

The streaming mel emits f32 (src/audio/nemotron_mel_stream.rs:218) and nothing in src/audio/fastconformer/ casts it, while the encoder weights are bf16. Every matmul and conv therefore runs f32 x bf16, and MLX inserts a bf16-to-f32 astype of the weight on every call, every frame: about 24 layers x (two FeedForward, q/k/v/out/pos, two pointwise convs, depthwise conv) plus proj (1024 to 4480). That is on the order of 1 GB of weights cast per frame before the f32 matmuls read them again.

Language (16.1 ms 4-bit, 22.4 ms 8-bit)

The 6.3 ms gap between checkpoints tracks weight bytes, and a C++ full-forward decode path was already tried with no gain ("the bottleneck is GPU ops, not graph build CPU overhead", src/models/nemotron_h.rs:1747-1750). This is the lowest-yield stage for this issue. fused_input promotes to f32 through the perception output (llm.rs:147-164), so the Nemotron-H residual stream runs in f32; changing that is a numerics change against the reference and is out of scope.

Codec (7.4 ms)

ChannelLayerNorm reduces through einsum on purpose (audio/nemotron_codec/layers.rs:40-45, 257-265; audio/native_reduce.rs:25-31) because the widened sum_axis flips RVQ codes. Replacing it with fast_layer_norm is a numerics change and out of scope.

Cross-frame overlap (last resort)

realtime_factor is the mean of per-frame sequential stage sums, each taken after forcing evaluation, so any cross-frame overlap makes that number meaningless. X1 may be adopted only if the profiler gains a separately reported completion-interval metric (wall time between consecutive frame completions over a stream fed faster than real time), documented in docs/nemotron-voicechat.md. It also adds up to one frame (80 ms) of audio output latency, which must be documented.

Proposed Solution

Do the work in this order, measuring after each step, and stop adding risk once the target is met:

1. Profiling first. Behind the existing profile flag, add sub-stage timers inside RvqEarTtsModel::step (embed_codes, guided_condition, run_backbone, and per MaskGIT pass the MoG head and RVQ encode) and inside StreamingPerception::step (mel push, pre_encode, stream_layers, project, materialize), plus a per-frame count of host syncs (eval, eval_all, try_eval, item_i32, array_to_raw_bytes). Extend FrameTiming / ProfileSummary (streaming/profile.rs:28-64) with an optional sub-stage map so existing --profile output stays backward compatible. The sub-stage forced evals serialize the graph, so they run only when profiling is on and their numbers are for attribution, not for the headline factor. Post the 4-bit breakdown as a comment on this issue before optimizing.
2. Exact changes, highest estimated gain first: P1, T2, T1, T3, P3, P2, C1, L1, C3, C2, T4, C4.
3. If still at or above 1.0: P4, then T5, each only with the parity tests passing. Any change that alters --stream --seed 0 output needs a paragraph in the PR explaining the numerics change and why the parity tests still hold.
4. X1 only as described above.

Implement everything in a single PR, run the local gate once at the end, and keep pushes to a minimum.

Scope

In scope: src/models/nemotron_voicechat/ (streaming session, profiler, llm.rs, tts/), src/audio/fastconformer/, src/audio/nemotron_codec/, src/audio/nemotron_mel_stream.rs, the CLI --profile summary output, and docs/nemotron-voicechat.md. New compiled wrappers in src/lib/mlxcel-core only if T5 is reached.

Out of scope: quantizing EAR-TTS, the encoder or the codec (changes the model, not its execution); changing the f32 promotion of the LLM input; batch size above 1; server changes beyond what the profiler output requires.

Implementation Notes

• Reuse: mlxcel_core::eval_all / async_eval_pair for merged syncs; the existing host_i32 / host_f32 helpers (session.rs:70-77); the load-time norms precomputation pattern in tts/rvq.rs:67-86.
• Constraints: MLX evaluation is thread-affine (streaming/session.rs:26-27), so keep all work on the session thread. Memoization caches must be bounded: T1 by the text ids actually seen (store only the [2, 1, hidden] result), P2 and P3 by the saturating window length (at most 71 entries).
• RNG order: generate_codes draws one uniform and one normal per non-zero pass from MLX's global RNG (tts/model.rs:280-283). No change may add, remove or reorder random draws, or --seed 0 output changes.
• Edge cases: the first timeline position uses the silence frame and skips TTS (streaming/session.rs:238-240); eos resets previous_code (streaming/session.rs:241-243); the uncached perception and language modes (use_perception_cache: false, use_language_cache: false) must keep working but need not get faster.
• Error handling: C1 must keep rejecting out-of-range codes with the same code id {bad} outside codebook range error, now raised from the host copy.

Acceptance Criteria

• Sub-stage profiling and the per-frame host-sync count are available behind --profile, and the 4-bit breakdown is posted on this issue before the first optimization lands.
• 4-bit realtime_factor below 1.0 on the same idle M1 Ultra with the same command and methodology (1 warm-up plus 3 measured runs, all three reported with stage p50s).
• 8-bit realtime_factor reported with all three runs; below 1.0 is a stretch goal.
• One additional 4-bit run longer than 71 frames (past perception window saturation) reported once, so the steady state is on record.
• --stream --seed 0 output (text, function ids, audio codes, waveform) is byte-identical before and after, unless a deliberate numerics change is documented in the PR.
• All existing offline and streaming parity unit tests pass, and the env-gated real-weight tests pass: tests/nemotron_voicechat_{streaming,stream,tts,codec,offline,front,llm}_real.rs with MLXCEL_VOICECHAT_MODEL and the matching MLXCEL_VOICECHAT_*_REF variables set. (PR update(nemotron_voicechat): bring streaming frame time under 80 ms #2054: streaming, tts, codec, offline, front pass; llm and stream skip because the PADDED / STREAM_REF dumps were not available.)
• Each exact change has a unit test comparing the new path against the old one on fixed inputs with exact equality (for example T2 stacked gather vs 31 takes, T1 memoized vs fresh condition, P2 and P3 cached vs recomputed position terms, P1 f32-cast weights vs promoted matmul).
• docs/nemotron-voicechat.md Performance table and per-stage sentence updated with the new numbers.
• If shared code is touched, no regression for other models: Qwen3-4B 4-bit greedy output identical to main in every case; additionally a Gemma 3 text checkpoint greedy-identical to main if src/models/gemma3_backbone.rs or src/models/gemma3.rs changes, and a Nemotron-H text checkpoint greedy-identical if src/models/nemotron_h.rs changes.
• The speedup reaches the real /v1/realtime path and mlxcel generate --stream, not a separate benchmark-only path.

Verification

Measurements must be taken alone on the machine: no concurrent builds, tests or model runs. The M1 Ultra has shown GPU throttling after about 100 s of sustained load, so keep each measured run short, pause between runs, and report it if timings drift across the three runs.

make verify
make test-fast FILTER=nemotron_voicechat
MLXCEL_VOICECHAT_MODEL=<ckpt> MLXCEL_VOICECHAT_REF=<ref> cargo test --release --test nemotron_voicechat_streaming_real
mlxcel generate -m <4bit-ckpt> --audio question.wav --stream --profile --seed 0   # 1 warm-up + 3 measured
mlxcel generate -m <8bit-ckpt> --audio question.wav --stream --profile --seed 0   # 1 warm-up + 3 measured

For byte identity, capture --stream --seed 0 output (text, codes, WAV) on main and on the branch and compare with cmp.