Skip to content

feat(tdt): add opt-in N-best beam decoding - #56

Merged
mudler merged 4 commits into
mudler:masterfrom
inlanger:feat/tdt-nbest
Jul 28, 2026
Merged

mudler merged 4 commits into
mudler:masterfrom
inlanger:feat/tdt-nbest

Conversation

@inlanger

@inlanger inlanger commented Jul 26, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Add an opt-in offline TDT N-best decoder without changing the existing greedy
path.

  • implements sequence-level TDT beam search in C++
  • returns ranked text, raw/normalized scores, token ids, emission frames, and durations
  • exposes the result through pk::Model, additive flat C API symbols, and CLI JSON
  • reuses the existing frontend, encoder, predictor, and joint implementations
  • adds model-independent tests, real-model CPU/CUDA tests, and a closed-loop NeMo parity gate

This is ready for design and API feedback before treating the new surface as
final.

Reference behavior

The implementation independently follows NVIDIA NeMo 2.7.3's
BeamTDTInfer.default_beam_search:

  • token and duration logits receive independent log_softmax
  • top non-blank token-duration pairs are expanded jointly
  • blank is paired only with non-zero durations
  • duplicate (token sequence, last frame) paths merge with logaddexp
  • final ranking uses score / len(y_sequence) by default, including NeMo's leading SOS/blank sentinel

One edge case intentionally differs from NeMo: when the duration beam contains only duration 0, blank uses the highest-scoring positive duration rather than the smallest positive duration. This preserves score ordering for the advertised beam_size=1 surface and is covered by a model-independent regression.

The C++ implementation additionally rejects malformed duration tables,
non-finite log probabilities, and a zero-duration expansion that does not
strictly reduce accumulated score. These checks fail explicitly; they do not
truncate valid hypotheses or apply NeMo's unsupported max_symbols_per_step
parameter to default beam.

NeMo labels this default TDT beam strategy experimental and recommends
malsd_batch
.
This PR deliberately starts with the smaller directly comparable reference
algorithm; it does not claim that default beam is the final high-performance
decoder design.

Surface

parakeet-cli transcribe \
  --model model.gguf \
  --input audio.wav \
  --decoder tdt \
  --beam-size 4 \
  --nbest 4

Additive C API symbols:

parakeet_capi_transcribe_path_nbest_json(...);
parakeet_capi_transcribe_pcm_nbest_json(...);

The existing ABI version and all greedy entry points are unchanged. Details and
the JSON shape are in
docs/tdt-nbest.md.

Validation

  • fork CI passed build, server E2E, strict NeMo N-best parity, and the existing greedy closed-loop assertion on 9f6c147
  • test_tdt_beam passed on CPU and CUDA with tdt-0.6b-v3-q8_0.gguf
  • post-review Q4 output on b9abbbc is byte-identical to 9f6c147, including all hypotheses, scores, and token metadata
  • post-review real-model test passed in regular and ASAN-only builds
  • 25b8a3a fixes the beam-size-1 legal blank fallback and adds a model-independent regression
  • model-independent suite passed 18/18 available tests on 25b8a3a; the regression also passed 100 consecutive runs
  • both C API paths, raw and normalized ranking, early validation, deterministic repeat decoding, and exported C symbols are covered

A full manually dispatched hosted workflow passed against 2f37c2c: Linux build and 18 model-independent tests, server E2E, strict NeMo N-best parity, and the existing greedy transcript assertion.

CUDA cold-process sanity check on an RTX PRO 2000 Blackwell, 7.4 s fixture,
with existing unrelated GPU services left running:

  • greedy median: 2.15 s
  • beam-4 median: 2.49 s

Those numbers include container startup, model load, mel, and encoder; they are
not presented as a decoder-only benchmark.

Open questions

  1. Is default sequence beam a useful first N-best surface, or should this target malsd_batch before merge?
  2. Is JSON-only additive C API appropriate initially, or should the library expose owned typed hypothesis arrays too?
  3. Should CUDA optimization wait for or build on #46, which accelerates the same repeated predictor/joint steps?

AI assistance disclosure

This change was implemented with Codex assistance. The commits use the
repository's required Assisted-by: trailer and intentionally have no AI
Signed-off-by or Co-Authored-By trailer. Human review and ownership are
required before merge.

Implement NeMo-compatible default TDT beam search, expose ranked hypotheses through Model, C API, and CLI surfaces, and add CPU/CUDA plus closed-loop NeMo parity coverage.

Assisted-by: Codex:gpt-5 [Codex]
@inlanger
inlanger marked this pull request as ready for review July 26, 2026 17:54
@inlanger

Copy link
Copy Markdown
Contributor Author

Fork CI completed successfully: build, server E2E, and strict closed-loop NeMo TDT N-best parity. Marking this ready for design/API feedback; the open questions in the description are intentionally unresolved.

inlanger and others added 3 commits July 26, 2026 20:26
Validate requests before inference, reject malformed duration tables and non-progressing zero-duration expansions, and broaden real-model API coverage.

[skip ci]

Assisted-by: Codex:gpt-5 [Codex]
Select the highest-scoring positive duration when zero is the only top-k blank duration, and cover the beam-1 fallback independently of a model.

Assisted-by: Codex:gpt-5 [Codex]
Keep the TDT beam test portable across compilers instead of relying on a transitive standard-library include.

Assisted-by: Codex:gpt-5 [Codex]
@inlanger

Copy link
Copy Markdown
Contributor Author

Addressed the requested beam_size=1 edge case in 25b8a3a (plus the portable test include in 2f37c2c).

  • The blank fallback now scans the full duration distribution and selects the highest-scoring positive duration.
  • test_tdt_beam_core covers duration scores where 0 ranks first, duration 1 is the smallest legal value, and duration 2 is the best-scoring legal blank transition.
  • The docs call out this intentional improvement over NeMo's min_non_zero_duration_idx fallback.

The current head passed the full hosted workflow: Linux build/model-independent tests, server E2E, strict NeMo N-best parity, and the greedy transcript assertion: https://github.com/inlanger/parakeet.cpp/actions/runs/30248394029

@mudler
mudler merged commit e747acd into mudler:master Jul 28, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants