feat: Nemotron-3-Diarization support + speaker-attributed ASR by localai-org-maint-bot · Pull Request #71 · mudler/parakeet.cpp

localai-org-maint-bot · 2026-09-26T11:20:23Z

What

Add offline speaker diarization and speaker-attributed ASR (SAS) to
parakeet.cpp, porting the Nemotron-3-Diarization (Sortformer) model to ggml.

This is Phase 1 (standalone diarization) and Phase 3.1-3.2 (SAS merge +
combined C-API) from the plan in docs/diarization-plan.md.

Architecture

Nemotron-3-Diarization is a 31-layer pre-LN RoPE Transformer
(d_model=512, 8 heads, head_dim=64, ff_dim=2048) with FeatureStacking
8x subsampling and a subpixel Conv1D 8x upsampling speaker head.

The model takes mel spectrograms as input and outputs per-frame
speaker activity probabilities for up to 8 speakers. Hysteresis
thresholding (onset=0.5, offset=0.5) produces speaker segments.

The implementation is a separate DiarizationModel class (not the ASR
Model), with its own GGUF arch string "diarization". The C-API
takes separate ASR and diarization contexts so callers can mix and
match model sizes.

Parity vs PyTorch reference

Metric	Value
probs max_diff	0.0097 (F16 noise)
probs mean_diff	0.000118
segments	584 (ours) vs 581 (reference) - 0.5%
up (upsampled) max_diff	0.0149 (F16 noise)
enc_out max_diff	0.009

Speaker-attributed ASR (SAS)

parakeet_capi_transcribe_and_diarize() runs ASR + diarization on the
same audio and merges the results: each ASR word is assigned to the
speaker with the largest temporal overlap, then consecutive
same-speaker words are grouped into utterances.

parakeet_ctx* asr  = parakeet_capi_load("parakeet-tdt-0.6b-v3.q8_0.gguf");
parakeet_ctx* diar = parakeet_capi_load("diarization.gguf");

int n = 0;
parakeet_sas_result* r = parakeet_capi_transcribe_and_diarize(
    asr, diar, samples, n_samples, sample_rate, &n);
// r[0].speaker, r[0].text, r[0].start, r[0].end, r[0].conf

char* json = parakeet_capi_transcribe_and_diarize_json(
    asr, diar, samples, n_samples, sample_rate);
// {"speakers":8,"utterances":[...],"words":[...]}

ggml bugs discovered and worked around

ggml_flash_attn_ext output layout: Returns [hd, H, T, 1],
not [hd, T, H, 1]. Workaround: no permute, just cont +
reshape_2d.
3D ggml_cont after ggml_permute is broken: It creates a
view instead of copying. Workaround: use 2D transpose + cont.
Conv1d weight layout: Stored as [OC, IC, k] in GGUF, read as
[k, IC, OC] (reversed). No permute needed - conv output matches
within F16 noise.

test_3d_permute.cpp documents bug #2 and is expected to FAIL.
test_cont_check.cpp and test_subpixel_check.cpp verify the
workarounds.

C-API changes (ABI v7)

New functions:

parakeet_capi_diarize_path / parakeet_capi_diarize_pcm
parakeet_capi_diarize_json
parakeet_capi_free_diar_results
parakeet_capi_transcribe_and_diarize (struct variant)
parakeet_capi_transcribe_and_diarize_json (JSON variant)
parakeet_capi_free_sas_results

New struct: parakeet_sas_result (speaker, text, start, end, conf).

Commits

feat(converter): add diarization model support - converter,
ModelLoader, Model::load rejection
feat(diarization): add encoder and speaker head - 31-layer
RoPE transformer, subpixel Conv1D upsampling head
feat(diarization): add DiarizationModel and SAS merge layer -
model class, hysteresis thresholding, SAS merge logic
feat(capi): add diarization and SAS C-API (ABI v7) - C-API
functions for diarization and speaker-attributed ASR
feat(diarization): add CLI, tests, build integration, and plan doc

How to verify

# Build
cmake -B build -DCMAKE_BUILD_TYPE=Release && cmake --build build -j

# Unit tests (no model needed)
./build/tests/test_sas_merge
./build/tests/test_cont_check
./build/tests/test_subpixel_check

# End-to-end (needs ASR + diarization GGUF + audio)
PARAKEET_TEST_GGUF=asr.gguf \
PARAKEET_TEST_DIAR_GGUF=diarization.gguf \
PARAKEET_TEST_COMBINED_WAV=test.wav \
./build/tests/test_combined_offline

# CLI
./build/examples/cli/diarize diarization.gguf test.wav

Not in this PR

Phase 2 (streaming diarization): researched, not yet implemented
Phase 3.3 (mel sharing optimization)
Quantized GGUF publication to HuggingFace
LocalAI integration

The full 3-phase plan is in docs/diarization-plan.md.

Assisted-by: Codex:GPT-5 [Codex]

Convert Nemotron-3-Diarization (.nemo) to GGUF with arch="diarization". The converter loads state_dict + config directly from the .nemo tar, bypassing SortformerEncLabelModel.restore_from() (which fails on NeMo versions that don't support self_attention_model='rope'). ModelLoader reads diarization-specific KV pairs (parakeet.diar.*, parakeet.encoder.self_attention_model, rope_base, rotary_fraction) and accepts models with no vocabulary (arch="diarization"). Model::load rejects diarization arch so the C-API falls through to DiarizationModel::load. test_model_loader skips vocab checks for diarization models.

DiarizationEncoder: 31-layer pre-LN RoPE Transformer (d_model=512, 8 heads, head_dim=64, ff_dim=2048). Uses ggml_flash_attn_ext for attention (CPU). FeatureStacking 8x subsampling in pre_encode. Key ggml workarounds: - flash_attn_ext returns [hd,H,T,1] not [hd,T,H,1]: no permute, just cont + reshape_2d - 3D ggml_cont after ggml_permute is broken (creates view not copy): use 2D transpose + cont instead DiarizationHead: encoder_proj (512->192), subpixel Conv1D 8x upsampling, sigmoid speaker head (192->8). Subpixel reshape uses reshape_3d -> reshape_2d -> 2D transpose+cont to avoid the 3D cont bug. Output is [n_spk, T_out] row-major via cont(transpose(sig)). Parity vs PyTorch reference: probs max_diff=0.0097 (F16 noise).

DiarizationModel: loads diarization GGUF, runs encoder + head, applies hysteresis thresholding (onset=0.5, offset=0.5) to produce SpeakerSegments with speaker labels and time ranges. SAS merge (sas_merge.hpp/.cpp): merge_asr_diarization assigns each ASR word to the speaker with the largest temporal overlap. Dominant speaker wins ties (first by start time, then by speaker index). group_speaker_words groups consecutive same-speaker words into SpeakerUtterances, splitting on speaker change or >0.5s gap. Words with no overlapping segment get speaker=-1. Parity: 584 segments (ours) vs 581 (reference) — 0.5% difference.

Bump ABI to 7. Add diarization C-API (Phase 1): - parakeet_capi_diarize_path / _pcm: run diarization on audio - parakeet_capi_diarize_json: JSON output with segments array - parakeet_capi_free_diar_results: free result array Add speaker-attributed ASR (SAS) C-API (Phase 3): - parakeet_sas_result struct (speaker, text, start, end, conf) - parakeet_capi_transcribe_and_diarize: run ASR + diarization, merge into speaker-attributed utterances (struct variant) - parakeet_capi_transcribe_and_diarize_json: same, JSON output with speakers/utterances/words arrays - parakeet_capi_free_sas_results: free result array SAS functions take separate ASR and diarization contexts so the caller can mix and match model types and sizes.

Build: - CMakeLists.txt: add diarization sources to parakeet library - examples/cli/CMakeLists.txt: add diarize CLI target - tests/CMakeLists.txt: register all diarization/SAS tests CLI: - examples/cli/diarize.cpp: minimal CLI (path -> segments JSON) Tests: - test_diarization.cpp: basic diarization smoke test - test_diarization_parity.cpp: full parity test vs PyTorch reference - test_diar_head_bisect.cpp: head stage-by-stage bisection - test_diar_layer0.cpp / _bisect.cpp: encoder layer 0 bisection - test_cont_check.cpp: verifies 2D cont works (PASS) - test_subpixel_check.cpp: verifies subpixel reshape (PASS) - test_3d_permute.cpp: documents 3D cont bug (expected FAIL) - test_sas_merge.cpp: 10 unit tests for SAS merge logic (PASS) - test_combined_offline.cpp: end-to-end SAS test (ASR + diarization) docs/diarization-plan.md: 3-phase plan with dependency graph.

Split the DiarizationEncoder into pre_encode() and transformer_forward() so the streaming path can run the transformer over [spkcache | chunk] without re-doing the mel front-end. Add forward_range() to DiarizationHead to compute probs for a sub-range of encoder frames (the chunk portion only). StreamingDiarization implements the AOSC mechanism from NeMo's Sortformer: chunked processing with a speaker cache, FIFO overflow, and per-speaker top-K compression of the spkcache. When fifo_len=0 (Nemotron-3 default) every chunk immediately overflows into the spkcache, which is then compressed back to spkcache_len frames. C-API ABI v8 adds: - parakeet_capi_diarize_stream_begin / _feed / _free (streaming diarization) - parakeet_capi_sas_stream_begin / _feed / _free (streaming SAS) Streaming SAS (Phase 3.4) buffers incoming PCM, feeds complete chunks to both ASR and diarization, and merges the results with the same SAS merge logic as the offline path. ModelLoader reads streaming config from GGUF KVs (parakeet.diar.chunk_len, spkcache_len, fifo_len, etc.). When absent, Nemotron-3-Diarization defaults are used so the streaming C-API works without converter changes. test_streaming_diarization verifies the streaming path produces segments in chronological order within the audio duration, and detects speakers. On the 2-minute test audio: 115 streaming segments vs 583 offline (the gap is expected — streaming uses chunk-local context, not full audio).

Write parakeet.diar.chunk_len, spkcache_len, fifo_len, and the AOSC compression tuning parameters as GGUF KV pairs. Values are read from the NeMo config 'streaming' section, with Nemotron-3-Diarization defaults (chunk_len=264, spkcache_len=264, fifo_len=0) when absent. The model loader already falls back to these defaults when the KVs are missing, so this is not a breaking change — but writing them explicitly ensures non-default models are handled correctly.

…3324) ## Summary Adds speaker diarization and speaker-attributed ASR (SAS) to vllm.cpp by linking [parakeet.cpp](https://github.com/mudler/parakeet.cpp) as a FetchContent dependency. The diarization stack (Nemotron-3-Diarization Sortformer encoder, speaker head, AOSC streaming, and SAS merge layer) lives in parakeet.cpp; vllm.cpp wraps it through a clean C++ seam. ## What's new ### C++ seam (`include/vllm/multimodal/diarization.h`) - **`Diarizer`** class: loads a diarization GGUF, exposes `Diarize()` and `DiarizeWavFile()` returning speaker segments - **`SpeakerAttributedASR`** struct + `TranscribeAndDiarize()` helpers: run both ASR and diarization on the same audio and merge results ### C ABI (`include/vllm.h`, ABI v30) - `vllm_diarization_load`: load a diarization GGUF as a separate engine handle - `vllm_diarize_path` / `vllm_diarize_pcm`: diarize a WAV file or raw PCM - `vllm_transcribe_and_diarize` / `_pcm`: combined ASR + diarization (speaker-attributed utterances) - `vllm_diarization_free` / `vllm_sas_result_free`: cleanup ### Server endpoints (registered only when a diarizer/SAS is attached) - **`POST /v1/audio/diarizations`**: returns speaker segments as JSON - **`POST /v1/audio/sas`**: returns speaker-attributed utterances as JSON ### CMake - `parakeet.cpp` is fetched from GitHub (or pointed at a local tree via `VLLM_CPP_PARAKEET_CPP_DIR`) - Gated by `VLLM_CPP_WITH_DIARIZATION` (ON by default). When disabled, every diarization function returns an error and the endpoints are not registered — the library is byte-identical to before. ## Dependencies - [parakeet.cpp](https://github.com/mudler/parakeet.cpp) — the diarization engine, GGUF model loader, and SAS merge layer - [Nemotron-3-Diarization-GGUF](https://huggingface.co/mudler/Nemotron-3-Diarization-GGUF) — the quantized model checkpoint ## Parakeet.cpp PR The parakeet.cpp diarization implementation is in [PR #71](mudler/parakeet.cpp#71) (Phases 1–3.4: offline diarization, streaming diarization, SAS merge, streaming SAS). ## Test plan - [ ] Build with `VLLM_CPP_WITH_DIARIZATION=ON` — verify FetchContent populates parakeet.cpp and the library links - [ ] Build with `VLLM_CPP_WITH_DIARIZATION=OFF` — verify the library is byte-identical to `main` - [ ] Load a diarization GGUF via `vllm_diarization_load` — verify segments are returned - [ ] Run `vllm_transcribe_and_diarize` with a Parakeet ASR checkpoint + diarization GGUF — verify speaker-attributed utterances - [ ] `POST /v1/audio/diarizations` — verify JSON response with speaker segments - [ ] `POST /v1/audio/sas` — verify JSON response with speaker-attributed utterances

The offline diarization output disagreed with NeMo on most frames and almost never detected a second speaker. Three causes: - The subpixel upsample used the wrong interleave. NeMo reshapes the conv output as [T, up, hidden], so output frame t*up+u reads conv channels u*tf..u*tf+tf at frame t. The old code read channel h*up+u and wrote frame u*T+t, which scrambled time across the whole clip. The conv now runs as im2col + mul_mat in F32, and its time-major output reshapes straight into the upsampled sequence. - NeMo peak-normalizes the waveform in offline mode (x / (max(x) + 1e-3)). Without it every log-mel bin was shifted. - The mel was not trimmed to floor(S / hop), and the upsampled output kept the FeatureStacking pad frames. The pipeline up to the probabilities is now DiarizationModel:: speaker_probs, which returns the same tensor as NeMo's offline forward(). On a 2-speaker clip the F32 max prob diff vs NeMo drops from 0.96 to 0.004 and the segments match exactly. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

scripts/gen_diar_baseline.py runs NeMo on a clip and stores the audio, the offline and streaming speaker probabilities and the diarize() segments in a baseline GGUF. test_diarization_accuracy checks the C++ offline pipeline against it: probability max/mean diff, segment count, speakers and boundaries (20 ms), and frame-level agreement (>= 99.5%). tests/fixtures/two_speakers.wav is 23.6 s of LibriSpeech (CC BY 4.0): speaker 1272 (1272-128104-0000, -0001, from the dev-clean set) and speaker 2086 (the existing speech.wav) alternating A-B-A-B with 0.5 s gaps. NeMo finds 5 segments across the 2 speakers. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

test_diar_layer0, test_diar_layer0_bisect, test_diar_head_bisect and test_diarization_parity were one-off debugging dumps with no assertions. test_3d_permute, test_cont_check and test_subpixel_check documented "ggml bugs" that are not bugs: ggml_permute takes the destination of each source axis (the inverse of torch.permute), and flash_attn_ext is documented to return [head_dim, n_head, T]. The subpixel check also asserted the wrong upsample layout. test_diarization_accuracy and test_streaming_diarization replace them with checks against NeMo. test_diarization is now labelled "model". Assisted-by: Claude:claude-opus-5-5 [Claude Code]

Streaming diarization did not follow NeMo's Sortformer streaming: - chunk_len, spkcache_len, fifo_len and spkcache_update_period are encoder frames in NeMo (264 = 21.12 s), not mel frames, and the speaker cache was 8x too small. - The cache stored channels-first embeddings but was concatenated as if time-major, so [cache | chunk] was scrambled after one chunk. - The cache must hold pre-norm embeddings: NeMo applies embed_norm inside the encoder, after the bypassed pre-encoder. - NeMo runs the head over the whole [cache | fifo | chunk] sequence and feeds the downsampled predictions to the cache update. The head ran on the chunk alone. - The converter wrote no streaming keys for .nemo input, and read them from a "streaming" section that does not exist. spkcache_sil_frames_per_spk is 1 for Nemotron-3, not 3. StreamingDiarization now mirrors streaming_update and _compress_spkcache (sync mode). Top-k ties go to the earlier frame so the result is deterministic. Segments carry across chunk boundaries. On the fixture and on 31 s and 68 s clips (the last compresses the cache twice) the streaming segments match NeMo streaming exactly. The encoder builds one shared graph for the offline and streaming entry points, works time-major end to end, reads RoPE base, rotary fraction and pre_block_norm from the GGUF, and throws on graph failure instead of asserting. The converter drops the dead ASR-path diarization code and accepts a HF id. C-API (ABI 7, nothing released used 8): - SAS JSON uses the shared JSON writers: text is escaped and long utterances are no longer cut at 512 bytes. - transcribe_and_diarize returns a status. Zero utterances is success, not NULL. free_sas_results takes the count and frees every .text. - Streaming diarization takes 16 kHz PCM and computes the mel with StreamingMel (bit-identical to the whole-clip mel), instead of asking the caller for mel chunks. - Streaming SAS transcribes the uncommitted audio when a diarization chunk completes and commits only words that end 1 s before the cut, so words at chunk edges are not split. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

ASR word and diarization segment boundaries can disagree by a frame or two, and a word that overlaps no segment got speaker -1. It now takes the nearest segment's speaker when that segment is within 0.5 s. Words farther from any segment are still unassigned. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

docs/superpowers/ holds local planning notes and stays out of the tree. The file is kept locally. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

NeMo's diarize() for Nemotron-3-Diarization runs cache-aware streaming inference (streaming_mode: true in the model config). diarize_path and diarize_pcm always ran the offline path, which attends over the whole clip: on a 12.3 min, 3 speaker recording it agreed with NeMo on 22% of speech frames and put almost everything on one speaker. The streaming path agrees on 100% and runs in 6.8 s instead of 11.8 s, since each chunk attends over at most 528 steps. The converter writes parakeet.diar.streaming_mode; GGUFs without it default to streaming. run_offline, run_streaming and segments_from_probs are public so tests can check both modes. test_diarization_accuracy now also checks the default diarize_pcm against NeMo's diarize(). PARAKEET_TEST_DIAR_VERBOSE prints the streaming probability diff per 10 s window. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

Feeds the audio through parakeet_capi_diarize_stream_* in 1 s pieces, the live path, and prints the same JSON as the offline mode. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

docs/diarization.md covers the model, conversion, offline and streaming inference, parity with NeMo, speaker-attributed ASR, speed and tests. AGENTS.md lists the new sources, tests, fixture, baseline script and C-API symbols. The implementation plan was a working note and now lives with the other plans outside the tree. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

…ASR (#71) Adds speaker diarization with nvidia/Nemotron-3-Diarization (Sortformer, up to 8 speakers) and speaker-attributed ASR that combines it with any Parakeet ASR model. Offline and live streaming paths are both exposed through the C-API (ABI 7, additive). diarize_* follows the checkpoint's streaming_mode like NeMo's diarize(): cache-aware streaming in 21.12 s chunks with a speaker cache. Against NeMo main the segments match exactly on a 23.6 s and a 68.5 s two-speaker clip, offline and streaming, and a 12.3 min three-speaker clip agrees on 100% of speech frames. F16 matches too; Q8_0 keeps the same segments. The 12.3 min clip runs in 6.8 s on CPU. tests/fixtures/two_speakers.wav (LibriSpeech, CC BY 4.0) and scripts/gen_diar_baseline.py back test_diarization_accuracy and test_streaming_diarization. See docs/diarization.md. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

mudler-agent · 2026-09-28T06:40:02Z

Squash-merged to master as ba693fa (pushed directly: the bot account cannot merge PRs). All CI checks were green on 78dbb4c; the squashed tree on top of #69 was rebuilt and re-tested locally before pushing.

mudler-agent · 2026-09-28T06:40:08Z

Squash-merged to master as ba693fa (pushed directly: the bot account cannot merge PRs). All CI checks were green on 78dbb4c; the squashed tree on top of #69 was rebuilt and re-tested locally before pushing.

mudler added 6 commits September 26, 2026 11:54

docs: design ROCm release and container support

59cf0fb

Assisted-by: Codex:GPT-5 [Codex]

mudler force-pushed the feat/diarization-sas branch from 880b8a3 to e497e08 Compare September 26, 2026 12:26

mudler-agent mentioned this pull request Sep 26, 2026

feat(diarization): add diarization and SAS support via parakeet.cpp mudler/vllm.cpp#3324

Merged

6 tasks

mudler added 9 commits September 27, 2026 19:23

chore: stop tracking the ROCm design spec

5633c78

docs/superpowers/ holds local planning notes and stays out of the tree. The file is kept locally. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

feat(cli): add --stream to the diarize example

ace15e7

Feeds the audio through parakeet_capi_diarize_stream_* in 1 s pieces, the live path, and prints the same JSON as the offline mode. Assisted-by: Claude:claude-opus-5-5 [Claude Code]

mudler closed this Sep 28, 2026

mudler-agent mentioned this pull request Sep 28, 2026

feat(parakeet-cpp): speaker diarization, including streaming and live transcription mudler/LocalAI#12323

Closed

Navigation Menu

feat: Nemotron-3-Diarization support + speaker-attributed ASR - #71

feat: Nemotron-3-Diarization support + speaker-attributed ASR#71
localai-org-maint-bot wants to merge 17 commits into
masterfrom
feat/diarization-sas

localai-org-maint-bot commented Sep 26, 2026

Uh oh!

mudler-agent commented Sep 28, 2026

Uh oh!

mudler-agent commented Sep 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

localai-org-maint-bot commented Sep 26, 2026

What

Architecture

Parity vs PyTorch reference

Speaker-attributed ASR (SAS)

ggml bugs discovered and worked around

C-API changes (ABI v7)

Commits

How to verify

Not in this PR

Uh oh!

mudler-agent commented Sep 28, 2026

Uh oh!

mudler-agent commented Sep 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants