nemotron_duplex: per-frame codec in the talker's step (full duplex 6/6)

configs/nemotron_duplex.yaml

@@ -1,18 +1,18 @@
 # NVIDIA NemotronLabs VoiceChat-11B — single-GPU full-duplex layout.
 #
-# All four stages (conformer_encoder, nano_llm, eartts_talker, audio_codec)
-# run on rank 0: the encoder + LLM decode loop, the talker loop, and the codec
-# all share one worker/GPU thread. This is the simplest working deployment and
+# All three nodes (conformer_encoder, nano_llm, eartts_talker, which decodes
+# its codes to PCM in the same step) run on rank 0: the encoder + LLM decode
+# loop and the talker loop share one worker/GPU thread. This is the simplest working deployment and
 # produces the same text + audio as the disaggregated layout
 # (configs/nemotron_duplex_disagg.yaml), which instead puts each
 # stream-consuming loop on its own GPU thread for pipeline parallelism.
 model: "nemotron_duplex"
 max_seq_len: 12288  # config: voicechat streaming_max_len / vllm_max_model_len
 
 # Sessions at once; the overflow queues. Each session holds a slot in
-# resources mamba_state, encoder_state, encoder_front and encoder_kv and its
-# pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise
-# the cap, raise their max_slots (cap + 1 for the recurrent pools, for the
+# resources mamba_state, encoder_state, encoder_front, codec_state and encoder_kv
+# and its pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to
+# raise the cap, raise their max_slots (cap + 1 for the recurrent pools, for the
 # sink; cap for encoder_kv) and the KV pools' max_num_pages with it.
 # The server refuses to start with a cap one of them cannot hold.
 max_concurrent_requests: 32
@@ -47,6 +47,5 @@ node_groups:
   - node_names:
       - conformer_encoder
       - eartts_talker
-      - audio_codec
     ranks:
       - 0

configs/nemotron_duplex_disagg.yaml

@@ -1,23 +1,24 @@
 # NVIDIA NemotronLabs VoiceChat-11B — disaggregated duplex layout.
 #
 # Each streaming consumer runs on its own GPU / worker thread so the
-# frame-synchronous decode loop, the talker loop, and the codec run in a
-# pipeline instead of serializing on one GPU thread. The colocated single-GPU
-# layout (configs/nemotron_duplex.yaml) also works and produces identical
-# output; this layout trades GPUs for pipeline parallelism / throughput.
+# frame-synchronous decode loop and the talker loop run in a pipeline instead
+# of serializing on one GPU thread. The colocated single-GPU layout
+# (configs/nemotron_duplex.yaml) also works and produces identical output; this
+# layout trades GPUs for pipeline parallelism / throughput.
 # Encoder + LLM share rank 0: encoder steps (one per chunk of the user's audio,
-# all live calls batched) interleave with the LLM's frame steps.
+# all live calls batched) interleave with the LLM's frame steps. The talker
+# decodes its codes to PCM in its own step (a few GFLOP per frame), so the
+# codec needs no GPU of its own.
 #
 #   rank 0: conformer_encoder (Encoder) + nano_llm (LLM)
-#   rank 1: eartts_talker     (Talker)
-#   rank 2: audio_codec       (Codec)
+#   rank 1: eartts_talker     (Talker, with the codec)
 model: "nemotron_duplex"
 max_seq_len: 12288
 
 # Sessions at once; the overflow queues. Each session holds a slot in
-# resources mamba_state, encoder_state, encoder_front and encoder_kv and its
-# pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise
-# the cap, raise their max_slots (cap + 1 for the recurrent pools, for the
+# resources mamba_state, encoder_state, encoder_front, codec_state and encoder_kv
+# and its pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to
+# raise the cap, raise their max_slots (cap + 1 for the recurrent pools, for the
 # sink; cap for encoder_kv) and the KV pools' max_num_pages with it.
 # The server refuses to start with a cap one of them cannot hold.
 max_concurrent_requests: 32
@@ -54,8 +55,3 @@ node_groups:
       - eartts_talker
     ranks:
       - 1
-
-  - node_names:
-      - audio_codec
-    ranks:
-      - 2

docs/models.rst

@@ -486,7 +486,7 @@ Fish Audio S2-Pro notes
 Nemotron VoiceChat (``nemotron_duplex``) notes
 ----------------------------------------------
 
-- Four nodes in four async partitions, joined by streaming edges, with the
+- Three nodes in three async partitions, joined by streaming edges, with the
   user's audio a client stream (``audio_in``) into the first:
   ``conformer_encoder`` (16 kHz audio → one LLM-space embedding per 80 ms frame,
   streamed: each chunk runs the layers over its new frames only, against the
@@ -495,12 +495,16 @@ Nemotron VoiceChat (``nemotron_duplex``) notes
   of input)
   → ``nano_llm`` (frame-synchronous decode loop: each step fuses one audio frame
   with the previous agent-text and tool-call tokens, emits one agent-text token)
-  → ``eartts_talker`` (one text token → 31 RVQ codes) → ``audio_codec`` (codes →
-  PCM: each 5-frame chunk is decoded with the previous 3 frames as left context,
-  which covers the decoder's receptive field, and a window's last 12 samples,
-  the inverse-STFT overlap, are held back and emitted with the next chunk, so
-  the stream equals the full-history decode sample for sample; that steady
-  window of 8 frames replays a CUDA graph, other windows decode eagerly).
+  → ``eartts_talker`` (one text token → 31 RVQ codes, decoded by the RVQ codec
+  in the same step to that frame's 80 ms of PCM, which goes straight to the
+  client: no codec step of its own, so a frame's audio leaves with the step
+  that made its codes. The decoder is causal, so each call keeps its stream
+  state in a recurrent-state pool, ``codec_state``: each ConvNeXt block's last
+  6 inputs and the inverse STFT's 12-sample overlap, ~95 KiB a slot in bf16. A
+  step decodes only the call's new frame against that state, and a fresh,
+  zeroed slot is the decoder's causal padding, so a call's PCM equals the full
+  decode of its frames from its first frame; at 32 sessions the decode is
+  ~100 GMAC per 80 ms tick, inside the talker's CUDA graph).
   The decode loops are stream-terminated: they end when the upstream stream
   closes, not on EOS, which is an ordinary per-frame token in duplex speech.
 - ``nano_llm`` declares a paged KV cache over its four attention layers with
@@ -512,7 +516,7 @@ Nemotron VoiceChat (``nemotron_duplex``) notes
   (``mamba``) and the agent-text sampler (``nano_sampler``). Its decode step is
   captured as a CUDA graph for batch sizes 1 to 32.
 - Every per-session resource is sized for 32 sessions: a slot in each of
-  ``mamba_state``, ``encoder_state`` and ``encoder_front`` (``max_slots``, one
+  ``mamba_state``, ``encoder_state``, ``encoder_front`` and ``codec_state`` (``max_slots``, one
   more than the sessions, for the sink) and in ``encoder_kv`` (``max_slots``,
   no sink), and its pages in ``nano_kv`` and
   ``talker_kv`` (``max_num_pages``; both sized from ``max_session_s``, and the
@@ -521,7 +525,7 @@ Nemotron VoiceChat (``nemotron_duplex``) notes
   all six with it. The server
   refuses to start with a cap one of them cannot hold, since a session past
   any of them would wait on admit, and its frames with it, until another ends.
-  The nano, the talker, the encoder and the codec each run at most 32 sessions
+  The nano, the talker and the encoder each run at most 32 sessions
   a step, their widest CUDA graph; more are split over several steps rather
   than run eagerly.
 - ``eartts_talker`` advances every live session in one backbone pass per frame
@@ -554,7 +558,11 @@ Nemotron VoiceChat (``nemotron_duplex``) notes
   time at 32 calls from 23.6 to 8.4 ms and cut live lag 2-6x at 16-32 calls
   (paced runs interleaved against the uncaptured encoder, ``MSTAR_SHM_ARENA=1``);
   32 calls staggered over 2 s all finish, heard latency p50 2.1-2.3 s (from the
-  user stopping to agent audio playable without gaps). Served audio is checked for intelligibility by
+  user stopping to agent audio playable without gaps). With the codec in the
+  talker's step, on one H200: 16 live calls keep real time whether they start
+  together or 2 s apart; 32 all finish but sit at the edge of real time (one run
+  kept up with 32 started together, another did not), so the shipped cap of 32
+  is an admission limit for live calls, not a real-time promise. Served audio is checked for intelligibility by
   transcribing it (``test/nemotron_duplex/asr_check.py``, Whisper
   large-v3-turbo): MaskGIT sampling is knife-edge, so attention-backend
   numerics change the waveform but not the words.
@@ -626,9 +634,9 @@ Nemotron VoiceChat (``nemotron_duplex``) notes
   lag; with ``--compare-one-shot`` it checks that a live call sending the whole
   clip as one append matches the upload save the last frame, and reports where
   the paced call's text first differs from it.
-- Default deployment: ``configs/nemotron_duplex.yaml`` (all four nodes on one
-  GPU); ``configs/nemotron_duplex_disagg.yaml`` puts the encoder+LLM, the talker
-  and the codec on three ranks. Launch with ``--tensor-comm-protocol SHM`` on a
+- Default deployment: ``configs/nemotron_duplex.yaml`` (all three nodes on one
+  GPU); ``configs/nemotron_duplex_disagg.yaml`` puts the encoder+LLM and the
+  talker (with its codec) on two ranks. Launch with ``--tensor-comm-protocol SHM`` on a
   single node. Text samples at ``temperature: 0.7, top_p: 0.9`` by default;
   send ``temperature: 0`` for greedy. The talker's mixture-of-Gaussians sampling
   is stochastic by design and seeded per request.

mstar/model/nemotron_duplex/components/audio_codec.py

@@ -23,11 +23,13 @@
 from __future__ import annotations
 
 import math
+from typing import NamedTuple
 
 import torch
 from torch import nn
 from torch.nn import functional as F
 
+from mstar.engine.resources.recurrent import RecurrentBlockConfig
 from mstar.model.nemotron_duplex.components._util import RawWeight, param
 from mstar.model.nemotron_duplex.components.dft import irfft_dim1
 from mstar.model.nemotron_duplex.config import CodecConfig
@@ -54,9 +56,12 @@ def __init__(self, dim: int, expand: int = 4, kernel: int = 7):
         self.pwconv1 = nn.Conv1d(dim, expand * dim, kernel_size=1)
         self.pwconv2 = nn.Conv1d(expand * dim, dim, kernel_size=1)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, tail: torch.Tensor | None = None) -> torch.Tensor:
+        """``tail``, for a stream, is the block's previous ``kernel - 1`` inputs
+        ``[B, C, kernel - 1]``; without one the block pads with zeros, as at the
+        stream's start."""
         residual = x
-        x = F.pad(x, [self.kernel - 1, 0])  # causal left pad
+        x = F.pad(x, [self.kernel - 1, 0]) if tail is None else torch.cat([tail.to(x.dtype), x], dim=-1)
         x = self.dwconv(x)
         # Channel-dim LayerNorm: normalize over C (== nn.LayerNorm over the
         # transposed last dim), eps 1e-6, biased variance.
@@ -84,11 +89,23 @@ def __init__(self, spec, kernel: int):
         super().__init__()
         self.layers, self._meta = _build_stack(spec, kernel)
 
-    def run(self, x: torch.Tensor) -> torch.Tensor:
-        """Apply the conv/block stack to a ``[B, C, T]`` tensor."""
+    def run(self, x: torch.Tensor, tails: list[torch.Tensor] | None = None) -> torch.Tensor:
+        """Apply the conv/block stack to a ``[B, C, T]`` tensor.
+
+        For a stream, ``tails`` holds each block's previous ``kernel - 1`` inputs,
+        in order, and the run replaces them with its own last ones. The resample
+        convs have kernel == stride, so the blocks hold all the state."""
+        block = 0
         for layer, (kind, stride, transpose) in zip(self.layers, self._meta, strict=True):
             if kind == "block":
-                x = layer(x)
+                if tails is None:
+                    x = layer(x)
+                    continue
+                tail, keep = tails[block], tails[block].shape[-1]
+                tails[block] = x[..., -keep:] if x.shape[-1] >= keep else (
+                    torch.cat([tail.to(x.dtype), x], dim=-1)[..., -keep:])
+                x = layer(x, tail)
+                block += 1
             elif transpose:
                 x = F.conv_transpose1d(x, layer.weight, stride=stride)
             else:
@@ -219,6 +236,74 @@ def check_window_envelope(window: torch.Tensor, hop_length: int) -> None:
                              "its inverse STFT cannot normalize")
 
 
