feat(QVAC-18625): Add better support for long-term streaming diarization with v2.1

[pratiknarola-t]

feat(QVAC-18625): Sortformer v2.1 streaming with NeMo Audio-Online Speaker Cache (AOSC)

TL;DR

Adds v2.1 streaming Sortformer to parakeet-cpp via a port of NeMo's Audio-Online Speaker Cache (AOSC) algorithm, plus a slot-continuity fix on the existing v1 streaming path, plus a regression test suite. Speakers now keep their original speaker_id even after long silences. The v1 path is bit-identical apart from the slot-continuity fix on its own streaming code.

The original problem

In v1 streaming, the session only carries a rolling history_ms window of audio context. If a speaker goes silent for longer than that window and then returns, the model has no memory of them — they get whatever speaker_id slot happens to be free, which may or may not match their original slot. The user sees what looks like a speaker swap mid-conversation.

The fix (model-side)

The newer v2 / v2.1 NeMo Sortformer introduced a speaker cache: a small set of acoustic embeddings retained per speaker across the entire session, plus a FIFO of recent embeddings, plus a per-session "silence profile" embedding. The model then processes a concatenated row

[ cache_rows | fifo_rows | (mel → subsampling) ]

through the FastConformer blocks every chunk, so it always sees a fresh attention window over every speaker it has ever heard in this session. After arbitrarily long silences, the same voice rebinds to its original slot.

For v1 (which doesn't have the cache architecture), this PR also adds a separate overlap-based slot-continuity remap so cross-chunk drift inside the history window is suppressed.

What's added

• Full AOSC port of NeMo's sortformer_modules.py + sortformer_diar_models.py into parakeet-cpp. Helpers in C++ carry // matches NeMo <fn> at <line> comments: _compress_spkcache, _get_silence_profile, _disable_low_scores / _boost_topk_scores, streaming_update, forward_streaming_step, and run_encoder_bypass_pre_encode (feeds pre-subsampled embeddings straight into the conformer stack — required because the cache lives in post-subsampling space).
• Encoder context windowing (chunk_left_context_ms=80, chunk_right_context_ms=560), defaults from NeMo's e2e_diarize_speech.py inference YAML.
• v2.1 detection by encoder shape (17 layers / 128 mels); cache_active=false for v1, so the v1 code path is bit-identical.
• Overlap-based slot-continuity remap on the v1 streaming path (separate fix, additive).
• New ctest test-sortformer-aosc-speakers-{abcba,abcdba} asserting three invariants against an RTTM ground truth: (a) every reference speaker is covered, (b) re-entering speakers land in the same hyp_<id> they were first assigned to (the AOSC contract), (c) frame-level DER ≤ 30 %.
• Two redistributable ElevenLabs-generated fixtures with hand-built ground-truth RTTMs (LIFO re-entry patterns).
• SortformerStreamSession::aosc_active() — runtime introspection getter so callers can tell v1 from v2.1+AOSC.
• examples/live-mic: banner branches on aosc_active(); v2.1 prints spkcache/fifo/lc/rc, v1 keeps the legacy history-window banner unchanged.
• README / PROGRESS docs updated (new v2.1 row in the model table; "Shipped / Not in-repo" status corrected; new Phase 17 closing §11.11.2's reservation).

The biggest blocker: mean_sil_emb

NeMo's algorithm uses a special "what does silence sound like in this room" baseline embedding to fill cache slots that aren't yet bound to a real speaker. Without it, those slots stay at zeros — out-of-distribution input for the model — and it panics and collapses every chunk to a single slot. Once the runtime EMA update rule from sortformer_modules.py was matched in C++, the cache-aware forward started behaving.

Results

Complete 5-way DER table (synthetic fixtures, q8_0 GGUFs, Apple M-series CPU):

fixture	mode	DER%	speakers tracked
abcba (3-speaker, LIFO re-entry)	v1 streaming	24.31	2 (no C)
abcba	v2.1+AOSC	27.29	3 (all)
abcba	v2.1 no-cache	23.74	2 (no C)
abcdba (4-speaker, LIFO re-entry)	v1 streaming	66.28	2 (collapsed)
abcdba	v2.1+AOSC	22.22	4 (all)
abcdba	v2.1 no-cache	65.72	2 (collapsed)

Voices used (free elevenlabs api key for testing):
A=Sarah (♀US), B=Adam (♂US), C=Alice (♀UK), D=George (♂UK). 1 s silence between turns. Ground-truth RTTM auto-generated from clip durations.

Per-fixture analysis

• three-english-speakers — no re-entry, no long silence; v1 offline almost perfect (2.14 %). v2.1 conflates the two female-adjacent voices and gets 26 %. v1 vs v2.1 streaming track their offline numbers closely because there's no silence-return event to trigger the cache.
• abcab_real (distinct voices) — every mode handles this fine. Confirms the new code doesn't regress on easy material.
• absaman (A→B→C→Aman→B→A, 4-speaker with long silences) — the failure mode this PR targets.
  • v1 offline handles it (0.81 %) because the model sees the whole clip at once.
  • v1 streaming drops to 51 % — by the time A returns, A's voice has aged out of the history window; A gets whatever slot was free.
  • v2.1 + AOSC holds at 21.6 % — A re-binds to its original slot via the cache. Most of the residual error is the v2.1-side acoustic conflation of Aman with Alex (two male voices, the cache thinks "close enough" and reuses Alex's slot).
  • v2.1 offline and v2.1 streaming nearly match — because v2.1's offline pipeline also runs the cache-aware forward, so streaming is consistent with offline.

Known limitation (and the trade-off)

v2.1 is not strictly better than v1. When two speakers are acoustically similar (notably two female voices in a row, or two male voices in a row), v2.1's cache + FIFO + attention layer thinks they're the same speaker and assigns them the same ID. v1 — without the cache — happens to separate them better in that exact scenario.

So in practice:

• For ~80–90 % of use cases (distinct voices, normal conversation), v1 wins.
• When the long-silence re-entry mismatch hits (≥ history_ms of silence between same-speaker turns), v1 catastrophically fails (~50 % DER) and v2.1 + AOSC catches it (~22 %).

This PR makes v2.1+AOSC available as an option that works correctly per the NeMo algorithm. It does not change the default model selection; downstream consumers keep choosing whichever variant fits their workload.

API additions (additive only)

• SortformerStreamingOptions::spkcache_enable (default true), spkcache_len, fifo_len, chunk_left_context_ms, chunk_right_context_ms, spkcache_update_period — public knobs matching NeMo's inference YAML.
• SortformerStreamSession::aosc_active() — runtime introspection.

No signature changes on existing methods. v1 GGUFs take the unchanged code path.

Validation

• ctest -R sortformer-aosc-speakers passes both -abcba and -abcdba with the v2.1 GGUF in parakeet-cpp/models/.
• v1 GGUF on the same harness fails with exit code 21 (continuity broken) — proving the test discriminates AOSC from non-AOSC.
• live-mic on a v2.1 model prints [live-mic] listening at 16 kHz mono (v2.1 diarization, AOSC). chunk=2000 ms spkcache_len=188 fifo_len=188 lc=80 ms rc=560 ms. — verified end-to-end with a live mic capture that emitted multiple speaker IDs.
• v1 path: existing test-sortformer-streaming stays bit-identical (4.17 % DER on the English fixture).

[GustavoA1604]

I checked locally and noticed

1. need to update download-all-models.sh to include the new model used
2. need to update sortformer-streaming test to use v2.1 model instead of v2 model

Besides that worked in windows and mac

Also Cursor flagged the following, not sure if relevant but could be solved by pasting this message in a new thread:

---

Critical / Must-Fix

1. spkcache_len_per_spk can go negative with small spkcache_len

In compress_speaker_cache (parakeet_sortformer.cpp):

const int spkcache_len_per_spk = spkcache_len / num_spks - A_sil;
const int strong_boost = (int) std::floor((float) spkcache_len_per_spk * cfg.strong_boost_rate);
const int weak_boost = (int) std::floor((float) spkcache_len_per_spk * cfg.weak_boost_rate);
With the defaults spkcache_len=188, num_spks=4, A_sil=3 this is 47 - 3 = 44 — fine. But if a caller passes spkcache_len < num_spks * A_sil (e.g. spkcache_len=8, num_spks=4, A_sil=3 → -1), the nth_element calls in boost_topk_scores receive a negative k after std::min(n_boost_per_spk, n_frames), and std::nth_element with a negative distance is UB. Add a guard at the top of the function:

if (spkcache_len_per_spk <= 0) {
// degenerate config; fill with silence profile and return
...
}
4. run_encoder_bypass_pre_encode cache invalidation — every chunk misses until FIFO reaches steady-state

The bypass encoder graph is cached by (bypass_pre_encode, T_enc, n_run_layers). T_enc = spkcache_n + fifo_n + T_chunk_pre grows chunk-by-chunk as the FIFO fills (0 → fifo_len), so for the first fifo_len / chunk_len ≈ 188 / 6 ≈ 31 chunks a new ggml graph is built from scratch on every call. With k_encoder_graph_cache_max = 4, those graphs evict each other immediately and zero reuse occurs.

This is a performance bug, not a correctness bug, but it could make the first ~60 seconds of a session noticeably slower on slower hardware. Consider either:

Caching by (bypass_pre_encode, T_enc_max) and passing a mask / sequence-length argument, or
Pre-building the graph at diarize_start for the known steady-state size spkcache_len + fifo_len + max_chunk_pre_frames and always feeding that size (padding with silence rows when the FIFO isn't full yet).
Medium / Should-Fix
5. v2.1 detection by encoder shape is fragile

const bool model_is_v2_1 =
pimpl_->model.encoder_cfg.n_layers == 17 &&
pimpl_->model.mel_cfg.n_mels == 128;
If NeMo ships a v3 or a v2.2 variant that happens to share {17 layers, 128 mels} but was not trained with the cache-aware concat forward, enabling AOSC on it will produce garbage silently. A GGUF metadata key (e.g. parakeet.model_variant = "sortformer-v2.1-aosc") set by the converter would be more robust. At minimum, document this assumption in diarization.h next to the detection logic and add a note that it must be revisited when a new variant is converted.

6. streaming_update parameter name chunk_pre_encode_lc is misleading

The function signature says chunk_pre_encode_lc but the call site passes the committed chunk slice (already offset past the left context):

const float * chunk_pre_committed = chunk_pre_encode_embs + (size_t) lc * D;
streaming_update(cache, chunk_pre_committed, chunk_len_eff, ...);
The name implies it includes the left context, which it does not. Rename to committed_chunk_pre_encode to match the call-site variable name and the comment in the function body.

7. load_wav_pcm16le_mono duplicated verbatim from test_sortformer_streaming.cpp

The comment in the new test file acknowledges this ("duplicated here on purpose"). For a 60-line helper this is borderline, but the two copies will drift. A shared test/test_utils.h header in the test/ directory would be the right solution. Not a blocker, but worth a TODO at minimum.

8. WAV fixtures committed as binary blobs (~11 MB total)

abcba.wav (~5.0 MB) and abcdba.wav (~5.9 MB) are committed directly into the repo. Git LFS would be the cleaner long-term approach, consistent with how the project will likely handle future audio test fixtures. If the project doesn't use LFS yet, at least leave a comment in CMakeLists.txt pointing to where the fixtures can be regenerated.

Low / Nice-to-Have
9. -std::numeric_limits::infinity() is not UB

The comment /* very-negative sentinel; -inf is UB with current FP flags */ appears three times. IEEE 754 infinity is a well-defined value; the UB concern applies to operations like inf - inf, not to storing or comparing the value. Using std::numeric_limits::lowest() (which is approximately -3.4e38) or -std::numeric_limits::infinity() directly would both be more readable than the magic −1.0e30f sentinel. Not a bug, just misleading documentation.

10. ring.erase is O(n) — pre-existing, but AOSC retention differs

AOSC retains only chunk_left_context_samples behind emit_end, which is much smaller than history_ms. So ring trims happen more aggressively and ring.erase is called more frequently (every chunk vs. lazily on v1). This amplifies the pre-existing O(n) cost. No action required now, but worth a note for a future std::deque refactor.

11. prev_chunk_full_segments populated on AOSC path unnecessarily

On the AOSC path slot_remap is always identity, but cur_full is still moved into prev_chunk_full_segments every chunk. This is harmless (just a small memory/copy overhead) but a if (!cache_active) guard around those two lines would clarify intent.

12. encoder_ms attribution is slightly surprising

const double encoder_ms = ms_since(t_enc) - dres.decode_ms;
t_enc is set before run_subsampling, so ms_since(t_enc) covers subsampling + bypass-encode + diarize. Subtracting decode_ms leaves the "everything except the diarizer head" time, which is actually the subsampling + conformer-bypass time. The field name encoder_ms matches the existing convention in the non-AOSC path, so this is consistent — just worth a comment explaining what the subtraction is doing.