feat(nemotron_voicechat): cache-aware online session (80 ms frame clock, persistent caches, aligned events, profiler)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
The offline loop from #1374 recomputes the whole timeline from full history. A conversation needs bounded work per 80 ms frame: a stateful session that accepts PCM chunks of any size, advances every network exactly one frame at a time from persistent caches, and emits events aligned to a frame index. This sub-issue adds `VoiceChatStreamingSession` with the five cache families (streaming log-mel, FastConformer attention/conv/subsampling caches, Nemotron-H sequence state, EAR-TTS KV caches, codec causal-conv and iSTFT overlap caches), the system-prompt prefill, `push_audio` / `flush` / `cancel`, the event types, and a per-frame stage profiler. It is a library surface (used by the CLI for a `--stream` mode and by the WebSocket endpoint in the next sub-issue).
## Current behavior
- After #1374, `NemotronVoiceChatModel::generate_offline` (`src/models/nemotron_voicechat/session.rs`) recomputes the log-mel and the full FastConformer pass over the whole waveform, runs the LLM with model-owned caches, and decodes the codec from full history. There is no incremental frontend, no per-frame encoder state, and no codec overlap state.
- `src/models/nemotron_h.rs:1521` `make_caches() -> Vec` and the `ModelOwnedSequenceState` machinery already give the LLM a persistent per-session state; `src/lib/mlxcel-core/src/cache.rs:4094` `RotatingKVCache` gives the TTS sliding layers a bounded cache. Nothing comparable exists for the Conformer or the codec.
- `src/server/audio_worker.rs:15-36` documents the thread-affinity rule every MLX model follows (load and evaluate on one owned thread); the session must be created and driven on one thread for the same reason.
## Expected behavior
```rust
let session = model.create_streaming_session(StreamingOptions {
system_prompt: Some("Be concise and answer in one sentence."),
seed: 0,
max_streaming_seconds: None,
use_language_cache: true,
use_perception_cache: true,
profile: false,
})?;
for chunk in pcm_chunks { for ev in session.push_audio(chunk, 16_000)? { handle(ev) } }
for ev in session.flush(true)? { handle(ev) }
```
`push_audio` buffers samples and, for every complete 1280-sample frame, runs one timeline step and returns the events it produced; a 300-sample chunk returns nothing and a 3000-sample chunk returns the events of two frames. Events:
| kind | fields |
|---|---|
| `AssistantTextDelta` | `frame_index`, `token_id`, `delta: String`, `text: String` (cumulative) |
| `FunctionDelta` | same shape, function channel |
| `UserTranscriptDelta` | `frame_index`, `delta`, `text` |
| `Audio` | `frame_index`, `samples: Vec` (exactly 1764 at 22050 Hz), `audio_codes: [i32; 31]` |
| `Done` | `frame_index` |
| `Cancelled` | `frame_index` |
Text deltas skip ids `{12, 11, 1, 2}`; `delta` is `cumulative[len(previous_cumulative)..]` when the new cumulative decode extends the old one, else the single-token decode (tokenizers may revise a partial multibyte sequence; the cumulative `text` stays authoritative).
### Per-frame algorithm (one 1280-sample frame)
1. **Perception (cached path).** `mel_frames = mel_stream.push(frame)` where the streaming log-mel keeps a sample buffer, emits only hop-aligned frames whose centers trail the input edge by `lookahead_samples = 1280`, and retains `ceil((256 + 1) / 160) = 2` frames of look-behind plus the preemphasis predecessor sample. `encoded = conformer_stream.push(mel_frames, emit_partial = true)` must return exactly one `[1, 1, 1024]` frame; error otherwise (the 1280-sample frame is 8 mel hops, which is one encoder frame after 8x subsampling). `projected = perception.proj(encoded)`.
The Conformer stream keeps, per layer, `attn_cache[layer]` = the last `left_context = 70` attention-input frames and `conv_cache[layer]` = the last `conv_kernel - 1 = 8` GLU-output frames, plus a `mel_cache` of the last 16 mel frames for the causal subsampling stack. Per layer: `r = x + 0.5 ff1(LN(x))`; `xn = LN_att(r)`; `kv = concat(attn_cache, xn)`; `r += attn.stream(q = xn, kv, pos_emb_for(len(kv)))` (no mask: the window is the allowed context; `pos_emb_for(L)` is the centered `2L - 1` slice); `attn_cache = kv[-70:]`; `g = GLU(pointwise_conv1(LN_conv(r)))`; `din = concat(conv_cache or zeros[8], g)`; `dw = depthwise_conv(din)` (valid, no padding); `conv_cache = din[-8:]`; `r += pointwise_conv2(SiLU(LN(dw)))`; `r += 0.5 ff2(LN(r))`; `x = LN_out(r)`. Subsampling is incremental: run `pre_encode` over `concat(mel_cache, new_mel)` and keep only the frames not yet emitted (`emitted` / `consumed` counters in mel-frame units, `base = (consumed - cache_len) / 8`).
Uncached fallback (`use_perception_cache = false`): keep the last `max(2, 70 + 0 + 1) * 1280` samples, recompute the full log-mel and encoder over that window, and take the second-to-last encoder frame (the last one contains right-edge padding).
2. **Transcript.** `rnnt.step(encoded)` with persistent `(last_token, lstm_state, tokens, text)`: loop up to `max_symbols = 10`, appending non-special pieces; return a delta when the token list grew.
3. **Language.** `fused = E(prev_text) + projected + 2 E(prev_function)` (`prev_* = 12` at timeline index 0); `h = llm.forward_embeds_to_hidden(fused, &mut language_caches)`; `text_id = argmax(lm_head(h))`, `function_id = argmax(function_head(h))` when `generate_channels` is true, else both `12` (prefill). With `use_language_cache = false` append `fused` to an input history and recompute from scratch every step (diagnostic mode).
4. **TTS.** At timeline index 0 emit `codec_silence_tokens` without running the TTS; otherwise `if text_id == 2 { previous_code = silence }`, `previous_code = tts.step(previous_code, text_id, &mut tts_caches)`. Timeline index increments here.
5. **Codec.** Only for frames with `generate_channels = true` (audio frames, not the prompt prefix): `clean = replace_control_codes(code)`; `samples = codec.decode_step(clean, &mut codec_cache)` where every ConvNeXt block prepends its cached `kernel - 1 = 6` latent frames (zeros on the first call) and re-caches the last 6, and the iSTFT stage prepends the cached `2 * ceil(6 / 4) = 4` spectrogram frames (real and imaginary separately), then, after the usual 6-sample window-edge trim on each side, drops `2 * 4 = 8` further samples from each end (445 spectrogram frames give 1792 raw samples, 1780 after the edge trim, 1764 after the overlap trim); the result is exactly 1764 sample-continuous samples. Error if the count differs.
6. Events are assembled with the current `frame_index`; `frame_index += 1`.
### Session lifecycle
- `new`: seed the RNG, run the TTS Aria warmup (from the model sub-issue), set `previous_code`, then prefill the system prompt: `ids = [1] + encode(prompt) + [2]`, and for each id run steps 3 and 4 with `audio_embedding = E(id)`, `generate_channels = false`, no codec (`timeline_index` advances; `frame_index` does not).
- `push_audio(samples, sample_rate)`: reject `sample_rate != 16000` and non-mono input; append to `pending`; while `pending.len() >= 1280` pop a frame and run the per-frame algorithm; raise `ContextLimit` when `frame_index >= max_frames` (`max_streaming_seconds / 0.08`) if a limit was set.
- `flush(pad_partial)`: if `pad_partial` and `pending` is non-empty, zero-pad to 1280 and run one more frame; clear the codec cache; mark closed; emit `Done`.
- `cancel()`: clear pending and the codec cache; mark closed; emit `Cancelled`. Pushing into a closed session is an error; `flush` / `cancel` on a closed session return nothing.
### Profiler
When `profile = true`, record per audio frame `perception_ms`, `rnnt_ms`, `language_ms`, `tts_ms`, `codec_ms`, `total_ms` (wall clock around each stage, after forcing evaluation of that stage's outputs). `summary(drop_first)` returns per-stage mean / p50 / p95 / max, `frames`, `processing_frames_per_second = 1000 / mean_total`, and `realtime_factor = mean_total / 80`.
## Implementation plan
1. **Streaming log-mel (`src/audio/nemotron_mel.rs`).** `StreamingLogMel { args, samples: Vec, buffer_start, total_samples, next_frame, lookahead_samples, lookbehind_frames, closed }` with `push(&[f32]) -> Vec /*[n, 128]*/` and `flush()`. Frame `i` is centered at sample `i * 160`; it is emitted once `total_samples - lookahead >= i * 160`; the buffer is trimmed to `keep_frame = next_frame - lookbehind_frames` hops. Test: `streaming_mel_joined_outputs_equal_offline` (push a 1 s signal in 300-sample chunks plus `flush`, compare to `log_mel_spectrogram` within 1e-5).
2. **Conformer streaming state (`src/audio/fastconformer.rs`).** `ConformerStreamingState::new(&FastConformerEncoder, chunk_frames = 1, att_context = [70, 0])` with `push(mel, emit_partial) -> Vec>`; `RelPosMultiHeadAttention::stream(q_in, kv_in, pos_emb)` and `RelPositionalEncoding::pos_emb_for(len)`; `FastConformerBlock::stream(...)` as in step 1 above; `materialize()` evaluating the cache slices at each frame so lazy graphs do not grow. Tests: `streamed_encoder_matches_offline_chunked_limited` (random weights, 64 mel frames, per-frame streaming vs one offline call within 1e-3), `push_of_eight_mel_frames_emits_one_encoder_frame`.
3. **RNNT state (`src/audio/rnnt.rs`).** `RnntStreamState { last_token, hidden: Option<(h, c)>, tokens, text }` with `step(&encoded) -> Option<(delta, text)>`.
4. **Codec streaming (`src/audio/nemotron_codec.rs`).** `CausalConv1dCache` keyed by block index and by `"istft_real"` / `"istft_imag"`, `update(states, key, padding, flush) -> padded` (prepend the cached tail or zeros, store the last `padding` frames, drop on flush); `NemotronCodec::decode_step(codes [1, 31, 1], &mut cache) -> [1, 1, 1764]`. Test: `decode_step_concatenation_equals_full_decode` (random weights, 6 frames, max abs diff < 1e-4; the final-frame tail is compared after `flush`).
5. **TTS state.** `RvqEarTtsModel::make_caches()` (from the model sub-issue) owned by the session; the Aria warmup runs in `new`.
6. **Session (`src/models/nemotron_voicechat/streaming.rs`).** `VoiceChatStreamingSession`, `VoiceChatEvent`, `StreamingOptions`, `VoiceChatProfile` / `FrameTiming`, `TokenAccumulator` (the delta rule), `VoiceChatError::ContextLimit`. `NemotronVoiceChatModel::create_streaming_session(&self, options) -> Result>`. Because the LLM state lives in `ModelOwnedSequenceState`, use a fresh `SequenceId` per session (`prepare_sequence_state`) and release it on close so two sessions created back to back on one model do not share state.
7. **CLI.** `mlxcel generate --audio in.wav --output-audio out.wav --stream` drives the session in 1280-sample pushes, prints deltas as they arrive, and with `--profile` prints the profiler summary (drop the first 5 cold frames). `--max-streaming-seconds` maps to `max_streaming_seconds`.
8. **Docs.** Extend `docs/nemotron-voicechat.md` with the session API, the cache-disable switches, and the measured real-time factor table.
## Validation
(a) Unit tests from steps 1 to 4, plus `src/models/nemotron_voicechat/streaming_tests.rs`: `push_audio_buffers_arbitrary_chunk_boundaries` (a fake model counting frames: 300 + 1000 + 1280 + 2000 samples gives 0, 1, 1, 1 frames and flush(pad) gives 1), `flush_without_padding_drops_partial_frame`, `closed_session_rejects_push`, `token_accumulator_delta_rule`, `profile_summary_percentiles`.
(b) Real checkpoint: `mlx-community/NemotronLabs-VoiceChat-11B-4bit` and `-8bit`.
```
./target/release/mlxcel generate -m models/NemotronLabs-VoiceChat-11B-4bit --audio /tmp/question.wav \
--output-audio /tmp/response_stream.wav -p "Be concise and answer in one sentence." --stream --profile
```
Acceptance: the first audio frame's text id, function id, and 31 codes are identical to the offline path for the same WAV and prompt (parity test `streaming_first_frame_matches_offline`, run with the real checkpoint behind an env-gated test); the streamed assistant text equals the offline text for a 5 s input; `response_stream.wav` has `frames * 1764` samples and is audibly the same answer; the profiler prints finite numbers and the 8-bit checkpoint's `realtime_factor` is recorded in the docs (a value above 1.0 is reported, not hidden).
## Acceptance criteria
- [ ] Streaming log-mel, Conformer, RNNT, TTS, and codec states advance one 80 ms frame per push with bounded memory (unit tests for each cache against the offline computation)
- [ ] `push_audio` / `flush` / `cancel` semantics and the event schema as specified, with frame-aligned indices
- [ ] System-prompt prefill advances the timeline without emitting events or audio
- [ ] First-frame parity with the offline path on the real checkpoint; full-utterance text equality on a 5 s input
- [ ] Profiler reports per-stage p50/p95 and the real-time factor
- [ ] docs/supported-models.md updated
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Dependencies
Blocked by #1374
- Part of epic #1372
Contributor guide
Research direction
Start by reading the streaming implementation plan and the thread-affinity rules in src/server/audio_worker.rs:15-36, then inspect src/audio/nemotron_mel.rs, src/audio/fastconformer.rs, src/audio/rnnt.rs, src/audio/nemotron_codec.rs, and src/models/nemotron_voicechat/session.rs. Run the named unit tests, streaming_tests.rs, and the CLI validation command. Done means the streaming session, caches, aligned events, profiler, CLI mode, and documentation meet the listed parity and continuity checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- audio-video-rtc, backend, cli, documentation, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100