+class CodecStreamState(NamedTuple):
+    """What a decode stream carries from one step to the next."""
+
+    # each ConvNeXt block's last ``kernel - 1`` inputs ``[B, C, kernel - 1]``, in decoder order
+    tails: list[torch.Tensor]
+    # ``[B, 2, n_fft - hop]`` float32: the overlap-added signal and squared window
+    # of the samples past the last final one, which the next frames still add to
+    overlap: torch.Tensor
+
+
+class CodecStreamAudio(NamedTuple):
+    """One decode step's samples, float32. ``end_before`` / ``end_after`` are the
+    stream's last ``stream_edge`` samples had it ended before / after this step's
+    frames: the full decode normalizes its right edge by the partial window sum."""
+
+    final: torch.Tensor        # [B, frames * wav_to_token_ratio]
+    end_before: torch.Tensor   # [B, stream_edge]
+    end_after: torch.Tensor    # [B, stream_edge]
+
+
+def codec_state_blocks(config: CodecConfig, dtype: torch.dtype | None = None) -> dict[str, RecurrentBlockConfig]:
+    """A decode stream's state as recurrent-pool blocks: per decoder stage, its
+    blocks' conv tails ``(num_blocks, channels, kernel - 1)`` (by default in the
+    engine's compute dtype, which the blocks compute in), and the inverse STFT's
+    overlap (``CodecStreamState.overlap``, float32)."""
+    blocks = {
+        f"stage{i}": RecurrentBlockConfig(shape=(config.num_blocks, ch, config.kernel_size - 1), dtype=dtype)
+        for i, ch in enumerate(reversed(config.channels))
+    }
+    blocks["overlap"] = RecurrentBlockConfig(shape=(2, config.n_fft - config.hop_length), dtype=torch.float32)
+    return blocks
+
+
+def _overlap_add_stream(
+    spec: torch.Tensor, overlap: torch.Tensor, n_fft: int, hop_length: int,
+) -> tuple[CodecStreamAudio, torch.Tensor]:
+    """``_spec_to_wav`` (window == n_fft, value range constrained, by matmul) for
+    a stream: ``spec``'s frames overlap-add onto the carried ``overlap``. The
+    samples no later frame reaches are final; the next ``n_fft - hop`` become
+    the new overlap."""
+    with torch.autocast(device_type=spec.device.type, enabled=False):
+        frames, held = spec.size(-1), n_fft - hop_length
+        size = frames * hop_length + held
+        window = torch.hann_window(n_fft, device=spec.device).unsqueeze(-1)
+        ifft = irfft_dim1(spec, n_fft)
+        ifft = torch.where(ifft >= 0, torch.minimum(ifft, window), torch.maximum(ifft, -window)) * window
+
+        def fold(x: torch.Tensor) -> torch.Tensor:
+            return F.fold(x, output_size=(1, size), kernel_size=(1, n_fft), stride=(1, hop_length))
+
+        carried = F.pad(overlap, (0, size - held))
+        sig = fold(ifft)[:, 0, 0] + carried[:, 0]
+        env = fold(window.square().expand(-1, frames)).reshape(1, size) + carried[:, 1]
+        out = _normalize(sig, env)
+        done, edge = frames * hop_length, held // 2
+        audio = CodecStreamAudio(
+            final=out[:, :done], end_before=_normalize(overlap[:, 0, :edge], overlap[:, 1, :edge]),
+            end_after=out[:, done:done + edge],
+        )
+        return audio, torch.stack([sig[:, done:], env[:, done:]], dim=1)
+
+
+def _normalize(sig: torch.Tensor, env: torch.Tensor) -> torch.Tensor:
+    """The overlap-added signal over its squared-window sum, 0 where that is ~0
+    (as ``_spec_to_wav``; a fresh stream's empty overlap)."""
+    return torch.where(env > _ENVELOPE_FLOOR, sig / env, torch.zeros_like(sig))
+
+
 class AudioCodec(nn.Module):
     def __init__(self, config: CodecConfig | None = None):
         super().__init__()
@@ -233,6 +318,11 @@ def __init__(self, config: CodecConfig | None = None):
         self.max_magnitude = cfg.max_magnitude
         # on the host: the model builds the codec under torch.device("meta")
         check_window_envelope(torch.hann_window(cfg.n_fft, device="cpu"), cfg.hop_length)
+        if (cfg.n_fft - cfg.hop_length) % 2:
+            # the stream trims half the overlap at each end, as the full decode does
+            raise ValueError(
+                f"the codec's stream needs an even overlap; n_fft {cfg.n_fft} - hop {cfg.hop_length} is odd"
+            )
 
     # ---- autoencoder halves -------------------------------------------------
     def ae_encode(self, x: torch.Tensor) -> torch.Tensor:
@@ -247,8 +337,16 @@ def ae_encode(self, x: torch.Tensor) -> torch.Tensor:
     def ae_decode(self, x: torch.Tensor, constrain_value_range: bool = True, by_matmul: bool = False) -> torch.Tensor:
         """Continuous latents ``[B, T', latent]`` -> waveform ``[B, 1, T]``; ``by_matmul``
         as ``_spec_to_wav`` takes it."""
-        x = x.transpose(-1, -2)
-        x = self.decoder.run(x)
+        spec = self._spectrum(self.decoder.run(x.transpose(-1, -2)))
+        with torch.autocast(device_type=spec.device.type, enabled=False):
+            return _spec_to_wav(
+                spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range,
+                by_matmul=by_matmul,
+            ).unsqueeze(1)
+
+    def _spectrum(self, x: torch.Tensor) -> torch.Tensor:
+        """The decoder stack's output ``[B, 2 * bins, frames]`` (log-magnitude,
+        phase) -> the complex one-sided spectrum ``[B, bins, frames]``, float32."""
         with torch.autocast(device_type=x.device.type, enabled=False):
             mag, ph = x.float().chunk(2, dim=1)
             mag = self.max_magnitude * torch.exp(-F.softplus(-mag + math.log(self.max_magnitude)))
@@ -259,13 +357,29 @@ def ae_decode(self, x: torch.Tensor, constrain_value_range: bool = True, by_matm
 
             spec_real = mag_mid * ph_mid
             spec_imag = mag_mid * ph_imag
-            spec = torch.cat([mag_dc * ph_dc, spec_real + 1j * spec_imag, mag_nyquist * ph_nyquist], 1)
-
-            x = _spec_to_wav(
-                spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range,
-                by_matmul=by_matmul,
-            ).unsqueeze(1)
-        return x
+            return torch.cat([mag_dc * ph_dc, spec_real + 1j * spec_imag, mag_nyquist * ph_nyquist], 1)
+
+    # ---- streaming decode ---------------------------------------------------
+    @property
+    def stream_edge(self) -> int:
+        """Samples the full decode trims off each end of its overlap-add: a
+        stream's first frame emits that many fewer, and its end that many more."""
+        return (self.n_fft - self.hop_length) // 2
+
+    def decode_stream(self, code: torch.Tensor, state: CodecStreamState) -> tuple[CodecStreamAudio, CodecStreamState]:
+        """Decode a stream's next code frames ``[B, T, num_quantizers]`` given its
+        ``state`` (zeros for a stream's first frames, which is the causal zero
+        padding of the full decode): the samples these frames make final, and
+        the state after them. Over a stream, ``audio.final`` (less
+        ``stream_edge`` samples at the first frame) followed by the last step's
+        ``audio.end_after`` is the full decode of all its frames.
+
+        Every step is the same fixed-shape work, so a CUDA graph captures it; the
+        inverse DFT is a matmul (see ``dft``)."""
+        tails = list(state.tails)
+        spec = self._spectrum(self.decoder.run(self.dequantize(code).transpose(-1, -2), tails))
+        audio, overlap = _overlap_add_stream(spec, state.overlap, self.n_fft, self.hop_length)
+        return audio, CodecStreamState(tails=tails, overlap=overlap)
 
     # ---- quantizer wrappers -------------------------------------------------
     def quantize(self, z: torch.Tensor) -> torch.Tensor: