Add speaker-tagged transcription to Nemotron ASR

[LysanderdeJong]

Summary

• add speaker-tagged transcription to nemotron_asr, driven by a nemotron_3_diar speaker-activity file: NeMo's masked multitalker ASR (default; one cache-aware stream per active speaker, so overlapping speech is transcribed per speaker) and word attribution (one transcript, cheaper)
• output speaker turns with text, a SegLST artifact, and, in streaming, finished turns as events

Third of three stacked PRs split from the original #681; #685 (spec-backed nemotron_asr options) and #686 (diarizer speaker_probabilities.safetensors and asr_laN profiles) are merged.

There is no pipeline or composite family: the diarizer and the ASR run as two commands, documented in docs/asr.md.

With NeMo's own diarizer output as input, masked mode reproduces NeMo's SegLST segment for segment on CPU F32: 100% on two full AMI meetings (14 and 17.5 min) and on every tested lookahead.

Changes and rationale

1. Attribution mode

One ASR pass; each word takes the speaker with the highest mean activity over a window around it (NeMo Python get_word_dict_content_offline: 80 ms frames [first - 1, last + 1), clamped at 0.01, lowest index wins ties). Words are formed from tokens (a leading space starts a word). Offline runs it after decoding; streaming attributes each word once the next word starts.

Evidence: on the AMI 20 s clip all 42 words get the same speaker as NeMo's function on the same diarizer file.

2. Masked mode

A port of SpeakerTaggedASR.perform_parallel_streaming_stt_spk with masked_asr=true (the Hugging Face guide's "Option 2"):

• Chunking. NeMo forces pad_and_drop_preencoded for this diarizer: each step feeds [9 pre-encode cache frames | 8·(lookahead + 1) new frames], runs the causal subsampling statelessly, and drops the first 2 encoder frames. This is a new raw-cache variant of the stream graph; the existing streaming path is untouched and byte-identical.
• Per-speaker streams. Encoder K/V and convolution caches live in per-speaker backend tensors; the decoder's LSTM state moved into NemotronDecoderStreamState. A speaker's stream starts at its first active chunk from zero caches and is frozen while inactive.
• Gating and masks. A speaker is active when its 80 ms activity exceeds 0.5 in the last two chunks. The mask is the last lookahead + 1 frames of the diarizer stream, repeated 8×, left-padded with zeros (so the cache frames are always masked) or cut from the right on the final chunk; masked frames become 0.0 and frames that were 0.0 become −16.6355, exactly as mask_features.
• Transcript. A port of update_sessionwise_seglsts_for_parallel: per-speaker segments that break on a pause over speaker_segment_gap_sec (1.0, NeMo sent_break_sec), with word continuations and leading punctuation kept on the previous segment. The hypothesis text follows NeMo exactly (space before punctuation removed, language tags stripped, SentencePiece's ⁇ for <unk>).
• Features equal one full-file frontend pass (with NeMo's zeroed padding frame) but are computed incrementally, so offline and streaming share the same step function and give byte-identical SegLST.
• Batching. Active speakers run in one encoder graph (batch = active speakers, per-row caches and attention masks). Caches stay in their graph row between steps and are written back inside the graph, so a stable speaker set needs no copies.

3. Lookahead resolution

Masked mode reads the lookahead from the file: chunk_len - 1 when chunk_right_context == 0 and the value is supported. An explicit lookahead_tokens wins with a warning; a file without a matching geometry (for example very_high) uses 13 with a warning that the masks are not NeMo-exact.

4. Outputs and options

• speaker_turns with text (--turns-out), labels speaker_k as in the diarizer
• seglst artifact in NeMo's schema (--out-dir writes seglst.json)
• masked mode: the transcript is speaker_k: words lines, in start time order offline and in finish order when streaming
• streaming: finished segments arrive as speaker-turn events and are never revised; segments over speaker_segment_max_sec (10 s) arrive as several events, the final result keeps full segments
• masked streaming also sends each finished segment's lines as partial text, so the server's transcript.text.delta events are append-only and concatenate to transcript.text.done; ordinary streaming and the server endpoints are unchanged
• new request options: speaker_probabilities, speaker_mode, speaker_mask, speaker_segment_gap_sec, speaker_segment_max_sec

Validation

Environment: Windows 11, MSVC (VS 2026), CUDA 13.3, 2× RTX 4070 Ti SUPER. References: NVIDIA-NeMo/Speech cf724ac, CPU FP32, the guide's Option 2 (max_num_of_spks=8 masked_asr=true parallel_speaker_strategy=true cache_gating=true binary_diar_preds=true att_context_size=[56,13] fifo_len=264 spkcache_update_period=222 target_lang=auto, strip_lang_tags=true), checkpoints nvidia/nemotron-3.5-asr-streaming-0.6b and nvidia/Nemotron-3-Diarization. NeMo runs are reproducible byte for byte.

Single-speaker baseline (unchanged paths)

Lookahead 13, CPU F32: offline tokens (id and 80 ms frame) equal NeMo on AMI 20 s (79/79) and a 120 s file (527/527). The existing streaming path matches too, except that it never computes the last centered feature frame, so NeMo's final . at that frame is missing (pre-existing, unchanged).

Masked mode vs NeMo, NeMo's diarizer output as input

To test the ASR side alone, NeMo's in-pipeline diarizer stream was written to our file format. Identical = same speaker, start, end and text.

Audio	Lookahead	CPU F32	CUDA F32
AMI EN2002d 20 s	0 / 3 / 6 / 13	4/4, 6/6, 5/5, 6/6	13: 5/6
120 s meeting (Dutch)	13	28/28	28/28
AMI IS1009a, 14 min	13	204/204	202/204
AMI ES2004a, 17.5 min	13	266/266	256/266

The CUDA differences are single tokens (uh, a repeated word) from float noise; running the speakers one at a time instead of batched gives the same differences.

End to end on AMI (our diarizer, asr_la13; attribution uses very_high)

Four AMI test meetings (ES2004a, IS1009a, TS3003a, EN2002a Mix-Headset; 92 min, 14,837 reference words from the AMI manual annotations), cpWER with meeteval, lowercase, [^a-z0-9' ] removed, CUDA F32:

System	ES2004a	IS1009a	TS3003a	EN2002a	Total
NeMo Python, masked (its own diarizer)	24.43%	36.62%	23.52%	37.41%	32.61% (4838)
This PR, masked (default)	24.54%	36.42%	23.24%	37.58%	32.64% (4843)
This PR, masked, speaker_mask=audio	24.88%	33.89%	24.66%	37.98%	32.81% (4868)
This PR, attribution	26.27%	33.99%	24.55%	43.20%	35.73% (5302)

Masked mode is within 5 errors of NeMo over 14,837 words. During development two alternatives were measured and dropped: filling masked frames with the log floor instead of NeMo's 0.0 (33.18% total) and NeMo-Speech.cpp's 160 ms onset window for attribution (35.43%, within noise of NeMo Python's span window, which is kept because it matches NeMo word for word). Attribution loses mostly deletions: one transcript cannot hold two speakers talking at once.

Performance

Masked mode, 120 s meeting, CUDA F32, k speakers active all the time. "One at a time" is the unbatched variant measured during development:

Active speakers	One at a time	Batched	Batched encoder	Peak VRAM
1	1.44 s	1.23 s	0.87 s	3630 MiB
2	3.18 s	1.62 s	0.98 s	3678 MiB
4	5.98 s	2.32 s	1.11 s	3741 MiB
8	12.47 s	3.80 s	1.46 s	3883 MiB

Peak VRAM includes 870 MiB used by other processes; going from 1 to 8 active speakers adds about 250 MiB. At 8 speakers the RNNT decoder (2.2 s) dominates. On real AMI audio masked mode runs at about 75x real time (17.5 min in 14 s).

Other checks

• unit tests nemotron_asr_speaker_tagging_test (ported NeMo segment, gating and mask cases, lookahead resolution, file validation) and nemotron_3_diar_streaming_test pass
• streaming events concatenate to the final result: on ES2004a, 10 of 2,357 words differ (5 ignoring punctuation and case), from the case below
• tools/check_loader_catalog_sync.py passes

Reproduce:

cmake -S . -B build/nemotron3 -G "Visual Studio 18 2026" -T "cuda=<CUDA 13.3 root>" -DGGML_CUDA=ON `
  -DAUDIOCPP_MODEL_SET=custom "-DAUDIOCPP_MODELS=nemotron_3_diar,nemotron_asr" -DENGINE_BUILD_TESTS=ON
cmake --build build/nemotron3 --config Release --target audiocpp_cli nemotron_asr_speaker_tagging_test
ctest --test-dir build/nemotron3 -C Release -R "nemotron"

build/nemotron3/bin/Release/audiocpp_cli.exe --task diar --family nemotron_3_diar --mode streaming `
  --model nemotron-3-diarization-bf16.gguf --audio meeting.wav `
  --session-option nemotron_3_diar.latency_profile=asr_la13 `
  --request-option return_frame_probabilities=true --out-dir diar
build/nemotron3/bin/Release/audiocpp_cli.exe --task asr --family nemotron_asr `
  --model nemotron-3.5-asr-streaming-0.6b --audio meeting.wav `
  --request-option speaker_probabilities=diar/speaker_probabilities.safetensors `
  --request-option speaker_mode=masked --turns-out turns.json --out-dir asr

Known limitations

• End of file: NeMo's multitalker diarizer sees one extra zero-valued feature frame; our file has floor(samples / 160) rows. The final chunk's mask can differ in the last 1.12 s.
• Streaming events and text: a speaker's stream is paused while they are silent, so a segment's last token (often punctuation or a word ending) can arrive only when they talk again. SegLST and the final speaker turns add it to the earlier segment; the already-sent event and delta do not have it (2% of words on a 1 h 48 min recording). Holding segments until then would add a median 9 s of latency.
• Offline attribution uses the full-file offline encoder graph (about 20–30 GB for 15+ minutes on CUDA, pre-existing); long recordings should use --mode streaming. Masked mode is chunked in both modes.
• speaker_mask=audio recomputes features per speaker and chunk and is about 15× slower than mel.
• Live microphone diarization is out of scope: streaming reads a precomputed diarizer file.

[0xShug0]

@LysanderdeJong Could you split the PR into smaller stacked PRs?

1. ASR spec migration.
2. Diarizer output and safetensors support.
3. Speaker-tagged ASR.

[LysanderdeJong]

Split into three stacked PRs:

1. nemotron_asr: move options to the model spec #685 ASR spec migration
2. nemotron_3_diar: write speaker probabilities as safetensors #686 Diarizer safetensors output and asr_laN profiles (stacked on nemotron_asr: move options to the model spec #685)
3. Add speaker-tagged transcription to Nemotron ASR #681 Speaker-tagged ASR (this PR, stacked on nemotron_3_diar: write speaker probabilities as safetensors #686)

Each PR's last commit is its own change. The final tree is identical to the previous #681.

[0xShug0]

@LysanderdeJong When speaker_probabilities + speaker_mode=masked + stream=true, the model emits incremental speaker-turn events, but the ASR server endpoint does not forward them.
For example, ordinary streaming returned:

data: {"type":"transcript.text.delta","delta":"Okay, so"}
data: {"type":"transcript.text.delta","delta":" even"}
data: {"type":"transcript.text.delta","delta":" the old"}

On the same audio, masked mode returned only the final result (text shortened here):

data: {"type":"transcript.text.done","text":"speaker_0: Okay, so ...","timing":{"ttft_ms":389.025}}
data: [DONE]

Is this expected behavior? If so, please document that masked mode returns only the final transcript through the server API. If not, my concern is that adding streaming support for this feature to the existing endpoints could break the boundaries and contracts of ASR and diarization.

I prefer reusing the existing transcript.text.delta events if the output can remain append-only and concatenating the deltas reproduces the final transcript. Clients using this feature can then parse the speaker labels from the text. Otherwise, consider a separate endpoint for structured speaker-segment streaming, without changing existing endpoint contracts.

Please also rebase.

[LysanderdeJong]

Not intended, thanks for testing it through the server. Masked streaming now also publishes each finished segment as partial text lines (speaker_k: words\n). They reach clients as ordinary transcript.text.delta events: append-only, and their concatenation is exactly transcript.text.done. No server or endpoint contract changes. The streamed text lists segments in the order they finish. For start-time order and exact times, use SegLST or speaker_turns via /details.

One documented limit: a speaker's stream is paused while they're silent, so a segment's last token (often punctuation) can arrive only when they speak again. The delta has already gone out without it. SegLST keeps it. That affects about 2% of words on a 1 h 48 min recording.

Rebased onto main; the fix is the new commit 4af4841.

[0xShug0]

@LysanderdeJong Thanks. PR merged!