diff --git a/CMakeLists.txt b/CMakeLists.txt index 39c4f95..a8afde8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -76,6 +76,7 @@ set(PARAKEET_SRC src/joint.cpp src/tdt.cpp src/rnnt.cpp + src/transducer_batch.cpp src/tokenizer.cpp src/search.cpp src/transcription.cpp) diff --git a/README.md b/README.md index d726b90..edc1821 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,29 @@ The `parakeet-cli` binary lands at `build/examples/cli/parakeet-cli`. --- +## Batching + +Single-clip transcription is the default and needs no flags: every `transcribe` call runs one clip at a time, byte-for-byte identical to before. Batching is an opt-in path for decoding several clips together, which matters when you serve many concurrent requests on a GPU. + +The win is on the **decode** side. A transducer (TDT/RNN-T) decodes autoregressively with tiny per-step prediction-LSTM and joint GEMMs; one clip launches hundreds of these matvec-sized kernels and leaves the GPU mostly idle between launches. Decoding N clips together coalesces each step into one batched GEMM, so the device stays busy. On the NVIDIA GB10 this reaches about **10-12x** at batch size 16 (CPU about 3-5x); the encoder is already compute-bound, so batching it gives no throughput win. CTC has no autoregressive decode, so batching does not apply to standalone CTC models. The batched path is bit-identical to running the clips one by one (greedy decode is deterministic). Full numbers and per-model tables are in [`benchmarks/BENCHMARK.md`](benchmarks/BENCHMARK.md#batched-decode-throughput). + +Measure it yourself: + +```bash +# Decode-only: serial vs batched decode of one clip replicated B times (the win in isolation). +parakeet-cli bench-decode --model --audio [--batch-sizes 1,4,8,16] [--threads N] [--reps R] [--json ] + +# Full transcribe (encoder + decode) over a manifest at several batch sizes. +parakeet-cli bench-batch --model --manifest [--decoder ctc|tdt] [--threads N] [--batch-sizes 1,4,8] [--json ] +``` + +To batch from code, use the batched entry points (single-clip B=1 is just N=1): + +- C++ (`src/model.hpp`): `Model::transcribe_16k_batch(pcms16k, decoder)` and `transcribe_16k_batch_with_timestamps(...)` take N clips of 16 kHz mono float PCM and return N results. +- C-API (`include/parakeet_capi.h`): `parakeet_capi_transcribe_pcm_batch(...)` (N transcripts) and `parakeet_capi_transcribe_pcm_batch_json(...)` (one JSON array of N `{text,words,tokens}` objects). These are what LocalAI's `parakeet-cpp` backend calls to coalesce concurrent requests; it leaves batching off by default and exposes a `batch_max_size` option to opt in. + +--- + ## C-API (`libparakeet.so`) `include/parakeet_capi.h` defines a flat, exception-free C-API meant for `dlopen` / FFI / LocalAI integration. Build the shared library with `-DPARAKEET_SHARED=ON`: diff --git a/benchmarks/BENCHMARK.md b/benchmarks/BENCHMARK.md index 352945b..3017b5c 100644 --- a/benchmarks/BENCHMARK.md +++ b/benchmarks/BENCHMARK.md @@ -84,6 +84,38 @@ Averaged over all models (LibriSpeech). Size is the mean GGUF size as a fraction > f32 is the faithful reference (agreement ≈ 0). q8_0 is near-lossless; K-quants (q6_k→q4_k) shrink the model further at a small, monotonic accuracy cost. See the per-model quant plots below. +## Batched decode throughput + +Decode batching coalesces the per-step prediction-LSTM and joint-network GEMMs across several utterances into single batched ops, so one decode loop advances B clips at once. It applies to **transducer (TDT/RNN-T) models only** — CTC has no autoregressive decode to batch. Speedup is `serial_ms / batched_ms`: the wall-clock of B independent single-clip decodes divided by one batched decode over the same B copies (encoder cost is paid once and excluded). The `clips/s @B=16` column is the batched decode throughput at B=16. + +**CPU** (cpu, q5_k, 8 threads), best-of-5, one clip replicated B times. + + +| Model | B=1 | B=4 | B=8 | B=16 | clips/s @B=16 | +|---|---|---|---|---|---| +| rnnt-0.6b | 0.98× | 2.81× | 3.91× | 5.01× | 378.9 | +| rnnt-1.1b | 1.01× | 2.27× | 3.54× | 4.19× | 366.1 | +| rt-eou-120m-v1 | 1.00× | 2.32× | 3.00× | 3.36× | 638.1 | +| tdt-0.6b-v2 | 0.90× | 2.78× | 4.12× | 5.19× | 458.8 | +| tdt-0.6b-v3 | 0.82× | 3.26× | 5.90× | 4.88× | 253.2 | +| tdt-1.1b | 0.98× | 2.55× | 3.74× | 4.57× | 483.2 | +| tdt_ctc-1.1b | 1.00× | 2.59× | 3.80× | 4.73× | 458.8 | +| tdt_ctc-110m | 1.03× | 2.46× | 3.12× | 3.61× | 834.5 | + +**GPU** (CUDA0, f16, 8 threads), best-of-5, one clip replicated B times. + + +| Model | B=1 | B=4 | B=8 | B=16 | clips/s @B=16 | +|---|---|---|---|---|---| +| rnnt-0.6b | 1.00× | 3.32× | 6.46× | 11.44× | 585.8 | +| rnnt-1.1b | 1.00× | 3.31× | 6.30× | 11.17× | 594.1 | +| rt-eou-120m-v1 | 0.98× | 3.11× | 5.87× | 10.25× | 951.1 | +| tdt-0.6b-v2 | 1.00× | 3.61× | 6.88× | 12.42× | 712.0 | +| tdt-0.6b-v3 | 1.00× | 3.45× | 6.49× | 11.45× | 614.5 | +| tdt-1.1b | 1.00× | 3.55× | 6.72× | 12.17× | 817.9 | +| tdt_ctc-1.1b | 1.00× | 3.45× | 6.50× | 11.47× | 746.9 | +| tdt_ctc-110m | 1.01× | 3.29× | 6.33× | 10.89× | 1259.7 | + ## Plots ### RTFx per model — NeMo vs ours (all dtypes), LibriSpeech diff --git a/benchmarks/results/decode_batch/cpu/rnnt-0.6b.json b/benchmarks/results/decode_batch/cpu/rnnt-0.6b.json new file mode 100644 index 0000000..ee23d2d --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/rnnt-0.6b.json @@ -0,0 +1,50 @@ +{ + "model": "rnnt-0.6b-q5_k.gguf", + "decoder": "ctc", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 13.32, + "batched_ms": 13.65, + "speedup": 0.98, + "serial_cps": 75.1, + "batched_cps": 73.3 + }, + { + "B": 4, + "serial_ms": 54.99, + "batched_ms": 19.56, + "speedup": 2.81, + "serial_cps": 72.7, + "batched_cps": 204.5 + }, + { + "B": 8, + "serial_ms": 103.03, + "batched_ms": 26.35, + "speedup": 3.91, + "serial_cps": 77.7, + "batched_cps": 303.6 + }, + { + "B": 16, + "serial_ms": 211.7, + "batched_ms": 42.23, + "speedup": 5.01, + "serial_cps": 75.6, + "batched_cps": 378.9 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/rnnt-1.1b.json b/benchmarks/results/decode_batch/cpu/rnnt-1.1b.json new file mode 100644 index 0000000..72f22ab --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/rnnt-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "rnnt-1.1b-q5_k.gguf", + "decoder": "ctc", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 11.36, + "batched_ms": 11.26, + "speedup": 1.01, + "serial_cps": 88.0, + "batched_cps": 88.8 + }, + { + "B": 4, + "serial_ms": 45.26, + "batched_ms": 19.92, + "speedup": 2.27, + "serial_cps": 88.4, + "batched_cps": 200.8 + }, + { + "B": 8, + "serial_ms": 94.43, + "batched_ms": 26.68, + "speedup": 3.54, + "serial_cps": 84.7, + "batched_cps": 299.9 + }, + { + "B": 16, + "serial_ms": 183.17, + "batched_ms": 43.7, + "speedup": 4.19, + "serial_cps": 87.4, + "batched_cps": 366.1 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/rt-eou-120m-v1.json b/benchmarks/results/decode_batch/cpu/rt-eou-120m-v1.json new file mode 100644 index 0000000..7abac20 --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/rt-eou-120m-v1.json @@ -0,0 +1,50 @@ +{ + "model": "realtime_eou_120m-v1-q5_k.gguf", + "decoder": "ctc", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 132, + "d_model": 512, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 5.08, + "batched_ms": 5.07, + "speedup": 1.0, + "serial_cps": 196.9, + "batched_cps": 197.3 + }, + { + "B": 4, + "serial_ms": 20.31, + "batched_ms": 8.75, + "speedup": 2.32, + "serial_cps": 196.9, + "batched_cps": 457.4 + }, + { + "B": 8, + "serial_ms": 42.16, + "batched_ms": 14.07, + "speedup": 3.0, + "serial_cps": 189.8, + "batched_cps": 568.6 + }, + { + "B": 16, + "serial_ms": 84.35, + "batched_ms": 25.07, + "speedup": 3.36, + "serial_cps": 189.7, + "batched_cps": 638.1 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/tdt-0.6b-v2.json b/benchmarks/results/decode_batch/cpu/tdt-0.6b-v2.json new file mode 100644 index 0000000..220c44a --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/tdt-0.6b-v2.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-0.6b-v2-q5_k.gguf", + "decoder": "tdt", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 10.89, + "batched_ms": 12.12, + "speedup": 0.9, + "serial_cps": 91.9, + "batched_cps": 82.5 + }, + { + "B": 4, + "serial_ms": 43.97, + "batched_ms": 15.83, + "speedup": 2.78, + "serial_cps": 91.0, + "batched_cps": 252.8 + }, + { + "B": 8, + "serial_ms": 89.54, + "batched_ms": 21.71, + "speedup": 4.12, + "serial_cps": 89.3, + "batched_cps": 368.5 + }, + { + "B": 16, + "serial_ms": 181.13, + "batched_ms": 34.88, + "speedup": 5.19, + "serial_cps": 88.3, + "batched_cps": 458.8 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/tdt-0.6b-v3.json b/benchmarks/results/decode_batch/cpu/tdt-0.6b-v3.json new file mode 100644 index 0000000..3850805 --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/tdt-0.6b-v3.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-0.6b-v3-q5_k.gguf", + "decoder": "tdt", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 17.43, + "batched_ms": 21.17, + "speedup": 0.82, + "serial_cps": 57.4, + "batched_cps": 47.2 + }, + { + "B": 4, + "serial_ms": 126.46, + "batched_ms": 38.81, + "speedup": 3.26, + "serial_cps": 31.6, + "batched_cps": 103.1 + }, + { + "B": 8, + "serial_ms": 262.68, + "batched_ms": 44.54, + "speedup": 5.9, + "serial_cps": 30.5, + "batched_cps": 179.6 + }, + { + "B": 16, + "serial_ms": 308.27, + "batched_ms": 63.18, + "speedup": 4.88, + "serial_cps": 51.9, + "batched_cps": 253.2 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/tdt-1.1b.json b/benchmarks/results/decode_batch/cpu/tdt-1.1b.json new file mode 100644 index 0000000..1e232f1 --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/tdt-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-1.1b-q5_k.gguf", + "decoder": "tdt", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 9.2, + "batched_ms": 9.42, + "speedup": 0.98, + "serial_cps": 108.7, + "batched_cps": 106.2 + }, + { + "B": 4, + "serial_ms": 37.35, + "batched_ms": 14.62, + "speedup": 2.55, + "serial_cps": 107.1, + "batched_cps": 273.6 + }, + { + "B": 8, + "serial_ms": 76.22, + "batched_ms": 20.36, + "speedup": 3.74, + "serial_cps": 105.0, + "batched_cps": 393.0 + }, + { + "B": 16, + "serial_ms": 151.37, + "batched_ms": 33.11, + "speedup": 4.57, + "serial_cps": 105.7, + "batched_cps": 483.2 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/tdt_ctc-1.1b.json b/benchmarks/results/decode_batch/cpu/tdt_ctc-1.1b.json new file mode 100644 index 0000000..02be6fc --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/tdt_ctc-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "tdt_ctc-1.1b-q5_k.gguf", + "decoder": "tdt", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 10.13, + "batched_ms": 10.15, + "speedup": 1.0, + "serial_cps": 98.7, + "batched_cps": 98.5 + }, + { + "B": 4, + "serial_ms": 42.21, + "batched_ms": 16.27, + "speedup": 2.59, + "serial_cps": 94.8, + "batched_cps": 245.8 + }, + { + "B": 8, + "serial_ms": 82.49, + "batched_ms": 21.72, + "speedup": 3.8, + "serial_cps": 97.0, + "batched_cps": 368.3 + }, + { + "B": 16, + "serial_ms": 164.85, + "batched_ms": 34.88, + "speedup": 4.73, + "serial_cps": 97.1, + "batched_cps": 458.8 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/cpu/tdt_ctc-110m.json b/benchmarks/results/decode_batch/cpu/tdt_ctc-110m.json new file mode 100644 index 0000000..1f6f8a8 --- /dev/null +++ b/benchmarks/results/decode_batch/cpu/tdt_ctc-110m.json @@ -0,0 +1,50 @@ +{ + "model": "tdt_ctc-110m-q5_k.gguf", + "decoder": "tdt", + "backend": "cpu", + "dtype": "q5_k", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 512, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 5.66, + "batched_ms": 5.5, + "speedup": 1.03, + "serial_cps": 176.8, + "batched_cps": 182.0 + }, + { + "B": 4, + "serial_ms": 17.7, + "batched_ms": 7.2, + "speedup": 2.46, + "serial_cps": 225.9, + "batched_cps": 555.6 + }, + { + "B": 8, + "serial_ms": 35.67, + "batched_ms": 11.42, + "speedup": 3.12, + "serial_cps": 224.3, + "batched_cps": 700.3 + }, + { + "B": 16, + "serial_ms": 69.27, + "batched_ms": 19.17, + "speedup": 3.61, + "serial_cps": 231.0, + "batched_cps": 834.5 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/rnnt-0.6b.json b/benchmarks/results/decode_batch/gpu/rnnt-0.6b.json new file mode 100644 index 0000000..b8e644a --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/rnnt-0.6b.json @@ -0,0 +1,50 @@ +{ + "model": "rnnt-0.6b", + "decoder": "rnnt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 19.42, + "batched_ms": 19.48, + "speedup": 1.0, + "serial_cps": 51.5, + "batched_cps": 51.3 + }, + { + "B": 4, + "serial_ms": 77.73, + "batched_ms": 23.43, + "speedup": 3.32, + "serial_cps": 51.5, + "batched_cps": 170.7 + }, + { + "B": 8, + "serial_ms": 156.65, + "batched_ms": 24.25, + "speedup": 6.46, + "serial_cps": 51.1, + "batched_cps": 329.9 + }, + { + "B": 16, + "serial_ms": 312.46, + "batched_ms": 27.32, + "speedup": 11.44, + "serial_cps": 51.2, + "batched_cps": 585.8 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/rnnt-1.1b.json b/benchmarks/results/decode_batch/gpu/rnnt-1.1b.json new file mode 100644 index 0000000..4db1b0b --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/rnnt-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "rnnt-1.1b", + "decoder": "rnnt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 18.73, + "batched_ms": 18.73, + "speedup": 1.0, + "serial_cps": 53.4, + "batched_cps": 53.4 + }, + { + "B": 4, + "serial_ms": 74.74, + "batched_ms": 22.55, + "speedup": 3.31, + "serial_cps": 53.5, + "batched_cps": 177.4 + }, + { + "B": 8, + "serial_ms": 150.25, + "batched_ms": 23.83, + "speedup": 6.3, + "serial_cps": 53.2, + "batched_cps": 335.7 + }, + { + "B": 16, + "serial_ms": 300.95, + "batched_ms": 26.93, + "speedup": 11.17, + "serial_cps": 53.2, + "batched_cps": 594.1 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/rt-eou-120m-v1.json b/benchmarks/results/decode_batch/gpu/rt-eou-120m-v1.json new file mode 100644 index 0000000..cdbcdcd --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/rt-eou-120m-v1.json @@ -0,0 +1,50 @@ +{ + "model": "rt-eou-120m-v1", + "decoder": "rnnt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 132, + "d_model": 512, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 10.44, + "batched_ms": 10.67, + "speedup": 0.98, + "serial_cps": 95.7, + "batched_cps": 93.7 + }, + { + "B": 4, + "serial_ms": 42.6, + "batched_ms": 13.68, + "speedup": 3.11, + "serial_cps": 93.9, + "batched_cps": 292.4 + }, + { + "B": 8, + "serial_ms": 86.26, + "batched_ms": 14.69, + "speedup": 5.87, + "serial_cps": 92.7, + "batched_cps": 544.8 + }, + { + "B": 16, + "serial_ms": 172.51, + "batched_ms": 16.82, + "speedup": 10.25, + "serial_cps": 92.7, + "batched_cps": 951.1 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/tdt-0.6b-v2.json b/benchmarks/results/decode_batch/gpu/tdt-0.6b-v2.json new file mode 100644 index 0000000..66f9c3b --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/tdt-0.6b-v2.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-0.6b-v2", + "decoder": "tdt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 17.33, + "batched_ms": 17.37, + "speedup": 1.0, + "serial_cps": 57.7, + "batched_cps": 57.6 + }, + { + "B": 4, + "serial_ms": 69.34, + "batched_ms": 19.22, + "speedup": 3.61, + "serial_cps": 57.7, + "batched_cps": 208.1 + }, + { + "B": 8, + "serial_ms": 139.06, + "batched_ms": 20.21, + "speedup": 6.88, + "serial_cps": 57.5, + "batched_cps": 395.9 + }, + { + "B": 16, + "serial_ms": 279.11, + "batched_ms": 22.47, + "speedup": 12.42, + "serial_cps": 57.3, + "batched_cps": 712.0 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/tdt-0.6b-v3.json b/benchmarks/results/decode_batch/gpu/tdt-0.6b-v3.json new file mode 100644 index 0000000..8446945 --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/tdt-0.6b-v3.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-0.6b-v3", + "decoder": "tdt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 18.59, + "batched_ms": 18.58, + "speedup": 1.0, + "serial_cps": 53.8, + "batched_cps": 53.8 + }, + { + "B": 4, + "serial_ms": 74.34, + "batched_ms": 21.53, + "speedup": 3.45, + "serial_cps": 53.8, + "batched_cps": 185.8 + }, + { + "B": 8, + "serial_ms": 148.99, + "batched_ms": 22.96, + "speedup": 6.49, + "serial_cps": 53.7, + "batched_cps": 348.4 + }, + { + "B": 16, + "serial_ms": 298.14, + "batched_ms": 26.04, + "speedup": 11.45, + "serial_cps": 53.7, + "batched_cps": 614.5 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/tdt-1.1b.json b/benchmarks/results/decode_batch/gpu/tdt-1.1b.json new file mode 100644 index 0000000..f207e83 --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/tdt-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "tdt-1.1b", + "decoder": "tdt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 14.79, + "batched_ms": 14.73, + "speedup": 1.0, + "serial_cps": 67.6, + "batched_cps": 67.9 + }, + { + "B": 4, + "serial_ms": 58.96, + "batched_ms": 16.62, + "speedup": 3.55, + "serial_cps": 67.8, + "batched_cps": 240.7 + }, + { + "B": 8, + "serial_ms": 118.73, + "batched_ms": 17.67, + "speedup": 6.72, + "serial_cps": 67.4, + "batched_cps": 452.8 + }, + { + "B": 16, + "serial_ms": 238.01, + "batched_ms": 19.56, + "speedup": 12.17, + "serial_cps": 67.2, + "batched_cps": 817.9 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/tdt_ctc-1.1b.json b/benchmarks/results/decode_batch/gpu/tdt_ctc-1.1b.json new file mode 100644 index 0000000..0798241 --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/tdt_ctc-1.1b.json @@ -0,0 +1,50 @@ +{ + "model": "tdt_ctc-1.1b", + "decoder": "tdt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 1024, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 15.18, + "batched_ms": 15.18, + "speedup": 1.0, + "serial_cps": 65.9, + "batched_cps": 65.9 + }, + { + "B": 4, + "serial_ms": 60.48, + "batched_ms": 17.51, + "speedup": 3.45, + "serial_cps": 66.1, + "batched_cps": 228.5 + }, + { + "B": 8, + "serial_ms": 122.27, + "batched_ms": 18.81, + "speedup": 6.5, + "serial_cps": 65.4, + "batched_cps": 425.2 + }, + { + "B": 16, + "serial_ms": 245.75, + "batched_ms": 21.42, + "speedup": 11.47, + "serial_cps": 65.1, + "batched_cps": 746.9 + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/decode_batch/gpu/tdt_ctc-110m.json b/benchmarks/results/decode_batch/gpu/tdt_ctc-110m.json new file mode 100644 index 0000000..c90a6a0 --- /dev/null +++ b/benchmarks/results/decode_batch/gpu/tdt_ctc-110m.json @@ -0,0 +1,50 @@ +{ + "model": "tdt_ctc-110m", + "decoder": "tdt", + "backend": "CUDA0", + "dtype": "f16", + "threads": 8, + "reps": 5, + "clip_frames": 131, + "d_model": 512, + "batch_sizes": [ + 1, + 4, + 8, + 16 + ], + "rows": [ + { + "B": 1, + "serial_ms": 8.53, + "batched_ms": 8.42, + "speedup": 1.01, + "serial_cps": 117.2, + "batched_cps": 118.7 + }, + { + "B": 4, + "serial_ms": 33.96, + "batched_ms": 10.33, + "speedup": 3.29, + "serial_cps": 117.8, + "batched_cps": 387.3 + }, + { + "B": 8, + "serial_ms": 68.86, + "batched_ms": 10.88, + "speedup": 6.33, + "serial_cps": 116.2, + "batched_cps": 735.1 + }, + { + "B": 16, + "serial_ms": 138.3, + "batched_ms": 12.7, + "speedup": 10.89, + "serial_cps": 115.7, + "batched_cps": 1259.7 + } + ] +} \ No newline at end of file diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index ed7d017..421bcff 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -5,7 +5,16 @@ #include "audio_io.hpp" #include "streaming.hpp" #include "transcription.hpp" -#include "ggml_graph.hpp" // pk::set_num_threads +#include "ggml_graph.hpp" // pk::set_num_threads, pk::global_backend +#include "backend.hpp" // pk::ensure_weights_realized +#include "encoder.hpp" +#include "prediction.hpp" +#include "joint.hpp" +#include "tdt.hpp" +#include "rnnt.hpp" +#include "transducer_batch.hpp" +#include "mel.hpp" +#include "mel_gpu.hpp" #include "ggml.h" #include "gguf.h" #include @@ -16,6 +25,7 @@ #include #include #include +#include #include #include @@ -637,6 +647,474 @@ static int cmd_bench(int argc, char** argv) { return 0; } +// Measures BATCHED encoder throughput at one or more batch sizes. Mirrors +// cmd_bench's arg parsing / model load / manifest read / warmup, but instead of +// timing one clip at a time it groups the clips into batches of size B and times +// the wall-clock cost of running every batch through transcribe_pcm_batch. +// +// The B=1 row goes through the SAME fused batched encoder path with one-clip +// batches, so it is the apples-to-apples baseline against which the batching win +// (B=4, B=8, ...) is read. +static int cmd_bench_batch(int argc, char** argv) { + std::string model, manifest, decoder_str, json_out; + std::string batch_sizes_str = "1,4,8"; + int threads = 0; // 0 == unset -> use the components' built-in default + for (int i = 0; i < argc; ++i) { + if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) { + model = argv[++i]; + } else if (std::strcmp(argv[i], "--manifest") == 0 && i + 1 < argc) { + manifest = argv[++i]; + } else if (std::strcmp(argv[i], "--decoder") == 0 && i + 1 < argc) { + decoder_str = argv[++i]; + } else if (std::strcmp(argv[i], "--threads") == 0 && i + 1 < argc) { + threads = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--batch-sizes") == 0 && i + 1 < argc) { + batch_sizes_str = argv[++i]; + } else if (std::strcmp(argv[i], "--json") == 0 && i + 1 < argc) { + json_out = argv[++i]; + } + } + if (model.empty() || manifest.empty()) { + std::fprintf(stderr, + "usage: parakeet-cli bench-batch --model --manifest " + "[--decoder ctc|tdt] [--threads N] [--batch-sizes 1,4,8] [--json ]\n"); + return 2; + } + + // Parse --batch-sizes (comma-separated positive ints). + std::vector batch_sizes; + { + std::stringstream ss(batch_sizes_str); + std::string tok; + while (std::getline(ss, tok, ',')) { + // Trim surrounding whitespace. + size_t b = tok.find_first_not_of(" \t"); + if (b == std::string::npos) continue; + size_t e = tok.find_last_not_of(" \t"); + int v = std::atoi(tok.substr(b, e - b + 1).c_str()); + if (v > 0) batch_sizes.push_back(v); + } + } + if (batch_sizes.empty()) { + std::fprintf(stderr, + "parakeet-cli bench-batch: no valid --batch-sizes (want e.g. 1,4,8)\n"); + return 2; + } + + // Resolve the decoder selector (matches `transcribe` / `bench`). + pk::Decoder dec = pk::Decoder::kDefault; + if (!decoder_str.empty()) { + if (decoder_str == "ctc") { + dec = pk::Decoder::kCTC; + } else if (decoder_str == "tdt") { + dec = pk::Decoder::kTDT; + } else { + std::fprintf(stderr, + "parakeet-cli bench-batch: unknown --decoder '%s' (want ctc|tdt)\n", + decoder_str.c_str()); + return 2; + } + } + + // Apply the thread count to EVERY ggml graph computation. When --threads is + // omitted we report the components' built-in default. + int reported_threads = threads; + if (threads > 0) { + pk::set_num_threads(threads); + } else { + reported_threads = 8; // the persistent-backend default (kDefaultThreads) + } + + bool man_ok = false; + std::vector paths = read_manifest(manifest, man_ok); + if (!man_ok) { + std::fprintf(stderr, "parakeet-cli bench-batch: failed to read manifest %s\n", + manifest.c_str()); + return 1; + } + if (paths.empty()) { + std::fprintf(stderr, "parakeet-cli bench-batch: manifest %s has no audio paths\n", + manifest.c_str()); + return 1; + } + + using clock = std::chrono::steady_clock; + auto ms_since = [](clock::time_point t0) { + return std::chrono::duration(clock::now() - t0).count(); + }; + + // Load ALL clips into memory ONCE (untimed) so the per-batch loop only times + // the encoder/decoder work, not audio decode/IO. + std::vector> clips; + clips.reserve(paths.size()); + double total_audio_sec = 0.0; + for (const std::string& p : paths) { + pk::Audio audio; + if (!pk::load_audio_16k_mono(p, audio)) { + std::fprintf(stderr, "parakeet-cli bench-batch: failed to load audio %s\n", + p.c_str()); + return 1; + } + total_audio_sec += (double)audio.samples.size() / 16000.0; + clips.push_back(std::move(audio.samples)); + } + + // Load the model ONCE -- timed, and excluded from per-batch proc_ms. + auto t_load = clock::now(); + std::unique_ptr m = pk::Model::load(model); + double load_ms = ms_since(t_load); + if (!m) { + std::fprintf(stderr, "parakeet-cli bench-batch: failed to load model %s\n", + model.c_str()); + return 1; + } + + // Warm up once (untimed): pays the one-time lazy weight upload / kernel init + // so per-batch timings are steady-state. + { + std::vector> warm{clips[0]}; + try { + (void)m->transcribe_pcm_batch(warm, 16000, dec); + } catch (const std::exception& e) { + std::fprintf(stderr, "parakeet-cli bench-batch: warmup failed: %s\n", e.what()); + return 1; + } + } + + struct BatchResult { int batch_size; double proc_ms; size_t n_clips; + double clips_per_sec; double rtfx; }; + std::vector results; + results.reserve(batch_sizes.size()); + + for (int B : batch_sizes) { + auto t_proc = clock::now(); + for (size_t s = 0; s < clips.size(); s += (size_t)B) { + size_t end = std::min(clips.size(), s + (size_t)B); + std::vector> chunk(clips.begin() + (long)s, + clips.begin() + (long)end); + try { + (void)m->transcribe_pcm_batch(chunk, 16000, dec); + } catch (const std::exception& e) { + std::fprintf(stderr, + "parakeet-cli bench-batch: transcribe failed at batch_size=%d: %s\n", + B, e.what()); + return 1; + } + } + double proc_ms = ms_since(t_proc); + double secs = proc_ms / 1000.0; + double clips_per_sec = secs > 0.0 ? (double)clips.size() / secs : 0.0; + double rtfx = secs > 0.0 ? total_audio_sec / secs : 0.0; + results.push_back({B, proc_ms, clips.size(), clips_per_sec, rtfx}); + } + + // Human-readable summary table to stderr. + std::fprintf(stderr, + "\nbench-batch: %zu clips, %.2f s audio, decoder=%s, threads=%d, load_ms=%.1f\n", + clips.size(), total_audio_sec, + decoder_str.empty() ? "default" : decoder_str.c_str(), + reported_threads, load_ms); + std::fprintf(stderr, " %-12s %-14s %-14s %-10s\n", + "batch_size", "proc_ms", "clips/sec", "RTFx"); + for (const BatchResult& r : results) { + std::fprintf(stderr, " %-12d %-14.1f %-14.2f %-10.2f\n", + r.batch_size, r.proc_ms, r.clips_per_sec, r.rtfx); + } + + // Hand-roll the JSON document. + std::string out; + out.reserve(256 + results.size() * 96); + out += "{\"model\":"; + bench_json_string(out, model); + out += ",\"decoder\":"; + bench_json_string(out, decoder_str.empty() ? std::string("default") : decoder_str); + char numbuf[96]; + std::snprintf(numbuf, sizeof(numbuf), ",\"threads\":%d", reported_threads); + out += numbuf; + std::snprintf(numbuf, sizeof(numbuf), ",\"n_clips\":%zu", clips.size()); + out += numbuf; + std::snprintf(numbuf, sizeof(numbuf), ",\"total_audio_sec\":%.6f", total_audio_sec); + out += numbuf; + out += ",\"results\":["; + for (size_t i = 0; i < results.size(); ++i) { + if (i) out += ','; + std::snprintf(numbuf, sizeof(numbuf), + "{\"batch_size\":%d,\"proc_ms\":%.3f,\"clips_per_sec\":%.6f,\"rtfx\":%.6f}", + results[i].batch_size, results[i].proc_ms, + results[i].clips_per_sec, results[i].rtfx); + out += numbuf; + } + out += "]}"; + + if (!json_out.empty()) { + std::ofstream of(json_out, std::ios::binary | std::ios::trunc); + if (!of) { + std::fprintf(stderr, "parakeet-cli bench-batch: failed to write %s\n", + json_out.c_str()); + return 1; + } + of << out << '\n'; + } else { + std::printf("%s\n", out.c_str()); + } + return 0; +} + +// --------------------------------------------------------------------------- +// bench-decode: encode ONE clip once, then time DECODE only -- serial (N +// separate tdt_greedy/rnnt_greedy calls over N copies of the encoder output) +// vs batched (transducer_greedy_batch over the same N copies) -- at several +// batch sizes. Reports decode wall-clock and the batched/serial speedup so the +// GPU win from batched decode can be measured in isolation from the encoder. +// +// The encoder cost is paid once and excluded; only the transducer decode loop +// is timed. Each rep is averaged over R repetitions (best/min recorded). The +// b=0 batched ids are compared against the serial ids (same clip, decode is +// deterministic) as a correctness sanity check. +// --------------------------------------------------------------------------- +static int cmd_bench_decode(int argc, char** argv) { + std::string model, audio, json_out; + std::string batch_sizes_str = "1,4,8,16"; + int threads = 0; // 0 == unset -> use the persistent-backend default + int reps = 5; + for (int i = 0; i < argc; ++i) { + if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) { + model = argv[++i]; + } else if (std::strcmp(argv[i], "--audio") == 0 && i + 1 < argc) { + audio = argv[++i]; + } else if (std::strcmp(argv[i], "--batch-sizes") == 0 && i + 1 < argc) { + batch_sizes_str = argv[++i]; + } else if (std::strcmp(argv[i], "--threads") == 0 && i + 1 < argc) { + threads = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--reps") == 0 && i + 1 < argc) { + reps = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--json") == 0 && i + 1 < argc) { + json_out = argv[++i]; + } + } + if (model.empty() || audio.empty()) { + std::fprintf(stderr, + "usage: parakeet-cli bench-decode --model --audio " + "[--batch-sizes 1,4,8,16] [--threads N] [--reps R] [--json ]\n"); + return 2; + } + if (reps < 1) reps = 1; + + // Parse --batch-sizes (comma-separated positive ints). + std::vector batch_sizes; + { + std::stringstream ss(batch_sizes_str); + std::string tok; + while (std::getline(ss, tok, ',')) { + size_t b = tok.find_first_not_of(" \t"); + if (b == std::string::npos) continue; + size_t e = tok.find_last_not_of(" \t"); + int v = std::atoi(tok.substr(b, e - b + 1).c_str()); + if (v > 0) batch_sizes.push_back(v); + } + } + if (batch_sizes.empty()) { + std::fprintf(stderr, + "parakeet-cli bench-decode: no valid --batch-sizes (want e.g. 1,4,8,16)\n"); + return 2; + } + + if (threads > 0) pk::set_num_threads(threads); + int reported_threads = threads > 0 ? threads : 8; // kDefaultThreads + + // Load the model components over the lower-level loader (we need the encoder + // / prediction / joint pieces, not the high-level Model::transcribe path). + pk::ModelLoader ml; + if (!ml.load(model)) { + std::fprintf(stderr, "parakeet-cli bench-decode: failed to load model %s\n", + model.c_str()); + return 1; + } + pk::ensure_weights_realized(ml); + pk::Encoder enc(ml); + pk::PredictionNet pred(ml); + pk::Joint joint(ml); + const auto& cfg = ml.config(); + const int blank = (int)cfg.blank_id; + const int maxs = (int)cfg.max_symbols; + const std::vector durations = cfg.tdt_durations; + + // Load the WAV. + pk::Audio a; + if (!pk::load_audio_16k_mono(audio, a)) { + std::fprintf(stderr, "parakeet-cli bench-decode: failed to load audio %s\n", + audio.c_str()); + return 1; + } + + // Mel front end (GpuMel on a non-CPU backend, else FFT MelFrontend), exactly + // as model.cpp's transcribe path does. + std::vector feats; + int n_mels = 0, T = 0; + if (std::string(pk::global_backend().device_name()) != "cpu") { + pk::GpuMel gmel(ml); + gmel.compute(a.samples, feats, n_mels, T); + } else { + pk::MelFrontend mel(ml); + mel.compute(a.samples, feats, n_mels, T); + } + + // Encoder -> enc_out [d_model, Tout] (channels-first); transpose to row-major + // enc_row [Tout, d_model] as the decoders expect. + std::vector enc_out; + int dm = 0, Tout = 0; + enc.forward(feats, n_mels, T, enc_out, dm, Tout); + std::vector enc_row((size_t)Tout * dm); + for (int t = 0; t < Tout; ++t) + for (int c = 0; c < dm; ++c) + enc_row[(size_t)t * dm + c] = enc_out[(size_t)c * Tout + t]; + + const bool use_tdt = !durations.empty(); + auto decode_serial_one = [&]() -> std::vector { + return use_tdt + ? pk::tdt_greedy(pred, joint, enc_row, Tout, dm, durations, blank, maxs, nullptr) + : pk::rnnt_greedy(pred, joint, enc_row, Tout, dm, blank, maxs, nullptr); + }; + + using clock = std::chrono::steady_clock; + auto ms_since = [](clock::time_point t0) { + return std::chrono::duration(clock::now() - t0).count(); + }; + + // Warm up (untimed): realize weights / CUDA kernels for both paths. + std::vector serial_ref = decode_serial_one(); + { + std::vector> encs1{enc_row}; + std::vector Ts1{Tout}; + std::vector> ids1; + pk::transducer_greedy_batch(pred, joint, encs1, Ts1, dm, durations, + blank, maxs, ids1, nullptr); + } + + struct Row { int B; double serial_ms; double batched_ms; double speedup; + double serial_cps; double batched_cps; }; + std::vector rows; + rows.reserve(batch_sizes.size()); + bool sanity_ok = true; + + for (int B : batch_sizes) { + std::vector> encs((size_t)B, enc_row); + std::vector Ts((size_t)B, Tout); + + // SERIAL: B separate single-clip decodes, best of R reps. + double serial_ms = 1e300; + for (int r = 0; r < reps; ++r) { + auto t0 = clock::now(); + for (int b = 0; b < B; ++b) (void)decode_serial_one(); + serial_ms = std::min(serial_ms, ms_since(t0)); + } + + // BATCHED: one transducer_greedy_batch over the B copies, best of R reps. + double batched_ms = 1e300; + std::vector> ids_last; + for (int r = 0; r < reps; ++r) { + std::vector> ids; + auto t0 = clock::now(); + pk::transducer_greedy_batch(pred, joint, encs, Ts, dm, durations, + blank, maxs, ids, nullptr); + batched_ms = std::min(batched_ms, ms_since(t0)); + ids_last = std::move(ids); + } + + // Sanity: batched ids[0] must equal the serial decode of the same clip. + if (!ids_last.empty() && ids_last[0] != serial_ref) { + sanity_ok = false; + std::fprintf(stderr, + "parakeet-cli bench-decode: WARN B=%d batched ids[0] != serial " + "(%zu vs %zu tokens) -- batched decode may be buggy\n", + B, ids_last[0].size(), serial_ref.size()); + } + + double speedup = batched_ms > 0.0 ? serial_ms / batched_ms : 0.0; + double serial_cps = serial_ms > 0.0 ? (double)B / (serial_ms / 1000.0) : 0.0; + double batched_cps = batched_ms > 0.0 ? (double)B / (batched_ms / 1000.0) : 0.0; + rows.push_back({B, serial_ms, batched_ms, speedup, serial_cps, batched_cps}); + } + + // Human-readable table to stderr. + std::fprintf(stderr, + "\nbench-decode: clip Tout=%d frames, d_model=%d, decoder=%s, threads=%d, " + "reps=%d (best-of), backend=%s\n", + Tout, dm, use_tdt ? "tdt" : "rnnt", reported_threads, reps, + pk::global_backend().device_name()); + std::fprintf(stderr, " %-6s %-12s %-12s %-10s %-14s %-14s\n", + "B", "serial_ms", "batched_ms", "speedup", "serial_cps", "batched_cps"); + for (const Row& r : rows) { + std::fprintf(stderr, " %-6d %-12.2f %-12.2f %-10.2f %-14.1f %-14.1f\n", + r.B, r.serial_ms, r.batched_ms, r.speedup, + r.serial_cps, r.batched_cps); + } + std::fprintf(stderr, " sanity (batched ids[0]==serial): %s\n", + sanity_ok ? "OK" : "MISMATCH (see WARN above)"); + + // Optional machine-readable JSON document (hand-rolled, same style as + // cmd_bench). Written ONLY when --json is given; the human table above + // always prints regardless. + if (!json_out.empty()) { + // basename of the model gguf path. + std::string model_base = model; + size_t slash = model_base.find_last_of("/\\"); + if (slash != std::string::npos) model_base = model_base.substr(slash + 1); + + std::string out; + out.reserve(512 + rows.size() * 96); + out += "{\"model\":"; + bench_json_string(out, model_base); + out += ",\"decoder\":"; + bench_json_string(out, use_tdt ? std::string("tdt") : std::string("ctc")); + out += ",\"backend\":"; + bench_json_string(out, std::string(pk::global_backend().device_name())); + char nb[64]; + std::snprintf(nb, sizeof(nb), ",\"threads\":%d", reported_threads); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"reps\":%d", reps); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"clip_frames\":%d", Tout); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"d_model\":%d", dm); + out += nb; + out += ",\"batch_sizes\":["; + for (size_t i = 0; i < batch_sizes.size(); ++i) { + if (i) out += ','; + std::snprintf(nb, sizeof(nb), "%d", batch_sizes[i]); + out += nb; + } + out += "],\"rows\":["; + for (size_t i = 0; i < rows.size(); ++i) { + if (i) out += ','; + const Row& r = rows[i]; + std::snprintf(nb, sizeof(nb), "{\"B\":%d", r.B); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"serial_ms\":%.2f", r.serial_ms); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"batched_ms\":%.2f", r.batched_ms); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"speedup\":%.2f", r.speedup); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"serial_cps\":%.1f", r.serial_cps); + out += nb; + std::snprintf(nb, sizeof(nb), ",\"batched_cps\":%.1f", r.batched_cps); + out += nb; + out += '}'; + } + out += "]}"; + + std::ofstream of(json_out, std::ios::binary | std::ios::trunc); + if (!of) { + std::fprintf(stderr, "parakeet-cli bench-decode: failed to write %s\n", + json_out.c_str()); + return 1; + } + of << out << '\n'; + } + return 0; +} + // Run a subcommand, then free the process-global backend while the GPU driver is // still alive (the subcommand's local Model is already destroyed by the time it // returns, releasing its device weight buffer). Avoids the CUDA "driver shutting @@ -654,6 +1132,10 @@ int main(int argc, char** argv) { return run_and_shutdown(cmd_transcribe, argc - 2, argv + 2); if (argc >= 2 && std::strcmp(argv[1], "quantize") == 0) return run_and_shutdown(cmd_quantize, argc - 2, argv + 2); + if (argc >= 2 && std::strcmp(argv[1], "bench-batch") == 0) + return run_and_shutdown(cmd_bench_batch, argc - 2, argv + 2); + if (argc >= 2 && std::strcmp(argv[1], "bench-decode") == 0) + return run_and_shutdown(cmd_bench_decode, argc - 2, argv + 2); if (argc >= 2 && std::strcmp(argv[1], "bench") == 0) return run_and_shutdown(cmd_bench, argc - 2, argv + 2); std::fprintf(stderr, @@ -664,6 +1146,10 @@ int main(int argc, char** argv) { " parakeet-cli quantize " "\n" " parakeet-cli bench --model --manifest " - "[--decoder ctc|tdt] [--threads N] [--json ]\n"); + "[--decoder ctc|tdt] [--threads N] [--json ]\n" + " parakeet-cli bench-batch --model --manifest " + "[--decoder ctc|tdt] [--threads N] [--batch-sizes 1,4,8] [--json ]\n" + " parakeet-cli bench-decode --model --audio " + "[--batch-sizes 1,4,8,16] [--threads N] [--reps R] [--json ]\n"); return 2; } diff --git a/include/parakeet_capi.h b/include/parakeet_capi.h index 7da7105..98cd4a5 100644 --- a/include/parakeet_capi.h +++ b/include/parakeet_capi.h @@ -44,6 +44,21 @@ char* parakeet_capi_transcribe_path(parakeet_ctx* ctx, const char* wav_path, char* parakeet_capi_transcribe_pcm(parakeet_ctx* ctx, const float* samples, int n_samples, int sample_rate, int decoder); +// Transcribe a batch of in-memory mono float PCM clips. `samples` is an array of +// `n_clips` pointers and `n_samples` an array of `n_clips` per-clip lengths; each +// clip is resampled to 16 kHz if `sample_rate != 16000`. `decoder` is as in +// parakeet_capi_transcribe_path (0=default,1=ctc,2=tdt/rnnt). On success returns +// 0 and fills `out` (a caller-allocated array of `n_clips` char*) with malloc'd +// NUL-terminated UTF-8 transcripts; release each with parakeet_capi_free_string. +// On error returns nonzero, sets the context's last error (see +// parakeet_capi_last_error), and leaves every out[] entry NULL: the caller owns +// nothing and has nothing to free. +int parakeet_capi_transcribe_pcm_batch(parakeet_ctx* ctx, + const float* const* samples, + const int* n_samples, int n_clips, + int sample_rate, int decoder, + char** out); + // Transcribe a WAV file returning a malloc'd UTF-8 JSON document with per-word // and per-token timestamps + confidence (matching NeMo timestamps=True and the // 'max_prob' confidence method). `decoder` is as in @@ -61,6 +76,21 @@ char* parakeet_capi_transcribe_pcm(parakeet_ctx* ctx, const float* samples, char* parakeet_capi_transcribe_path_json(parakeet_ctx* ctx, const char* wav_path, int decoder); +// Batched transcription with timestamps, returning ONE malloc'd JSON string that +// is a JSON ARRAY of n_clips objects, each identical in shape to +// parakeet_capi_transcribe_path_json's document ({"text","words","tokens"}). +// samples_concat holds all clips' 16 kHz mono float samples concatenated; +// n_samples gives each clip's sample count; n_clips is the array length. +// decoder: 0=default,1=ctc,2=tdt. PRECONDITION (caller MUST uphold, not +// validated here): the sum of n_samples[0..n_clips) equals the number of floats +// in samples_concat. A larger sum reads out of bounds. +// Returns the JSON string on success (free with parakeet_capi_free_string), or +// NULL on error (see parakeet_capi_last_error). +char* parakeet_capi_transcribe_pcm_batch_json(parakeet_ctx* ctx, + const float* samples_concat, + const int* n_samples, int n_clips, + int sample_rate, int decoder); + // --------------------------------------------------------------------------- // Streaming API (cache-aware streaming RNN-T, e.g. the EOU model // nvidia/parakeet_realtime_eou_120m-v1). The stream session buffers incoming diff --git a/scripts/gen_benchmark_md.py b/scripts/gen_benchmark_md.py index 3b1e674..659a906 100644 --- a/scripts/gen_benchmark_md.py +++ b/scripts/gen_benchmark_md.py @@ -215,6 +215,83 @@ def build_diverse_section(models: list[dict]) -> str: return "\n".join(lines) +# ── Batched decode throughput ────────────────────────────── + +def _decode_batch_subtable(json_dir: Path, kind: str) -> str: + """Render one machine sub-table from decode-batch JSON files in `json_dir`. + + `kind` is "CPU" or "GPU" (for the caption). Returns "" when the directory is + absent or holds no JSON files. + """ + if not json_dir.is_dir(): + return "" + files = sorted(json_dir.glob("*.json")) + if not files: + return "" + + docs = [] + for fp in files: + with open(fp) as f: + d = json.load(f) + d["_label"] = fp.stem + docs.append(d) + docs.sort(key=lambda d: d["_label"]) + + # One-line caption from the first doc's metadata. + meta = docs[0] + threads = meta.get("threads", "?") + backend = meta.get("backend", "?") + reps = meta.get("reps", "?") + dtype = meta.get("dtype", "?") + caption = ( + f"**{kind}** ({backend}, {dtype}, {threads} threads), " + f"best-of-{reps}, one clip replicated B times." + ) + + # Column set: the B values to show (fixed display order, only those present). + show_bs = [1, 4, 8, 16] + header = "| Model | " + " | ".join(f"B={b}" for b in show_bs) + " | clips/s @B=16 |" + sep = "|" + "---|" * (len(show_bs) + 2) + + lines = [caption + "\n", "", header, sep] + for d in docs: + by_b = {int(r["B"]): r for r in d.get("rows", [])} + cells = [] + for b in show_bs: + r = by_b.get(b) + cells.append(f"{r['speedup']:.2f}\u00d7" if r else "\u2014") + r16 = by_b.get(16) + cps16 = f"{r16['batched_cps']:.1f}" if r16 else "\u2014" + lines.append(f"| {d['_label']} | " + " | ".join(cells) + f" | {cps16} |") + return "\n".join(lines) + "\n" + + +def build_batched_decode_section(results_dir: Path) -> str: + """Decode-batching speedup tables (CPU + GPU). Empty string when no data.""" + cpu = _decode_batch_subtable(results_dir / "decode_batch" / "cpu", "CPU") + gpu = _decode_batch_subtable(results_dir / "decode_batch" / "gpu", "GPU") + if not cpu and not gpu: + return "" + + intro = ( + "## Batched decode throughput\n\n" + "Decode batching coalesces the per-step prediction-LSTM and joint-network " + "GEMMs across several utterances into single batched ops, so one decode " + "loop advances B clips at once. It applies to **transducer (TDT/RNN-T) " + "models only** \u2014 CTC has no autoregressive decode to batch. Speedup is " + "`serial_ms / batched_ms`: the wall-clock of B independent single-clip " + "decodes divided by one batched decode over the same B copies (encoder " + "cost is paid once and excluded). The `clips/s @B=16` column is the " + "batched decode throughput at B=16.\n" + ) + parts = [intro] + if cpu: + parts.append("\n" + cpu) + if gpu: + parts.append("\n" + gpu) + return "".join(parts) + + # ── Findings ───────────────────────────────────────────────────────────────────── def build_findings(models: list[dict], dtypes: list[str]) -> str: @@ -341,6 +418,7 @@ def main(): build_methodology(models, dtypes) + "\n", build_headline_table(models) + "\n", build_quant_table(models, dtypes) + "\n", + build_batched_decode_section(results_dir) + "\n", build_plots_section(plots_dir, out_path) + "\n", build_diverse_section(models) + "\n", build_findings(models, dtypes) + "\n", diff --git a/src/conformer.cpp b/src/conformer.cpp index 2bc3038..8ef6645 100644 --- a/src/conformer.cpp +++ b/src/conformer.cpp @@ -27,10 +27,148 @@ static ggml_tensor* clone_weight_opt(ggml_context* ctx, const ModelLoader& ml, } // Build the ConformerConvolution sub-graph (everything AFTER norm_conv) on the -// conv input `c` (= norm_conv(residual)), ne [D, T] (channels fastest). Returns -// the conv output tensor, ne [D, T] (row-major [T, D]). See header for the NeMo -// `ConformerConvolution.forward` mapping. `pool` keeps host-side mask/scale/ -// shift buffers alive until the enclosing Backend::compute finishes. +// conv input `c` (= norm_conv(residual)), ne [D, T, B] (channels fastest, batch +// on ne2). Returns the conv output tensor, ne [D, T, B]. See header for the NeMo +// `ConformerConvolution.forward` mapping. `valid_len[b]` is the per-item valid +// time length; the pre-depthwise pad mask is built per item so trailing-pad +// frames cannot leak through the time-mixing depthwise conv. All other ops are +// per-frame (1x1 convs == matmul, GLU, norm, silu) and broadcast over ne2. +// `pool` keeps host-side mask/scale/shift buffers alive until the enclosing +// Backend::compute finishes. +static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, + const std::string& pre, ggml_tensor* c, + int D, int T, int K, int B, + const std::vector& valid_len, + const std::string& conv_norm_type, + bool conv_causal, GraphInputPool& pool) { + const float ln_eps = 1e-5f; + const float bn_eps = 1e-5f; + const int pad = (K - 1) / 2; // symmetric padding (offline model) + + // -- pointwise_conv1 (Conv1d d->2d, k=1): 1x1 conv == linear over channels. + ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight"); + pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d] + ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias"); + ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2d, T, B] + if (pw1b) y = ggml_add(ctx, y, pw1b); + + // -- GLU over channel dim (NeMo F.glu(x, dim=1)). y is [2D, T, B]; each half + // is D wide along ne0, stepping T via nb[1] and B via nb[2]. + ggml_tensor* a = ggml_view_3d(ctx, y, D, T, B, y->nb[1], y->nb[2], 0); + ggml_tensor* b = ggml_view_3d(ctx, y, D, T, B, y->nb[1], y->nb[2], + (size_t)D * y->nb[0]); + ggml_tensor* glu = ggml_mul(ctx, ggml_cont(ctx, a), + ggml_sigmoid(ctx, ggml_cont(ctx, b))); // [d, T, B] + + // -- pad_mask: zero padded time positions before depthwise conv, per item. + // The depthwise conv is the only time-mixing op here; without this mask the + // trailing-pad region (no longer zero after pointwise bias/GLU) would leak + // into the last valid output frame of a batched item. Mask is [1, T, B] and + // broadcasts over ne0 (D). Emit only if some item has valid_len < T. + bool need_mask = false; + for (int bi = 0; bi < B; ++bi) if (valid_len[bi] < T) { need_mask = true; break; } + if (need_mask) { + std::vector& md = pool.alloc_f32((size_t)T * B); + for (int bi = 0; bi < B; ++bi) + for (int t = 0; t < T; ++t) + md[(size_t)bi * T + t] = (t < valid_len[bi]) ? 1.0f : 0.0f; + int64_t tm_ne[3] = {1, T, B}; + ggml_tensor* tmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, tm_ne, + md.data(), md.size() * sizeof(float)); + glu = ggml_mul(ctx, glu, tmask); + } + + // -- depthwise_conv (Conv1d d->d, k=K, groups=d). F32 im2col throughout. + // Transpose ne0<->ne1 (keep batch on ne2): [D,T,B] -> [T,C,B]. + ggml_tensor* glu_tcb = ggml_cont(ctx, ggml_permute(ctx, glu, 1, 0, 2, 3)); // [T, C, B] + ggml_tensor* dww = clone_weight(ctx, ml, pre + "conv.depthwise_conv.weight"); // [K,1,C] + ggml_tensor* dw = nullptr; + { + // ggml's 1D im2col asserts b->ne[3] == 1, so the depthwise conv cannot + // take a batch dim. Run it per item (each slice has ne[3]==1, exactly the + // B=1 path) and reassemble along the batch axis. For B==1 the loop runs + // once and skips the concat, so the result is byte-identical to before. + const int Tn = (int)glu_tcb->ne[0]; + const int Cn = (int)glu_tcb->ne[1]; + for (int bi = 0; bi < B; ++bi) { + // item bi's [T, C, 1] slice out of glu_tcb [T, C, B]; cont for im2col. + ggml_tensor* item = ggml_view_3d(ctx, glu_tcb, Tn, Cn, 1, + glu_tcb->nb[1], glu_tcb->nb[2], + (size_t)bi * glu_tcb->nb[2]); // [T,C,1] + item = ggml_cont(ctx, item); + // im2col layout: [W=T, H=1, C, N=1]. + ggml_tensor* nb = ggml_reshape_4d(ctx, item, Tn, 1, Cn, 1); // [T,1,C,1] + ggml_tensor* ic; + if (conv_causal) { + ggml_tensor* nbp = ggml_pad_ext(ctx, nb, /*lp0*/K - 1, /*rp0*/0, + 0, 0, 0, 0, 0, 0); // [T+K-1,1,C,1] + ic = ggml_im2col(ctx, dww, nbp, /*s0*/1, /*s1*/0, + /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/0, + /*is_2D*/false, GGML_TYPE_F32); + } else { + ic = ggml_im2col(ctx, dww, nb, /*s0*/1, /*s1*/0, + /*p0*/pad, /*p1*/0, /*d0*/1, /*d1*/0, + /*is_2D*/false, GGML_TYPE_F32); + } + // mul_mat result r2 ne = [OW=T, 1, C, 1]; drop the unit ne1 -> [T, C, 1]. + ggml_tensor* r2 = ggml_mul_mat(ctx, ic, dww); + ggml_tensor* item_dw = ggml_reshape_3d(ctx, r2, r2->ne[0], r2->ne[2], 1); // [T,C,1] + // reassemble along ne2 (batch). For B==1 there is no concat. + dw = (dw == nullptr) ? item_dw : ggml_concat(ctx, dw, item_dw, /*dim*/2); + } + // dw is [OW=T, C, B]. + } + ggml_tensor* dwb = clone_weight_opt(ctx, ml, pre + "conv.depthwise_conv.bias"); // [C] + // transpose ne0<->ne1 (keep batch on ne2): [T,C,B] -> [C,T,B]. + ggml_tensor* dwt = ggml_cont(ctx, ggml_permute(ctx, dw, 1, 0, 2, 3)); // [C, T, B] + if (dwb) dwt = ggml_add(ctx, dwt, dwb); // broadcast [C] over T,B + + // -- norm (between depthwise conv and SiLU). + ggml_tensor* normed; + if (conv_norm_type == "layer_norm") { + ggml_tensor* g = clone_weight(ctx, ml, pre + "conv.batch_norm.weight"); // [C] + ggml_tensor* bb = clone_weight(ctx, ml, pre + "conv.batch_norm.bias"); // [C] + normed = ggml_norm(ctx, dwt, ln_eps); // normalize over ne0=C + normed = ggml_mul(ctx, normed, g); // *gamma (broadcast [C] over T,B) + normed = ggml_add(ctx, normed, bb); // +beta + } else { + // batch_norm (inference): fold into per-channel scale/shift constants: + // scale = g / sqrt(var+eps); shift = b - mean*scale. Computed host-side. + std::vector& sc = pool.alloc_f32(D); + std::vector& sh = pool.alloc_f32(D); + std::vector g, bb, m, var; + pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.weight").c_str(), g); + pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.bias").c_str(), bb); + pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_mean").c_str(), m); + pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_var").c_str(), var); + for (int cc = 0; cc < D; ++cc) { + sc[cc] = g[cc] / std::sqrt(var[cc] + bn_eps); + sh[cc] = bb[cc] - m[cc] * sc[cc]; + } + int64_t d_ne[1] = {D}; + ggml_tensor* scale = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne, + sc.data(), sc.size() * sizeof(float)); + ggml_tensor* shift = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne, + sh.data(), sh.size() * sizeof(float)); + normed = ggml_add(ctx, ggml_mul(ctx, dwt, scale), shift); // [C, T, B] + } + + // -- SiLU (Swish), then pointwise_conv2 (Conv1d d->d, k=1). + normed = ggml_silu(ctx, normed); + ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight"); + pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d] + ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias"); + ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T, B] + if (pw2b) cout = ggml_add(ctx, cout, pw2b); + return cout; // [D, T, B]; this is layers[i].conv output +} + +// Scalar (B=1) overload: the verbatim v1 2-D ConformerConvolution sub-graph on +// conv input `c`, ne [D, T] (channels fastest). Used by the single-clip / +// forward path (build_graph below, forward_with_conv localization, and +// conv_module_forward) so B=1 runs the lean 2-D graph and is bit-exact with v1. +// Distinguished from the batched overload by signature: this takes `int +// valid_len`, the batched one takes `int B, const std::vector&`. static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, const std::string& pre, ggml_tensor* c, int D, int T, int K, int valid_len, @@ -144,6 +282,80 @@ ConformerLayer::ConformerLayer(const ModelLoader& ml, int layer_idx) assert((conv_kernel_ - 1) % 2 == 0 && "depthwise kernel must be odd"); } +ggml_tensor* ConformerLayer::build_graph_batched(ggml_context* ctx, + ggml_tensor* xt, int T, int B, + ggml_tensor* pe, int pos_len, + const std::vector& valid_len, + GraphInputPool& pool) const { + const int D = d_model_; + const int K = conv_kernel_; + const float ln_eps = 1e-5f; // LayerNorm eps (NeMo nn.LayerNorm default) + assert(pos_len == 2 * T - 1); + + const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + "."; + const ModelLoader& ml = ml_; + + // LayerNorm over the channel dim (ne0 = D), affine. Input ne [D, T, B]; + // ggml_norm normalizes ne0 and the affine [D] broadcasts over T and B. + auto layer_norm = [&](ggml_tensor* in, const std::string& nm) { + ggml_tensor* g = clone_weight(ctx, ml, pre + nm + ".weight"); // [D] + ggml_tensor* b = clone_weight(ctx, ml, pre + nm + ".bias"); // [D] + ggml_tensor* y = ggml_norm(ctx, in, ln_eps); // normalize over ne0 + y = ggml_mul(ctx, y, g); // broadcast [D] over T,B + y = ggml_add(ctx, y, b); + return y; + }; + // nn.Linear: ggml weight ne = [in, out]. in ne [in, T, B] -> [out, T, B]. + auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) { + ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight"); + ggml_tensor* y = ggml_mul_mat(ctx, W, in); + if (bias) { + ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias"); + if (B) y = ggml_add(ctx, y, B); + } + return y; + }; + // ConformerFeedForward: linear1(d->ff) -> SiLU -> linear2(ff->d). in [D, T, B]. + auto feed_forward = [&](ggml_tensor* in, const std::string& ff) { + ggml_tensor* h = linear(in, ff + ".linear1", /*bias*/true); // [FF, T, B] + h = ggml_silu(ctx, h); // Swish == SiLU + h = linear(h, ff + ".linear2", /*bias*/true); // [D, T, B] + return h; + }; + + // === Stage A: r = x + 0.5 * FFN1(norm_ff1(x)). === + ggml_tensor* h1 = layer_norm(xt, "norm_feed_forward1"); + h1 = feed_forward(h1, "feed_forward1"); + h1 = ggml_scale(ctx, h1, 0.5f); // fc_factor + ggml_tensor* r = ggml_add(ctx, xt, h1); // [D, T, B] + + // === Stage B: r = r + self_attn(norm_self_att(r)). === + ggml_tensor* attn_in = layer_norm(r, "norm_self_att"); + RelPosAttention attn(ml_, layer_idx_); + ggml_tensor* attn_out = attn.build_graph_batched(ctx, attn_in, T, B, pe, + pos_len, valid_len, pool); // [D, T, B] + r = ggml_add(ctx, r, attn_out); + + // === Stage C: r = r + conv(norm_conv(r)). === + ggml_tensor* c = layer_norm(r, "norm_conv"); // [D, T, B] + ggml_tensor* conv_out = build_conv_module(ctx, ml, pre, c, D, T, K, B, + valid_len, + conv_norm_type_, conv_causal_, pool); + r = ggml_add(ctx, r, conv_out); + + // === Stage D: r = r + 0.5 * FFN2(norm_ff2(r)); out = norm_out(r). === + ggml_tensor* h2 = layer_norm(r, "norm_feed_forward2"); + h2 = feed_forward(h2, "feed_forward2"); + h2 = ggml_scale(ctx, h2, 0.5f); + r = ggml_add(ctx, r, h2); + r = layer_norm(r, "norm_out"); + return r; // [D, T, B] -> per item row-major [T, D] +} + +// Scalar (B=1) builder: the verbatim v1 2-D conformer layer graph. The fused +// single-clip encoder and the scalar test entry points (forward / +// forward_with_conv) route through here so B=1 runs the lean 2-D graph and is +// bit-exact with v1. The batched builder above is used by forward_batch (B>1). ggml_tensor* ConformerLayer::build_graph(ggml_context* ctx, ggml_tensor* xt, int T, ggml_tensor* pe, int pos_len, int valid_len, diff --git a/src/conformer.hpp b/src/conformer.hpp index dbbbe22..23402fb 100644 --- a/src/conformer.hpp +++ b/src/conformer.hpp @@ -59,6 +59,13 @@ class ConformerLayer { ggml_tensor* pe, int pos_len, int valid_len, GraphInputPool& pool) const; + // Batched GRAPH-BUILDER. `xt` is [D, T, B]; `pe` is [D, pos_len] (shared + // across the batch). `valid_len` is per item (size B). Returns [D, T, B]. + ggml_tensor* build_graph_batched(ggml_context* ctx, ggml_tensor* xt, int T, + int B, ggml_tensor* pe, int pos_len, + const std::vector& valid_len, + GraphInputPool& pool) const; + // x: [T, d_model]; pos_emb: [pos_len=2T-1, d_model]; out: [T, d_model]. void forward(const std::vector& x, int T, const std::vector& pos_emb, int pos_len, diff --git a/src/decode_common.hpp b/src/decode_common.hpp new file mode 100644 index 0000000..5bae7e0 --- /dev/null +++ b/src/decode_common.hpp @@ -0,0 +1,21 @@ +#pragma once +#include +namespace pk { +// argmax over a[0..n): first index of the max (matches torch.max tie-break). +inline int decode_argmax(const float* a, int n) { + int best = 0; float bv = a[0]; + for (int i = 1; i < n; ++i) if (a[i] > bv) { bv = a[i]; best = i; } + return best; +} +// NeMo rescaled max_prob confidence over a[0..n) at index k: +// conf = (N*p_max - 1)/(N - 1), p_max = softmax(a)[k]. Stable softmax. +inline float decode_max_prob_conf(const float* a, int n, int k) { + float mx = a[0]; + for (int i = 1; i < n; ++i) if (a[i] > mx) mx = a[i]; + double denom = 0.0; + for (int i = 0; i < n; ++i) denom += std::exp((double)a[i] - (double)mx); + const double p_max = std::exp((double)a[k] - (double)mx) / denom; + const double N = (double)n; + return (float)((N * p_max - 1.0) / (N - 1.0)); +} +} // namespace pk diff --git a/src/encoder.cpp b/src/encoder.cpp index c87af2d..871971b 100644 --- a/src/encoder.cpp +++ b/src/encoder.cpp @@ -94,4 +94,76 @@ void Encoder::forward_capture(const std::vector& mel, int n_mels, int T, Tout = Tp; } +void Encoder::forward_batch(const MelBatch& mels, + std::vector>& enc_outs, + int& d_model, int& Tout, + std::vector& valid_Tout) const { + // Phase 5: the WHOLE batched encoder is ONE fused ggml graph, mirroring + // forward_capture but at B>1: subsampling -> xscaling -> pos_emb -> N + // conformer layers, all [d_model, T', B]. We return the raw [d_model, Tp, B] + // tensor and do the channels-first transpose host-side while splitting per + // item (see the index mapping below). + GraphInputPool pool; + Subsampling sub(ml_); + int Tp = 0; + std::vector vout; + // Per-item entry valid frames: offline convention is T-1 per clip. + std::vector vin(mels.B); + for (int b = 0; b < mels.B; ++b) vin[b] = mels.valid_T[b] - 1; + + std::vector flat; // receives [d_model, Tp, B] (ne0=d_model fastest) + bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0, + [&](ggml_context* ctx) -> ggml_tensor* { + // ---- 1. Subsampling (batched): mel -> x [d_model, T', B] (+ valid). ---- + ggml_tensor* x = sub.build_graph_batched(ctx, mels.data.data(), + mels.n_mels, mels.T_max, mels.B, pool, Tp, vout, vin); + assert((int)x->ne[0] == d_model_); + + // ---- 2. xscaling (gated; off for this model). ---- + if (xscaling_) x = ggml_scale(ctx, x, std::sqrt((float)d_model_)); + + // ---- 3. Relative positional encoding pos_emb [d_model, 2T'-1]. ---- + const int pos_len = 2 * Tp - 1; + std::vector& pe_host = pool.alloc_f32(); + rel_pos_encoding(Tp, d_model_, pe_host); // row-major [pos_len, d_model] + int64_t pe_ne[2] = {d_model_, pos_len}; + ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne, + pe_host.data(), pe_host.size() * sizeof(float)); + + // ---- 4. Conformer layer stack (all in-graph, shared pe). ---- + for (int i = 0; i < n_layers_; ++i) { + ConformerLayer layer(ml_, i); + x = layer.build_graph_batched(ctx, x, Tp, mels.B, pe, pos_len, vout, pool); + } + return x; // [d_model, Tp, B] + }, flat); + + assert(ok && "batched encoder graph failed"); + (void)ok; + + // flat is the [d_model, Tp, B] tensor flattened in ggml order (ne0=d_model + // fastest): element (c, t, b) sits at index ((size_t)b*Tp + t)*d_model_ + c. + // Split per item and transpose to channels-first [d_model, Tp] + // (enc_out[c*Tp + t]), matching what forward()/forward_capture return. + d_model = d_model_; + Tout = Tp; + valid_Tout = vout; + // Each enc_outs[b] is channels-first [d_model, valid_Tout[b]]: compact to the + // per-item non-pad frame count so the row stride equals valid_Tout[b]. The + // fused graph runs every item at the padded width Tp, but the trailing + // (Tp - vout[b]) columns are pad-derived; emitting them would (a) make the + // row stride differ from valid_Tout[b] (decoders index enc_out[c*Tout + t] + // with Tout = valid_Tout[b]) and (b) feed pad frames into the decoder. Both + // corrupt a padded (shorter) item's decode. + enc_outs.assign(mels.B, std::vector()); + for (int b = 0; b < mels.B; ++b) { + const int tv = vout[b]; + enc_outs[b].resize((size_t)d_model_ * tv); + for (int t = 0; t < tv; ++t) + for (int c = 0; c < d_model_; ++c) + enc_outs[b][(size_t)c * tv + t] = + flat[(((size_t)b * Tp) + t) * d_model_ + c]; + } +} + } // namespace pk diff --git a/src/encoder.hpp b/src/encoder.hpp index 540aa78..9da6c87 100644 --- a/src/encoder.hpp +++ b/src/encoder.hpp @@ -13,6 +13,15 @@ namespace pk { // // The `valid_len` (number of non-pad output frames) is derived from Subsampling // and threaded into every ConformerLayer (attention + conv pad masking). +// A batch of clips' mel features, pre-stacked to T_max. +struct MelBatch { + std::vector data; // contiguous [B][n_mels][T_max], row-major: data[(b*n_mels + m)*T_max + t] + int n_mels = 0; + int T_max = 0; // padded time length (max over clips) + int B = 0; + std::vector valid_T; // per-item true mel frame count, size B (<= T_max) +}; + class Encoder { public: explicit Encoder(const ModelLoader& ml); @@ -23,6 +32,19 @@ class Encoder { void forward(const std::vector& mel, int n_mels, int T, std::vector& enc_out, int& d_model, int& Tout) const; + // Batched encoder. Runs all B clips through ONE fused ggml graph + // (subsampling -> conformer stack -> output). Returns per-item encoder + // outputs, each row-major [d_model, valid_Tout[b]] (channels-first), the same + // orientation as forward(): each enc_outs[b] is compacted to that item's + // non-pad frame count so its row stride equals valid_Tout[b] (pad-derived + // trailing frames are dropped). `d_model` and `Tout` (the padded T') are + // filled; `valid_Tout` holds each item's non-pad output-frame count (size B). + // For B=1 this reduces to forward(). + void forward_batch(const MelBatch& mels, + std::vector>& enc_outs, + int& d_model, int& Tout, + std::vector& valid_Tout) const; + // Same as forward(), but also captures the per-layer outputs at indices // `capture_layers` (each row-major [T', d_model]) into `layer_outs` (parallel // to capture_layers). Used by the parity test to localize divergence. diff --git a/src/joint.cpp b/src/joint.cpp index 8e20df2..90cc24e 100644 --- a/src/joint.cpp +++ b/src/joint.cpp @@ -103,6 +103,45 @@ void Joint::step_logits(const float* enc_proj_t, assert(ok && "step_logits graph failed"); } +void Joint::step_logits_batch(const float* enc_proj_gathered, + const float* g, int pred_hidden, int n, + std::vector& logits) const { + assert(pred_hidden == pred_hidden_ && "pred_hidden mismatch"); + const int H = joint_hidden_; + + // Batched per-step joint over N items on the PERSISTENT backend. Mirrors + // step_logits with a batch axis (ggml ne1 = N): each of the two matmuls is + // applied across all N columns at once, and the biases broadcast over N. + // N=1 reduces exactly to step_logits. The gathered enc_proj input holds one + // joint_hidden row per item (item k at offset k*H), and g holds one + // pred_hidden vector per item (item k at offset k*pred_hidden). + bool ok = pk::run_graph(0, 0, + [&](ggml_context* ctx) -> ggml_tensor* { + // Gathered enc_proj rows: [H, N]. + int64_t ep_ne[2] = { H, n }; + ggml_tensor* ep = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ep_ne, + enc_proj_gathered, (size_t)H * n * sizeof(float)); + // Batched pred-net output g: [P, N]. + int64_t g_ne[2] = { pred_hidden_, n }; + ggml_tensor* gv = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, g_ne, + g, (size_t)pred_hidden_ * n * sizeof(float)); + // pred_proj = pred.weight·g + pred.bias (P->H). Weight ne=[P,H]. + ggml_tensor* Wp = pk::clone_weight(ctx, ml_, "joint.pred.weight"); + ggml_tensor* pp = ggml_mul_mat(ctx, Wp, gv); // [H, N] + ggml_tensor* bp = pk::clone_weight(ctx, ml_, "joint.pred.bias"); + pp = ggml_add(ctx, pp, bp); // bp [H] broadcasts over N + // f = ReLU(enc_proj + pred_proj) + ggml_tensor* f = ggml_relu(ctx, ggml_add(ctx, ep, pp)); // [H, N] + // logits = joint_net.2.weight·f + joint_net.2.bias (H->V). Weight ne=[H,V]. + ggml_tensor* Wo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.weight"); + ggml_tensor* y = ggml_mul_mat(ctx, Wo, f); // [V, N] + ggml_tensor* bo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.bias"); + y = ggml_add(ctx, y, bo); // bo [V] broadcasts over N + return y; // [V_plus, N] + }, logits); + assert(ok && "step_logits_batch graph failed"); +} + void Joint::forward(const std::vector& enc, int T, int enc_hidden, const std::vector& pred, int U, int pred_hidden, std::vector& logits, int& V_plus_out) const { diff --git a/src/joint.hpp b/src/joint.hpp index 2163c01..8369cb2 100644 --- a/src/joint.hpp +++ b/src/joint.hpp @@ -66,6 +66,15 @@ class Joint { const float* g, int pred_hidden, std::vector& logits) const; + // Batched per-step joint for N items. + // enc_proj_gathered: [joint_hidden * N], each item's enc_proj row for its + // current frame (item n at offset n*joint_hidden). + // g: [pred_hidden * N], batched pred output (item n at n*pred_hidden). + // logits out: [V_plus * N], item n at offset n*V_plus. + void step_logits_batch(const float* enc_proj_gathered, + const float* g, int pred_hidden, int n, + std::vector& logits) const; + int joint_hidden() const { return joint_hidden_; } // V_plus = vocab_size + 1 + num_durations diff --git a/src/model.cpp b/src/model.cpp index a408b0f..fcaa594 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -11,11 +11,13 @@ #include "joint.hpp" #include "tdt.hpp" #include "rnnt.hpp" +#include "transducer_batch.hpp" #include "transcription.hpp" #include "decode_types.hpp" #include "backend.hpp" #include "ggml_graph.hpp" +#include #include #include @@ -47,6 +49,38 @@ std::unique_ptr Model::load(const std::string& gguf_path) { return m; } +// Decode one item's encoder output (row-major [d_model, Tout], channels-first) +// into a transcript. Mirrors the tail of transcribe_16k exactly. +static std::string decode_enc_out(const ModelLoader& loader, + const std::vector& enc_out, + int d_model, int Tout, bool use_tdt) { + const ParakeetConfig& cfg = loader.config(); + if (use_tdt) { + std::vector enc_row((size_t)Tout * d_model); + for (int t = 0; t < Tout; ++t) + for (int c = 0; c < d_model; ++c) + enc_row[(size_t)t * d_model + c] = enc_out[(size_t)c * Tout + t]; + PredictionNet pred(loader); + Joint joint(loader); + const int max_symbols = static_cast(cfg.max_symbols); + std::vector ids; + if (!cfg.tdt_durations.empty()) + ids = tdt_greedy(pred, joint, enc_row, Tout, d_model, + cfg.tdt_durations, (int)cfg.blank_id, max_symbols); + else + ids = rnnt_greedy(pred, joint, enc_row, Tout, d_model, + (int)cfg.blank_id, max_symbols); + return detokenize(loader.tokenizer_pieces(), ids); + } else { + CTCDecoder ctc(loader); + std::vector logits; int vocab_plus_1 = 0; + ctc.forward(enc_out, d_model, Tout, logits, vocab_plus_1); + std::vector ids = ctc_greedy(logits, Tout, vocab_plus_1, + (int)cfg.blank_id); + return detokenize(loader.tokenizer_pieces(), ids); + } +} + std::string Model::transcribe_16k(const std::vector& pcm16k, Decoder decoder) const { const ParakeetConfig& cfg = loader_.config(); @@ -74,53 +108,155 @@ std::string Model::transcribe_16k(const std::vector& pcm16k, const bool use_tdt = (decoder == Decoder::kTDT) || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); - if (use_tdt) { - // 3a. TDT path: transpose encoder output to row-major [Tout, d_model]. - // enc_out from Encoder is [d_model, Tout] (channels-first). - std::vector enc_row(static_cast(Tout) * d_model); - for (int t = 0; t < Tout; ++t) - for (int c = 0; c < d_model; ++c) - enc_row[t * d_model + c] = enc_out[c * Tout + t]; + return decode_enc_out(loader_, enc_out, d_model, Tout, use_tdt); +} - // 3b. Prediction net + Joint. - PredictionNet pred(loader_); - Joint joint(loader_); +// Stage a batch of 16 kHz mono clips into a MelBatch: per-clip log-mel +// (GpuMel on a non-CPU backend, else the byte-identical FFT MelFrontend), +// zero-padded and stacked to the batch's longest clip (T_max). data layout is +// [B][n_mels][T_max], index (b*n_mels+m)*T_max+t; valid_T[b] is each clip's +// true frame count. +static MelBatch build_mel_batch(const ModelLoader& loader, + const std::vector>& pcms16k) { + const bool gpu = std::string(pk::global_backend().device_name()) != "cpu"; + MelBatch mb; + mb.B = (int)pcms16k.size(); + std::vector> feats(mb.B); + std::vector Ts(mb.B, 0); + int n_mels = 0; + for (int b = 0; b < mb.B; ++b) { + int nm = 0, T = 0; + if (gpu) { GpuMel g(loader); g.compute(pcms16k[b], feats[b], nm, T); } + else { MelFrontend m(loader); m.compute(pcms16k[b], feats[b], nm, T); } + n_mels = nm; Ts[b] = T; + } + mb.n_mels = n_mels; + mb.T_max = 0; + for (int b = 0; b < mb.B; ++b) mb.T_max = std::max(mb.T_max, Ts[b]); + mb.valid_T = Ts; + mb.data.assign((size_t)mb.B * n_mels * mb.T_max, 0.0f); + for (int b = 0; b < mb.B; ++b) + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < Ts[b]; ++t) + mb.data[((size_t)b * n_mels + m) * mb.T_max + t] = + feats[b][(size_t)m * Ts[b] + t]; + return mb; +} - // max_symbols: greedy max symbols emitted per frame, read from the model - // metadata (parakeet.decoding.max_symbols; NeMo default 10). - const int max_symbols = static_cast(cfg.max_symbols); +// Transpose the batched encoder outputs (channels-first [d_model, valid_Tout[b]]) +// into per-item row-major [valid_Tout[b], d_model] for the transducer decoder. +static void batch_enc_to_row_major(const std::vector>& enc_outs, + const std::vector& valid_Tout, int d_model, + std::vector>& encs, + std::vector& Ts) { + const int B = (int)enc_outs.size(); + encs.assign(B, {}); Ts.assign(B, 0); + for (int b = 0; b < B; ++b) { + const int tb = valid_Tout[b]; + Ts[b] = tb; + encs[b].resize((size_t)tb * d_model); + for (int t = 0; t < tb; ++t) + for (int c = 0; c < d_model; ++c) + encs[b][(size_t)t * d_model + c] = enc_outs[b][(size_t)c * tb + t]; + } +} - // Branch on the duration table: TDT (durations present) uses the - // duration-aware greedy loop; a pure RNNT transducer (no durations, e.g. - // arch ∈ {rnnt, hybrid_rnnt_ctc}) uses the standard RNNT greedy loop. - std::vector ids; - if (!cfg.tdt_durations.empty()) { - ids = tdt_greedy( - pred, joint, enc_row, Tout, d_model, - cfg.tdt_durations, static_cast(cfg.blank_id), max_symbols); - } else { - ids = rnnt_greedy( - pred, joint, enc_row, Tout, d_model, - static_cast(cfg.blank_id), max_symbols); - } +std::vector Model::transcribe_16k_batch( + const std::vector>& pcms16k, Decoder decoder) const { + const ParakeetConfig& cfg = loader_.config(); + const bool use_tdt = (decoder == Decoder::kTDT) + || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); - // 4a. Detokenize. - return detokenize(loader_.tokenizer_pieces(), ids); + // 1. Per-clip mel, then stack to T_max. + MelBatch mb = build_mel_batch(loader_, pcms16k); + // 2. Batched encoder. + Encoder encoder(loader_); + std::vector> enc_outs; int d_model = 0, Tout = 0; + std::vector valid_Tout; + encoder.forward_batch(mb, enc_outs, d_model, Tout, valid_Tout); + + // 3. Decode (each enc_out is [d_model, valid_Tout[b]]). + std::vector outs(mb.B); + if (use_tdt) { + // Batched transducer (TDT/RNNT) greedy decode: build per-item row-major + // [T, d_model] from the channels-first [d_model, T] encoder outputs. + std::vector> encs; + std::vector Ts; + batch_enc_to_row_major(enc_outs, valid_Tout, d_model, encs, Ts); + PredictionNet pred(loader_); + Joint joint(loader_); + std::vector> ids; + pk::transducer_greedy_batch(pred, joint, encs, Ts, d_model, + cfg.tdt_durations, (int)cfg.blank_id, + (int)cfg.max_symbols, ids, nullptr); + for (int b = 0; b < mb.B; ++b) + outs[b] = detokenize(loader_.tokenizer_pieces(), ids[b]); } else { - // 3b. CTC path: head -> log-probs [Tout, vocab+1]. - CTCDecoder ctc(loader_); - std::vector logits; - int vocab_plus_1 = 0; - ctc.forward(enc_out, d_model, Tout, logits, vocab_plus_1); + // CTC stays per-item (no autoregressive decode to batch). + for (int b = 0; b < mb.B; ++b) + outs[b] = decode_enc_out(loader_, enc_outs[b], d_model, valid_Tout[b], use_tdt); + } + return outs; +} - // 4b. CTC greedy collapse -> token ids. Blank is the last column. - const int blank_id = static_cast(cfg.blank_id); - std::vector ids = ctc_greedy(logits, Tout, vocab_plus_1, blank_id); +std::vector Model::transcribe_pcm_batch( + const std::vector>& pcms, int sample_rate, + Decoder decoder) const { + if (sample_rate <= 0) { + throw std::runtime_error("parakeet: invalid sample_rate"); + } + std::vector> r(pcms.size()); + for (size_t i = 0; i < pcms.size(); ++i) + r[i] = (sample_rate == 16000) ? pcms[i] + : resample_linear(pcms[i], sample_rate, 16000); + return transcribe_16k_batch(r, decoder); +} - // 5b. Detokenize. - return detokenize(loader_.tokenizer_pieces(), ids); +// Decode one item's encoder output (channels-first [d_model, Tout]) into a +// Transcription (text + per-word timestamps + tokens). Mirrors the decode tail +// of transcribe_16k_with_timestamps exactly. +static Transcription decode_enc_out_with_timestamps( + const ModelLoader& loader, const std::vector& enc_out, + int d_model, int Tout, bool use_tdt, float frame_sec) { + const ParakeetConfig& cfg = loader.config(); + Transcription result; + std::vector toks; + if (use_tdt) { + std::vector enc_row((size_t)Tout * d_model); + for (int t = 0; t < Tout; ++t) + for (int c = 0; c < d_model; ++c) + enc_row[(size_t)t * d_model + c] = enc_out[(size_t)c * Tout + t]; + PredictionNet pred(loader); + Joint joint(loader); + const int max_symbols = (int)cfg.max_symbols; + if (!cfg.tdt_durations.empty()) + tdt_greedy(pred, joint, enc_row, Tout, d_model, cfg.tdt_durations, + (int)cfg.blank_id, max_symbols, &toks); + else + rnnt_greedy(pred, joint, enc_row, Tout, d_model, + (int)cfg.blank_id, max_symbols, &toks); + } else { + CTCDecoder ctc(loader); + std::vector logits; int vocab_plus_1 = 0; + ctc.forward(enc_out, d_model, Tout, logits, vocab_plus_1); + ctc_greedy(logits, Tout, vocab_plus_1, (int)cfg.blank_id, &toks); + // NeMo CTC word end_offset is the NEXT collapsed token's start frame + // (cumulative run lengths), not start+1. ctc_greedy emits span == 1; + // rewrite each token's span to (next_frame - frame) so group_words' + // frame+span rule reproduces NeMo's end_offset. The final token keeps + // span == 1 (its true run length is unknown to the collapse, and within + // the 1-frame word-end tolerance). + for (size_t i = 0; i + 1 < toks.size(); ++i) + toks[i].span = toks[i + 1].frame - toks[i].frame; } + std::vector ids; + ids.reserve(toks.size()); + for (const TokenInfo& ti : toks) ids.push_back(ti.id); + result.text = detokenize(loader.tokenizer_pieces(), ids); + result.words = group_words(toks, loader.tokenizer_pieces(), frame_sec); + result.tokens = std::move(toks); + return result; } Transcription Model::transcribe_16k_with_timestamps( @@ -155,60 +291,68 @@ Transcription Model::transcribe_16k_with_timestamps( const bool use_tdt = (decoder == Decoder::kTDT) || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); - Transcription result; - std::vector toks; + Transcription result = decode_enc_out_with_timestamps( + loader_, enc_out, d_model, Tout, use_tdt, frame_sec); + return result; +} - if (use_tdt) { - // Transpose channels-first [d_model, Tout] -> row-major [Tout, d_model]. - std::vector enc_row(static_cast(Tout) * d_model); - for (int t = 0; t < Tout; ++t) - for (int c = 0; c < d_model; ++c) - enc_row[static_cast(t) * d_model + c] = enc_out[static_cast(c) * Tout + t]; +std::vector Model::transcribe_16k_batch_with_timestamps( + const std::vector>& pcms16k, Decoder decoder) const { + const ParakeetConfig& cfg = loader_.config(); + const float frame_sec = + (float)cfg.hop_length * (float)cfg.subsampling_factor / (float)cfg.sample_rate; + const bool use_tdt = (decoder == Decoder::kTDT) + || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); + MelBatch mb = build_mel_batch(loader_, pcms16k); + + Encoder encoder(loader_); + std::vector> enc_outs; int d_model = 0, Tout = 0; + std::vector valid_Tout; + encoder.forward_batch(mb, enc_outs, d_model, Tout, valid_Tout); + + std::vector outs(mb.B); + if (use_tdt) { + // Batched transducer (TDT/RNNT) greedy decode with timestamps. Build + // per-item row-major [T, d_model] from channels-first [d_model, T]. + std::vector> encs; + std::vector Ts; + batch_enc_to_row_major(enc_outs, valid_Tout, d_model, encs, Ts); PredictionNet pred(loader_); Joint joint(loader_); - const int max_symbols = static_cast(cfg.max_symbols); - - if (!cfg.tdt_durations.empty()) { - tdt_greedy(pred, joint, enc_row, Tout, d_model, - cfg.tdt_durations, static_cast(cfg.blank_id), - max_symbols, &toks); - } else { - rnnt_greedy(pred, joint, enc_row, Tout, d_model, - static_cast(cfg.blank_id), max_symbols, &toks); + std::vector> ids; + std::vector> toks; + pk::transducer_greedy_batch(pred, joint, encs, Ts, d_model, + cfg.tdt_durations, (int)cfg.blank_id, + (int)cfg.max_symbols, ids, &toks); + // Assemble each Transcription exactly as decode_enc_out_with_timestamps' + // transducer tail does. + for (int b = 0; b < mb.B; ++b) { + Transcription& result = outs[b]; + result.text = detokenize(loader_.tokenizer_pieces(), ids[b]); + result.words = group_words(toks[b], loader_.tokenizer_pieces(), frame_sec); + result.tokens = std::move(toks[b]); } - // TDT/RNNT TokenInfo.span is already the per-token end-offset extent - // (duration for TDT, 1 for RNNT) -> group_words' frame+span rule is - // correct as-is. } else { - // CTC path: head -> log-probs [Tout, vocab+1]. - CTCDecoder ctc(loader_); - std::vector logits; - int vocab_plus_1 = 0; - ctc.forward(enc_out, d_model, Tout, logits, vocab_plus_1); - - const int blank_id = static_cast(cfg.blank_id); - ctc_greedy(logits, Tout, vocab_plus_1, blank_id, &toks); - - // NeMo CTC word end_offset = the NEXT collapsed token's start frame - // (cumulative run lengths), not start+1. ctc_greedy documents span == 1; - // rewrite each token's span to (next_frame - frame) so group_words' - // `frame + span` rule reproduces NeMo's end_offset exactly. The final - // token keeps span == 1 (its true run-length is unknown to the collapse, - // and within the 1-frame word-end tolerance). - for (size_t i = 0; i + 1 < toks.size(); ++i) { - toks[i].span = toks[i + 1].frame - toks[i].frame; - } + // CTC stays per-item (not a transducer; no autoregressive decode). + for (int b = 0; b < mb.B; ++b) + outs[b] = decode_enc_out_with_timestamps( + loader_, enc_outs[b], d_model, valid_Tout[b], use_tdt, frame_sec); } + return outs; +} - // Detokenize the flat text from the emitted ids. - std::vector ids; - ids.reserve(toks.size()); - for (const TokenInfo& ti : toks) ids.push_back(ti.id); - result.text = detokenize(loader_.tokenizer_pieces(), ids); - result.words = group_words(toks, loader_.tokenizer_pieces(), frame_sec); - result.tokens = std::move(toks); - return result; +std::vector Model::transcribe_pcm_batch_with_timestamps( + const std::vector>& pcms, int sample_rate, + Decoder decoder) const { + if (sample_rate <= 0) { + throw std::runtime_error("parakeet: invalid sample_rate"); + } + std::vector> r(pcms.size()); + for (size_t i = 0; i < pcms.size(); ++i) + r[i] = (sample_rate == 16000) ? pcms[i] + : resample_linear(pcms[i], sample_rate, 16000); + return transcribe_16k_batch_with_timestamps(r, decoder); } std::string Model::transcribe_pcm(const std::vector& pcm, int sample_rate, diff --git a/src/model.hpp b/src/model.hpp index 76e997e..8b7d95c 100644 --- a/src/model.hpp +++ b/src/model.hpp @@ -35,6 +35,13 @@ class Model { std::string transcribe_path(const std::string& wav_path, Decoder decoder = Decoder::kDefault) const; + // Transcribe a batch of mono float PCM clips. Each is resampled to 16 kHz if + // needed, then all run through the batched encoder; decode is per item. + // Returns one transcript per input, in order. + std::vector transcribe_pcm_batch( + const std::vector>& pcms, int sample_rate, + Decoder decoder = Decoder::kDefault) const; + // Transcribe raw mono float PCM, returning the flat text plus per-word and // per-token timestamps + confidence (matching NeMo timestamps=True + // 'max_prob' confidence). If `sample_rate != 16000` the audio is linearly @@ -48,6 +55,13 @@ class Model { const std::string& wav_path, Decoder decoder = Decoder::kDefault) const; + // Batched timestamped transcription. Each clip is resampled to 16 kHz if + // needed, all run through the batched encoder; decode + timestamp extraction + // are per item. Returns one Transcription per input, in order. + std::vector transcribe_pcm_batch_with_timestamps( + const std::vector>& pcms, int sample_rate, + Decoder decoder = Decoder::kDefault) const; + const ParakeetConfig& config() const { return loader_.config(); } // The underlying loaded GGUF. Exposed so the streaming C-API can build a @@ -67,6 +81,15 @@ class Model { std::string transcribe_16k(const std::vector& pcm16k, Decoder decoder) const; + // Core batched orchestration: N 16 kHz clips -> N transcripts. Stacks mels, + // runs forward_batch, decodes each item with the existing greedy decoders. + std::vector transcribe_16k_batch( + const std::vector>& pcms16k, Decoder decoder) const; + + // Core batched timestamped orchestration: N 16 kHz clips -> N Transcriptions. + std::vector transcribe_16k_batch_with_timestamps( + const std::vector>& pcms16k, Decoder decoder) const; + // Core orchestration for the timestamps path: 16 kHz mono PCM -> full // Transcription (text + per-token TokenInfo + grouped words). Shared by the // two timestamp entry points. diff --git a/src/parakeet_capi.cpp b/src/parakeet_capi.cpp index 66737d7..a11ffca 100644 --- a/src/parakeet_capi.cpp +++ b/src/parakeet_capi.cpp @@ -16,7 +16,7 @@ #include // ABI version. Bump on breaking changes. -#define PARAKEET_CAPI_ABI_VERSION 1 +#define PARAKEET_CAPI_ABI_VERSION 2 // The opaque context: a loaded model plus a buffer for the last error message. struct parakeet_ctx { @@ -246,6 +246,53 @@ extern "C" char* parakeet_capi_transcribe_pcm(parakeet_ctx* ctx, const float* sa } } +extern "C" int parakeet_capi_transcribe_pcm_batch(parakeet_ctx* ctx, + const float* const* samples, + const int* n_samples, int n_clips, + int sample_rate, int decoder, + char** out) { + if (!ctx) return 1; + if (!ctx->model) { ctx->last_error = "context has no loaded model"; return 1; } + if (!samples || !n_samples || !out || n_clips < 0) { + ctx->last_error = "invalid batch arguments"; + return 1; + } + // Contract: on any error path (validation, exception, OOM) every out[] + // entry is left NULL, so the caller owns nothing and frees nothing. + for (int i = 0; i < n_clips; ++i) out[i] = nullptr; + try { + std::vector> pcms(n_clips); + for (int i = 0; i < n_clips; ++i) { + if (!samples[i] || n_samples[i] < 0) { + ctx->last_error = "invalid samples buffer in batch"; + return 1; + } + pcms[i].assign(samples[i], samples[i] + n_samples[i]); + } + std::vector texts = + ctx->model->transcribe_pcm_batch(pcms, sample_rate, to_decoder(decoder)); + ctx->last_error.clear(); + for (int i = 0; i < n_clips; ++i) { + char* s = dup_to_c(texts[i]); + if (!s) { + // Roll back the strings already allocated this call so every + // out[] entry is NULL on return (out[i..] are already NULL). + for (int j = 0; j < i; ++j) { std::free(out[j]); out[j] = nullptr; } + ctx->last_error = "out of memory"; + return 2; + } + out[i] = s; + } + return 0; + } catch (const std::exception& e) { + ctx->last_error = e.what(); + return 3; + } catch (...) { + ctx->last_error = "unknown error"; + return 3; + } +} + extern "C" char* parakeet_capi_transcribe_path_json(parakeet_ctx* ctx, const char* wav_path, int decoder) { @@ -273,6 +320,45 @@ extern "C" char* parakeet_capi_transcribe_path_json(parakeet_ctx* ctx, } } +extern "C" char* parakeet_capi_transcribe_pcm_batch_json(parakeet_ctx* ctx, + const float* samples_concat, const int* n_samples, int n_clips, + int sample_rate, int decoder) { + if (!ctx) return nullptr; + if (!ctx->model) { ctx->last_error = "context has no loaded model"; return nullptr; } + if (!samples_concat || !n_samples || n_clips < 0) { + ctx->last_error = "invalid batch arguments"; return nullptr; + } + try { + std::vector> pcms(n_clips); + size_t off = 0; + for (int i = 0; i < n_clips; ++i) { + if (n_samples[i] < 0) { ctx->last_error = "invalid clip length"; return nullptr; } + pcms[i].assign(samples_concat + off, samples_concat + off + n_samples[i]); + off += (size_t)n_samples[i]; + } + std::vector trs = + ctx->model->transcribe_pcm_batch_with_timestamps(pcms, sample_rate, + to_decoder(decoder)); + const pk::ParakeetConfig& cfg = ctx->model->config(); + const float frame_sec = + (float)cfg.hop_length * (float)cfg.subsampling_factor / (float)cfg.sample_rate; + std::string json = "["; + for (size_t i = 0; i < trs.size(); ++i) { + if (i) json += ','; + json += transcription_to_json(trs[i], frame_sec); + } + json += "]"; + ctx->last_error.clear(); + char* out = dup_to_c(json); + if (!out) { ctx->last_error = "out of memory"; return nullptr; } + return out; + } catch (const std::exception& e) { + ctx->last_error = e.what(); return nullptr; + } catch (...) { + ctx->last_error = "unknown error"; return nullptr; + } +} + // --------------------------------------------------------------------------- // Streaming API // --------------------------------------------------------------------------- diff --git a/src/prediction.cpp b/src/prediction.cpp index d21ec04..3ef8066 100644 --- a/src/prediction.cpp +++ b/src/prediction.cpp @@ -112,6 +112,90 @@ void PredictionNet::step(int32_t token_id, bool is_sos, assert(ok && "pred-net step graph failed"); } +// --------------------------------------------------------------------------- +// Batched single-step advance: the same LSTM math as step(), but with a batch +// axis N. Inputs and state are laid out [H, N] in ggml (item n is column n, +// offset n*H in the flat host buffer). Gate slices become [H, N] views into the +// [4H, N] z, using z->nb[1] as the column stride. N=1 reduces to step(). +// --------------------------------------------------------------------------- +void PredictionNet::step_batch(const std::vector& token_ids, + const std::vector& is_sos, + const BatchedPredState& in, + std::vector& g, + BatchedPredState& out_state) const { + const int H = H_; + const int L = n_layers_; + const int N = (int)token_ids.size(); + assert(N > 0 && (int)is_sos.size() == N && "batch size mismatch"); + + // Lazily fetch the embedding table to the host (device-safe), exactly as + // step() does. + if (embed_host_.empty()) { + pk::ensure_weights_realized(ml_); + ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight"); + assert(emb && "missing decoder.prediction.embed.weight"); + embed_host_.resize((size_t)vocab_p1_ * H); + ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb)); + } + + // Layer-0 input [H*N]: zeros for SOS items, else the embedding row. + std::vector x0((size_t)H * N, 0.0f); + for (int n = 0; n < N; ++n) { + if (!is_sos[n]) { + assert(token_ids[n] >= 0 && token_ids[n] < vocab_p1_ && "embedding id out of range"); + std::memcpy(&x0[(size_t)n * H], &embed_host_[(size_t)token_ids[n] * H], + (size_t)H * sizeof(float)); + } + } + + out_state.h.assign((size_t)L, std::vector((size_t)H * N)); + out_state.c.assign((size_t)L, std::vector((size_t)H * N)); + + bool ok = pk::run_graph(0, 0, [&](ggml_context* ctx) -> ggml_tensor* { + int64_t ne2[2] = { H, N }; + ggml_tensor* layer_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2, + x0.data(), (size_t)H * N * sizeof(float)); + ggml_tensor* top_h = nullptr; + for (int l = 0; l < L; ++l) { + const std::string s = "_l" + std::to_string(l); + ggml_tensor* Wih = pk::clone_weight(ctx, ml_, + ("decoder.prediction.dec_rnn.lstm.weight_ih" + s).c_str()); + ggml_tensor* Whh = pk::clone_weight(ctx, ml_, + ("decoder.prediction.dec_rnn.lstm.weight_hh" + s).c_str()); + ggml_tensor* bih = pk::clone_weight(ctx, ml_, + ("decoder.prediction.dec_rnn.lstm.bias_ih" + s).c_str()); + ggml_tensor* bhh = pk::clone_weight(ctx, ml_, + ("decoder.prediction.dec_rnn.lstm.bias_hh" + s).c_str()); + ggml_tensor* h_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2, + in.h[l].data(), (size_t)H * N * sizeof(float)); + ggml_tensor* c_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2, + in.c[l].data(), (size_t)H * N * sizeof(float)); + // z = W_ih·x + b_ih + W_hh·h_in + b_hh [4H, N] + // (bias [4H] broadcasts over the N columns). + ggml_tensor* z = ggml_add(ctx, + ggml_add(ctx, ggml_mul_mat(ctx, Wih, layer_in), bih), + ggml_add(ctx, ggml_mul_mat(ctx, Whh, h_in), bhh)); + // Gate slices (i, f, g, o), each [H, N]. The view keeps z's FULL + // column stride (z->nb[1] = 4H elems), reading only H contiguous + // elements per column, so consecutive columns skip the other three + // gate blocks. (Do NOT change the stride to H*sizeof(float).) + ggml_tensor* i = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], 0))); + ggml_tensor* f = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)H * sizeof(float)))); + ggml_tensor* gg = ggml_tanh (ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)2 * H * sizeof(float)))); + ggml_tensor* o = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)3 * H * sizeof(float)))); + // c' = f*c_in + i*g ; h' = o*tanh(c') + ggml_tensor* c_out = ggml_add(ctx, ggml_mul(ctx, f, c_in), ggml_mul(ctx, i, gg)); + ggml_tensor* h_out = ggml_mul(ctx, o, ggml_tanh(ctx, c_out)); + pk::capture_graph_output(c_out, &out_state.c[l]); + pk::capture_graph_output(h_out, &out_state.h[l]); + layer_in = h_out; + top_h = h_out; + } + return top_h; + }, g); + assert(ok && "pred-net step_batch graph failed"); +} + // --------------------------------------------------------------------------- // Full-sequence forward pass (unchanged API; now driven by step() so there is a // single LSTM implementation). Carries (h, c) state across timesteps; the diff --git a/src/prediction.hpp b/src/prediction.hpp index 54cb6ca..6e5b8ed 100644 --- a/src/prediction.hpp +++ b/src/prediction.hpp @@ -14,6 +14,13 @@ struct PredState { std::vector> c; // c[layer] = [hidden] }; +// Batched LSTM state: one (h,c) per layer, each holding N items' columns laid +// out [H*N] (item n at offset n*H). Generalizes PredState to a batch. +struct BatchedPredState { + std::vector> h; // h[layer] size H*N + std::vector> c; // c[layer] size H*N +}; + // RNN-Transducer prediction network — NeMo RNNTDecoder prediction net. // // Architecture: @@ -67,6 +74,18 @@ class PredictionNet { std::vector& g, PredState& out_state) const; + // Advance the LSTM one token for N items at once. + // token_ids[n]: embedding index for item n (ignored where is_sos[n]). + // is_sos[n]: use the zero SOS input for item n (1=true). + // in: batched prior state (h[L],c[L] each [H*N]). + // g: OUT, top-layer h' for all items [H*N] (item n at n*H). + // out_state: OUT, new batched (h',c'). + void step_batch(const std::vector& token_ids, + const std::vector& is_sos, + const BatchedPredState& in, + std::vector& g, + BatchedPredState& out_state) const; + int hidden_size() const { return H_; } int num_layers() const { return n_layers_; } diff --git a/src/relpos_attention.cpp b/src/relpos_attention.cpp index b9cb5f6..3fb1a81 100644 --- a/src/relpos_attention.cpp +++ b/src/relpos_attention.cpp @@ -41,6 +41,9 @@ ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt, int T, ggml_tensor* pe, int pos_len, int valid_len, GraphInputPool& pool) const { + // Scalar (B=1) builder: the verbatim v1 2-D/3D relative-position attention + // graph. The single-clip conformer layer routes here so B=1 runs the lean + // graph and is bit-exact with v1. build_graph_batched below serves B>1. const int D = d_model_; const int H = n_heads_; const int dk = d_head_; @@ -156,6 +159,152 @@ ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt, return linear("linear_out.weight", "linear_out.bias", merged); // [D, T] } +ggml_tensor* RelPosAttention::build_graph_batched( + ggml_context* ctx, ggml_tensor* xt, int T, int B, ggml_tensor* pe, + int pos_len, const std::vector& valid_len, + GraphInputPool& pool) const { + const int D = d_model_; + const int H = n_heads_; + const int dk = d_head_; + const float scale = 1.0f / std::sqrt((float)dk); + assert(pos_len == 2 * T - 1); + assert((int)valid_len.size() == B); + + const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn."; + const ModelLoader& ml = ml_; + + // ---- linear projections (nn.Linear: ggml W ne=[in,out]) ---- + // The bias is added only when requested AND present: NeMo configures the + // attention linears with bias=False in some checkpoints + // (parakeet-tdt-0.6b-v2/-v3) and bias=True in others (110m). + auto linear = [&](const char* w, const char* b, ggml_tensor* in) { + ggml_tensor* W = clone_weight(ctx, ml, pre + w); + ggml_tensor* y = ggml_mul_mat(ctx, W, in); // [out, *] + if (b && ml.tensor(pre + b)) { + ggml_tensor* B = clone_weight(ctx, ml, pre + b); + y = ggml_add(ctx, y, B); // broadcast [out] over cols + } + return y; + }; + // xt is [D, T, B]; mul_mat batches over ne2 -> q/k/v are [D, T, B]. pe is + // shared [D, P] (NO batch) -> p is [D, P]. + ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T, B] + ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); // [D, T, B] + ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt); // [D, T, B] + ggml_tensor* p = linear("linear_pos.weight", nullptr, pe); // [D, P] + + // ---- split into heads (batched): [D, n, B] -> [dk, H, n, B] -> [dk, n, H, B] ---- + auto to_heads_b = [&](ggml_tensor* t, int n) { + t = ggml_reshape_4d(ctx, t, dk, H, n, B); // [dk, H, n, B] + t = ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H, B] + return t; + }; + // p is shared (no batch) -> keep the 3D head-split: [dk, P, H]. + auto to_heads = [&](ggml_tensor* t, int n) { + t = ggml_reshape_3d(ctx, t, dk, H, n); // [dk, H, n] + t = ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H] + return t; + }; + ggml_tensor* qh = to_heads_b(q, T); // [dk, T, H, B] + ggml_tensor* kh = to_heads_b(k, T); // [dk, T, H, B] + ggml_tensor* vh = to_heads_b(v, T); // [dk, T, H, B] + ggml_tensor* ph = to_heads(p, pos_len); // [dk, P, H] (ne3=1, broadcast over B) + + // ---- pos_bias_u/v: ne [dk, H] -> [dk, 1, H, 1] to broadcast over T and B ---- + ggml_tensor* bu = clone_weight(ctx, ml, pre + "pos_bias_u"); // [dk, H] + ggml_tensor* bv = clone_weight(ctx, ml, pre + "pos_bias_v"); // [dk, H] + bu = ggml_reshape_4d(ctx, bu, dk, 1, H, 1); + bv = ggml_reshape_4d(ctx, bv, dk, 1, H, 1); + ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H, B] + ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H, B] + + // ---- ac = q_u @ k^T : mul_mat([dk,T,H,B],[dk,T,H,B]) -> [T_k, T_q, H, B] ---- + ggml_tensor* ac = ggml_mul_mat(ctx, kh, qu); // [T(key), T(query), H, B] + + // ---- bd = q_v @ p^T -> [P, T_q, H, B], then rel_shift -> [T, T, H, B] ---- + // ph is [dk, P, H] (ne3=1) and broadcasts over the batch (ne3=B) of qv. + ggml_tensor* bd = ggml_mul_mat(ctx, ph, qv); // [P, T, H, B] + // 4D rel-shift: identical ne0/ne1 stride+offset arithmetic as the 3D path, + // with the batch axis threaded through every reshape/view via ne3=B and + // nb[3]. ne2/ne3 are passive (no offset on them). + bd = ggml_pad_ext(ctx, bd, /*lp0*/1, /*rp0*/0, 0,0, 0,0, 0,0); // [2T, T, H, B] + bd = ggml_reshape_4d(ctx, bd, T, 2 * T, H, B); // [T, 2T, H, B] + bd = ggml_view_4d(ctx, bd, T, 2 * T - 1, H, B, + bd->nb[1], bd->nb[2], bd->nb[3], bd->nb[1]); // [T, 2T-1, H, B] + bd = ggml_cont(ctx, bd); + bd = ggml_reshape_4d(ctx, bd, 2 * T - 1, T, H, B); // [2T-1, T, H, B] + bd = ggml_view_4d(ctx, bd, T, T, H, B, + bd->nb[1], bd->nb[2], bd->nb[3], 0); // [T, T, H, B] + bd = ggml_cont(ctx, bd); + + // ---- scores = ac + bd ; softmax(scores*scale + mask) ---- + ggml_tensor* scores = ggml_add(ctx, ac, bd); // [T_k, T_q, H, B] + + // Per-item additive mask [T_k, T_q, 1, B]: 0 where query qi may attend to key + // kj, -inf otherwise. (1) pad mask: key kj valid iff kj < valid_len[b]. + // (2) chunked-limited window for streaming models. See NeMo _create_masks. + // Mask shape is [T, T, 1, B] (ne2=1) so soft_max_ext broadcasts it over the + // head axis (ne2=H) while indexing per item on ne3=B. Verified against the + // ggml CPU kernel ggml_compute_forward_soft_max_f32: it reads the mask at + // i12 = i02 % ne12 (head, here ne12=1 -> always 0) and i13 = i03 % ne13 + // (batch, here ne13=B -> exact per-item), and ggml_soft_max_impl asserts + // a->ne[2] % mask->ne[2] == 0 and a->ne[3] % mask->ne[3] == 0. + const int chunk_size = chunked_limited_ ? (att_right_ + 1) : 0; + const int left_chunks = (chunked_limited_ && chunk_size > 0) + ? (att_left_ / chunk_size) : 0; + std::vector& mask_host = pool.alloc_f32((size_t)B * T * T); + { + float* md = mask_host.data(); + const float ninf = -INFINITY; + for (int b = 0; b < B; ++b) { + const int vl = valid_len[b]; + for (int qi = 0; qi < T; ++qi) { + const int cq = chunked_limited_ ? (qi / chunk_size) : 0; + for (int kj = 0; kj < T; ++kj) { + bool ok = (kj < vl); + if (ok && chunked_limited_) { + const int ck = kj / chunk_size; + const int diff = cq - ck; + ok = (diff >= 0 && diff <= left_chunks); + } + md[(size_t)b * T * T + (size_t)qi * T + kj] = ok ? 0.0f : ninf; + } + } + } + } + int64_t mask_ne[4] = {T, T, 1, B}; + ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, mask_ne, + mask_host.data(), mask_host.size() * sizeof(float)); + ggml_tensor* attn = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // [T_k, T_q, H, B] + + // ---- context = attn @ v -> [dk, T_q, H, B] ---- + ggml_tensor* vtk = ggml_cont(ctx, ggml_permute(ctx, vh, 1, 0, 2, 3)); // [T_k, dk, H, B] + ggml_tensor* ctxh = ggml_mul_mat(ctx, vtk, attn); // [dk, T_q, H, B] + + // ---- concat heads: [dk, T, H, B] -> [dk, H, T, B] -> [D, T, B] ---- + ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, ctxh, 0, 2, 1, 3)); // [dk, H, T, B] + merged = ggml_reshape_3d(ctx, merged, D, T, B); // [D, T, B] + + // Zero the context for PADDED query rows (NeMo masks padded query rows fully + // -> output reduces to linear_out.bias). Apply a per-item query-row mask + // [1, T, B] (broadcast over D). Emit only when some item has valid_len < T. + bool any_pad = false; + for (int b = 0; b < B; ++b) any_pad = any_pad || (valid_len[b] < T); + if (any_pad) { + std::vector& qmask_host = pool.alloc_f32((size_t)B * T); + for (int b = 0; b < B; ++b) + for (int qi = 0; qi < T; ++qi) + qmask_host[(size_t)b * T + qi] = (qi < valid_len[b]) ? 1.0f : 0.0f; + int64_t qm_ne[3] = {1, T, B}; + ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, qm_ne, + qmask_host.data(), qmask_host.size() * sizeof(float)); + merged = ggml_mul(ctx, merged, qmask); // broadcast over D + } + + // ---- output projection ---- + return linear("linear_out.weight", "linear_out.bias", merged); // [D, T, B] +} + void RelPosAttention::forward(const std::vector& x, int T, const std::vector& pos_emb, int pos_len, int valid_len, diff --git a/src/relpos_attention.hpp b/src/relpos_attention.hpp index 67b03b3..91d2534 100644 --- a/src/relpos_attention.hpp +++ b/src/relpos_attention.hpp @@ -41,6 +41,13 @@ class RelPosAttention { ggml_tensor* pe, int pos_len, int valid_len, GraphInputPool& pool) const; + // Batched GRAPH-BUILDER. `xt` is [D, T, B]; `pe` is [D, pos_len] (shared + // across the batch). `valid_len` is per item (size B). Returns [D, T, B]. + ggml_tensor* build_graph_batched(ggml_context* ctx, ggml_tensor* xt, int T, + int B, ggml_tensor* pe, int pos_len, + const std::vector& valid_len, + GraphInputPool& pool) const; + // x: [T, d_model]; pos_emb: [2T-1, d_model]; out: [T, d_model]. void forward(const std::vector& x, int T, const std::vector& pos_emb, int pos_len, diff --git a/src/rnnt.cpp b/src/rnnt.cpp index e607720..3c32517 100644 --- a/src/rnnt.cpp +++ b/src/rnnt.cpp @@ -1,37 +1,10 @@ #include "rnnt.hpp" +#include "decode_common.hpp" #include #include namespace pk { -namespace { -// argmax over a[0..n) returning the first index of the maximum value. -// torch.max(dim) returns the FIRST max index on ties; match that. -int argmax(const float* a, int n) { - int best = 0; - float bv = a[0]; - for (int i = 1; i < n; ++i) { - if (a[i] > bv) { bv = a[i]; best = i; } - } - return best; -} - -// NeMo's rescaled `max_prob` confidence (method 'max_prob', alpha==1.0): -// conf = (N * p_max - 1) / (N - 1), p_max = softmax(logits)[k]. -// For RNN-T the confidence slice is the FULL joint output vector (V_plus = -// vocab + 1, blank included; no durations) — NeMo log_softmaxes the whole -// joint output. N == n (the slice size). Stable softmax (subtract the max). -float max_prob_conf_logits(const float* a, int n, int k) { - float mx = a[0]; - for (int i = 1; i < n; ++i) if (a[i] > mx) mx = a[i]; - double denom = 0.0; - for (int i = 0; i < n; ++i) denom += std::exp((double)a[i] - (double)mx); - const double p_max = std::exp((double)a[k] - (double)mx) / denom; - const double N = (double)n; - return (float)((N * p_max - 1.0) / (N - 1.0)); -} -} // namespace - RnntDecodeState rnnt_decode_init(const PredictionNet& pred) { RnntDecodeState st; st.state = pred.zero_state(); @@ -103,7 +76,7 @@ std::vector rnnt_decode_frames(const PredictionNet& pred, const Joint& joint.step_logits(enc_proj.data() + (size_t)t * H, g.data(), (int)g.size(), logits); - const int k = argmax(logits.data(), token_count); + const int k = decode_argmax(logits.data(), token_count); // Blank -> stop emitting at this frame and advance time. if (k == blank_id) break; @@ -117,7 +90,7 @@ std::vector rnnt_decode_frames(const PredictionNet& pred, const Joint& // max_prob confidence): frame = the (local) encoder frame t at // emission, conf = max_prob over the full joint output vector // (N = V_plus = vocab+1), span = 1 (RNN-T advances one frame). - const float conf = max_prob_conf_logits(logits.data(), token_count, k); + const float conf = decode_max_prob_conf(logits.data(), token_count, k); tokens->push_back(TokenInfo{ (int32_t)k, (int32_t)t, conf, 1 }); } st.last_token = (int32_t)k; diff --git a/src/subsampling.cpp b/src/subsampling.cpp index 0a2658b..f06b881 100644 --- a/src/subsampling.cpp +++ b/src/subsampling.cpp @@ -45,6 +45,200 @@ int Subsampling::valid_out_len(int T, int in_valid_frames) const { return valid; } +ggml_tensor* Subsampling::build_graph_batched(ggml_context* ctx, + const float* mel, + int n_mels, int T, int B, GraphInputPool& pool, + int& out_Tp, std::vector& out_valid, + const std::vector& valid_in) const { + const int C = conv_channels_; + const int F = n_mels; // feature dim (80) + const ModelLoader& ml = ml_; + + // This task targets the NON-causal (offline) model only. Batched causal + // subsampling (per-stage time masking with a batch axis) is out of scope: + // the causal branch below still operates on the single-item assumption. + GGML_ASSERT(!(causal_ && B > 1) && "batched causal subsampling not supported"); + + // --- Input (host-side): ggml conv data layout is [W=feat, H=T, IC=1, N=B]. + // NeMo conv input is [B,1,T,feat] (H=T, W=feat). We must feed + // x[(b*T + t)*F + f] = mel(item=b, feat=f, time=t). mel is per-item + // feat-major [F,T] (mel[(b*F + f)*T + t]); transpose into time-major per + // item in pool-owned storage (extra b*T block offset), feed as input. + std::vector& x_host = pool.alloc_f32((size_t)B * T * F); + for (int b = 0; b < B; ++b) + for (int t = 0; t < T; ++t) + for (int f = 0; f < F; ++f) + x_host[((size_t)b * T + t) * F + f] = + mel[((size_t)b * n_mels + f) * T + t]; + + int64_t x_ne[4] = {F, T, 1, B}; + ggml_tensor* x = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, x_ne, + x_host.data(), + x_host.size() * sizeof(float)); + + // Subsampling conv padding. NeMo dw_striding uses k=3, s=2 on each stage; the + // padding differs by model: + // non-causal (offline): symmetric (k-1)/2 = 1 on every side, applied + // directly via the conv's p0/p1 (byte-identical to the old path). + // causal (causal_downsampling=True, e.g. parakeet_realtime_eou_120m): + // NeMo CausalConv2D pads BOTH spatial axes (time H and feature W) with + // left = k-1 = 2, right = stride-1 = 1 (F.pad order (W_l,W_r,H_l,H_r)). + // ggml conv takes one symmetric p per axis, so for the causal case we pad + // explicitly with ggml_pad_ext (lp0/rp0 = W=feature, lp1/rp1 = H=time) + // and run the conv with p=0. + const bool causal = causal_; + auto pad_causal = [&](ggml_tensor* t) -> ggml_tensor* { + return ggml_pad_ext(ctx, t, /*lp0*/2, /*rp0*/1, /*lp1*/2, /*rp1*/1, + 0, 0, 0, 0); + }; + + // NeMo's MaskedConvSequential zeros the trailing (pad) time frames of the + // conv input BEFORE every stage. We replicate this per-item, per-stage in + // BOTH paths: + // - Causal: the right pad is +1, so the last valid output frame DOES read + // the trailing pad input frame; per-stage input masking is required for + // correctness even at B=1. + // - Non-causal (offline), B>1: a shorter clip is zero-padded to T_max, but + // after every conv stage bias+ReLU make the padded time region NON-ZERO, + // so the last valid output frame of a short item reads contaminated + // values instead of the clean conv zero-edge a standalone clip sees. + // Zeroing the trailing pad time frames before each stage reproduces the + // standalone boundary (the conv's own symmetric pad supplies clean zeros). + // The mask is per-item: [1, H, 1, B], md[b*H + h] = (h < vt[b]) ? 1 : 0, + // broadcasting over ne0 (W=feat) and ne2 (C). + auto mask_time = [&](ggml_tensor* t, const std::vector& vt) -> ggml_tensor* { + const int H = (int)t->ne[1]; + const int Bx = (int)t->ne[3]; + bool any = false; + for (int b = 0; b < Bx; ++b) { + int v = (b < (int)vt.size()) ? vt[b] : H; + if (v < H) { any = true; break; } + } + if (!any) return t; + std::vector& md = pool.alloc_f32((size_t)Bx * H); + for (int b = 0; b < Bx; ++b) { + int v = (b < (int)vt.size()) ? vt[b] : H; + for (int h = 0; h < H; ++h) + md[(size_t)b * H + h] = (h < v) ? 1.0f : 0.0f; + } + int64_t m_ne[4] = {1, H, 1, Bx}; + ggml_tensor* tm = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, m_ne, + md.data(), md.size() * sizeof(float)); + return ggml_mul(ctx, t, tm); // broadcast over ne0(W), ne2(C) + }; + // Per-item per-stage valid TIME lengths at the INPUT of each conv stage, + // mirroring valid_out_len's recurrence (and the old single-item valid_t0/1/2). + const int all_paddings = causal_ ? 3 : 2; + std::vector vt_stage0(B), vt_stage1(B), vt_stage2(B); // input of stage0/1/2 + for (int b = 0; b < B; ++b) { + int vi = (b < (int)valid_in.size()) ? valid_in[b] : -1; + int v0 = (vi >= 0) ? vi : (T - 1); // before stage 0 + int v1 = (v0 + all_paddings - 3) / 2 + 1; // before stage 1 (after stage 0) + int v2 = (v1 + all_paddings - 3) / 2 + 1; // before stage 2 (after stage 1) + vt_stage0[b] = v0; + vt_stage1[b] = v1; + vt_stage2[b] = v2; + } + + // ---- Stage 1: full Conv2d(1 -> C, k=3, s=2) + ReLU ---- + // kernel conv.0.weight: torch [C,1,3,3] -> ggml ne [3,3,1,C] = [KW,KH,IC,OC]. + ggml_tensor* w0 = clone_weight(ctx, ml, "encoder.pre_encode.conv.0.weight"); + ggml_tensor* b0 = clone_weight(ctx, ml, "encoder.pre_encode.conv.0.bias"); + x = mask_time(x, vt_stage0); // zero trailing pad time frames (both paths) + if (causal) { + x = pad_causal(x); + x = ggml_conv_2d(ctx, w0, x, /*s0*/2, /*s1*/2, /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/1); + } else { + x = ggml_conv_2d(ctx, w0, x, /*s0*/2, /*s1*/2, /*p0*/1, /*p1*/1, /*d0*/1, /*d1*/1); + } + // x: ne [OW=F/2, OH=T/2, OC=C, 1]. Add bias broadcast over channels: + // reshape bias to [1,1,C,1] so it broadcasts across W,H. + x = ggml_add(ctx, x, ggml_reshape_4d(ctx, b0, 1, 1, C, 1)); + x = ggml_relu(ctx, x); + + // ---- Stages 2 & 3: depthwise(k=3,s=2,p=1,groups=C) + pointwise(k=1) + ReLU ---- + struct StageW { const char* dw_w; const char* dw_b; const char* pw_w; const char* pw_b; }; + const StageW stages[2] = { + { "encoder.pre_encode.conv.2.weight", "encoder.pre_encode.conv.2.bias", + "encoder.pre_encode.conv.3.weight", "encoder.pre_encode.conv.3.bias" }, + { "encoder.pre_encode.conv.5.weight", "encoder.pre_encode.conv.5.bias", + "encoder.pre_encode.conv.6.weight", "encoder.pre_encode.conv.6.bias" }, + }; + const std::vector* stage_valid_t[2] = {&vt_stage1, &vt_stage2}; + for (int si = 0; si < 2; ++si) { + const StageW& s = stages[si]; + // Depthwise: weight torch [C,1,3,3] -> ggml ne [3,3,1,C] = [KW,KH,1,C]. + // ggml_conv_2d_dw_direct expects a:[KW,KH,1,C], b:[W,H,C,N]. + ggml_tensor* dww = clone_weight(ctx, ml, s.dw_w); + ggml_tensor* dwb = clone_weight(ctx, ml, s.dw_b); + x = mask_time(x, *stage_valid_t[si]); // zero trailing pad time frames (both paths) + if (causal) { + x = pad_causal(x); + x = ggml_conv_2d_dw_direct(ctx, dww, x, /*s0*/2, /*s1*/2, /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/1); + } else { + x = ggml_conv_2d_dw_direct(ctx, dww, x, /*s0*/2, /*s1*/2, /*p0*/1, /*p1*/1, /*d0*/1, /*d1*/1); + } + // x: ne [OW, OH, C, 1]. dw_direct keeps WHCN; make it contiguous so the + // bias add and following ops see a standard layout. + x = ggml_cont(ctx, x); + x = ggml_add(ctx, x, ggml_reshape_4d(ctx, dwb, 1, 1, C, 1)); + + // Pointwise: weight torch [C,C,1,1] -> ggml ne [1,1,C,C] = [KW,KH,IC,OC]. + ggml_tensor* pww = clone_weight(ctx, ml, s.pw_w); + ggml_tensor* pwb = clone_weight(ctx, ml, s.pw_b); + x = ggml_conv_2d(ctx, pww, x, /*s0*/1, /*s1*/1, /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/1); + x = ggml_add(ctx, x, ggml_reshape_4d(ctx, pwb, 1, 1, C, 1)); + x = ggml_relu(ctx, x); + } + + // x: ne [F'=OW, T'=OH, C, B]. NeMo flatten (per item): + // [B,C,T',F'].transpose(1,2).reshape(B,T',C*F') + // -> per time t, vector is channel-major: idx = c*F' + f. + const int Fp = (int)x->ne[0]; // F' + const int Tp = (int)x->ne[1]; // T' + // Want contiguous [F', C, T', B] so flat[b] = t*(C*F') + c*F' + f. + // current dims (0,1,2,3) = (F', T', C, B); permute to (F', C, T', B). + ggml_tensor* xp = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); + ggml_tensor* flat = ggml_reshape_3d(ctx, xp, (int64_t)C * Fp, Tp, B); // [C*F', T', B] + + // --- Length masking (faithful to NeMo MaskedConvSequential) --- + // Valid output frames never read masked input frames (kernel reach stays + // inside the valid region), so we can run the conv stack spatially and zero + // the flattened conv output at frames >= valid_out[b] before the Linear. + out_valid.assign(B, 0); + bool any_masked = false; + for (int b = 0; b < B; ++b) { + int vi = (b < (int)valid_in.size()) ? valid_in[b] : -1; + int vo = valid_out_len(T, vi); + out_valid[b] = (vo > Tp) ? Tp : vo; + if (vo < Tp) any_masked = true; + } + if (any_masked) { + // [1, Tp, B] mask: md[b*Tp + t] = (t < valid_out[b]) ? 1 : 0; broadcasts + // over ne0 (the C*F' feature axis). + std::vector& outmask = pool.alloc_f32((size_t)B * Tp); + for (int b = 0; b < B; ++b) { + // out_valid[b] == min(valid_out_len(T, vi), Tp); since this loop is + // bounded by Tp, "t < out_valid[b]" matches the unclamped "t < vo". + for (int t = 0; t < Tp; ++t) + outmask[(size_t)b * Tp + t] = (t < out_valid[b]) ? 1.0f : 0.0f; + } + int64_t mk_ne[3] = {1, Tp, B}; + ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, mk_ne, + outmask.data(), outmask.size() * sizeof(float)); + flat = ggml_mul(ctx, flat, mask); + } + + // ---- Linear out: torch [d_model, C*F'] -> ggml ne [C*F', d_model]. ---- + ggml_tensor* ow = clone_weight(ctx, ml, "encoder.pre_encode.out.weight"); + ggml_tensor* ob = clone_weight(ctx, ml, "encoder.pre_encode.out.bias"); + ggml_tensor* y = ggml_mul_mat(ctx, ow, flat); // [d_model, T', B] + y = ggml_add(ctx, y, ob); // broadcast bias [d_model] over T',B + + out_Tp = Tp; + return y; // ne [d_model, T', B] contiguous. +} + ggml_tensor* Subsampling::build_graph(ggml_context* ctx, const std::vector& mel, int n_mels, int T, GraphInputPool& pool, diff --git a/src/subsampling.hpp b/src/subsampling.hpp index 02aa20e..b2ae256 100644 --- a/src/subsampling.hpp +++ b/src/subsampling.hpp @@ -27,6 +27,15 @@ class Subsampling { int n_mels, int T, GraphInputPool& pool, int& out_Tp, int& out_valid, int in_valid_frames = -1) const; + // Batched GRAPH-BUILDER. `mel` is contiguous [B][n_mels][T] (mel[(b*n_mels+m)*T+t]), + // `valid_in` holds per-item valid mel frame counts (size B; element <0 means use + // the offline T-1 convention for that item). Returns [d_model, T', B] + // (ne0=d_model, ne1=T', ne2=B). `out_valid` (size B) receives each item's + // non-pad output-frame count. + ggml_tensor* build_graph_batched(ggml_context* ctx, const float* mel, + int n_mels, int T, int B, GraphInputPool& pool, + int& out_Tp, std::vector& out_valid, + const std::vector& valid_in) const; // mel: row-major [n_mels, T] (feat-major inner = T) — i.e. mel[m*T + t]. // out: row-major [Tout, d_model] (time-major) matching baseline subsampling_out. void forward(const std::vector& mel, int n_mels, int T, diff --git a/src/tdt.cpp b/src/tdt.cpp index f731ce7..72b5f71 100644 --- a/src/tdt.cpp +++ b/src/tdt.cpp @@ -1,38 +1,10 @@ #include "tdt.hpp" +#include "decode_common.hpp" #include #include namespace pk { -namespace { -// argmax over a[0..n) returning the first index of the maximum value. -// torch.max(dim) returns the FIRST max index on ties; match that. -int argmax(const float* a, int n) { - int best = 0; - float bv = a[0]; - for (int i = 1; i < n; ++i) { - if (a[i] > bv) { bv = a[i]; best = i; } - } - return best; -} - -// NeMo's rescaled `max_prob` confidence (method 'max_prob', alpha==1.0): -// conf = (N * p_max - 1) / (N - 1), p_max = softmax(slice)[argmax]. -// Computed numerically from the RAW logit slice a[0..n): p_max is the softmax -// probability of the argmax over the slice (equivalently exp of the max -// log_softmax value), and N == n (the slice size = num token classes incl. -// blank). Stable softmax (subtract the max). -float max_prob_conf_logits(const float* a, int n, int k) { - float mx = a[0]; - for (int i = 1; i < n; ++i) if (a[i] > mx) mx = a[i]; - double denom = 0.0; - for (int i = 0; i < n; ++i) denom += std::exp((double)a[i] - (double)mx); - const double p_max = std::exp((double)a[k] - (double)mx) / denom; - const double N = (double)n; - return (float)((N * p_max - 1.0) / (N - 1.0)); -} -} // namespace - std::vector tdt_greedy(const PredictionNet& pred, const Joint& joint, const std::vector& enc, int T, int enc_hidden, const std::vector& durations, @@ -102,8 +74,8 @@ std::vector tdt_greedy(const PredictionNet& pred, const Joint& joint, g.data(), (int)g.size(), logits); // Split: token logits [0, token_count), duration logits [token_count, V_plus). - const int k = argmax(logits.data(), token_count); - const int d_k = argmax(logits.data() + token_count, num_dur); + const int k = decode_argmax(logits.data(), token_count); + const int d_k = decode_argmax(logits.data() + token_count, num_dur); skip = durations[d_k]; // Commit state + last_token ONLY when k != blank. @@ -117,7 +89,7 @@ std::vector tdt_greedy(const PredictionNet& pred, const Joint& joint, // (NeMo log_softmaxes that slice; exclude the duration // logits). N = token_count = vocab + 1. // span = durations[d_k] (the duration/skip applied to the token). - const float conf = max_prob_conf_logits(logits.data(), token_count, k); + const float conf = decode_max_prob_conf(logits.data(), token_count, k); tokens->push_back(TokenInfo{ (int32_t)k, (int32_t)t, conf, (int32_t)skip }); } diff --git a/src/transducer_batch.cpp b/src/transducer_batch.cpp new file mode 100644 index 0000000..4320c7b --- /dev/null +++ b/src/transducer_batch.cpp @@ -0,0 +1,215 @@ +#include "transducer_batch.hpp" +#include "decode_common.hpp" +#include +#include + +namespace pk { + +// Batched greedy transducer decode. This is a faithful transposition of the +// per-item loops in tdt.cpp (tdt_greedy) and rnnt.cpp (rnnt_decode_frames), +// run for N items in lockstep "rounds". Each round runs ONE batched prediction +// step and ONE batched joint step over all active items, then applies the +// EXACT per-item rule from the oracle to each item independently. +// +// Parity rationale: +// - The per-item `g_valid` cache (skip the batched LSTM forward on rounds where +// no active item emitted) is a pure speed optimization: recomputing g from the +// SAME committed state yields an identical g, so reusing the cached g column is +// byte-identical. Mirrors the per-item `g_valid` in tdt.cpp/rnnt.cpp. +// - State recovery / masking: we only copy out_state columns into `committed` +// for items that emitted this round; non-emitting items keep their prior +// committed columns, exactly as the per-item loop leaves committed unchanged +// on a blank. +void transducer_greedy_batch( + const PredictionNet& pred, const Joint& joint, + const std::vector>& encs, + const std::vector& T, + int enc_hidden, + const std::vector& durations, + int blank_id, int max_symbols, + std::vector>& ids, + std::vector>* toks) { + + const bool is_tdt = !durations.empty(); + const int N = (int)encs.size(); + assert((int)T.size() == N); + + const int Hj = joint.joint_hidden(); + const int Hp = pred.hidden_size(); + const int L = pred.num_layers(); + const int Vp = joint.V_plus(); + const int num_dur = (int)durations.size(); + // RNNT: argmax over the full V_plus. TDT: argmax over the token slice + // (vocab+1), durations live in [token_count, V_plus). Mirrors the oracle: + // tdt.cpp token_count = V_plus - num_dur + // rnnt.cpp token_count = V_plus + const int token_count = is_tdt ? (Vp - num_dur) : Vp; + + // Per-item precomputed encoder projection [T[n], Hj]. + std::vector> ep(N); + for (int n = 0; n < N; ++n) { + joint.precompute_enc_proj(encs[n], T[n], enc_hidden, ep[n]); + } + + // Outputs. + ids.assign(N, {}); + if (toks) toks->assign(N, {}); + + // Per-item host state. + std::vector t(N, 0); + std::vector active(N, 0); + std::vector last_token(N, -1); + std::vector have_token(N, 0); + // Per-frame symbol counter (TDT symbols_added / RNNT emitted). + std::vector sym_at_frame(N, 0); + for (int n = 0; n < N; ++n) active[n] = (T[n] > 0) ? 1 : 0; + + // Per-item prediction-net cache validity. 0 = stale (g column must be + // recomputed before the joint), 1 = fresh (the persistent `g` buffer still + // holds this item's correct column from a prior round). All stale initially + // so the first round runs pred from SOS. Mirrors tdt.cpp/rnnt.cpp `g_valid`: + // set false on emit (committed state advanced), reused otherwise. Bit-exact: + // recomputing g from an UNCHANGED committed state yields an identical g, so + // skipping the pred step on all-valid rounds reuses the same values. + std::vector g_valid(N, 0); + + // Committed batched LSTM state, zero-initialized [L][Hp*N]. + BatchedPredState committed; + committed.h.assign((size_t)L, std::vector((size_t)Hp * N, 0.0f)); + committed.c.assign((size_t)L, std::vector((size_t)Hp * N, 0.0f)); + + // Scratch reused across rounds. + std::vector token_ids(N); + std::vector is_sos(N); + std::vector g; // [Hp*N] + BatchedPredState out_state; + std::vector enc_proj_gathered((size_t)Hj * N); + std::vector logits; // [Vp*N] + + auto any_active = [&]() { + for (int n = 0; n < N; ++n) if (active[n]) return true; + return false; + }; + + // Commit item n's state columns (offset n*Hp, all L layers, both h and c) + // from out_state into committed. Used on every emit in both branches. + auto commit_state = [&](int n) { + for (int l = 0; l < L; ++l) { + std::memcpy(&committed.h[l][(size_t)n * Hp], + &out_state.h[l][(size_t)n * Hp], (size_t)Hp * sizeof(float)); + std::memcpy(&committed.c[l][(size_t)n * Hp], + &out_state.c[l][(size_t)n * Hp], (size_t)Hp * sizeof(float)); + } + }; + + // Assumes max_symbols >= 1 (NeMo default 10). The per-round emit rule below + // mirrors the oracle inner loops, which never run at max_symbols==0. + while (any_active()) { + // (1) Batched prediction step from the committed state, but only when + // some active item's cache is stale. If every active item already has a + // fresh `g` column (no emit since it was last computed), the persistent + // `g` buffer still holds the correct values and we skip the LSTM forward + // entirely — the same per-item caching as tdt.cpp/rnnt.cpp, batched. + bool any_stale = false; + for (int n = 0; n < N; ++n) { + if (active[n] && !g_valid[n]) { any_stale = true; break; } + } + if (any_stale) { + // Build inputs from committed last_token/have_token. Inactive items + // still need valid inputs (their output is ignored): SOS / blank_id. + for (int n = 0; n < N; ++n) { + is_sos[n] = have_token[n] ? 0 : 1; + token_ids[n] = have_token[n] ? last_token[n] : (int32_t)blank_id; + } + pred.step_batch(token_ids, is_sos, committed, g, out_state); + // Every active item's g column is now fresh. (Recomputing a + // non-emitter's g from its unchanged committed state reproduces its + // cached value bit-for-bit, so marking it valid is exact.) + for (int n = 0; n < N; ++n) if (active[n]) g_valid[n] = 1; + } + + // (2) Gather each active item's enc_proj row for its current frame and + // run ONE batched joint step -> logits[Vp*N]. + for (int n = 0; n < N; ++n) { + int tf = t[n]; + if (tf < 0) tf = 0; // t[n] only grows from 0; lower clamp is defensive. Upper clamp matters for inactive/boundary columns. + if (tf > T[n] - 1) tf = T[n] - 1; // clamp (boundary / inactive) + if (T[n] <= 0) tf = 0; // no frames: harmless, ignored + const float* src = (T[n] > 0) ? (ep[n].data() + (size_t)tf * Hj) : ep[n].data(); + if (T[n] > 0) { + std::memcpy(&enc_proj_gathered[(size_t)n * Hj], src, (size_t)Hj * sizeof(float)); + } else { + std::memset(&enc_proj_gathered[(size_t)n * Hj], 0, (size_t)Hj * sizeof(float)); + } + } + joint.step_logits_batch(enc_proj_gathered.data(), g.data(), Hp, N, logits); + + // (3) Per-item rule from the oracle. + for (int n = 0; n < N; ++n) { + if (!active[n]) continue; + const float* lz = logits.data() + (size_t)n * Vp; + const int k = decode_argmax(lz, token_count); + + if (is_tdt) { + // --- tdt.cpp inner iteration --- + const int d_k = decode_argmax(lz + token_count, num_dur); + int skip = durations[d_k]; + + if (k != blank_id) { + ids[n].push_back((int32_t)k); + if (toks) { + const float conf = decode_max_prob_conf(lz, token_count, k); + (*toks)[n].push_back(TokenInfo{ (int32_t)k, (int32_t)t[n], conf, + (int32_t)skip }); + } + last_token[n] = (int32_t)k; + have_token[n] = 1; + commit_state(n); + g_valid[n] = 0; // committed state advanced -> g stale next round + } + // ALWAYS: symbols_added += 1; t += skip; need_loop = (skip == 0). + sym_at_frame[n] += 1; + t[n] += skip; + + // Inner loop in tdt.cpp continues iff (need_loop && symbols_added + // < max_symbols), i.e. (skip == 0 && sym_at_frame < max_symbols). + // Otherwise the frame is done. Post-inner: tdt.cpp does + // if (skip == 0) skip = 1; // dead for t (skip is local) + // if (symbols_added == max_symbols) t += 1; + const bool frame_done = !(skip == 0 && sym_at_frame[n] < max_symbols); + if (frame_done) { + if (sym_at_frame[n] == max_symbols) t[n] += 1; + sym_at_frame[n] = 0; + } + } else { + // --- rnnt.cpp inner iteration --- + if (k == blank_id) { + // Blank -> stop emitting at this frame, advance time. + t[n] += 1; + sym_at_frame[n] = 0; + } else { + ids[n].push_back((int32_t)k); + if (toks) { + const float conf = decode_max_prob_conf(lz, token_count, k); + (*toks)[n].push_back(TokenInfo{ (int32_t)k, (int32_t)t[n], conf, 1 }); + } + last_token[n] = (int32_t)k; + have_token[n] = 1; + commit_state(n); + g_valid[n] = 0; // committed state advanced -> g stale next round + sym_at_frame[n] += 1; + // emitted == max_symbols exits the inner while -> advance frame. + if (sym_at_frame[n] >= max_symbols) { + t[n] += 1; + sym_at_frame[n] = 0; + } + } + } + + // Recompute activity after time update. + active[n] = (t[n] < T[n]) ? 1 : 0; + } + } +} + +} // namespace pk diff --git a/src/transducer_batch.hpp b/src/transducer_batch.hpp new file mode 100644 index 0000000..a207236 --- /dev/null +++ b/src/transducer_batch.hpp @@ -0,0 +1,21 @@ +#pragma once +#include "prediction.hpp" +#include "joint.hpp" +#include "decode_types.hpp" +#include +#include +namespace pk { +// Batched greedy decode for N utterances. encs[n]: row-major [T[n], enc_hidden]. +// durations empty -> RNNT (advance-by-1); non-empty -> TDT (advance-by-duration). +// Outputs per item: ids[n], and (if toks != nullptr) TokenInfo[n]. Produces +// output bit-identical to per-item rnnt_greedy / tdt_greedy. +void transducer_greedy_batch( + const PredictionNet& pred, const Joint& joint, + const std::vector>& encs, + const std::vector& T, // [N] per-item frame counts + int enc_hidden, + const std::vector& durations, // empty=RNNT + int blank_id, int max_symbols, + std::vector>& ids, // OUT [N][.] + std::vector>* toks); // OUT [N][.] or nullptr +} // namespace pk diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fa54a71..95fe32c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -13,20 +13,29 @@ pk_add_test(test_fft) pk_add_test(test_mel) pk_add_test(test_mel_gpu) pk_add_test(test_subsampling) +pk_add_test(test_subsampling_batch) pk_add_test(test_relpos_attention) +pk_add_test(test_relpos_attention_batch) pk_add_test(test_conformer) +pk_add_test(test_conformer_batch) pk_add_test(test_conv_eou) pk_add_test(test_encoder) +pk_add_test(test_encoder_batch) pk_add_test(test_encoder_eou) pk_add_test(test_streaming_encoder) pk_add_test(test_ctc) pk_add_test(test_prediction) pk_add_test(test_prediction_step) +pk_add_test(test_prediction_step_batch) pk_add_test(test_joint) +pk_add_test(test_joint_step_batch) pk_add_test(test_transducer_core) pk_add_test(test_tdt_greedy) +pk_add_test(test_transducer_greedy_batch) +pk_add_test(test_transducer_greedy_batch_rnnt) pk_add_test(test_timestamps_tokens) pk_add_test(test_timestamps) +pk_add_test(test_transcribe_batch_ts) pk_add_test(test_tokenizer) pk_add_test(test_transcribe) pk_add_test(test_transcribe_speech) @@ -38,31 +47,36 @@ pk_add_test(test_transcribe_eou) pk_add_test(test_streaming_decode) pk_add_test(test_streaming_mel) pk_add_test(test_capi) +pk_add_test(test_capi_batch) pk_add_test(test_capi_stream) pk_add_test(test_capi_timestamps) -set_tests_properties(test_model_loader test_mel test_mel_gpu test_subsampling test_relpos_attention - test_conformer test_conv_eou test_encoder test_encoder_eou +pk_add_test(test_capi_batch_json) +set_tests_properties(test_model_loader test_mel test_mel_gpu test_subsampling test_subsampling_batch test_relpos_attention test_relpos_attention_batch + test_conformer test_conformer_batch test_conv_eou test_encoder test_encoder_batch test_encoder_eou test_streaming_encoder test_ctc test_prediction - test_prediction_step - test_joint test_transducer_core test_tdt_greedy - test_timestamps_tokens test_timestamps test_tokenizer test_transcribe + test_prediction_step test_prediction_step_batch + test_joint test_joint_step_batch test_transducer_core test_tdt_greedy + test_transducer_greedy_batch test_transducer_greedy_batch_rnnt + test_timestamps_tokens test_timestamps test_transcribe_batch_ts test_tokenizer test_transcribe test_transcribe_speech test_transcribe_tdt test_transcribe_0_6b test_transcribe_ctc test_transcribe_rnnt test_transcribe_eou - test_streaming_decode test_streaming_mel test_capi test_capi_stream - test_capi_timestamps + test_streaming_decode test_streaming_mel test_capi test_capi_batch test_capi_stream + test_capi_timestamps test_capi_batch_json PROPERTIES LABELS "model") # These tests read fixtures/baselines via paths relative to the project root. -set_tests_properties(test_mel test_mel_gpu test_subsampling test_relpos_attention test_conformer - test_conv_eou test_encoder test_encoder_eou test_streaming_encoder - test_ctc test_prediction test_prediction_step - test_joint - test_transducer_core test_tdt_greedy test_timestamps_tokens - test_timestamps +set_tests_properties(test_mel test_mel_gpu test_subsampling test_subsampling_batch test_relpos_attention test_relpos_attention_batch test_conformer test_conformer_batch + test_conv_eou test_encoder test_encoder_batch test_encoder_eou test_streaming_encoder + test_ctc test_prediction test_prediction_step test_prediction_step_batch + test_joint test_joint_step_batch + test_transducer_core test_tdt_greedy + test_transducer_greedy_batch test_transducer_greedy_batch_rnnt + test_timestamps_tokens + test_timestamps test_transcribe_batch_ts test_tokenizer test_transcribe test_transcribe_speech test_transcribe_tdt test_transcribe_0_6b test_transcribe_ctc test_transcribe_rnnt test_transcribe_eou - test_streaming_decode test_streaming_mel test_capi test_capi_stream - test_capi_timestamps + test_streaming_decode test_streaming_mel test_capi test_capi_batch test_capi_stream + test_capi_timestamps test_capi_batch_json PROPERTIES WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}) # Python converter check (skips with exit 77 when the venv/model are absent). diff --git a/tests/test_capi_batch.cpp b/tests/test_capi_batch.cpp new file mode 100644 index 0000000..228447c --- /dev/null +++ b/tests/test_capi_batch.cpp @@ -0,0 +1,38 @@ +#include "model.hpp" +#include "audio_io.hpp" +#include +#include +#include +#include + +// Batch-API B=1 equivalence smoke test. +// +// Loads the real-speech fixture, transcribes it single-clip, then runs it as a +// 2-clip batch of the same audio. Asserts both batch results equal the single +// result, proving the batch path is byte-identical to the single path. This is +// self-consistency (our own code on both sides), so it only needs +// PARAKEET_TEST_GGUF. +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + if (!gguf) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + auto model = pk::Model::load(gguf); + if (!model) { std::fprintf(stderr, "load failed\n"); return 1; } + + pk::Audio audio; + if (!pk::load_audio_16k_mono("tests/fixtures/speech.wav", audio) || audio.samples.empty()) { + std::fprintf(stderr, "wav load failed\n"); + return 1; + } + const std::vector& pcm = audio.samples; + + std::string single = model->transcribe_pcm(pcm, 16000); + auto batch = model->transcribe_pcm_batch({pcm, pcm}, 16000); + if (batch.size() != 2) { std::fprintf(stderr, "batch size %zu\n", batch.size()); return 1; } + if (batch[0] != single || batch[1] != single) { + std::fprintf(stderr, "MISMATCH single='%s' b0='%s' b1='%s'\n", + single.c_str(), batch[0].c_str(), batch[1].c_str()); + return 1; + } + std::fprintf(stderr, "OK batch==single: '%s'\n", single.c_str()); + return 0; +} diff --git a/tests/test_capi_batch_json.cpp b/tests/test_capi_batch_json.cpp new file mode 100644 index 0000000..0ccb38c --- /dev/null +++ b/tests/test_capi_batch_json.cpp @@ -0,0 +1,45 @@ +#include "parakeet_capi.h" +#include "audio_io.hpp" +#include +#include +#include +#include +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + if (!gguf) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + parakeet_ctx* ctx = parakeet_capi_load(gguf); + if (!ctx) { std::fprintf(stderr, "load failed\n"); return 1; } + pk::Audio a; + if (!pk::load_audio_16k_mono("tests/fixtures/speech.wav", a) || a.samples.empty()) { + std::fprintf(stderr, "wav load failed\n"); parakeet_capi_free(ctx); return 1; + } + char* single = parakeet_capi_transcribe_path_json(ctx, "tests/fixtures/speech.wav", 0); + if (!single) { std::fprintf(stderr, "single failed: %s\n", parakeet_capi_last_error(ctx)); parakeet_capi_free(ctx); return 1; } + std::string single_doc(single); + parakeet_capi_free_string(single); + + std::vector concat; + concat.insert(concat.end(), a.samples.begin(), a.samples.end()); + concat.insert(concat.end(), a.samples.begin(), a.samples.end()); + int n_samples[2] = { (int)a.samples.size(), (int)a.samples.size() }; + char* batch = parakeet_capi_transcribe_pcm_batch_json(ctx, concat.data(), n_samples, 2, 16000, 0); + if (!batch) { std::fprintf(stderr, "batch failed: %s\n", parakeet_capi_last_error(ctx)); parakeet_capi_free(ctx); return 1; } + std::string doc(batch); + parakeet_capi_free_string(batch); + + auto text_field = [](const std::string& s) -> std::string { + size_t p = s.find("\"text\":\""); + if (p == std::string::npos) return ""; + p += 8; size_t q = s.find('"', p); + return s.substr(p, q - p); + }; + std::string t = text_field(single_doc); + bool is_array = !doc.empty() && doc.front() == '[' && doc.back() == ']'; + size_t cnt = 0, pos = 0; + std::string needle = "\"text\":\"" + t + "\""; + while ((pos = doc.find(needle, pos)) != std::string::npos) { ++cnt; pos += needle.size(); } + bool ok = is_array && !t.empty() && cnt == 2; + std::fprintf(stderr, "array=%d text='%s' count=%zu -> %s\n", is_array, t.c_str(), cnt, ok?"OK":"FAIL"); + parakeet_capi_free(ctx); + return ok ? 0 : 1; +} diff --git a/tests/test_conformer_batch.cpp b/tests/test_conformer_batch.cpp new file mode 100644 index 0000000..e9d33b9 --- /dev/null +++ b/tests/test_conformer_batch.cpp @@ -0,0 +1,155 @@ +#include "conformer.hpp" +#include "model_loader.hpp" +#include "graph_builder.hpp" +#include "ggml_graph.hpp" +#include "backend.hpp" +#include "ggml.h" +#include "parity.hpp" +#include +#include +#include +#include +#include + +// Batched-layer isolation test for ConformerLayer::build_graph_batched (T3.2), +// which also exercises the conv module's B>1 pad-masking (T3.1) for the first +// time. It empirically validates the assembled batched conformer layer by +// comparing per-item batched output slices against the scalar +// ConformerLayer::forward reference. +// +// Input source: a real baseline tensor ("l0_attn_in" = [T, d_model] row-major) +// is reused as a deterministic [T,D] layer input, with baseline "pos_emb" +// ([2T-1, d_model] row-major) as the positional encoding. This test validates +// batched==scalar equivalence of the SAME layer on the SAME input; it does NOT +// need to match NeMo's layer output, so any consistent [T,D] input works as +// long as both paths get the same data. +// +// Layout mapping (load-bearing): +// scalar x : row-major [T, D] -> x[t*D + c] +// batched xt : ggml [D, T, B] -> data index (b*T + t)*D + c +// pe : ggml [D, pos_len] -> pe[p*D + c] (shared across batch) +// out : ggml [D, T, B] -> out[(b*T + t)*D + c] +// +// Case A (equal-length B=2, identical clips): item0 == item1 == scalar ref. +// Proves the assembled batched layer is correct under batching. +// Case B (ragged + padding invariance): item0 full (valid=T), item1 short +// (valid=T/2). item1's valid rows match the scalar short-clip ref; item0 is +// UNCHANGED vs the scalar full ref (no leakage from the shorter neighbor, +// exercising the conv module's B>1 pad-masking). + +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + const char* base = std::getenv("PARAKEET_TEST_BASELINE"); + if (!gguf || !base) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + + pk::ModelLoader ml; + if (!ml.load(gguf)) return 1; + + // Deterministic [T, d_model] layer input from the baseline. + std::vector x; std::vector xshape; + if (!pktest::load_baseline(base, "l0_attn_in", x, xshape)) return 1; + if (xshape.size() != 2) { std::fprintf(stderr, "l0_attn_in rank=%zu\n", xshape.size()); return 1; } + const int T = (int)xshape[0]; + const int D = (int)xshape[1]; + + // Relative positional encoding: baseline "pos_emb" is [2T-1, d_model] row-major. + std::vector pe; std::vector pshape; + if (!pktest::load_baseline(base, "pos_emb", pe, pshape)) return 1; + if (pshape.size() != 2) { std::fprintf(stderr, "pos_emb rank=%zu\n", pshape.size()); return 1; } + const int pos_len = (int)pshape[0]; + if (pos_len != 2 * T - 1 || (int)pshape[1] != D) { + std::fprintf(stderr, "pos_emb shape=[%d,%lld] expected [%d,%d]\n", + pos_len, (long long)pshape[1], 2 * T - 1, D); + return 1; + } + + pk::ConformerLayer layer(ml, /*layer_idx*/0); + + // Helper: run build_graph_batched for B items, all sharing the same input x + // (row-major [T,D]) and pe, with per-item valid_len. Returns [D,T,B] row-major. + auto run_batched = [&](int B, const std::vector& valid_len, + std::vector& out) -> bool { + pk::GraphInputPool pool; + // Build the [D,T,B] host buffer: xb[(b*T+t)*D + c] = x[t*D + c]. + std::vector xb((size_t)B * T * D); + for (int b = 0; b < B; ++b) + for (int t = 0; t < T; ++t) + for (int c = 0; c < D; ++c) + xb[((size_t)(b * T + t)) * D + c] = x[(size_t)t * D + c]; + return pk::run_graph(/*mem_bytes*/0, /*n_threads*/4, + [&](ggml_context* ctx) -> ggml_tensor* { + int64_t xt_ne[3] = {D, T, B}; + ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, xt_ne, + xb.data(), (size_t)B * T * D * sizeof(float)); + int64_t pe_ne[2] = {D, pos_len}; + ggml_tensor* pet = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne, + pe.data(), (size_t)pos_len * D * sizeof(float)); + return layer.build_graph_batched(ctx, xt, T, B, pet, pos_len, + valid_len, pool); + }, out); + }; + + // Extract item b's [T,D] (row-major) from a [D,T,B] row-major buffer. + auto slice_item = [&](const std::vector& out, int b) { + std::vector item((size_t)T * D); + for (int t = 0; t < T; ++t) + for (int c = 0; c < D; ++c) + item[(size_t)t * D + c] = out[((size_t)(b * T + t)) * D + c]; + return item; + }; + + bool ok = true; + + // ---- Case A: equal-length B=2, identical clips ---- + { + std::vector ref; + layer.forward(x, T, pe, pos_len, /*valid_len*/T, ref); + + std::vector out; + if (!run_batched(2, {T, T}, out)) { + std::fprintf(stderr, "[caseA] run_graph failed\n"); + return 1; + } + std::vector item0 = slice_item(out, 0); + std::vector item1 = slice_item(out, 1); + bool a0 = pktest::compare(item0, ref, "caseA_item0", 1e-3f, 1e-3f); + bool a1 = pktest::compare(item1, ref, "caseA_item1", 1e-3f, 1e-3f); + ok = ok && a0 && a1; + } + + // ---- Case B: ragged + padding invariance ---- + { + const int v0 = T; + const int v1 = T / 2; + + std::vector ref0; + layer.forward(x, T, pe, pos_len, /*valid_len*/v0, ref0); + std::vector ref1; + layer.forward(x, T, pe, pos_len, /*valid_len*/v1, ref1); + + std::vector out; + if (!run_batched(2, {v0, v1}, out)) { + std::fprintf(stderr, "[caseB] run_graph failed\n"); + return 1; + } + std::vector item0 = slice_item(out, 0); + std::vector item1 = slice_item(out, 1); + + // item0 (full) must be UNCHANGED vs the full scalar ref (padding + // invariance): the shorter neighbor must not perturb it. + bool b0 = pktest::compare(item0, ref0, "caseB_item0", 1e-3f, 1e-3f); + + // item1's first v1 valid query rows must match ref1's first v1 rows. + std::vector item1_valid((size_t)v1 * D); + std::vector ref1_valid((size_t)v1 * D); + for (int t = 0; t < v1; ++t) + for (int c = 0; c < D; ++c) { + item1_valid[(size_t)t * D + c] = item1[(size_t)t * D + c]; + ref1_valid[(size_t)t * D + c] = ref1[(size_t)t * D + c]; + } + bool b1 = pktest::compare(item1_valid, ref1_valid, "caseB_item1", 1e-3f, 1e-3f); + ok = ok && b0 && b1; + } + + return ok ? 0 : 1; +} diff --git a/tests/test_encoder_batch.cpp b/tests/test_encoder_batch.cpp new file mode 100644 index 0000000..489ae6e --- /dev/null +++ b/tests/test_encoder_batch.cpp @@ -0,0 +1,105 @@ +#include "encoder.hpp" +#include "model_loader.hpp" +#include "parity.hpp" +#include +#include +#include +#include +#include + +// T5.2: end-to-end equivalence + padding invariance for the fused batched +// encoder (forward_batch). +// +// Build a ragged B=2 batch where item0 is the full baseline mel (length T0) +// and item1 is the first T1=(3*T0)/4 frames, zero-padded up to T0. Run both +// clips standalone via Encoder::forward, run the batch via forward_batch, and +// assert each item's VALID output region matches its standalone result. A +// divergence on item1 (the shorter, zero-padded clip) or any perturbation of +// item0 by the shorter neighbor signals pad leakage in the batched encoder, +// NOT a test bug. Tolerance 5e-2/5e-2 mirrors test_encoder.cpp (error +// accumulates over the 17 conformer layers). +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + const char* base = std::getenv("PARAKEET_TEST_BASELINE"); + if (!gguf || !base) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + + pk::ModelLoader ml; + if (!ml.load(gguf)) { std::fprintf(stderr, "model load failed\n"); return 1; } + + // Baseline "mel" is [n_mels, T0] row-major (feat-major inner = T0). + std::vector mel; std::vector ms; + if (!pktest::load_baseline(base, "mel", mel, ms)) return 1; + if (ms.size() != 2) { std::fprintf(stderr, "mel rank=%zu\n", ms.size()); return 1; } + const int n_mels = (int)ms[0]; + const int T0 = (int)ms[1]; + const int T1 = (T0 * 3) / 4; + + pk::Encoder enc(ml); + + // --- Standalone references --------------------------------------------- + std::vector e0, e1; int dm = 0, to0 = 0, to1 = 0; + enc.forward(mel, n_mels, T0, e0, dm, to0); + + std::vector mel1((size_t)n_mels * T1); + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < T1; ++t) + mel1[(size_t)m * T1 + t] = mel[(size_t)m * T0 + t]; + int dm1 = 0; + enc.forward(mel1, n_mels, T1, e1, dm1, to1); + + // --- Batched (T_max=T0; item1 zero-padded) ----------------------------- + pk::MelBatch mb; + mb.B = 2; mb.n_mels = n_mels; mb.T_max = T0; mb.valid_T = { T0, T1 }; + mb.data.assign((size_t)2 * n_mels * T0, 0.0f); + for (int m = 0; m < n_mels; ++m) { + for (int t = 0; t < T0; ++t) mb.data[((size_t)0 * n_mels + m) * T0 + t] = mel[(size_t)m * T0 + t]; + for (int t = 0; t < T1; ++t) mb.data[((size_t)1 * n_mels + m) * T0 + t] = mel1[(size_t)m * T1 + t]; + } + std::vector> eo; int dmb = 0, tob = 0; std::vector vt; + enc.forward_batch(mb, eo, dmb, tob, vt); + + if (dmb <= 0 || tob <= 0 || eo.size() != 2 || vt.size() != 2) { + std::fprintf(stderr, "bad batched output dmb=%d tob=%d |eo|=%zu |vt|=%zu\n", + dmb, tob, eo.size(), vt.size()); + return 1; + } + + // Valid-length plumbing sanity. forward() reports Tout = each clip's OWN Tp, + // so to0 (full clip, T0) and to1 (shorter clip, T1) generally differ. The + // batch is padded to T_max=T0, so the batched tob == to0 (item0's Tp). + // vt[b] is the per-item NON-PAD valid frame count (= subsampling + // valid_out_len). For a clip whose length is not a clean power-of-two + // multiple the offline "T-1" convention can leave one trailing pad output + // frame, so vt[b] <= tob. The comparison below slices to vt[b] columns so + // only genuine (non-pad) frames are checked. + std::fprintf(stderr, + "[encbatch] dm=%d to0=%d to1=%d | dmb=%d tob=%d vt={%d,%d}\n", + dm, to0, to1, dmb, tob, vt[0], vt[1]); + if (dmb != dm || tob != to0) + std::fprintf(stderr, "[encbatch] WARN shape mismatch: dmb=%d(dm=%d) tob=%d(to0=%d)\n", + dmb, dm, tob, to0); + if (vt[0] < 1 || vt[0] > tob || vt[1] < 1 || vt[1] > to1) + std::fprintf(stderr, "[encbatch] WARN vt out of range: vt={%d,%d} tob=%d to1=%d\n", + vt[0], vt[1], tob, to1); + + // enc_out is channels-first [d_model, Tout] (full[c*Tout + t]); slice the + // first Tvalid columns from each row. + auto slice_cols = [&](const std::vector& full, int Tfull, int Tvalid) { + std::vector s((size_t)dmb * Tvalid); + for (int c = 0; c < dmb; ++c) + for (int t = 0; t < Tvalid; ++t) + s[(size_t)c * Tvalid + t] = full[(size_t)c * Tfull + t]; + return s; + }; + // forward_batch compacts each eo[b] to its own valid_Tout[b] columns, so its + // full row width is vt[b] (not the padded tob). The standalone references + // keep their own full Tout (to0 / to1). + std::vector b0 = slice_cols(eo[0], vt[0], vt[0]); + std::vector b1 = slice_cols(eo[1], vt[1], vt[1]); + std::vector ref0 = slice_cols(e0, to0, vt[0]); + std::vector ref1 = slice_cols(e1, to1, vt[1]); + + bool a = pktest::compare(b0, ref0, "encbatch.item0", 5e-2f, 5e-2f); + bool b = pktest::compare(b1, ref1, "encbatch.item1", 5e-2f, 5e-2f); + return (a && b) ? 0 : 1; +} diff --git a/tests/test_joint_step_batch.cpp b/tests/test_joint_step_batch.cpp new file mode 100644 index 0000000..a76b421 --- /dev/null +++ b/tests/test_joint_step_batch.cpp @@ -0,0 +1,46 @@ +#include "joint.hpp" +#include "prediction.hpp" +#include "encoder.hpp" +#include "model_loader.hpp" +#include "parity.hpp" +#include +#include +#include +#include +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + const char* base = std::getenv("PARAKEET_TEST_BASELINE"); + if (!gguf || !base) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + pk::ModelLoader ml; if (!ml.load(gguf)) return 1; + std::vector mel; std::vector ms; + if (!pktest::load_baseline(base, "mel", mel, ms)) return 1; + const int n_mels=(int)ms[0], T=(int)ms[1]; + pk::Encoder enc(ml); std::vector eo; int dm=0,Tout=0; + enc.forward(mel,n_mels,T,eo,dm,Tout); + std::vector encr((size_t)Tout*dm); + for (int t=0;t ep; joint.precompute_enc_proj(encr, Tout, dm, ep); + const int Hj = joint.joint_hidden(), Vp = joint.V_plus(); + pk::PredictionNet pred(ml); const int Hp = pred.hidden_size(); + std::vector frames = {0, Tout/2, Tout-1}; + std::vector toks = {7, 13, 21}; + const int N = 3; + std::vector> gs(N); + pk::PredState z = pred.zero_state(); + for (int n=0;n> ref(N); + for (int n=0;n epg((size_t)Hj*N), gg((size_t)Hp*N); + for (int n=0;n lb; joint.step_logits_batch(epg.data(), gg.data(), Hp, N, lb); + bool ok = (int)lb.size()==Vp*N; + for (int n=0;n col(lb.begin()+(size_t)n*Vp, lb.begin()+(size_t)(n+1)*Vp); + ok = pktest::compare(col, ref[n], "jointbatch", 1e-3f, 1e-3f) && ok; + } + return ok?0:1; +} diff --git a/tests/test_prediction_step_batch.cpp b/tests/test_prediction_step_batch.cpp new file mode 100644 index 0000000..bac4201 --- /dev/null +++ b/tests/test_prediction_step_batch.cpp @@ -0,0 +1,33 @@ +#include "prediction.hpp" +#include "model_loader.hpp" +#include "parity.hpp" +#include +#include +#include +#include +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + if (!gguf) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + pk::ModelLoader ml; if (!ml.load(gguf)) return 1; + pk::PredictionNet pred(ml); + const int H = pred.hidden_size(), L = pred.num_layers(); + std::vector toks = {5, 0, 42}; + std::vector sos = {0, 1, 0}; + const int N = 3; + std::vector> g_ref(N); + pk::PredState z = pred.zero_state(); + for (int n = 0; n < N; ++n) { pk::PredState os; pred.step(toks[n], sos[n], z, g_ref[n], os); } + pk::BatchedPredState bin; + bin.h.assign(L, std::vector((size_t)H*N, 0.0f)); + bin.c.assign(L, std::vector((size_t)H*N, 0.0f)); + std::vector gb; pk::BatchedPredState bout; + pred.step_batch(toks, sos, bin, gb, bout); + bool ok = (int)gb.size() == H*N; + for (int n = 0; n < N && ok; ++n) { + std::vector col(gb.begin()+(size_t)n*H, gb.begin()+(size_t)(n+1)*H); + ok = pktest::compare(col, g_ref[n], "predbatch.g", 1e-4f, 1e-4f) && ok; + } + // also check out_state top-layer h column matches g (sanity) and sizes + if (ok && ((int)bout.h.size()!=L || (int)bout.h[L-1].size()!=H*N)) { std::fprintf(stderr,"state shape\n"); ok=false; } + return ok ? 0 : 1; +} diff --git a/tests/test_relpos_attention_batch.cpp b/tests/test_relpos_attention_batch.cpp new file mode 100644 index 0000000..7cc50b6 --- /dev/null +++ b/tests/test_relpos_attention_batch.cpp @@ -0,0 +1,149 @@ +#include "relpos_attention.hpp" +#include "model_loader.hpp" +#include "graph_builder.hpp" +#include "ggml_graph.hpp" +#include "backend.hpp" +#include "ggml.h" +#include "parity.hpp" +#include +#include +#include +#include +#include + +// Batched-attention isolation test for RelPosAttention::build_graph_batched +// (the 4D rel-shift added in T4.1). It empirically validates the riskiest code +// path in the batched encoder by comparing per-item batched output slices +// against the scalar RelPosAttention::forward reference. +// +// Input source: the SAME real baseline tensors the scalar test uses +// ("l0_attn_in" = [T, d_model] row-major, "pos_emb" = [2T-1, d_model] row-major), +// so the batched path is exercised on the exact distribution NeMo produces. +// +// Layout mapping (load-bearing): +// scalar x : row-major [T, D] -> x[t*D + c] +// batched xt : ggml [D, T, B] -> data index (b*T + t)*D + c +// pe : ggml [D, pos_len] -> pe[p*D + c] (shared across batch) +// out : ggml [D, T, B] -> out[(b*T + t)*D + c] +// +// Case A (equal-length B=2, identical clips): item0 == item1 == scalar ref. +// Proves the 4D rel-shift is correct under batching. +// Case B (ragged + padding invariance): item0 full (valid=T), item1 short +// (valid=T/2). item1's valid rows match the scalar short-clip ref; item0 is +// UNCHANGED vs the scalar full ref (no leakage from the shorter neighbor). + +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + const char* base = std::getenv("PARAKEET_TEST_BASELINE"); + if (!gguf || !base) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + + pk::ModelLoader ml; + if (!ml.load(gguf)) return 1; + + // Attention input: baseline "l0_attn_in" is [T, d_model] row-major. + std::vector x; std::vector xshape; + if (!pktest::load_baseline(base, "l0_attn_in", x, xshape)) return 1; + if (xshape.size() != 2) { std::fprintf(stderr, "l0_attn_in rank=%zu\n", xshape.size()); return 1; } + const int T = (int)xshape[0]; + const int D = (int)xshape[1]; + + // Relative positional encoding: baseline "pos_emb" is [2T-1, d_model] row-major. + std::vector pe; std::vector pshape; + if (!pktest::load_baseline(base, "pos_emb", pe, pshape)) return 1; + if (pshape.size() != 2) { std::fprintf(stderr, "pos_emb rank=%zu\n", pshape.size()); return 1; } + const int pos_len = (int)pshape[0]; + if (pos_len != 2 * T - 1 || (int)pshape[1] != D) { + std::fprintf(stderr, "pos_emb shape=[%d,%lld] expected [%d,%d]\n", + pos_len, (long long)pshape[1], 2 * T - 1, D); + return 1; + } + + pk::RelPosAttention attn(ml, /*layer_idx*/0); + + // Helper: run build_graph_batched for B items, all sharing the same input x + // (row-major [T,D]) and pe, with per-item valid_len. Returns [D,T,B] row-major. + auto run_batched = [&](int B, const std::vector& valid_len, + std::vector& out) -> bool { + pk::GraphInputPool pool; + // Build the [D,T,B] host buffer: xb[(b*T+t)*D + c] = x[t*D + c]. + std::vector xb((size_t)B * T * D); + for (int b = 0; b < B; ++b) + for (int t = 0; t < T; ++t) + for (int c = 0; c < D; ++c) + xb[((size_t)(b * T + t)) * D + c] = x[(size_t)t * D + c]; + return pk::run_graph(/*mem_bytes*/0, /*n_threads*/4, + [&](ggml_context* ctx) -> ggml_tensor* { + int64_t xt_ne[3] = {D, T, B}; + ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, xt_ne, + xb.data(), (size_t)B * T * D * sizeof(float)); + int64_t pe_ne[2] = {D, pos_len}; + ggml_tensor* pet = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne, + pe.data(), (size_t)pos_len * D * sizeof(float)); + return attn.build_graph_batched(ctx, xt, T, B, pet, pos_len, + valid_len, pool); + }, out); + }; + + // Extract item b's [T,D] (row-major) from a [D,T,B] row-major buffer. + auto slice_item = [&](const std::vector& out, int b) { + std::vector item((size_t)T * D); + for (int t = 0; t < T; ++t) + for (int c = 0; c < D; ++c) + item[(size_t)t * D + c] = out[((size_t)(b * T + t)) * D + c]; + return item; + }; + + bool ok = true; + + // ---- Case A: equal-length B=2, identical clips ---- + { + std::vector ref; + attn.forward(x, T, pe, pos_len, /*valid_len*/T, ref); + + std::vector out; + if (!run_batched(2, {T, T}, out)) { + std::fprintf(stderr, "[caseA] run_graph failed\n"); + return 1; + } + std::vector item0 = slice_item(out, 0); + std::vector item1 = slice_item(out, 1); + bool a0 = pktest::compare(item0, ref, "caseA_item0", 1e-3f, 1e-3f); + bool a1 = pktest::compare(item1, ref, "caseA_item1", 1e-3f, 1e-3f); + ok = ok && a0 && a1; + } + + // ---- Case B: ragged + padding invariance ---- + { + const int v0 = T; + const int v1 = T / 2; + + std::vector ref0; + attn.forward(x, T, pe, pos_len, /*valid_len*/v0, ref0); + std::vector ref1; + attn.forward(x, T, pe, pos_len, /*valid_len*/v1, ref1); + + std::vector out; + if (!run_batched(2, {v0, v1}, out)) { + std::fprintf(stderr, "[caseB] run_graph failed\n"); + return 1; + } + std::vector item0 = slice_item(out, 0); + std::vector item1 = slice_item(out, 1); + + // item0 (full) must be UNCHANGED vs the full scalar ref (padding invariance). + bool b0 = pktest::compare(item0, ref0, "caseB_item0", 1e-3f, 1e-3f); + + // item1's first v1 query rows must match ref1's first v1 rows. + std::vector item1_valid((size_t)v1 * D); + std::vector ref1_valid((size_t)v1 * D); + for (int t = 0; t < v1; ++t) + for (int c = 0; c < D; ++c) { + item1_valid[(size_t)t * D + c] = item1[(size_t)t * D + c]; + ref1_valid[(size_t)t * D + c] = ref1[(size_t)t * D + c]; + } + bool b1 = pktest::compare(item1_valid, ref1_valid, "caseB_item1", 1e-3f, 1e-3f); + ok = ok && b0 && b1; + } + + return ok ? 0 : 1; +} diff --git a/tests/test_subsampling_batch.cpp b/tests/test_subsampling_batch.cpp new file mode 100644 index 0000000..becd8e1 --- /dev/null +++ b/tests/test_subsampling_batch.cpp @@ -0,0 +1,117 @@ +#include "subsampling.hpp" +#include "model_loader.hpp" +#include "ggml_graph.hpp" +#include "graph_builder.hpp" +#include "parity.hpp" +#include "ggml.h" +#include +#include +#include +#include + +// T2.2: batched-vs-standalone per-item equivalence for ConvSubsampling. +// +// Build a ragged B=2 batch where item0 is the full baseline mel (length T0) +// and item1 is the first T1=T0/2 frames of that same mel, zero-padded up to +// T0. Run build_graph_batched once, then assert each item's VALID output +// frames match Subsampling::forward run on that clip standalone. A larger diff +// on item1 (the shorter, zero-padded clip) would signal pad leakage in the +// batched builder, NOT a test bug. +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + const char* base = std::getenv("PARAKEET_TEST_BASELINE"); + if (!gguf || !base) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + + pk::ModelLoader ml; + if (!ml.load(gguf)) { std::fprintf(stderr, "model load failed\n"); return 1; } + + // Baseline "mel" is [n_mels, T0] row-major (feat-major inner = T0). + std::vector mel; std::vector ms; + if (!pktest::load_baseline(base, "mel", mel, ms)) return 1; + if (ms.size() != 2) { std::fprintf(stderr, "mel rank=%zu\n", ms.size()); return 1; } + const int n_mels = (int)ms[0]; + const int T0 = (int)ms[1]; + const int T1 = T0 / 2; + + pk::Subsampling sub(ml); + + // --- Standalone references --------------------------------------------- + // item0: full mel. + std::vector r0; int Tout0 = 0, dm = 0, vl0 = 0; + sub.forward(mel, n_mels, T0, r0, Tout0, dm, vl0); + + // item1: first T1 mel frames (feat-major [n_mels, T1]). + std::vector mel1((size_t)n_mels * T1); + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < T1; ++t) + mel1[(size_t)m * T1 + t] = mel[(size_t)m * T0 + t]; + std::vector r1; int Tout1 = 0, dm1 = 0, vl1 = 0; + sub.forward(mel1, n_mels, T1, r1, Tout1, dm1, vl1); + + // --- Batched build ------------------------------------------------------ + // Stack to T_max=T0; item1 is zero-padded for t >= T1. Layout the builder + // expects: contiguous [B][n_mels][Tmax], i.e. mel[(b*n_mels+m)*Tmax + t]. + const int Tmax = T0, B = 2; + std::vector stacked((size_t)B * n_mels * Tmax, 0.0f); + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < T0; ++t) + stacked[((size_t)0 * n_mels + m) * Tmax + t] = mel[(size_t)m * T0 + t]; + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < T1; ++t) + stacked[((size_t)1 * n_mels + m) * Tmax + t] = mel1[(size_t)m * T1 + t]; + + // Per-item entry valid counts. The offline convention's entry valid length + // for a clip of T mel frames is T-1; pass explicit positive counts here so + // item1 is masked at its true valid length (NOT Tmax-1). + const std::vector valid_in = { T0 - 1, T1 - 1 }; + + // GraphInputPool host buffers (mel transpose + masks) are registered by the + // builder and must outlive the run_graph compute. Declare the pool in this + // scope and capture it by reference (mirrors Subsampling::forward, which + // owns the pool in the enclosing scope of run_graph). + pk::GraphInputPool pool; + std::vector outflat; int Tp = 0; std::vector vout; + bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4, + [&](ggml_context* ctx) -> ggml_tensor* { + return sub.build_graph_batched(ctx, stacked.data(), n_mels, Tmax, B, + pool, Tp, vout, valid_in); + }, outflat); + if (!ok) { std::fprintf(stderr, "batched graph failed\n"); return 1; } + + if (Tp <= 0 || dm <= 0) { std::fprintf(stderr, "bad Tp=%d dm=%d\n", Tp, dm); return 1; } + std::fprintf(stderr, + "[subbatch] Tp=%d dm=%d Tout0=%d Tout1=%d valid_out={%d,%d}\n", + Tp, dm, Tout0, Tout1, + vout.size() > 0 ? vout[0] : -1, vout.size() > 1 ? vout[1] : -1); + + // outflat is [d_model, Tp, B] (ne0=d_model fastest): + // element (c, t, b) at index ((size_t)b*Tp + t)*dm + c. + auto slice = [&](int b, int Tvalid) { + std::vector s((size_t)Tvalid * dm); + for (int t = 0; t < Tvalid; ++t) + for (int c = 0; c < dm; ++c) + s[(size_t)t * dm + c] = outflat[(((size_t)b * Tp + t) * dm) + c]; + return s; + }; + // r_b is [Tout_b, d_model] row-major; compare against the first Tout_b + // valid frames of each batched item. + std::vector s0 = slice(0, Tout0); + std::vector s1 = slice(1, Tout1); + + bool a = pktest::compare(s0, r0, "subbatch.item0", 1e-3f, 1e-3f); + // item1 (the shorter, zero-padded clip) compares two DIFFERENT code paths by + // design: the batched builder (s1) masks the padded time region per conv + // stage, while the scalar forward (r1) now runs the lean 2-D graph on the + // standalone clip. They are numerically equivalent on interior frames (match + // to ~1e-3) but diverge at the SINGLE trailing valid frame: its 3x- + // downsampled receptive field straddles the clip boundary, where the batched + // per-stage time masking and the standalone conv's own zero-edge produce a + // different value (different op orderings by design). Compare the interior + // frames (skipping that last boundary frame) at a modest 5e-3 to absorb the + // fp rounding of the two distinct op orders. + const int Tcmp = (Tout1 > 1) ? Tout1 - 1 : Tout1; + std::vector s1i(s1.begin(), s1.begin() + (size_t)Tcmp * dm); + std::vector r1i(r1.begin(), r1.begin() + (size_t)Tcmp * dm); + bool b2 = pktest::compare(s1i, r1i, "subbatch.item1", 5e-3f, 5e-3f); + return (a && b2) ? 0 : 1; +} diff --git a/tests/test_transcribe_batch_ts.cpp b/tests/test_transcribe_batch_ts.cpp new file mode 100644 index 0000000..4fd5947 --- /dev/null +++ b/tests/test_transcribe_batch_ts.cpp @@ -0,0 +1,38 @@ +#include "model.hpp" +#include "audio_io.hpp" +#include +#include +#include +#include +#include +static bool words_close(const std::vector& a, const std::vector& b) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) { + if (a[i].text != b[i].text) return false; + if (std::fabs(a[i].start - b[i].start) > 1e-3f) return false; + if (std::fabs(a[i].end - b[i].end) > 1e-3f) return false; + } + return true; +} +int main() { + const char* gguf = std::getenv("PARAKEET_TEST_GGUF"); + if (!gguf) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + auto model = pk::Model::load(gguf); + if (!model) { std::fprintf(stderr, "load failed\n"); return 1; } + pk::Audio a; + if (!pk::load_audio_16k_mono("tests/fixtures/speech.wav", a) || a.samples.empty()) { + std::fprintf(stderr, "wav load failed\n"); return 1; + } + std::vector half(a.samples.begin(), a.samples.begin() + (a.samples.size()*3)/4); + pk::Transcription r0 = model->transcribe_with_timestamps(a.samples, 16000); + pk::Transcription r1 = model->transcribe_with_timestamps(half, 16000); + auto batch = model->transcribe_pcm_batch_with_timestamps({a.samples, half}, 16000); + if (batch.size() != 2) { std::fprintf(stderr, "size %zu\n", batch.size()); return 1; } + bool ok = batch[0].text == r0.text && batch[1].text == r1.text + && words_close(batch[0].words, r0.words) && words_close(batch[1].words, r1.words); + std::fprintf(stderr, "item0 text %s words %zu/%zu; item1 text %s words %zu/%zu -> %s\n", + (batch[0].text==r0.text?"OK":"DIFF"), batch[0].words.size(), r0.words.size(), + (batch[1].text==r1.text?"OK":"DIFF"), batch[1].words.size(), r1.words.size(), + ok?"OK":"FAIL"); + return ok ? 0 : 1; +} diff --git a/tests/test_transducer_greedy_batch.cpp b/tests/test_transducer_greedy_batch.cpp new file mode 100644 index 0000000..8438d74 --- /dev/null +++ b/tests/test_transducer_greedy_batch.cpp @@ -0,0 +1,54 @@ +#include "transducer_batch.hpp" +#include "tdt.hpp" +#include "prediction.hpp" +#include "joint.hpp" +#include "encoder.hpp" +#include "mel.hpp" +#include "audio_io.hpp" +#include "model_loader.hpp" +#include "parity.hpp" +#include +#include +#include +#include +#include +static bool toks_equal(const std::vector& a, const std::vector& b){ + if (a.size()!=b.size()) return false; + for (size_t i=0;i1e-4f) return false; + } + return true; +} +int main(){ + const char* gguf=std::getenv("PARAKEET_TEST_GGUF"); const char* base=std::getenv("PARAKEET_TEST_BASELINE"); + if(!gguf||!base){ std::fprintf(stderr,"env not set; skip\n"); return 77; } + pk::ModelLoader ml; if(!ml.load(gguf)) return 1; + const auto& cfg = ml.config(); + if (cfg.tdt_durations.empty()){ std::fprintf(stderr,"no TDT durations; skip\n"); return 77; } + // Compute a REAL speech mel so the decode actually emits tokens (the + // emit/commit/duration paths get exercised, not just all-blank). Falls back + // to nothing else: speech.wav ships in tests/fixtures. + pk::Audio audio; if(!pk::load_audio_16k_mono("tests/fixtures/speech.wav", audio)){ std::fprintf(stderr,"speech.wav load failed\n"); return 1; } + pk::MelFrontend melfe(ml); + std::vector mel; int n_mels=0, T0=0; + melfe.compute(audio.samples, mel, n_mels, T0); // row-major [n_mels, T0] + pk::Encoder enc(ml); + auto enc_row=[&](const std::vector& m,int Tn){ std::vector eo;int dm=0,to=0; enc.forward(m,n_mels,Tn,eo,dm,to); + std::vector r((size_t)to*dm); for(int t=0;t mel1((size_t)n_mels*T1); + for(int m=0;m r0,r1; + auto id0 = pk::tdt_greedy(pred,joint,e0,to0,dm,cfg.tdt_durations,blank,maxs,&r0); + auto id1 = pk::tdt_greedy(pred,joint,e1,to1,dm,cfg.tdt_durations,blank,maxs,&r1); + std::vector> encs={e0,e1}; std::vector Ts={to0,to1}; + std::vector> ids; std::vector> tk; + pk::transducer_greedy_batch(pred,joint,encs,Ts,dm,cfg.tdt_durations,blank,maxs,ids,&tk); + bool ok = ids.size()==2 && ids[0]==id0 && ids[1]==id1 && toks_equal(tk[0],r0) && toks_equal(tk[1],r1); + std::fprintf(stderr,"item0 ids %zu/%zu words; item1 ids %zu/%zu -> %s\n", ids[0].size(),id0.size(),ids[1].size(),id1.size(), ok?"OK":"FAIL"); + return ok?0:1; +} diff --git a/tests/test_transducer_greedy_batch_rnnt.cpp b/tests/test_transducer_greedy_batch_rnnt.cpp new file mode 100644 index 0000000..758e5bb --- /dev/null +++ b/tests/test_transducer_greedy_batch_rnnt.cpp @@ -0,0 +1,51 @@ +#include "transducer_batch.hpp" +#include "rnnt.hpp" +#include "prediction.hpp" +#include "joint.hpp" +#include "encoder.hpp" +#include "model_loader.hpp" +#include "parity.hpp" +#include +#include +#include +#include +#include +static bool toks_equal(const std::vector& a, const std::vector& b){ + if (a.size()!=b.size()) return false; + for (size_t i=0;i1e-4f) return false; + } + return true; +} +int main(){ + // RNNT-only model (the 110m anchor is TDT, not pure RNNT). Self-skip unless a + // dedicated RNNT GGUF is provided. + const char* gguf=std::getenv("PARAKEET_TEST_GGUF_RNNT"); const char* base=std::getenv("PARAKEET_TEST_BASELINE"); + if(!gguf||!base){ std::fprintf(stderr,"PARAKEET_TEST_GGUF_RNNT not set; skip\n"); return 77; } + pk::ModelLoader ml; if(!ml.load(gguf)) return 1; + const auto& cfg = ml.config(); + if (!cfg.tdt_durations.empty()){ std::fprintf(stderr,"model has TDT durations; not pure RNNT; skip\n"); return 77; } + std::vector mel; std::vector ms; + if(!pktest::load_baseline(base,"mel",mel,ms)) return 1; + const int n_mels=(int)ms[0], T0=(int)ms[1]; + pk::Encoder enc(ml); + auto enc_row=[&](const std::vector& m,int Tn){ std::vector eo;int dm=0,to=0; enc.forward(m,n_mels,Tn,eo,dm,to); + std::vector r((size_t)to*dm); for(int t=0;t mel1((size_t)n_mels*T1); + for(int m=0;m r0,r1; + auto id0 = pk::rnnt_greedy(pred,joint,e0,to0,dm,blank,maxs,&r0); + auto id1 = pk::rnnt_greedy(pred,joint,e1,to1,dm,blank,maxs,&r1); + std::vector> encs={e0,e1}; std::vector Ts={to0,to1}; + std::vector> ids; std::vector> tk; + std::vector no_dur{}; + pk::transducer_greedy_batch(pred,joint,encs,Ts,dm,no_dur,blank,maxs,ids,&tk); + bool ok = ids.size()==2 && ids[0]==id0 && ids[1]==id1 && toks_equal(tk[0],r0) && toks_equal(tk[1],r1); + std::fprintf(stderr,"item0 ids %zu/%zu words; item1 ids %zu/%zu -> %s\n", ids[0].size(),id0.size(),ids[1].size(),id1.size(), ok?"OK":"FAIL"); + return ok?0:1; +}