Skip to content

NemotronLabs VoiceChat-11B PR 1: streaming consumers inside loops, stream-terminated partition end - #336

Merged
stephen-dwq merged 17 commits into
mainfrom
engine/streaming-loop-fixes
Oct 3, 2026
Merged

stephen-dwq merged 17 commits into
mainfrom
engine/streaming-loop-fixes

Conversation

@merceod

@merceod merceod commented Sep 28, 2026

Copy link
Copy Markdown
Collaborator

Five small graph/worker/conductor fixes that a stream-consuming node inside a Loop needs. Split out of the Nemotron VoiceChat work (#215 follow-up) so PersonaPlex can use them too.

  1. Loop re-arms pure-streaming inner nodes every iteration (mstar/graph/base.py). A node whose inputs are all streams has no loop-back edge, so nothing re-injected into it after the first chunk and the loop stalled after one item.
  2. The worker finds stream consumers nested in Loop / Sequential / Parallel (mstar/worker/worker.py). The consumer lookup only matched a walk that was itself a single node, so a decode loop consuming a streamed edge got no stream buffer and a KeyError when routing.
  3. Stream-terminated loops stop on the final chunk (worker.py). A stream-consuming loop has no internal stop signal, so when it consumes the terminal chunk its innermost loop is stopped. Not triggered under continue_after_done, so Qwen3-Omni's talker is unaffected.
  4. The conductor ends a stream-terminated partition on the worker's final-chunk signal (mstar/conductor/conductor.py). Before, producer_done for the next partition only went out when the model returned request_done, so models had to guess the end from the connection counters, which race (VoiceChat sessions lost their last audio frame). Now the worker's signal also sets request_done.
  5. A request cancelled after its batch was scheduled is left out of the batch instead of failing the whole batch (Worker._build_executing_batch).

Tests: test_streaming_loop_rearm.py, test_streaming_consumer_nodes.py, test_stream_partition_done.py, test_worker_batch_skips_dropped_requests.py. ruff check . clean, test/modular green. Verified live on the VoiceChat branch (single GPU, text identical to the offline oracle, real speech, 64 concurrent sessions).

@merceod merceod changed the title engine: streaming consumers inside loops, stream-terminated partition end NemotronLabs VoiceChat-11B PR 1: streaming consumers inside loops, stream-terminated partition end Sep 28, 2026

@NSagan271 NSagan271 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good, but this will have to be rebased and the logic implemented with the rust graph (lmk if you want me to look into that @merceod)

Comment thread mstar/worker/worker.py Resolved
Comment thread mstar/graph/base.py Resolved
…if their tensors cannot be found: when speculation is cancelled, chunks are returned to the stream buffer while still remaining ingested in the graph, so a chunk can be ingested, dereferenced, and then ingested again
…equests if their tensors cannot be found: when speculation is cancelled, chunks are returned to the stream buffer while still remaining ingested in the graph, so a chunk can be ingested, dereferenced, and then ingested again"

This reverts commit 1fa63e8.
The helper it named went away with the revert above; describe the bug class
instead, since the backstop is not specific to that one cause.
@NSagan271

Copy link
Copy Markdown
Collaborator

Updated PR, merging in the latest version of main. The original bugs fixed by this PR changed as follows:

  1. "Loop re-arms pure-streaming inner nodes every iteration": this was fixed during the rust transition already and dropped from this PR
  2. Building the streaming node cache so that streaming edges can be ingested in a loop: taken unchanged
  3. "The conductor ends a stream-terminated partition on the worker's final-chunk signal": taken unchanged
  4. "Stream-terminated loops stop on the final chunk": ported over; now a node -> innermost loop cache is built at initialization time, which the worker uses to terminate loops on the final chunk. The original implementation used NestedLoopIdxs, which is now inside the runtime. This new version is runtime-agnostic.
  5. "A request cancelled after its batch was scheduled is left out of the batch instead of failing the whole batch": at least on the latest version of main, it seems like a request cannot be cancelled after the batch is scheduled, so a KeyError can't happen on the request info. A KeyError check in retrieving the batch tensors is kept.

Also fixed a few more issues found:

  1. When speculation is rolled back or a rid is dropped, the streaming chunks are returned to the corresponding stream buffers without being un-ingested from the graph node. Now the flag-setting call at submit is folded into a new commit_speculation, which in the same crossing marks the batch speculatively scheduled and un-ingests the staged streaming chunks of rolled-back or dropped rids. set_speculatively_scheduled remains for clearing the flag.
  2. The new token counting happened after dereferencing tensors, so it could raise a KeyError if a tensor is freed. Now, the order of the two calls is swapped.
  3. The _mark_stream_ingested docstring had incorrect rationale; updated it

@NSagan271 NSagan271 left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merged main in, probably worth a second set of eyes before merging

@stephen-dwq
stephen-dwq merged commit 7970852 into main Oct 3, 2026
4 checks passed
@stephen-dwq
stephen-dwq deleted the engine/streaming-loop-fixes branch October 3, 2026 03:14
NSagan271 added a commit that referenced this pull request Oct 3, 2026
PR 286 landed as a SQUASH (33a70b5), so engine/dit-scaffold is not an
ancestor of main -- its commits and main's squash are separate history with
the same content. Sourcing the DiT changes from main rather than from the
branch, which also picks up two things the branch does not have:

- Review changes made to the scaffold before #286 merged. Notably
  `_carried_inputs` now RAISES when a loop-back edge is missing past
  iteration 0 instead of silently reseeding, with tests. A branch-sourced
  merge would have reverted that.
- #336 (streaming consumers inside loops) and #302 (Qwen3-TTS 1.7B), plus
  the rest of main.

Resolution rule: main is authoritative for every shared file, because the
branch inherited those at 4d5bf44 and never edited them -- verified by
diffing 4d5bf44..369017e, which touches only three of the conflicted
files. So denoise_loop.py, lora.py, graph/base.py, worker.py, client.py,
adding_models.rst and four test files were taken from main wholesale, and
the three the branch did change (docs/clients.rst, examples/sdk_image_edit.py,
adapters.py's registry) took main's base with the branch's additions
re-applied on top.

The registry-style lists (README, cli/main, registry.py, pyproject) are
unions; rerere replayed my earlier resolutions and each was checked against
main rather than trusted -- all four are main's content plus the flux /
z_image entries.

Ported flux2_klein and z_image onto the renamed scaffold API, which a clean
textual merge leaves silently broken (`cg_key_info` calls `bucket_key_for`,
so a subclass still defining `shape_key_for` raises NotImplementedError on
the first forward):

  shape_key_for -> bucket_key_for
  capture_shapes -> capture_buckets   (ctor kwarg and the models' own method)
  shape_key      -> bucket_key        (parameters and locals)

Scoped to the two model packages: `get_capture_shapes` in cuda_graph_config /
cuda_graph_runner is the unrelated piecewise API and keeps its name.

CI core list green (759 passed, 1 skipped). test/rust/ plus the diffusion,
flux2_klein, z_image and openai suites: 578 passed, 16 skipped (weights not
in the local HF cache).

Pre-existing and NOT from this merge: test_z_image_native_equivalence.py::
test_caption_encoder_tap_is_the_second_to_last_hidden_state fails a 1-ULP
`torch.equal` batch-invariance assertion identically at 369017e.

Co-Authored-By: Claude <noreply@anthropic.com>
stephen-dwq added a commit that referenced this pull request Oct 14, 2026
… (supersedes #215) (#339)

* add recurrent state resource, mirroring the existing paged kv resource

* add sink slot for backends that don't ignore padding rows

* Add GDN linear attention resource

* Add gated delta net model components

* add qwen 3.5 layer stack, verified against HF

* model integration works

* add TP (untested)

* bug fix and add 27b

* add gdn tensor parallel test

* bug fixes

* bug fix for conductor new tokens under TP

* optimize

* refactor outputs to allow for batched d2h in check stop

* benchmark change

* add more worker timing + batched output for qwen

* batched pos ids

* refactor rope

* bug fixes

* batch text inputs into vision prefill

* config changes

* Optimize vision encoder

* ruff

* remove benchmark files that shouldn't be tracked

* remove comments from configs

* trim docstrings

* fix benchmark runner

* respond to PR comments

* ruff

* Qwen3.5 / GDN follow-ups and fixes: stop-check row order, padding rows, TP alignment, sampler (#285)

* step: is_padding_row on StepContext

* recurrent pool: padding and capture rows take no slot

* cuda graph runner: keep padding rows' pages resident again

* BatchedModelOutput: carry the forward's row order

* engine: stamp row_request_ids on collected outputs

* worker: map stop-check rows through the forward order

* gdn: 32-byte gate pad, zero the pad block

* gdn tp: zero the rank's gate pad

* sampler: split softmax survives -inf leading blocks

* qwen3.5: image-only prompts, 400s for bad input, rgba, top_k

* deps: flashinfer >= 0.6.14

* gdn wrappers: pass the pool's pad index to the conv kernels

* gdn manager: null slot id for the conv kernels

* rope: declare force_double_buffer

* numa: only pin when nobody narrowed us

* norm: honest RMSNormGated comment

* qwen3.5: fix stale capture comment

* benchmark: fix Qwen3.5 docstring

* docs: qwen3.5 slot sizing

* test: padding rows take no pool slots

* test: row mapping through the forward order

* test: gate pad alignment

* test: pad accounting for the 32-byte alignment

* test: split softmax with -inf regions

* test: conv kernels honour the pad index

* test: numa policy cases

* test: qwen3.5 prompt edge cases

* test: skip vision parity without qwen3_5 in transformers

* test: skip the gdn graph test without CUDA

* worker: call the row mapper through the class

* test: fix the reference shape in the masked-rows test

* kv: padding rows run against SINK_PAGE and hold no pages

* cuda graph runner: padding rows hold nothing, so free per capture and per step

* attn wrappers: room for one sink index per padding row

* test: kv padding rows take no pages

* padding rows for the GDN prefill kernel: newer flashinfer can't run GDN on cuda 12.8, and the flashinfer 0.6.14 prefill kernel fails when given length-0 rows

---------

Co-authored-by: NSagan271 <nsagan@stanford.edu>

* files leftover from merge

* add fused all-reduce + add + rmsnorm

* Perf fixes: fewer cuda streams + remove work involving cuda kernels from preplan by using H2D from buffers

* env var doc change

* Fix weird cuda slowdown issues

* default fp32 gdn state to match vllm

* avoid two h2ds for shm and check stop; check stop batched; cache sampler buffers; improvements to planning

* bug fix + limit preprocess torch threads

* prefill vision cap + add warmup requests

* capture ViT encoder

* Clean up PR: remove useless tests and features

* gdn perf

* improved rmsnorm

* kv: a padding row takes a sink page for every page its span covers

A padding row spanning more than one page got one SINK_PAGE, so the plan
addressed its later tokens through the next entries of the indices buffer,
the previous row's pages, and its writes landed there.

* model: replay the qwen3.5 vit block loop once per prompt, one segment per image

Replaying per image re-planned the one ragged wrapper while the previous
image's schedule copy could still be queued, so an image could attend with the
next one's lengths.

* worker: key pinned d2h buffers by size class, not shape

* test: read the ragged wrapper's planned layout, not the removed _cu_host

* fused qkv

* recurrent: add Mamba2Geometry next to DeltaNetGeometry

* recurrent: export Mamba2Geometry

* linear_attn: add the MAMBA2 variant and its dt clamp

* test: Mamba2Geometry blocks round-trip, sharding, family check

* linear_attn: slot-indexed Mamba-2 state update kernel and scan

* linear_attn: Mamba2Manager planned against the recurrent pool

* linear_attn: build the MAMBA2 variant

* resources: Mamba2Callable for the layer body

* test: Mamba-2 resource plan, step-vs-scan, prefill continuity, Triton kernel

* Add nemotron_duplex model scaffold + Nemotron-H backbone

Port of NVIDIA NemotronLabs VoiceChat-11B (half-duplex S2S) into M*.
This first cut lands the package, the M* contract for the text path, and
the fully-implemented Nemotron-H "nano" LLM backbone.

- config.py: NanoConfig / EarTTSConfig / ConformerSTTConfig from the
  published per-stage config.json (hybrid_override_pattern -> 56 layers =
  27 Mamba-2 + 25 MLP + 4 attention).
- components/nemotron_h.py: hybrid layer dispatch, NoPE attention,
  squared-ReLU MLP, and a pure-PyTorch Mamba-2 mixer with per-request
  conv/ssm state (Option A: PerRequestState, no mamba_ssm dependency).
  Module tree mirrors the checkpoint; verified against the real
  model.safetensors header -> 342/342 nano params match (0 missing /
  unexpected / shape-mismatch).
- submodules.py: NemotronHLLMSubmodule (text path, eager Option-A);
  Conformer / EarTTS / codec nodes stubbed (Phase 4/5).
- nemotron_duplex_model.py: 8 contract methods for the text path
  (prefill_text -> decode), KV config for the 4 attention layers, duplex
  design documented for Phase 6.
- registry + CLI + configs/nemotron_duplex.yaml wiring.

Remaining: Conformer STT encoder, RNN-T, EarTTS Gemma3 talker, RVQ codec
(mapped from the checkpoint; to be built to load-parity next). Mamba-2
scan numerics still to be verified against the HF reference.

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

* Implement weight loading for all nemotron_duplex components

Port every stage of the composite VoiceChat-11B checkpoint to exact
param parity so the single model.safetensors (1632 tensors) loads with
zero missing / unexpected / shape-mismatched tensors.

New components (all load-verified against the real checkpoint header):
- conformer.py: Fast-Conformer perception (mel preproc, dw-striding
  subsampling, 24 rel-pos Conformer layers, LLM projection) — 640 tensors.
- rnnt.py: RNN-T prediction (2-layer LSTM) + joint network — 15 tensors.
- eartts_talker.py: Gemma3 talker (28 layers, QK-norm), subword encoder,
  gated audio/text fusion, MoG head, RVQ embeddings + per-voice prompt
  latents — 421 tensors.
- audio_codec.py: DAC/Vocos-style ConvNeXt RVQ codec (enc/dec heterogeneous
  strided-conv + block stacks, projected RVQ) — 214 tensors.
- _util.py: exact-shape param helper (int64 checkpoint buffers registered
  as non-grad params so the params-only loader fills them) + RawWeight.

Each component exposes a static `remap` (checkpoint key -> param path) used
by both load_weights and the verification harness. get_submodule now builds
and loads every node (meta -> materialize -> load_weights).

Verification (meta-build vs safetensors header, no 44GB download): all six
stages PASS; whole checkpoint 1632/1632 tensors claimed, 0 uncovered, 0
dupes. Total params 11.095B (nano 9.48B), matching the advertised 11B.

Forward passes for the audio stages remain Phase 4/5 (need the NeMo/EarTTS
reference to verify numerics); this commit is loading only.

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

* Add offline forward for nano Nemotron-H backbone (+ 2 numeric fixes)

Implement engine-free full-sequence forward passes for offline inference,
validated against the reference NemotronHForCausalLM (NVIDIA-Nemotron-Nano-9B-v2
modeling_nemotron_h.py, torch_forward slow path):

- NemotronHForCausalLM.offline_forward(input_ids) -> logits: embed -> hybrid
  blocks -> final norm -> lm_head, no M* paged-cache dependency.
- NemotronHBlock.offline_forward / NemotronHLLM.offline_forward dispatch per
  layer kind; NemotronHAttention.offline_forward does eager causal SDPA (GQA,
  NoPE).
- The Mamba-2 recurrence (_ssd_scan) matches the reference; two fixes:
  * gated RMSNorm is now GROUPED (group_size = intermediate/n_groups = 1280),
    matching MambaRMSNormGated — was normalizing over the full 10240.
  * block/final norms use a manual fp32 RMSNorm on the offline path; the shared
    FlashInfer RMSNorm computes in bf16 regardless of input dtype (verified
    ~0.05 error on F32 input), which compounded across 56 layers.

Pre-fix validation already showed 100% next-token argmax agreement vs the HF
reference (structure correct); the norm fixes address the logit-magnitude
divergence (final re-validation in progress).

Only the nano backbone forward is implemented so far. Conformer / RNN-T /
EarTTS talker / codec forwards are still load-only (Phase 4/5).

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

* Implement + verify forward passes for all 5 modules; codec/rnnt config-driven

Forward passes ported from the NVIDIA-NeMo/Speech `nemotron-labs-voicechat`
reference and numerically verified against it (fp32, same checkpoint weights
loaded into both M* and reference):

- RVQ codec (audio_codec.py): STFT/iSTFT, causal ConvNeXt stacks, projected-RVQ.
  decode cosine 0.99999994; ae_encode 0.99999994; encode 100% in-distribution.
- RNN-T (rnnt.py): LSTM predictor + joint (broadcast-add -> joint_net, raw
  logits). predict/joint bit-exact (cosine 1.000000).
- Conformer perception (conformer.py): mel featurizer, causal dw-striding
  subsampling (factor 8), rel-pos MHA, causal conv module (conv_norm_type=
  layer_norm), Macaron 0.5, att_context [70,0]. cosine 1.00000000 vs NeMo.
- EarTTS talker (eartts_talker.py): Gemma3 backbone (QK-norm pre-rope, query
  scaling 256**-0.5, rope theta 1e6/1e4 by layer), gated audio/text fusion, MoG
  head, T5Gemma subword encoder. hidden/mog cosine > 0.9999.
- nano Nemotron-H (nemotron_h.py): offline_forward (embed -> hybrid blocks ->
  norm -> lm_head), grouped gated RMSNorm, fp32 offline norms. Mamba-2 group->
  head mapping uses contiguous/interleave (the mamba_ssm kernel + reference
  decode convention); the reference *naive prefill* uses tile (.repeat), an HF
  fallback quirk we intentionally do not copy. Kernel confirmation pending.

Config: codec + rnnt now read every dim from the checkpoint config.json via
NemotronDuplexConfig.from_pretrained (CodecConfig derives the conv/block stacks
from channel_mult/rates/num_blocks; RnntConfig from _rnnt_merge_info) instead of
module-level constants. Verified: config-built modules load 214/9/6 params and
forward at correct shapes. Conformer/talker/nano config migration still to come.

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

* Make Conformer + talker config-driven; nano grouping confirmed vs mamba kernel

Move all module-level architecture hyperparameters into the config dataclasses,
populated from the checkpoint's config.json.

- config.py: EarTTSConfig now holds the full Gemma3 backbone + MoG-head config
  (from tts_config.backbone_config / mog_head_config); ConformerSTTConfig holds
  the encoder + mel-preprocessor config (from perception.encoder / preprocessor,
  with derived head_dim/ff_dim/win_length/hop_length/subsampling_flat).
  NemotronDuplexConfig.from_pretrained now populates codec, rnnt, eartts and stt
  from config.json.
- conformer.py / eartts_talker.py: removed every architecture global (D, N_HEADS,
  FF, H, HEAD_DIM, rope thetas, MoG dims, ...); all read from config, threaded
  through submodules. proj_mus / subsampling flatten / stage counts are derived,
  not hardcoded. Only pure numerical constants (attention-mask fill, log epsilon,
  vocab-table cardinalities) remain as documented literals.
- get_submodule threads the config.json-sourced config into every component.

Re-verified after the refactor (numerics unchanged): Conformer encoder cosine
1.00000000 vs NeMo, full stack 0.99999982; talker hidden 0.9999999, mog_* >0.9999.
All config-built components load 214/640/9/6/421 params, 0 missing.

nano grouping: confirmed the Mamba-2 contiguous/interleave group->head mapping
against the real mamba_ssm kernel (mamba_chunk_scan_combined) -> cosine 1.000000,
settling that the reference naive-prefill .repeat (tile) is the HF fallback quirk.

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

* Phase A: standalone offline_inference orchestration (STT half + audio wiring)

Add NemotronDuplexModel.offline_inference — a standalone duplex S2S loop that
mirrors the reference NemotronVoiceChat.offline_inference and bypasses the M*
serving engine (walk-graph/engine integration is Phase B).

Pipeline (bs=1, fp32, greedy): perception -> per-frame [AddFusion(prev agent
text, audio, prev function) -> nano non-cache re-prefill step -> lm_head/
function_head -> argmax] -> talker RVQ codes -> codec waveform. The reference
itself uses Nemotron no-cache re-prefill, so this reuses the verified
offline_forward directly (no nano decode-with-cache / Mamba-state-in-engine).

- config.py: AddFusion is parameter-free (confirmed: 0 fusion weights in the
  checkpoint); channel weights (agent 1.0 / audio 1.0 / function 2.0) and the
  use_function_head / predict_user_text flags now read from config.json
  (model.stt.model). Text special-token ids added (TODO verify vs nano tokenizer).

STT half (audio -> agent text/function tokens) is fully implemented from the
verified perception + nano forwards. The talker autoregressive sampling loop
(EarTTSTalker.infer_codes_one_step) is being finalised; _decode_audio_path wires
it + the verified codec.decode and no-ops until it lands.

Reference for e2e verification is up: an isolated nemoref conda env (torch 2.11
+ nemo_toolkit + the Speech nemotron-labs-voicechat repo) imports NemotronVoiceChat.

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

* Talker: autoregressive RVQ code generation (KV cache + MoG sampler + MaskGIT)

Implement the offline audio-generation loop on EarTTSTalker, ported from NeMo
RVQEARTTSModel.generate_step / MoGHead.infer / depthsum_encoding_step and
DuplexEARTTS.infer_codes_one_step (branch nemotron-labs-voicechat):

- KV cache: _Attn/TalkerLayer/backbone_forward take an optional per-layer (k,v)
  cache + absolute start_pos; rope applied at absolute positions; causal mask
  offset for new-vs-cached. Non-cached path unchanged (earlier forward still verified).
- MoG sampler: MogHead.infer(x, temperature) — greedy argmax (or Gumbel-max) over
  mixture logits; low-rank mean up-projected by low_mat, mean = mu*exp(logs)+mu_res.
- RVQ: _rvq_quantize (nearest-entry residual) + generate_step (MaskGIT schedule,
  num_iter=8, greedy/deterministic default; CFG omitted).
- AR wrappers: init_state (warms KV from the pre-baked speaker latent + bos),
  initial_prev_codes (codec silence), infer_codes_one_step (one frame -> [B,31] codes).

Verified (fp32): KV-cache single-step vs full recompute cos 1.0000000; MoG infer
reconstructed from the verified head max-diff 0.0; end-to-end 6-frame run ->
codes [2,6,31] in [0,1024), greedy-deterministic. Earlier numerics unchanged.

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

* Talker: text conditioning (char-vocab) + generate_codes — audio path verified

Close the audio-out conditioning gap. context_hidden_size=None, so text
conditioning is subword_ids -> embed_subword (char path) + subword-flag +
bos/eos embeddings. Ports NeMo build_vocabs + the conditioning path:

- build_char_vocab / subword_to_char_ids: single-char tokenizer tokens densely
  re-indexed; each subword -> its in-vocab char ids (pad subword -> pad char).
- text_conditioning(subword_ids, ...) -> cond [B,T,H]: embed_subword (char
  encoder + mean-pool + proj) + continuation-flag + BOS/EOS embeddings.
- generate_codes(subword_stream, tokenizer, speaker, temperature): full AR loop
  (init_state warmup -> per-frame cond -> infer_codes_one_step greedy, EOS ->
  codec silence) -> codes [B, T, 31].

Verified against the LIVE NeMo reference (Gumbel-disabled deterministic, since
the reference MoG sampler is otherwise stochastic and can't reproduce its own
waveform): char vocab all 131073 maps exact; per-frame cond cosine 1.000 (26/26);
sampler 100% code agreement; seeded end-to-end 100% codes + decoded audio
cos 1.0000 / max-diff 0.0. Self-contained generate_codes is 92.7% codes / 0.94
audio, gap isolated to the frame-0 init_state warmup (prompt-state handling) —
a known follow-up; everything after frame 0 is exact.

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

* Wire full offline_inference audio path (tokenizer + talker generate_codes)

Complete Phase A: M*.offline_inference now runs the whole duplex S2S pipeline
standalone (audio in -> agent text -> RVQ codes -> waveform).

- tokenizer: source the nano text tokenizer from the Spark nano/ folder (base
  repo ships only rnnt_tokenizer/); 131072 vocab, 256 single-char tokens -> the
  talker char vocab. Added a small adapter exposing the NeMo tokenizer interface
  (.tokenizer.vocab) the talker's char-vocab builder expects.
- _decode_audio_path: use talker.generate_codes(text_stream, tokenizer,
  text_eos_id) -> verified codec.decode.

Standalone run (audio in -> audio out) vs the reference: gen_text 100% match,
tokens_audio [1,26,31] in-range, audio [1,45864] (26 frames x 1764), text
detokenizes. STT half bit-exact; audio path exact modulo the isolated frame-0
init_state warmup (fix in progress).

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

* A.5: text sampling (temp/top-p/rep-penalty) + batched offline inference

- _sample_text_token: exact port of the reference sampler (greedy at
  temperature==0; else repetition/presence penalty -> temperature -> top-p ->
  multinomial, special-id argmax bypass). Wired into the text channel;
  function/ASR channels stay greedy.
- offline_inference is now batched: a per-utterance loop (_infer_one) + _collate
  that pads/stacks tokens_text/function/audio across B. Correct B>1 results;
  true continuous batching is the Phase-B engine path.

Verified: bs=1 greedy still 100% vs reference; a batch of 2 different-length
utterances -> each row 100% identical to its single-run and to the reference;
sampling (temp 0.8 / top_p 0.9 / rep-pen 1.3) runs. Detokenized text returned.

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

* B.1: nano cached O(T) decode (attention KV + Mamba conv/ssm state)

Additive cached-decode path on the Nemotron-H nano (offline_forward untouched),
the "caching of Mamba layers" the streaming path needs:

- NemotronHCache: per-layer attn[i]=[k,v] (NoPE -> position-agnostic, no rope
  reconcile on append), conv[i]=(conv_dim, conv_kernel-1) fixed rolling buffer,
  ssm[i]=(nheads, head_dim, d_state); MLP layers stateless.
- prefill(input_embeds) / decode_step(input_embed, cache) threaded through
  attention (prefill + decode_forward, is_causal=False single query), Mamba2
  (prefill + single-step recurrence), block, LLM, CausalLM.
- Conv-state rolling fix: fixed conv_kernel-1 buffer, left-zero-padded for the
  first token; decode does an explicit windowed dot product matching the padded
  conv1d at the new position.

Verified vs offline_forward (L=24, fp32): pure-decode + prefill/decode splits
all give hidden cosine > 0.99998 and lm_head argmax 24/24 identical; prefilled
positions are bit-exact. Residual is fp32 reduction-order noise.

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

* A.5/frame-0: init_state warmup reproduces the reference exactly (audio bit-exact)

Close the last self-contained-audio gap. The old init_state warmup diverged from
the reference at the prompt boundary; dumping the live reference's
set_init_inputs('Aria') -> get_init_inputs -> warmup showed:
- the BOS frame is embed_code(depthsum(codec_silence)) + bos_emb (shifted last
  silence code), not latent[36];
- the last two warmup positions are conditioned on the shifted prompt text
  [pad*36, eos] via embed_subword (was zero cond);
- the first prev_codes is the speech-pad frame (all codebook_size -> zero
  embedding), not codec-silence;
- no codec encode is needed (code[:,35] == codec_silence exactly).

init_state now replicates the P=37 warmup (latent[0..P-2] + built BOS frame,
warmup text cond, returns speech-pad prev_codes); takes tokenizer ids + char
maps, falls back to the simplified warmup if omitted. generate_codes threads them.

Verified vs the live deterministic reference: init KV pos=37, max|K|diff 5.9e-3
(fp32), prev_codes exact; self-contained generate_codes 100% codes, exact-frame
100%, decoded audio cos 1.00000 / max-diff 0.0 (was 92.68% / 0.94).

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

* B.2 + C: online streaming duplex inference + OpenAI Realtime API

B.2 (online, DuplexStream on the model): audio fed in chunks via append_audio;
each new ~80ms perception frame runs one duplex step (fuse -> nano cached-decode
step [B.1] -> sample text/function -> talker frame -> codec) yielding agent
audio+text deltas. nano attention-KV + Mamba conv/ssm state and the talker KV
state persist across the stream (O(1)/frame). Perception is re-encoded over a
growing buffer (correct because the Conformer is causal); codec decodes the full
code history and emits the new tail (multi-frame receptive field). Talker warmup
uses the exact frame-0 init_state args. prepare_streaming/create_stream on the
model; offline _decode_audio_path also now passes the full warmup ids.

Verified: streaming the same wav in 160ms chunks -> gen_text 100% IDENTICAL to
offline_inference (26/26 frames), one audio chunk per frame. online == offline.

C (realtime_api.py): OpenAI Realtime-API-compatible WebSocket server (FastAPI) —
session.update / input_audio_buffer.append|commit -> response.audio.delta +
response.audio_transcript.delta, continuous full-duplex mode, one DuplexStream
per connection (so N connections = N independent streams = batched online).

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

* feat(nemotron-duplex): system-prompt support in streaming + Realtime API

Prime the nano cache with the system prompt before any audio, mirroring how
offline_inference prepends prompt_emb to the audio channel and holds gen_text at
PAD across the prompt region. DuplexStream._prime_prompt runs one cached
prefill/decode step per prompt token; prev_text is now derived from global
gen_text position so the first audio frame after a primed prompt uses the last
prompt PAD (not BOS).

- create_stream(prompt_tokens=... | system_prompt=...); encode_prompt() tokenizes
  [bos]+ids+[eos] like the reference encode_system_prompt.
- Realtime API: session.update `instructions` -> primed system prompt; re-prime on
  input_audio_buffer.clear / response.cancel; warn if changed mid-turn.

Verified: streaming-with-prompt == offline-with-prompt 100%; Realtime
instructions path == offline-with-prompt 100%; prompt changes the token stream;
no-prompt path unchanged (100%).

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

* fix(nemotron-duplex): real-speech audio (CFG+noise+top-p) + streaming lookahead

Generating outputs on real speech exposed two bugs the near-silent reference
clip had masked:

1. Talker produced silence codes for real text. The MoG head was missing
   classifier-free guidance, noise injection, and top-p mixture sampling. Add
   them per the reference _get_generation_config (guidance_scale=0.5,
   noise_scale=0.8, top_p=0.8, num_iter=8), config-driven on EarTTSConfig:
   - MogHead.infer: split [cond;uncond] after the trunk, x_cond+gs*(x_cond-x_uncond),
     top-p filter, Gumbel-max component sampling.
   - generate_step: dual cond/uncond hidden, per-iter guidance, z=mu+exp(logs)*eps*ns.
   - init_state / infer_codes_one_step: maintain a second (null_emb) KV cache for
     the unconditional stream.
   Result: codes go from ~3 unique/codebook (silence) to ~500 total; agent audio
   peak 0.0002 -> 0.22 (real speech). Now stochastic by design.

2. Streaming re-encode diverged from offline on real speech (text 78%). The
   Fast-Conformer subsampling has 1 frame of right-lookahead, so the newest frame
   per chunk was not final. Hold back streaming_lookahead_frames (=1) during
   append_audio; consume them at flush. Text parity 78% -> ~98-100%, audio now
   bit-exact between offline and streaming perception.

Also fix the module docstring (the model is full-duplex, not half-duplex) and
refresh the stale phasing note.

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

* fix(nemotron-duplex): seed talker sampling — stop agent-audio cutoff

The MoG/MaskGIT talker sampler is stochastic (noise + Gumbel). The offline and
streaming audio paths drew it from the uncontrolled global RNG, whose state after
the nano/perception ops occasionally lands on a collapse-to-constant-code attractor:
the talker emits one identical RVQ code frame for a whole utterance, which the codec
renders as silence — so the agent's speech cut off after the first utterance.

Diagnosed: with a dedicated seeded generator, 0/12 seeds collapse (independent of
CFG); only the unseeded path collapses. The reference does torch.manual_seed(0)
before generation for exactly this reason.

Fix: thread a dedicated seeded torch.Generator (config `inference_seed`, default 0)
through _decode_audio_path and DuplexStream's per-frame talker step. Verified: every
multi-frame utterance now vocalizes (turn_taking peaks [.,0.157,.,0.141];
interruptions [.,0.162,.,0.235,.]), longest stuck-code run 0–1, audio reproducible.

Also drops ENGINE_INTEGRATION.md (design notes for the in-progress engine work).

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

* feat(nemotron-duplex): E1 — batched Mamba state for true batched engine inference

Make the hybrid backbone's Mamba conv/ssm state batchable so B>1 requests advance
in a single fused engine forward (the core of true batched inference; attention
already batches via the paged-KV pool — only Mamba state needed it).

- MambaStateAccessor: per-request read/write by batch position + `seq_lens` (how the
  packed [sum(seq_lens), H] rows split into requests). Each request keeps its own
  Option-A conv/ssm slot; a CUDA-graph-capturable fixed-buffer pool (Option B) can
  later replace the dict without touching the mixer.
- Mamba2Mixer.forward dispatches: mamba_state=None -> single-sequence (offline/bs=1,
  unchanged); all-decode batch (1 tok/req) -> _decode_batched (one fused SSD step over
  B); varlen prefill -> per-request segments. Extracted _forward_seq (per-request
  full-seq math) + _decode_batched (batched rolling conv + batched single SSD step,
  contiguous group->head expansion matching _ssd_scan).

Verified: batched forward() decode == standalone prefill+decode_step per request
(max|Δ| 8.5e-4 over 4 reqs x 6 steps); offline inference unchanged (100% vs reference).

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

* feat(nemotron-duplex): E1 — batch the nano_llm engine node (forward_batched)

Wire the batched Mamba into the engine's nano_llm node so B>1 requests advance in
one fused forward instead of the sequential per-request fallback:

- can_batch -> True for multi-request batches (eager path; attention batches via the
  paged-KV pool, Mamba via the per-request MambaStateAccessor).
- preprocess now threads per-request seq_lens; _mamba_state passes them so the mixer
  segments the packed rows correctly (all-decode -> fused batch step; varlen prefill
  -> per-request segments).
- forward_batched runs the packed batch once and reads each request's next-token
  logits at its last packed position, returning {request_id: {"logits": [...]}}.

Single-request path unchanged (can_batch False for bs=1). Batched Mamba correctness
verified at the mixer level (prior commit); full serving-stack E2E is the E5 milestone.

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

* feat(nemotron-duplex): E2 — implement audio_codec engine forward

AudioCodecDecoderSubmodule.forward (STATELESS, audio_codec flavor): decode the
streamed RVQ code window to 22.05 kHz int16 PCM. Under a LeftContextChunkPolicy the
input carries `left_context` prior frames + the new `chunk`; decode the whole window
(the codec is causal / has a multi-frame receptive field) and emit only the new tail
(n_new frames), mirroring the standalone DuplexStream "decode history, emit tail".

Verified bit-exact vs the standalone codec.decode (max|Δ int16| = 0) and the
tail-emit windowing (20-frame window, n_new=5 -> 5*wav_to_token_ratio samples).

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

* feat(nemotron-duplex): E2 — implement conformer_encoder engine forward

ConformerEncoderSubmodule.forward (STATELESS): raw 16 kHz waveform chunk ->
per-frame LLM-space embeddings (T, H) via the Fast-Conformer perception (mel +
causal subsampling + encoder). These are the user-audio channel; the nano decode
node fuses them per frame with its own fed-back previous text/function tokens
(AddFusion) before stepping — that frame-synchronous fusion is the E4 walk work.

Verified bit-exact vs the standalone perception (max|Δ| = 0.0). RNN-T user-transcript
emission is a follow-up (predict_user_text is off by default).

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

* feat(nemotron-duplex): E3 — implement eartts_talker engine forward

EarTTSTalkerSubmodule: agent text token -> num_quantizers RVQ codes per frame. The
MoG/MaskGIT/CFG sampling is internal (not the engine's categorical sampler), and the
talker keeps its own KV + CFG-uncond state per request in PerRequestState (eager,
like the nano's Mamba state). Warms up from the speaker prompt on the first frame,
then steps via the verified infer_codes_one_step with a per-request seeded generator;
forward + forward_batched carry each request's state. _create_talker_submodule now
passes the char-vocab maps.

Verified bit-exact: per-frame _step codes == standalone generate_codes (0/806
mismatched). All four node forwards (nano_llm batched, conformer_encoder, eartts_talker,
audio_codec) now implemented + verified vs the standalone oracle; remaining engine work
is the walk graphs / async partitions (E4, incl. frame-synchronous AddFusion) and the
full serving-stack E2E check (E5).

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

* feat(nemotron-duplex): E4 — define the full-duplex engine walk graphs

Define everything the M* engine needs to serve the model as four async partitions
(frame-synchronous S2S), replacing the text-only placeholder walks:

  Encoder --[audio_frame, FixedChunk(1)]--> LLM
  LLM     --[new_token,   FixedChunk(1)]--> Talker
  Talker  --[codec_tokens, LeftContext ]--> Codec

- get_graph_walk_graphs: encode / prefill_text / decode / talker_decode / codec_chunk.
  The nano decode loop is frame-driven: each iter consumes one streamed audio frame,
  fuses it with the fed-back previous agent-text + function tokens (AddFusion), emits
  the agent text token (client + streamed to talker) and feeds both back.
- get_partitions (Encoder->LLM->Talker->Codec) + get_partition_topology (chunk policies)
  + get_aux_sampling_configs (nano "function" channel).
- nano submodule: _fuse_frame does the per-frame AddFusion in prepare_inputs; the decode
  forward samples the function token via the aux sampler; postprocess feeds prev_text /
  prev_func back into the loop. New config fields codec_chunk_frames / left_context_frames.

Verified structurally: get_worker_graphs() derives worker graphs for all five walks
(resolves every edge route + cross-partition streaming connection), partitions/topology/
aux/engine-types all well-formed. End-to-end batched serving through the worker stack
(driving the partitions via forward-pass-args) is the remaining E5 milestone.

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

* feat(nemotron-duplex): E4 — multi-partition forward-pass-args driver + audio prompt

Drive the four async partitions at runtime (mirrors qwen3_omni):
- get_initial_forward_pass_args routes each partition to its initial walk: Encoder->
  encode (audio_features), LLM->prefill_text if a system prompt is present else
  straight to decode, Talker->talker_decode, Codec->codec_chunk (self-triggered from
  their upstream streams; request_done when audio isn't requested).
- get_partition_forward_pass_args: Encoder runs once; LLM transitions prefill_text->
  decode and feeds prev_text/prev_func back each frame, finishing when the audio stream
  is exhausted; Talker/Codec finish when their upstream stream is done (_stream_exhausted).
- process_prompt now accepts raw user audio (audio_features) + optional system prompt
  instead of raising on audio input.

Verified structurally: get_worker_graphs derives all five walks, and per-partition
initial routing + the prefill->decode transition check out. This completes the engine
*definitions*; the end-to-end batched run through the worker stack is E5.

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

* test(nemotron-duplex): structural engine tests for the duplex walk graphs

test/modular/test_nemotron_duplex_model.py (8 tests, no weights/GPU): the four-
partition walks, engine types, producer chain, topology routing, nano function aux
channel, worker-graph derivation (resolves every edge + streaming connection),
per-partition initial forward-pass-args routing, and the LLM prefill->decode
transition. Mirrors test_qwen3_tts_model; makes the engine definitions CI-checked.

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

* wip(nemotron-duplex): E5 live-serving fixes + trace logging (mid-debug)

Progress from a live `mstar-serve` run of the full stack (conductor + worker + GPU)
with the model on cuda:0 — all four engines load, all five walks are served, and
requests execute through the pipeline. Fixes found live:

- load_audio: decode uploaded audio via soundfile (base uses torchcodec, whose native
  libs aren't available here).
- process_prompt: map the data worker's `audio_inputs` key to `audio_features` (the
  conformer_encoder node input).
- conformer_encoder.forward: emit under `audio_frame` (the LLM decode edge name) — was
  `audio_embeds`, so the LLM decode loop never received frames.

Known-remaining (this WIP): a request with audio input is accepted and routed but the
frame-synchronous Encoder->LLM streaming loop hangs (no crash) — the deepest runtime
piece. NDTRACE logging added to encoder.forward / nano.prepare_inputs /
get_initial+partition_forward_pass_args to localize it on the next run.

Note: serve with `--tensor-comm-protocol SHM` (single node, no RDMA/InfiniBand) and
PYTHONPATH set to this worktree so `mstar-serve` picks up this code.

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

* docs(nemotron-duplex): E5 resume/migration notes for handoff

Full engine-integration status (E1-E4 done+verified, E5 live-serving in progress),
the exact serve/request repro (PYTHONPATH + SHM transport gotchas), and a decision
tree for the open Encoder->LLM streaming-loop hang keyed off the NDTRACE logging.

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

* docs(nemotron-duplex): pin the E5 hang — LLM decode lacks consumer-gating

Static diagnosis vs qwen3_omni's Thinker->Talker (proven stream-consuming decode):
the LLM decode loop is the initial walk with no prefill/trigger gating, so on a single
worker it can start on an empty audio_frame StreamBuffer and deadlock the Encoder.
E5_RESUME.md now has the concrete fix plan (add an LLM consumer-gating walk mirroring
talker_prefill; terminate via the worker-side Loop's check_stop, not _stream_exhausted;
or fold the encoder into a prefill_audio Sequential).

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

* fix(nemotron-duplex): seed frame-0 prev_text/prev_func for the decode loop (UNVERIFIED)

Static root-cause of the E5 hang: the decode-loop node requires prev_text+prev_func in
its input_names, but frame 0 of an audio-only request has no prior sampled token, so the
node never becomes ready -> hang (qwen3's talker_prefill seeds the analogous first-decode
input). Fix: process_prompt emits prev_text=[BOS]/prev_func=[PAD]; get_initial_forward_pass_args
(LLM audio-only->decode) provides them as frame-0 inputs; iters >=1 use the loop-back edges.

UNVERIFIED — no GPU here (reclaimed by user). Structural tests (8) + ruff pass. Must be
run on a GPU (ptc) to confirm; see E5_RESUME.md. NDTRACE logging still present.

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

* feat(nemotron-duplex): E5 — full-duplex serving works E2E (text + audio)

The frame-synchronous Encoder→LLM→Talker→Codec pipeline (two chained
stream-consuming loops) now runs through the M* engine, matching the
verified offline_inference oracle on text and producing real agent speech.

Framework (generic — first model to exercise chained streaming into loops):
- worker: build the streaming consumer-node cache and per-partition node
  set by recursing via section.get_nodes(), so a consumer nested inside a
  Loop (nano_llm in decode_loop, eartts_talker in talker_decode_loop) is
  found — else streamed chunks get next_node="" / the StreamBuffer is never
  created.
- worker: stream-terminated loops — when a stream-consuming loop node
  consumes the terminal chunk (final_stream_rids), stop its enclosing loop
  (get_nested_loop_idxs_for_node). Prior loops self-stopped on EOS; a
  stream-driven loop had no stop signal and spun to max_iters.
- graph: re-arm a loop's pure-streaming inner nodes (no loop-back) into
  ready_for_streaming on each iteration (Loop._advance_one_iter →
  WorkerGraphStateRegistry.rearm_streaming); loop-back nodes are re-armed
  implicitly via re-injection. Without this the talker/codec stream loops
  stalled after one item.

Model (nemotron_duplex):
- encoder emits T per-frame stream items (was one (T,H) blob) so the LLM's
  FixedChunkPolicy(1) can pace them; runs eager (torch.compile inductor
  PyKernelArg error on variable audio-length shapes).
- seed prev_text(BOS)/prev_func(PAD) for decode iteration 0 so the
  frame-synchronous node is ready before any loop-back exists.
- NemotronHBlock.forward routes mamba_state only to mamba mixers (attention
  / MLP inherit the base Attention.forward signature — no mamba_state).
- check_stop no longer halts on EOS (a normal per-frame token in duplex);
  the decode loop ends on audio-frame stream exhaustion.
- talker prepare_inputs reads the streamed "new_token"; Encoder→LLM and
  LLM→Talker connections use continue_after_done=False (loops are
  stream-terminated, 1:1 frames).
- codec is stateful: keep a per-request rolling left-context of prior code
  frames, decode context+new, emit only the new frames (matches DuplexStream
  math); Talker→Codec uses a non-overlapping FixedChunkPolicy. Fixes audio
  ballooning (~window/chunk×) from re-emitted overlap.
- _fuse_frame tolerates the empty terminal audio_frame chunk (producer_done
  race) by running a no-audio step.

Add configs/nemotron_duplex_disagg.yaml (Encoder+LLM / Talker / Codec on
their own ranks) — each stream-consuming loop on its own worker thread.

Validated (disagg, VoiceChat-11B): text token-identical to offline_inference;
3 concurrent requests fuse into batched forwards; tiny clip audio == oracle
(37044 samples), 12s clip = 12.08s real speech (was 44.9s pre-codec-fix).

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

* docs(nemotron-duplex): single-GPU config is the full working duplex layout

The single-GPU layout (all four stages on rank 0) runs the full duplex
pipeline and produces identical text + audio to the disaggregated layout —
the earlier colocation "deadlock" was the pure-streaming-loop stall fixed in
E5, not a single-thread limit. Update both config comments accordingly
(single-GPU = simplest working deployment; disagg = pipeline parallelism).

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

* test(nemotron-duplex): unit + e2e tests for the full-duplex serving path

CPU unit tests (CI, no GPU / weights) — regression coverage for the E5 fixes:
- test_nemotron_duplex_model.py: + stream-connection chunk policies
  (1-per-step, continue_after_done=False; codec = non-overlapping
  codec_chunk_frames), frame-0 prev_text/prev_func seeding in process_prompt,
  and seeded initial decode inputs.
- test_nemotron_duplex_submodules.py (new): nano check_stop does NOT stop on
  EOS (a normal per-frame token) but does at max_tokens; _fuse_frame tolerates
  the empty terminal audio_frame chunk; the codec is stateful (per-request
  left-context, emits only the new frames, cleanup clears state, requests
  isolated).
- test_streaming_loop_rearm.py (new, framework): a Loop whose inner node's
  inputs are all streaming is re-armed into ready_for_streaming each iteration
  (verified to fail without Loop._advance_one_iter's rearm_streaming call).

End-to-end serving scripts (GPU + real weights) under test/nemotron_duplex/,
mirroring test/orpheus/: _env.py, .sample.env, launch_server.sh (single-GPU or
`disagg`), duplex_request.py (audio in -> agent text + speech WAV, -n for
concurrent batching), oracle_compare.py (offline_inference ground truth), README.

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

* test(nemotron-duplex): default duplex_request.py audio to a bundled sample

--audio is now optional; with none given, resolve a sample input
(turn_taking.wav / interruptions.wav / tool_call.wav) from the base
VoiceChat-11B HF repo via hf_hub_download (honoring NEMOTRON_DUPLEX_CACHE_DIR),
so `python duplex_request.py` works out of the box.

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

* test(nemotron-duplex): default input is a prepared user-only clip that gets a reply

The VoiceChat-11B demo wavs (turn_taking.wav, ...) are 2-channel recordings of
the WHOLE conversation (user left, agent right) meant for listening — feeding
one as-is lets the model hear its own side and it just stays silent. Instead,
default_audio() now takes the user (left) channel, isolates the first user turn,
and appends trailing silence (this is a frame-synchronous model: it replies in
the frames AFTER the user stops), then resamples to 16 kHz mono and caches it.

On the default clip the agent replies coherently ("Oh, I love cooking! How about
this chocolate chip cookie recipe? ..."). README updated with the user-audio
guidance (user-only + trailing silence).

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

* test(nemotron-duplex): add duplex/streaming request test

duplex_stream_request.py sends a user utterance + a window of no-sound frames
(the "keep talking after the user stops" case) and consumes the reply as it
streams back via /generate stream=True, reporting the turn boundary (start of
reply, when text+audio return to silence) and how much of the reply was voiced.

Findings it surfaces (verified against the server):
- with a longer --silence the agent emits MORE reply text — it doesn't wait for
  the user, it fills silence with plausible follow-up turns;
- the reply AUDIO is sparse (talker/codec voices ~1.5 s per reply of the dense
  text), matching the standalone offline_inference reference — a model/reference
  characteristic, not a serving bug. The engine has no incremental-audio-input
  path (fixed-buffer input, streamed output); that lives in realtime_api.py.

Factor the user-only input prep into _audio_prep.py (extract left/user channel,
isolate first turn, append trailing silence, 16 kHz mono) shared by both
duplex_request.py (6 s reply pad) and duplex_stream_request.py (configurable).

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

* test(nemotron-duplex): add a browser web UI (webui.py)

A tiny local web app to talk to the model from a browser. It's a small proxy:
the page records mic audio (Web Audio -> 16-bit WAV) or takes an uploaded wav and
posts it to /chat; the proxy resamples to 16 kHz mono, appends the silence reply
window, forwards to the model server's /generate with streaming, and streams the
agent's text + audio back over SSE for the page to display and play.

    python test/nemotron_duplex/webui.py --mstar-url http://127.0.0.1:8019 --port 8500
    # open http://127.0.0.1:8500/

Browser <-> proxy is same-origin (no CORS); proxy <-> model server is a plain
server-to-server call via MStarClient. getUserMedia needs a secure context
(localhost / 127.0.0.1, or HTTPS for remote). Verified end to end: /chat streams
back coherent agent text + audio.

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

* fix(nemotron-duplex): correct text-channel special token ids (eos=2, pad=12)

The text-channel special-token ids were wrong: text_eos_id=12 and
text_pad_id=0. Per the checkpoint config.json (model.speech_generation.model)
the tokenizer has bos_token=<s> (1), eos_token=</s> (2), and
pad_token=<SPECIAL_12> (12).

The eos mistake mattered most: <SPECIAL_12> (id 12) is the frequent
between-word / turn-taking "pad-pair" frame, not end-of-utterance. With
text_eos_id=12, `inference_force_speech_silence_on_eos` forced the talker
silent on the pad frames that fall *between* articulated words, so the spoken
reply collapsed after ~4 words. With eos=2 the talker only silences on the
true </s>, and the reply is voiced through the sentence (ASR-verified: old
"oh I love" ~0.2s -> now "Oh, I love cooking. How about this chocolate..."
~2.6s voiced). text_pad_id=12 also fixes the prev_func frame-0 seed and the
talker's char-warmup padding.

Short-clip engine output stays token-identical to the offline oracle; the
19 modular tests read the ids from config so they still pass; ruff clean.

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

* docs(nemotron-duplex): reply audio is real speech (token-id fix), not "sparse limitation"

The truncated/sparse-audio caveat in the test README and web UI was written
before the text_eos_id/text_pad_id fix (bc096dfc). With correct ids the agent
voices intelligible multi-sentence speech (long-context: coherent multi-turn
reply, ~13 s voiced, ASR-verified). Downgrade the caveat to the real residual:
very dense late sentences of a long reply can still thin out.

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

* Move E5 resume notes out of the model directory

* Move engine integration notes out of the model directory

* Drop the test README (folds into docs/models.rst)

* Add nano resource keys to the duplex config

* Bind Nemotron-H attention to the nano KV and attention resources

* Port duplex submodules to declared steps and resource sampling

* Declare nano resources and per-request sampling on the duplex model

* Keep the duplex package import lazy

* Describe the duplex components package

* Describe nemotron_duplex as full duplex in the CLI defaults

* Test duplex node resources and request configs

* Test nano step declaration, prompt priming and codec context

* Document the nemotron_duplex model family

* Drop oracle_compare.py in favour of parity_compare.py

* Add the standalone-vs-engine text parity script

* Point the duplex docs at parity_compare.py

* Add the Mamba-2 pool and resource keys to the duplex config

* Step Nemotron-H Mamba-2 layers through the recurrent pool

* Declare the Mamba-2 pool and resource steps on the nano node

* Declare the nano's recurrent-state pool and Mamba-2 resource

* Test the nano's Mamba-2 pool geometry and resource specs

* Test the nano step declares the Mamba-2 pool and resource

* Capture the nano decode step as a CUDA graph

* Test the nano decode capture config

* Add the nemotron_duplex benchmark package

* Benchmark concurrent duplex sessions: tick time, on-time ticks, response latency

* Talker: batched per-frame step over sessions with unequal caches

* Advance every live talker session in one backbone pass

* Test the batched talker step against the sequential one

* Decode equal-length codec windows of several sessions in one call

* Test codec batching across sessions

* Document the nano's recurrent pool, decode capture and the batched talker

* End stream-driven duplex partitions on the worker's final-chunk signal

* Add the talker KV, attention and position resource keys

* Talker: attention through the KV pool, warm-up inputs, host-side MaskGIT schedule, pre-drawn noise

* Serve the talker from the KV pool: two combined streams, warm-up prefill, captured decode

* Declare the talker's paged KV, attention and positions

* Test the talker resource specs and request labels

* Test the pooled talker: schedule, noise, warm-up, step declaration, packing

* Apply talker RoPE in torch from planned positions

* Lay out talker noise rows-first for graph replay

* Assert rows-first talker noise layout

* Declare the talker's KV head dim as FlashInfer's 128

* Size the talker KV pool by the padded head dim

* Zero-pad the talker's q/k/v to the KV head dim

* Pin the talker KV head dim to the padded size

* Test the padded pooled attention against the plain attention

* Keep warm-up steps off the talker's fixed-shape decode capture

* Test that warm-up batches run eager and steady-state batches replay

* Document the pooled talker

* Compute the talker warm-up under the serving autocast; zero the unused noise slots

* Size the talker KV pool by 128-position pages and reserve Mamba slots for padding rows

* Pin the Mamba pool size with the padding-row reserve

* Document the talker KV pool and Mamba padding-slot sizing

* Fuse frame-step batches and draw talker noise in two launches per row

* Test that the batched frame fusion equals the per-request fusion

* Fuse frame batches whose rows have no audio frame with a zero term

* Test the batched fusion with a row past the end of its audio

* Move each frame-step input to the device before fusing the batch

* Document the measured multi-session capacity and the ASR check

* Drop the Mamba padding-slot reserve: padding rows take no slot since #258's follow-up

* Pin the Mamba pool size to the session count

* Document the Mamba pool as session slots only

* Let the duplex demo client prime a system prompt

* test: a dropped request is left out of the batch, the others run

* Add the Whisper transcription check for served duplex audio

* Point the docs at the repo's transcription check

* Document the comparison with vLLM-Omni's native duplex serving

* Seed the first fed-back agent token as PAD after a system prompt

* Test the PAD seed after a system prompt

* Carry the system prompt into the LLM's first decode step instead of a prefill walk

* Fuse the prompt tokens ahead of the first frame and gate that step off the decode capture

* Test the single-walk prompt routing and seeds

* Test the prompt-plus-frame step, its capture gate and the sentinel

* Hand the nano and talker outputs over as batch rows with nothing for the stop check

* Test the row-addressed nano outputs and the codec's stacked-row chunks

* Keep the talker off speculative scheduling: its next step reads the codes its previous step stored

* Test that the talker node is not scheduled speculatively

* Cut the codec's left context to the decoder's receptive field

* Carry the codec's right edge into the next chunk and lift the LLM's prompt gate after its first step

* Test the exact chunked codec stream and the prompt gate lifting

* Document the exact chunked codec and the re-measured capacity

* Shorten the codec context comment

* Shorten the talker scheduling comment

* Tighten the nano and talker comments

* Drop the worker batch test, the fix and its tests landed with #336

* Declare the KV caches with PagedKVConfig

* broken from merge; fixed

* disable autocast in rvq, was causing auditory hallucinations during silent periods; fp32 insetad of bf16

* talker did not actually use guidance, noise and top-p from config

* we cannot index empty final chunks

* pad short clips to prevent index error on token stream ending with empty chunk

* reject audio shorter than a single encoder frame

* concurrency cap to avoid stalling

* final chunk handling for codec

* cap session length; we cannot serve more than this concurrently without oom stall

* stopgap to prevent long hangs; can be revisted once chunked prefill is in

* drop concurrency cap to safely faster-than-real-time; bump up session cap

* use model-suggested sampling params

* input validations

---------

Co-authored-by: NSagan271 <nsagan@stanford.edu>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Gaurav Shah <gshah2@andrew.cmu.edu>
Co-authored-by: Keisuke Kamahori <kamahori@cs.washington.edu>
Co-authored-by: Stephen Duan <s129652@cs.washington.edu>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants