mstar.model.nemotron_duplex.submodules

NodeSubmodules for NVIDIA NemotronLabs VoiceChat-11B (full duplex).

Nodes, and the resources each declares (NemotronDuplexModel.get_node_resources):

conformer_encoder (none) 16 kHz speech -> per-frame LLM embeds (+ RNN-T) nano_llm nano_kv, nano_attn, Nemotron-H hybrid Mamba-2 / attention / MLP (9B):

mamba_state, mamba, paged KV for the 4 attention layers, recurrent-pool nano_sampler slots for the 27 Mamba-2 layers, the text sampler

eartts_talker talker_kv, talker_attn, Gemma3 talker -> 31 RVQ codes per frame; two KV

talker_pos labels per session (cond / uncond) in one plan

audio_codec (none) RVQ codes -> 22.05 kHz PCM, per-request left context

Both autoregressive nodes keep their whole state in engine resources and capture their decode step as a CUDA graph; the encoder and the codec run eager.

Classes

AudioCodecDecoderSubmodule(codec, config)	RVQ codec decoder: talker codes -> 22.05 kHz int16 PCM (streaming).
ConformerEncoderSubmodule(perception, ...)	Fast-Conformer streaming STT encoder: raw 16 kHz audio -> per-frame LLM-space embeddings, emitted one frame per stream item so the LLM's FixedChunkPolicy(1) paces the frame-synchronous decode loop.
EarTTSTalkerSubmodule(talker, config[, ...])	Gemma3 talker: one agent text token -> num_quantizers RVQ codes per session per frame, all live sessions in one backbone pass.
NemotronHLLMSubmodule(language_model, config)	Nemotron-H backbone node.

class mstar.model.nemotron_duplex.submodules.AudioCodecDecoderSubmodule(codec, config)

Bases: NodeSubmodule

RVQ codec decoder: talker codes -> 22.05 kHz int16 PCM (streaming).

The codec is causal with a multi-frame receptive field, so decoding each chunk in isolation clicks at the boundaries. Each request keeps a rolling context of its previous codec_left_context_frames code frames (in its PerRequestState); a chunk of NEW frames is decoded as context + new and only the new frames’ samples are emitted. Same “decode with left context, emit the tail” math as the verified standalone path, but O(1) per chunk instead of re-decoding the whole history. The inverse STFT’s last n_fft - hop samples of a window still lack the following frames’ overlap, so a chunk holds those samples back and the next chunk emits them recomputed with its context: the emitted stream is the full-history decode sample for sample; the final chunk emits its held tail too. Requests whose windows have the same length are decoded in one batched call. Declares no resources.

Parameters:

• codec (nn.Module)

• config (NemotronDuplexConfig)

CONTEXT_KEY = 'codec_ctx'

TAIL_KEY = 'codec_tail'

can_batch(batch, model_inputs)

Return type:

bool

disable_autocast: bool = True

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, codes=None, flush=None, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, codes=None, flush=None, **kwargs)

Batched form of forward: a multi-request batch’s outputs, per request or as a BatchedModelOutput. Override when can_batch returns True.

Return type:

dict[str, dict[str, list[Tensor]]]

prepare_inputs(graph_walk, fwd_info, inputs, is_final_stream_chunk=False, **kwargs)

Return type:

NodeInputs | None

preprocess(graph_walk, engine_inputs, inputs)

class mstar.model.nemotron_duplex.submodules.ConformerEncoderSubmodule(perception, rnnt_decoder, rnnt_joint, config)

Bases: NodeSubmodule

Fast-Conformer streaming STT encoder: raw 16 kHz audio -> per-frame LLM-space embeddings, emitted one frame per stream item so the LLM’s FixedChunkPolicy(1) paces the frame-synchronous decode loop.

Holds the perception stack plus the RNN-T decoder/joint (loaded; the user-transcript channel is not emitted yet). Declares no resources.

Parameters:

• perception (nn.Module)

• rnnt_decoder (nn.Module)

• rnnt_joint (nn.Module)

• config (NemotronDuplexConfig)

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, wav=None, lens=None, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

NodeInputs

class mstar.model.nemotron_duplex.submodules.EarTTSTalkerSubmodule(talker, config, subword_to_char=None, char_pad_idx=None)

Bases: ARNodeSubmodule

Gemma3 talker: one agent text token -> num_quantizers RVQ codes per session per frame, all live sessions in one backbone pass.

The talker’s KV lives in the engine’s paged KV pool under two labels per request: main (text-conditioned stream) and uncond (null-conditioned stream for classifier-free guidance), combined into one attention plan whose rows are [cond rows; uncond rows] – the order mog_head.infer expects. A session’s first step prefills the 37 speaker warm-up positions together with its first frame (38 tokens per label); every later step is one token per label. The MoG / MaskGIT sampling is internal (not the engine’s categorical sampler); its noise is drawn per row, with the row’s own seeded generator, in preprocess – so a session’s speech does not depend on its batch mates and the whole step is a CUDA graph replay for all-decode batches.

Parameters:

• talker (nn.Module)

• config (NemotronDuplexConfig)

• subword_to_char (dict | None)

• char_pad_idx (int | None)

DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32]

DECODE_KEY = 'decode'

GEN_KEY = 'talker_gen'

PREV_CODES_KEY = 'talker_prev_codes'

SPEAKER = 'Aria'

can_batch(batch, model_inputs)

Return type:

bool

cg_key_info(graph_walk, per_request_info, **kwargs)

Which capture serves this batch: the steady-state decode capture (DECODE_KEY) when every session is past its first step, None (run eager) when any session prefills its warm-up this step. The engine leases the slot before prepare_inputs runs, so this answers from the request state, like declare_step does.

Parameters:

• graph_walk (str)

• per_request_info (Mapping[str, Any])

• kwargs (Any)

Return type:

str | None

declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs)

Declare this batch’s step for the runner to drive: which cache streams it touches, what spans they grow by, which plans back it, which streams fork, and what commits when it lands. The runner drives the declaration before preprocess and commits it after the forward, so a declaring submodule keeps no plan or advance calls of its own. None means the submodule still plans and advances through the facade itself.

request_ids pairs positionally with inputs. Under a captured graph the batch is padded to the bucket’s shape, so it carries the padding rows’ ids too — declare their segments like any other row.

slot_lease is the slot this step will replay on, or None for an eager step. A submodule whose declaration differs between the two (cosmos3 packs both guidance branches into one plan for the captured shape) must key off this, not off its own capture key: the key says the batch could be captured, the lease says it was.

piecewise_leases names the regions of this node that hold a slot for this step. Such a region declares, plans and commits its own work, so a resource it owns must be left out of this declaration.

per_request_info holds the real rows’ forward-pass info (capture passes dummy metadata); padding rows are absent, so read it with .get.

Parameters:

• graph_walk (str)

• request_ids (list[str])

• inputs (list[ARNodeInputs])

• slot_lease (SlotLease | None)

• piecewise_leases (Mapping[str, SlotLease] | None)

Return type:

SubmoduleStep

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, x=None, spans=None, noise_u=None, noise_eps=None, **kwargs)

Batched form of forward: a multi-request batch’s outputs, per request or as a BatchedModelOutput. Override when can_batch returns True.

Return type:

BatchedModelOutput

get_accelerator_graph_configs(device, tp_world_size=1)

Capture the steady-state step (one token per label per session). preprocess (text conditioning, noise, warm-up assembly) runs eagerly before the graph; the captured region is the 2N-row backbone over the pool and the MaskGIT sampling loop.

Parameters:

• device (device)

• tp_world_size (int)

Return type:

list[AcceleratorGraphConfig]

postprocess(request_id, request_info, outputs, **kwargs)

Per-request postprocessing on the submodule outputs.

Runs on the GPU thread inside execute_batch, after the forward (eager, batched, or CUDA-graph replay). inputs is the request’s prepare_inputs result for this step, so a submodule can finish a step the captured graph could not hold (e.g. combine guidance branches and run a Python multistep scheduler against the step’s input latents).

Keep it metadata-only where possible. Avoid reading tensor values — .item() / .cpu() sync here block the GPU thread and forfeit the worker’s async-scheduling overlap; stop-condition decisions that need token values (e.g. EOS) belong in check_stop. A captured-path tail that must read scheduler state is the sanctioned exception — it costs the same sync wherever it runs.

Typical uses:

• rebind output names for graph routing (outputs["text_inputs"] = outputs["new_token"]);

• drop keys on a per-request basis for static-capture submodules (e.g. Qwen3-Omni Thinker dropping thinker_states for requests that don’t need audio);

• finish a captured step from inputs (Cosmos3 denoise tail).

Modifies outputs in-place; returns nothing.

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)

warmup()

(cond [P, H], uncond [P, H], prev_codes [Q]) of the speaker warm-up. Computed once per speaker, under the serving autocast (bf16 on CUDA) no matter who asks first: the CUDA-graph capture’s first preprocess runs outside the engine’s autocast scope, and a warm-up computed there differs in the last bits from one computed under it – enough to move a session’s first frame (the talker’s frame-0 decision is a knife-edge).

Return type:

tuple[Tensor, Tensor, Tensor]

property warmup_len: int

class mstar.model.nemotron_duplex.submodules.NemotronHLLMSubmodule(language_model, config)

Bases: ARNodeSubmodule

Nemotron-H backbone node.

prefill_text primes the request with the system prompt exactly like the reference: every prompt token is fused with the agent channel (BOS on the first token, PAD afterwards) and a PAD function token, nothing is sampled, and the fed-back prev_text / prev_func leave the prompt region at PAD. decode is frame-synchronous: each step fuses one streamed audio frame with the previous agent-text and function tokens (AddFusion), then samples the agent text (sampler resource) and the tool-call token (greedy).

Parameters:

• language_model (nn.Module)

• config (NemotronDuplexConfig)

DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32]

DECODE_KEY = 'decode'

PROMPT_DONE_KEY = 'prompt_consumed'

can_batch(batch, model_inputs)

Return type:

bool

cg_key_info(graph_walk, per_request_info, **kwargs)

Decode capture unless a request still has its prompt pending (its first step is prompt + frame and runs eager). The routing sets that flag once for the whole loop, so the node also checks whether the request’s first step has already run.

Parameters:

• graph_walk (str)

• per_request_info (Mapping[str, Any])

• kwargs (Any)

Return type:

str | None

declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs)

Declare this batch’s step for the runner to drive: which cache streams it touches, what spans they grow by, which plans back it, which streams fork, and what commits when it lands. The runner drives the declaration before preprocess and commits it after the forward, so a declaring submodule keeps no plan or advance calls of its own. None means the submodule still plans and advances through the facade itself.

request_ids pairs positionally with inputs. Under a captured graph the batch is padded to the bucket’s shape, so it carries the padding rows’ ids too — declare their segments like any other row.

slot_lease is the slot this step will replay on, or None for an eager step. A submodule whose declaration differs between the two (cosmos3 packs both guidance branches into one plan for the captured shape) must key off this, not off its own capture key: the key says the batch could be captured, the lease says it was.

piecewise_leases names the regions of this node that hold a slot for this step. Such a region declares, plans and commits its own work, so a resource it owns must be left out of this declaration.

per_request_info holds the real rows’ forward-pass info (capture passes dummy metadata); padding rows are absent, so read it with .get.

Parameters:

• graph_walk (str)

• request_ids (list[str])

• inputs (list[ARNodeInputs])

• slot_lease (SlotLease | None)

• piecewise_leases (Mapping[str, SlotLease] | None)

Return type:

SubmoduleStep

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, input_embeds, seq_lens, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

• input_embeds (Tensor)

• seq_lens (list[int])

Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, input_embeds, seq_lens, **kwargs)

One fused forward over the packed batch: attention over the planned paged KV, Mamba-2 over each request’s slot in the recurrent pool.

The outputs are row-addressed (row i is request i): the engine copies each batch tensor out of the graph’s buffers once and hands every request a view, instead of one clone per request per name. Nothing is offered to check_stop (it would be copied to the host per step): the decode loop ends on the audio stream’s final chunk or the loop’s own max_iters, never on a token.

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

• input_embeds (Tensor)

• seq_lens (list[int])

Return type:

BatchedModelOutput

get_accelerator_graph_configs(device, tp_world_size=1)

Capture the frame-synchronous decode step.

prepare_inputs is host-only and preprocess (embedding + AddFusion) runs eagerly before the graph on both capture and replay; the captured region is forward_batched: the 56-layer backbone over the paged KV and the recurrent pool, the two heads and the sampler. The dummy row is one fused frame: a zero audio frame plus the BOS / PAD carry-ins.

Parameters:

• device (device)

• tp_world_size (int)

Return type:

list[AcceleratorGraphConfig]

postprocess(request_id, request_info, outputs, **kwargs)

Per-request postprocessing on the submodule outputs.

Runs on the GPU thread inside execute_batch, after the forward (eager, batched, or CUDA-graph replay). inputs is the request’s prepare_inputs result for this step, so a submodule can finish a step the captured graph could not hold (e.g. combine guidance branches and run a Python multistep scheduler against the step’s input latents).

Keep it metadata-only where possible. Avoid reading tensor values — .item() / .cpu() sync here block the GPU thread and forfeit the worker’s async-scheduling overlap; stop-condition decisions that need token values (e.g. EOS) belong in check_stop. A captured-path tail that must read scheduler state is the sanctioned exception — it costs the same sync wherever it runs.

Typical uses:

• rebind output names for graph routing (outputs["text_inputs"] = outputs["new_token"]);

• drop keys on a per-request basis for static-capture submodules (e.g. Qwen3-Omni Thinker dropping thinker_states for requests that don’t need audio);

• finish a captured step from inputs (Cosmos3 denoise tail).

Modifies outputs in-place; returns nothing.

Parameters:

• request_id (str)

• request_info (CurrentForwardPassInfo)

• outputs (dict[str, list[Tensor]])

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Host-side bookkeeping only; the embedding lookups and AddFusion run in preprocess on the GPU thread.

Parameters:

• graph_walk (str)

• fwd_info (CurrentForwardPassInfo)

• inputs (dict[str, list[Tensor]])

Return type:

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

• inputs (list[ARNodeInputs])

Return type:

dict[str, Tensor | Any]