diff --git a/.github/workflows/typescript-ci.yml b/.github/workflows/typescript-ci.yml index 3eea822e..7bed9b9e 100644 --- a/.github/workflows/typescript-ci.yml +++ b/.github/workflows/typescript-ci.yml @@ -166,16 +166,14 @@ jobs: # transcription/streaming/cancel/extension coverage. jfk.wav ships in-repo; # only the model paths need exporting (fetch-canary handles that). # - # Windows is no-model-only for now: that is the koffi DLL-resolution proof - # this leg was added for, and it is unaffected. The model tier is deferred - # because loading a model spins up the ggml CPU threadpool, and on MSVC with - # OpenMP (vcomp) the runtime crashes (0xC0000005) at process exit under - # node/koffi, after the run itself succeeds. Skipping the canary makes the - # model-gated tests and the (model-loading) examples self-skip cleanly. Drop - # this guard to re-enable the Windows model tier once that teardown crash in - # the MSVC OpenMP runtime is fixed. + # Windows now runs the full model tier too. It was previously no-model-only + # because loading a model spun up ggml's CPU threadpool, and on MSVC with + # OpenMP (vcomp) the runtime crashed (0xC0000005) at process exit under + # node/koffi. The build now defaults to ggml's native threadpool (no OpenMP + # — see the OpenMP CENTRAL POLICY in CMakeLists.txt), so there is no vcomp + # pool to crash at teardown, and the native barrier is MSVC-correct (the + # transcribe_threadpool_oversubscription test guards it). Re-enabled here. - uses: ./.github/actions/fetch-canary - if: runner.os != 'Windows' with: hf-token: ${{ secrets.HF_TOKEN }} - name: Conformance (no-model tier always; model tier when canary present) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9d64c434..d366e8ea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -125,12 +125,15 @@ option(TRANSCRIBE_LTO "Enable LTO in Release" OFF) # and bench workflow relies on. The Python wheel/provider build sets it ON to # produce a shared libtranscribe the FFI layer can dlopen. # -# TRANSCRIBE_USE_OPENMP / TRANSCRIBE_USE_SYSTEM_BLAS default ON to keep the -# developer build fast (ggml OpenMP, Accelerate/system BLAS host decoder). -# Official provider wheels pass them OFF so no libgomp/libiomp or OpenBLAS/MKL -# runtime is vendored into the wheel where it can collide with PyTorch/NumPy. +# TRANSCRIBE_USE_OPENMP defaults OFF: we use ggml's native CPU threadpool +# everywhere (see the OpenMP CENTRAL POLICY below — OpenMP's process-global pool +# is a teardown/coexistence liability for an embeddable library, and the native +# pool is now correct on MSVC and under oversubscription). Opt in with +# -DTRANSCRIBE_USE_OPENMP=ON. TRANSCRIBE_USE_SYSTEM_BLAS stays ON for the +# Accelerate/system-BLAS host decoder; official provider wheels pass it OFF so no +# OpenBLAS/MKL runtime is vendored where it can collide with PyTorch/NumPy. option(TRANSCRIBE_BUILD_SHARED "Build libtranscribe + ggml as shared libraries" OFF) -option(TRANSCRIBE_USE_OPENMP "Use OpenMP (ggml + Parakeet host-decoder TU)" ON) +option(TRANSCRIBE_USE_OPENMP "Use OpenMP for ggml's CPU threadpool" OFF) option(TRANSCRIBE_USE_SYSTEM_BLAS "Link non-Apple system BLAS for the host decoder" ON) # Dynamic ggml backends: each backend (CPU, Vulkan, CUDA, ...) becomes a @@ -265,6 +268,9 @@ set(GGML_METAL ${TRANSCRIBE_METAL} CACHE BOOL "" FORCE) set(GGML_VULKAN ${TRANSCRIBE_VULKAN} CACHE BOOL "" FORCE) set(GGML_CUDA ${TRANSCRIBE_CUDA} CACHE BOOL "" FORCE) set(GGML_BLAS OFF CACHE BOOL "" FORCE) +# tinyBLAS (Justine Tunney's llamafile_sgemm CPU kernels): ~29% faster encoder on +# CPU (q8_0 GEMM), numerically WER-equivalent. On by default; CPU-backend only. +set(GGML_LLAMAFILE ON CACHE BOOL "" FORCE) # Conservative x86 floor (see the option's comment above): one switch fans # out to GGML_NATIVE plus the full x86 SIMD tier list. Placed BEFORE @@ -283,34 +289,36 @@ if(TRANSCRIBE_X86_CONSERVATIVE) endforeach() endif() -# OpenMP — CENTRAL POLICY (read before changing the Windows/MSVC threading). +# OpenMP — CENTRAL POLICY. # -# ggml defaults GGML_OPENMP ON and probes for it gracefully, so leave ggml's own -# default in place when we want OpenMP. Only force it OFF (matching the -# GGML_OPENMP=OFF posture official wheels use) when OpenMP is switched off, so -# one TRANSCRIBE_USE_OPENMP knob covers ggml and our host-decoder TU. +# DEFAULT: OFF. We use ggml's native CPU threadpool everywhere; OpenMP is an +# opt-in (-DTRANSCRIBE_USE_OPENMP=ON), not the default. The TRANSCRIBE_USE_OPENMP +# knob simply drives ggml's GGML_OPENMP. # -# The catch: ggml's NON-OpenMP CPU threadpool barrier (ggml_barrier's custom -# spin path) DEADLOCKS under MSVC codegen on Windows — any multi-threaded CPU -# run wedges its workers in the barrier spin (it works on every other compiler). -# OpenMP's `#pragma omp barrier` is the only working multi-threaded CPU path on -# MSVC. But OpenMP is a per-CONSUMER tradeoff, not a global flag, so this is not -# one switch — it is a deliberate, documented split: +# Why native, not OpenMP — OpenMP's runtime pool is process-global and outlives +# any single compute, which is a liability for an embeddable/dlopen'd library: +# - Teardown crash: a binding (node/koffi, ctypes, ...) calls in on a worker +# thread, libgomp/vcomp spawns its pool there, and at process teardown the +# loader unmaps the runtime out from under those still-live pool threads -> +# SIGSEGV / 0xC0000005. The native pool joins its threads per graph_compute, +# so nothing outlives the call and there is nothing to unmap-race. +# - Coexistence: a vendored libgomp/libiomp can collide with numpy/torch/MKL's +# own OpenMP in one process. The native pool vendors no second runtime. # -# - Rust binding (CPU is its default backend, standalone process): MUST have -# OpenMP on Windows. Its build.rs passes -DGGML_OPENMP=ON; the guard below -# honors that explicit value. MSVC auto-links vcomp via the /openmp pragma in -# the ggml objects (no -fopenmp flag, so the link manifest is untouched); -# vcomp140.dll ships in the VC++ runtime the binary already needs. -# - Python wheels (default to the Vulkan GPU backend; deliberately vendor NO -# OpenMP so ggml's runtime cannot collide with numpy/torch's own OpenMP/MKL -# in one process — see python-wheels.yml): leave GGML_OPENMP unset and get -# the force-off below. KNOWN LIMITATION as a consequence: multi-threaded CPU -# compute on Windows is unsupported for the wheels (they use Vulkan). Lifting -# it needs a non-OpenMP fix (e.g. single-threaded CPU on Windows, or patching -# ggml's barrier) rather than forcing OpenMP and the coexistence hazard back. +# This used to be a per-consumer SPLIT because ggml's native barrier deadlocked +# under MSVC (its spin `relax` was a no-op on MSVC, so a waiter starved an +# un-arrived worker under oversubscription) — OpenMP was the only working +# multi-threaded CPU path on Windows. That barrier is now fixed (ggml-cpu.c: +# YieldProcessor() relax + a bounded-spin ggml_thread_yield fallback), so the +# native pool is correct on MSVC and under oversubscription. The split is gone: +# - Rust binding: no longer needs -DGGML_OPENMP=ON on Windows; native pool. +# - Python wheels: already vendored no OpenMP; multi-threaded CPU now works on +# every platform (the old "CPU-on-Windows unsupported" limitation is lifted). +# - perf: native persistent/ephemeral pool is on par with OpenMP (it is +# llama.cpp's default). Thread count is sized affinity-aware in +# transcribe-batch-util (default_n_threads) to avoid needless oversubscription. # -# So: honor an explicit -DGGML_OPENMP=... (the Rust opt-in); otherwise force OFF. +# So: honor an explicit -DGGML_OPENMP=... (opt-in); otherwise force OFF. if(NOT TRANSCRIBE_USE_OPENMP AND NOT DEFINED CACHE{GGML_OPENMP}) set(GGML_OPENMP OFF CACHE BOOL "" FORCE) endif() diff --git a/bindings/rust/sys/build.rs b/bindings/rust/sys/build.rs index 552cdfeb..398dc6aa 100644 --- a/bindings/rust/sys/build.rs +++ b/bindings/rust/sys/build.rs @@ -157,29 +157,27 @@ fn main() { if feature("CUDA") { cfg.define("TRANSCRIBE_CUDA", "ON"); } - // Force OpenMP OFF unless explicitly opted in. TRANSCRIBE_USE_OPENMP - // defaults ON and auto-detects, but its `-fopenmp` shows up only as a - // manifest link_flag → a `cargo:rustc-link-arg` that does NOT propagate to - // downstream binaries, so a static consumer link fails with undefined - // GOMP_*/omp_* symbols. A self-contained static build is the default; - // `--features openmp` opts in (and then owns providing the OpenMP runtime). + // Keep OpenMP OFF unless explicitly opted in. TRANSCRIBE_USE_OPENMP already + // defaults OFF in CMake (the native ggml threadpool is the default path); we + // set it explicitly here so `--features openmp` is the single switch. We keep + // it OFF by default because OpenMP's `-fopenmp` shows up only as a manifest + // link_flag → a `cargo:rustc-link-arg` that does NOT propagate to downstream + // binaries, so a static consumer link would fail with undefined GOMP_*/omp_* + // symbols. A self-contained static build is the default; `--features openmp` + // opts in (and then owns providing the OpenMP runtime). cfg.define( "TRANSCRIBE_USE_OPENMP", if feature("OPENMP") { "ON" } else { "OFF" }, ); - // ggml's non-OpenMP CPU threadpool barrier deadlocks under MSVC on Windows; - // OpenMP's `#pragma omp barrier` is the only working multi-threaded CPU path - // there. The Rust binding is CPU-default and standalone (no numpy/torch - // OpenMP-coexistence concern), so it opts ggml into OpenMP on Windows. - // GGML-internal only: TRANSCRIBE_USE_OPENMP stays off, so the link manifest - // emits no (GNU-only) -fopenmp and the Parakeet host-decoder TU is unchanged; - // MSVC auto-links vcomp via the ggml objects' /openmp pragma. CMakeLists - // honors this explicit GGML_OPENMP and skips its force-off. See the full - // per-consumer policy (incl. why the Python wheels stay OpenMP-free) in the - // "OpenMP — CENTRAL POLICY" block of the root CMakeLists.txt. - if target_os == "windows" { - cfg.define("GGML_OPENMP", "ON"); - } + // Windows no longer needs ggml's OpenMP. ggml's native CPU threadpool barrier + // used to deadlock under MSVC (its spin `relax` compiled to a no-op, starving + // an un-arrived worker), so OpenMP was the only multi-threaded CPU path there. + // That barrier is fixed upstream (ggml-cpu.c: YieldProcessor relax + a + // bounded-spin ggml_thread_yield fallback), so the native pool is correct on + // MSVC — and avoids the vcomp pool that crashes at teardown when a host + // dlopen/dlcloses this lib. We therefore leave GGML_OPENMP at its + // CMakeLists-forced OFF on every platform; `--features openmp` is the only + // opt-in. See the "OpenMP — CENTRAL POLICY" block in the root CMakeLists.txt. // Escape hatch: forward arbitrary configure args so the curated features are // never a hard ceiling. Anything CMake accepts (-DGGML_*, a -DTRANSCRIBE_* diff --git a/docs/build-windows.md b/docs/build-windows.md index da68e5b7..e1f90d4b 100644 --- a/docs/build-windows.md +++ b/docs/build-windows.md @@ -93,10 +93,12 @@ cmake --build build --target transcribe-cli --config Release Notes: - The Visual Studio generator is **multi-config**, so pass `--config Release` at build time (omitting it yields a Debug build). -- A successful configure prints `Found OpenMP`. **OpenMP matters on - Windows**: ggml's non-OpenMP CPU threadpool barrier deadlocks under MSVC, - so OpenMP is the working multi-threaded CPU path. MSVC auto-links it - (`vcomp`) — no action needed, just don't disable it. +- OpenMP is **not** required on Windows. ggml's native CPU threadpool is the + default everywhere (see the OpenMP CENTRAL POLICY in the top-level + `CMakeLists.txt`); the MSVC barrier bug that once forced OpenMP here is fixed, + so the native pool is correct under MSVC and oversubscription. Build with the + defaults; opt into OpenMP only with `-DTRANSCRIBE_USE_OPENMP=ON` if you have a + specific reason. - `no BLAS found — decoder uses scalar fallback` is fine; it only affects host-side decoder speed, not correctness. - `uv not found on PATH` is harmless — it only skips an optional test. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f1edfff3..c9fd52b1 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -124,27 +124,6 @@ add_library(transcribe third_party/miniz/miniz.c ) -# The parakeet host decoder's matmul fallback (linear()) parallelizes its -# large output projection across rows. Compile ONLY this TU with OpenMP so -# the `#pragma omp parallel for` in linear() activates; libgomp is already -# linked transitively via ggml. Scoped to one TU and numerically identical -# (rows summed in the same order). The primary out_w path is the ggml graph -# in joint_step; this row-parallelism covers the host fallback and the -# remaining host-side matmuls. -# -# Gated on TRANSCRIBE_USE_OPENMP so official provider wheels (which pass it OFF, -# alongside GGML_OPENMP=OFF) do not link an OpenMP runtime into this TU. When -# off, the `#pragma omp` is an ignored no-op and linear() runs serially — the -# same path taken when OpenMP simply is not found. -if(TRANSCRIBE_USE_OPENMP) - find_package(OpenMP QUIET) - if(OpenMP_CXX_FOUND) - set_source_files_properties(arch/parakeet/decoder.cpp PROPERTIES - COMPILE_OPTIONS "${OpenMP_CXX_FLAGS}") - target_link_libraries(transcribe PRIVATE OpenMP::OpenMP_CXX) - endif() -endif() - target_include_directories(transcribe PUBLIC $ diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index 4f69dc75..5634e7b6 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -902,22 +902,7 @@ transcribe_status run( } // Thread count. - { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, eb.graph); @@ -1570,20 +1555,7 @@ transcribe_status encode_one_to_host( pos_buf.size() * sizeof(float)); } - { - int n_threads = cc->n_threads; - if (n_threads <= 0) n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t0 = ggml_time_us(); if (ggml_backend_sched_graph_compute(cc->sched, eb.graph) != GGML_STATUS_SUCCESS) @@ -1672,8 +1644,7 @@ transcribe_status run_batch( std::vector> mel_bufs(n); std::vector mel_nf(n, 0); int n_threads = cc->n_threads; - if (n_threads <= 0) n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); + if (n_threads <= 0) n_threads = transcribe::default_n_threads(); int64_t mel_us = 0, enc_us = 0; const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_threads, [&](int b) { diff --git a/src/arch/canary_qwen/model.cpp b/src/arch/canary_qwen/model.cpp index 4e81ffab..b6464d97 100644 --- a/src/arch/canary_qwen/model.cpp +++ b/src/arch/canary_qwen/model.cpp @@ -830,20 +830,7 @@ transcribe_status init_context( // --------------------------------------------------------------------------- void apply_thread_policy(CanaryQwenSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } void build_relpos_emb_host(std::vector & pos_buf, @@ -1559,8 +1546,7 @@ transcribe_status run_batch( std::vector mel_nf(n, 0); int n_mel_threads = cc->n_threads; if (n_mel_threads <= 0) - n_mel_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); + n_mel_threads = transcribe::default_n_threads(); const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_mel_threads, [&](int b) { if (pcm[b] == nullptr || n_samples[b] <= 0) return true; diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index 38cf26a2..40e54d1d 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -941,24 +941,7 @@ transcribe_status run( } // Set thread count. - { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) { - fn(be, n_threads); - } - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); // Compute encoder graph. const int64_t t_enc_start = ggml_time_us(); @@ -1686,20 +1669,7 @@ transcribe_status encode_one_to_host( pos_buf.size() * sizeof(float)); } - { - int n_threads = cc->n_threads; - if (n_threads <= 0) n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t0 = ggml_time_us(); if (ggml_backend_sched_graph_compute(cc->sched, eb.graph) != GGML_STATUS_SUCCESS) @@ -1786,8 +1756,7 @@ transcribe_status run_batch( std::vector> mel_bufs(n); std::vector mel_nf(n, 0); int n_threads = cc->n_threads; - if (n_threads <= 0) n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); + if (n_threads <= 0) n_threads = transcribe::default_n_threads(); int64_t mel_us = 0, enc_us = 0; const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_threads, [&](int b) { diff --git a/src/arch/funasr_nano/model.cpp b/src/arch/funasr_nano/model.cpp index fd18d5ab..a802f1b8 100644 --- a/src/arch/funasr_nano/model.cpp +++ b/src/arch/funasr_nano/model.cpp @@ -283,20 +283,7 @@ transcribe_status build_funasr_nano_prompt(const transcribe::Tokenizer & tok, } void apply_thread_policy(FunAsrNanoSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } // --------------------------------------------------------------------------- @@ -1212,8 +1199,7 @@ transcribe_status run_batch( std::vector T_lfr(n, 0); int n_mel_threads = cc->n_threads; if (n_mel_threads <= 0) - n_mel_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); + n_mel_threads = transcribe::default_n_threads(); const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_mel_threads, [&](int b) { if (pcm[b] == nullptr || n_samples[b] <= 0) return true; diff --git a/src/arch/gigaam/model.cpp b/src/arch/gigaam/model.cpp index 60853af3..8c2092f2 100644 --- a/src/arch/gigaam/model.cpp +++ b/src/arch/gigaam/model.cpp @@ -701,22 +701,7 @@ transcribe_status run_batch_encode(GigaamSession * gc, // Set the per-backend thread count (mirrors single-shot through the // scheduler; the encoder is the dominant cost on CPU). - { - int n_threads = gc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(gc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(gc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(gc->sched, gc->n_threads); const int64_t t_enc_start = ggml_time_us(); if (ggml_backend_sched_graph_compute(gc->sched, eb.graph) != GGML_STATUS_SUCCESS) { diff --git a/src/arch/granite/model.cpp b/src/arch/granite/model.cpp index 2662c65b..4194ded0 100644 --- a/src/arch/granite/model.cpp +++ b/src/arch/granite/model.cpp @@ -785,22 +785,7 @@ transcribe_status run( } // Thread count. - { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); // ----- Compute ----- const int64_t t_enc_start = ggml_time_us(); @@ -1290,19 +1275,7 @@ transcribe_status reset_ctx_g(GraniteSession * cc, int mb) { } void apply_threads_g(GraniteSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } // encoder + projector for one utterance from a PRECOMPUTED mel buffer (the @@ -1497,8 +1470,7 @@ transcribe_status run_batch( std::vector mel_t_enc(n, 0); int n_mel_threads = cc->n_threads; if (n_mel_threads <= 0) - n_mel_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); + n_mel_threads = transcribe::default_n_threads(); const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_mel_threads, [&](int b) { if (pcm[b] == nullptr || n_samples[b] <= 0) return true; diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index bd180f7c..71527183 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -25,6 +25,7 @@ #include "weights.h" #include "transcribe-arch.h" +#include "transcribe-batch-util.h" #include "transcribe-debug.h" #include "transcribe-flash-policy.h" #include "transcribe-load-common.h" @@ -428,19 +429,7 @@ transcribe_status init_context( namespace { void apply_thread_count(ggml_backend_sched_t sched, int n_threads) { - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(sched, n_threads); } } // namespace diff --git a/src/arch/medasr/model.cpp b/src/arch/medasr/model.cpp index 33dc350d..ebd1af85 100644 --- a/src/arch/medasr/model.cpp +++ b/src/arch/medasr/model.cpp @@ -938,22 +938,7 @@ transcribe_status run_batch_encode(MedAsrSession * gc, } // Per-backend thread count (matching single-shot through the scheduler). - { - int n_threads = gc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(gc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(gc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(gc->sched, gc->n_threads); const int64_t t_enc_start = ggml_time_us(); if (ggml_backend_sched_graph_compute(gc->sched, eb.graph) diff --git a/src/arch/moonshine/model.cpp b/src/arch/moonshine/model.cpp index 3578884a..34fe9f57 100644 --- a/src/arch/moonshine/model.cpp +++ b/src/arch/moonshine/model.cpp @@ -332,20 +332,7 @@ transcribe_status init_context( // Set per-backend thread count from cc->n_threads (with a sensible default). void apply_thread_policy(MoonshineSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } bool ensure_compute_ctx(MoonshineSession * cc, size_t mem_size) { diff --git a/src/arch/moonshine_streaming/model.cpp b/src/arch/moonshine_streaming/model.cpp index f1853739..b1698243 100644 --- a/src/arch/moonshine_streaming/model.cpp +++ b/src/arch/moonshine_streaming/model.cpp @@ -39,6 +39,7 @@ #include "transcribe/moonshine_streaming.h" #include "transcribe-arch.h" +#include "transcribe-batch-util.h" #include "transcribe-debug.h" #include "transcribe-flash-policy.h" #include "transcribe-load-common.h" @@ -347,20 +348,7 @@ transcribe_status init_context( } void apply_thread_policy(MoonshineStreamingSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } bool ensure_compute_ctx(MoonshineStreamingSession * cc, size_t mem_size) { diff --git a/src/arch/parakeet/decoder.cpp b/src/arch/parakeet/decoder.cpp index a7904ec8..63aa5342 100644 --- a/src/arch/parakeet/decoder.cpp +++ b/src/arch/parakeet/decoder.cpp @@ -1,22 +1,24 @@ // arch/parakeet/decoder.cpp - Parakeet TDT decoder implementation. // -// See decoder.h for the API contract and the rationale for running -// the decoder on host instead of via a backend graph. The -// implementation is intentionally dependency-light: , , -// for the math; the only ggml include is for tensor reads -// at load time. +// See decoder.h for the API contract. The predictor LSTM, the encoder +// projection, and the joint network all run as ggml backend graphs on a +// single shared CPU backend + persistent threadpool (owned by PredGraph; +// see build_pred_graph). Decode weights are uploaded to resident ggml +// tensors at load time and the host-side fp32 mirrors are then freed. // -// Numerical-accuracy strategy: every step is a small enough operation -// (a few hundred dot products of length 640) that compounding error -// stays within ~1e-4 of the reference on CPU. The -// per-step dump points wired below feed scripts/compare_tensors.py -// for the bring-up loop and are gated on TRANSCRIBE_DUMP_DIR. +// Numerical-accuracy strategy: ggml's reduction order differs slightly +// from a naive host loop, so per-step values drift ~1e-7 vs the reference; +// over a full utterance this stays within ~1e-4 and the greedy transcript +// is unchanged. The per-step dump points wired below feed +// scripts/compare_tensors.py for the bring-up loop and are gated on +// TRANSCRIBE_DUMP_DIR. #include "decoder.h" #include "parakeet.h" #include "weights.h" +#include "transcribe-batch-util.h" #include "transcribe-debug.h" #include "transcribe-log.h" @@ -28,11 +30,10 @@ #include "ggml.h" #include "ggml-backend.h" #include "ggml-alloc.h" +#include "ggml-cpu.h" // ggml_threadpool_params_default (the rest of the CPU + // backend API is reached via the registry, see above) #include -#ifdef _OPENMP -#include -#endif #if TRANSCRIBE_HAS_BLAS # ifdef __APPLE__ @@ -59,11 +60,11 @@ namespace transcribe::parakeet { namespace { // Read all elements of a ggml_tensor into a host fp32 vector, -// dequantizing as needed via ggml_get_type_traits()->to_float. The -// host decoder runs entirely on fp32 (linear / lstm_step / joint_step -// in this file are all hand-rolled fp32 matmuls), so quantized weight -// tensors are dequantized exactly once at load time and the per-decode -// hot path pays no readback or dequant cost. +// dequantizing as needed via ggml_get_type_traits()->to_float. Decode +// weights are uploaded to resident fp32 ggml tensors once at load time +// (see build_pred_weights / build_joint_weight), so quantized weight +// tensors are dequantized exactly once here and the per-decode hot path +// pays no readback or dequant cost. // // Source dtype support is whatever ggml's type traits register a // to_float implementation for: F32, F16, BF16, all the Q* and IQ* @@ -136,38 +137,29 @@ bool read_tensor_to_f32(const ggml_tensor * t, return true; } -// Make the joint output-projection weight + bias resident as ggml tensors on -// the model (the immutable half). Built once at load; only ever read after -// that, so it is safe to share across every context. The per-step compute -// graph that consumes this is built per decode call (build_joint_graph below). -// On any failure it frees its partial state and returns false, so the caller -// can retry (native→fp32) or fall back to the host matmul (w_ready stays false). -// -// `use_native` selects the weight dtype: -// - true (default): keep the source quant dtype. ggml streams the quantized -// bytes (~4× less weight bandwidth) and quantizes the activation on the fly -// — the same regime the encoder runs in, WER-validated equal to fp32 on -// English. This is the faster path on bandwidth-bound CPUs/iGPUs. -// - false (fallback): dequantize the weight to fp32 once, so ggml accumulates -// fp32×fp32 with no activation down-conversion. Numerically faithful to the -// host reference (max rel logit diff ~3e-7 vs ~3e-3 for a Q4 native weight). -// The caller tries native first; TRANSCRIBE_JOINT_FP32=1 forces fp32. -bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w, - bool use_native) { +// Make the joint network's weights resident as fp32 ggml tensors on the model +// (the immutable half): the encoder projection (g_enc_w/g_enc_b), the predictor +// projection (g_pred_w/g_pred_b), and the output projection (gw_w/gw_b). Built +// once at load; only ever read after, so it is safe to share across every +// context. The per-decode compute graph that consumes these is built per call +// (build_joint_graph below). out_w is dequantized to fp32 from the model tensor +// `src_out_w`; the other weights come from the host mirrors, which are freed +// here once uploaded. On any failure it frees its partial state and returns +// false (w_ready stays false), so the decode reports a hard error. +bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w) { const int joint_h = j.joint_h; const int joint_n = j.joint_n; - // Free any partial state so the caller can cleanly retry (e.g. native→fp32). auto fail = [&]() -> bool { if (j.w_buf != nullptr) { ggml_backend_buffer_free(j.w_buf); j.w_buf = nullptr; } if (j.w_ctx != nullptr) { ggml_free(j.w_ctx); j.w_ctx = nullptr; } if (j.w_backend != nullptr) { ggml_backend_free(j.w_backend); j.w_backend = nullptr; } + j.g_enc_w = nullptr; j.g_enc_b = nullptr; + j.g_pred_w = nullptr; j.g_pred_b = nullptr; j.gw_w = nullptr; j.gw_b = nullptr; j.w_ready = false; return false; }; - const ggml_type w_type = use_native ? src_out_w->type : GGML_TYPE_F32; - // A CPU backend used ONLY to allocate the resident weight buffer — never // graph_compute'd (each decode call owns its own compute backend), so it // carries no per-step state and is safe to keep on the shared model. @@ -178,125 +170,371 @@ bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w, } ggml_init_params ip {}; - ip.mem_size = ggml_tensor_overhead() * 4; + ip.mem_size = ggml_tensor_overhead() * 8; ip.mem_buffer = nullptr; ip.no_alloc = true; j.w_ctx = ggml_init(ip); if (j.w_ctx == nullptr) return fail(); - j.gw_w = ggml_new_tensor_2d(j.w_ctx, w_type, joint_h, joint_n); - j.gw_b = ggml_new_tensor_1d(j.w_ctx, GGML_TYPE_F32, joint_n); + j.g_enc_w = ggml_new_tensor_2d(j.w_ctx, GGML_TYPE_F32, j.d_enc, joint_h); + j.g_enc_b = ggml_new_tensor_1d(j.w_ctx, GGML_TYPE_F32, joint_h); + j.g_pred_w = ggml_new_tensor_2d(j.w_ctx, GGML_TYPE_F32, j.pred_hidden, joint_h); + j.g_pred_b = ggml_new_tensor_1d(j.w_ctx, GGML_TYPE_F32, joint_h); + j.gw_w = ggml_new_tensor_2d(j.w_ctx, GGML_TYPE_F32, joint_h, joint_n); + j.gw_b = ggml_new_tensor_1d(j.w_ctx, GGML_TYPE_F32, joint_n); j.w_buf = ggml_backend_alloc_ctx_tensors(j.w_ctx, j.w_backend); if (j.w_buf == nullptr) return fail(); - // Upload the weight (verbatim native dtype, or dequantized to fp32 under - // TRANSCRIBE_JOINT_FP32) and the fp32 bias. - if (w_type == src_out_w->type) { - const size_t nb = ggml_nbytes(src_out_w); - std::vector tmp(nb); - ggml_backend_tensor_get(src_out_w, tmp.data(), 0, nb); - ggml_backend_tensor_set(j.gw_w, tmp.data(), 0, nb); - } else { - // fp32 graph weight from a non-fp32 source: dequantize once into fp32. + // out_w: dequantize the model tensor to fp32 once (faithful to the host + // reference: fp32×fp32 accumulation, no activation down-conversion). + { std::vector tmp; if (!read_tensor_to_f32(src_out_w, tmp)) return fail(); - ggml_backend_tensor_set(j.gw_w, tmp.data(), 0, - tmp.size() * sizeof(float)); + ggml_backend_tensor_set(j.gw_w, tmp.data(), 0, tmp.size() * sizeof(float)); } - ggml_backend_tensor_set(j.gw_b, j.out_b.data(), 0, - static_cast(joint_n) * sizeof(float)); + // The rest come from the host mirrors. + ggml_backend_tensor_set(j.g_enc_w, j.enc_w.data(), 0, j.enc_w.size() * sizeof(float)); + ggml_backend_tensor_set(j.g_enc_b, j.enc_b.data(), 0, j.enc_b.size() * sizeof(float)); + ggml_backend_tensor_set(j.g_pred_w, j.pred_w.data(), 0, j.pred_w.size() * sizeof(float)); + ggml_backend_tensor_set(j.g_pred_b, j.pred_b.data(), 0, j.pred_b.size() * sizeof(float)); + ggml_backend_tensor_set(j.gw_b, j.out_b.data(), 0, j.out_b.size() * sizeof(float)); + + // Host mirrors are now resident in ggml — release them (load-time scratch). + std::vector().swap(j.enc_w); + std::vector().swap(j.enc_b); + std::vector().swap(j.pred_w); + std::vector().swap(j.pred_b); + std::vector().swap(j.out_b); j.w_ready = true; return true; } +// Make the predictor LSTM weights resident as fp32 ggml tensors — the immutable +// half consumed by the per-call PredGraph. Built once at load. ne fast-to-slow +// is [pred_hidden, 4*pred_hidden] for Wx/Wh (the row-major [4*H, H] host bytes +// read as a ggml mul_mat operand) and [4*pred_hidden] for the bias. On any +// failure frees its partial state and returns false; the caller leaves +// lstm_ready false; building the pred graph then fails and the decode returns a hard error (no host fallback remains). +bool build_pred_weights(HostPredictor & p) { + const int H = p.pred_hidden; + const int four_H = 4 * H; + const int L = static_cast(p.lstm.size()); + if (L == 0 || H == 0) return false; + + auto fail = [&]() -> bool { + if (p.lstm_w_buf != nullptr) { ggml_backend_buffer_free(p.lstm_w_buf); p.lstm_w_buf = nullptr; } + if (p.lstm_w_ctx != nullptr) { ggml_free(p.lstm_w_ctx); p.lstm_w_ctx = nullptr; } + if (p.lstm_w_backend != nullptr) { ggml_backend_free(p.lstm_w_backend); p.lstm_w_backend = nullptr; } + for (auto & lh : p.lstm) { lh.g_Wx = nullptr; lh.g_Wh = nullptr; lh.g_b = nullptr; } + p.lstm_ready = false; + return false; + }; + + // Alloc-only backend (never graph_compute'd), mirroring HostJoint::w_backend. + p.lstm_w_backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + if (p.lstm_w_backend == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: ggml CPU backend init failed (pred)"); + return fail(); + } + + ggml_init_params ip {}; + ip.mem_size = ggml_tensor_overhead() * static_cast(L * 3 + 1); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + p.lstm_w_ctx = ggml_init(ip); + if (p.lstm_w_ctx == nullptr) return fail(); + + for (int l = 0; l < L; ++l) { + auto & lh = p.lstm[l]; + lh.g_Wx = ggml_new_tensor_2d(p.lstm_w_ctx, GGML_TYPE_F32, H, four_H); + lh.g_Wh = ggml_new_tensor_2d(p.lstm_w_ctx, GGML_TYPE_F32, H, four_H); + lh.g_b = ggml_new_tensor_1d(p.lstm_w_ctx, GGML_TYPE_F32, four_H); + } + + p.lstm_w_buf = ggml_backend_alloc_ctx_tensors(p.lstm_w_ctx, p.lstm_w_backend); + if (p.lstm_w_buf == nullptr) return fail(); + + for (int l = 0; l < L; ++l) { + auto & lh = p.lstm[l]; + ggml_backend_tensor_set(lh.g_Wx, lh.Wx.data(), 0, lh.Wx.size() * sizeof(float)); + ggml_backend_tensor_set(lh.g_Wh, lh.Wh.data(), 0, lh.Wh.size() * sizeof(float)); + ggml_backend_tensor_set(lh.g_b, lh.b.data(), 0, lh.b.size() * sizeof(float)); + // Host mirrors are now resident in ggml — release them (load-time scratch). + std::vector().swap(lh.Wx); + std::vector().swap(lh.Wh); + std::vector().swap(lh.b); + } + + p.lstm_ready = true; + return true; +} + } // namespace // --------------------------------------------------------------------------- // Per-call joint compute graph // --------------------------------------------------------------------------- // -// The MUTABLE half of the joint output projection. Built fresh on the stack in -// each decode call (cheap: it allocates only the [joint_h] activation input, -// the matmul node, and the [joint_n] logits output — the weight is the shared -// resident HostJoint::gw_w/gw_b, not re-uploaded), and torn down at function -// exit. Because every concurrent decode owns its own JointGraph (its own -// compute backend + I/O tensors), decode on contexts sharing one model is -// reentrant — nothing per-step is shared. +// The MUTABLE half of the joint network: a full per-call graph +// pred_proj = pred_w @ pred_in + pred_b [joint_h] +// summed = enc_proj + pred_proj [joint_h] +// activated = activation(summed) [joint_h] +// logits = out_w @ activated + out_b [joint_n] +// built on a BORROWED backend — the shared decoder pool owned by PredGraph — so +// the joint runs on the SAME single threadpool as the pred LSTM and enc_proj +// (no separate pool, no oversubscription). Inputs are the predictor output +// [pred_hidden] and this frame's precomputed encoder projection [joint_h]; +// output is the logits [joint_n]. Built fresh per decode call (reentrant: each +// decode owns its own I/O tensors); weights are the shared resident +// HostJoint::g_pred_w/g_pred_b/gw_w/gw_b. namespace { struct JointGraph { ggml_context * ctx = nullptr; - ggml_backend_t backend = nullptr; + ggml_backend_t backend = nullptr; // BORROWED (PredGraph's); NOT freed here ggml_backend_buffer_t buf = nullptr; ggml_cgraph * graph = nullptr; - ggml_tensor * act = nullptr; // [joint_h] fp32 input + ggml_tensor * pred_in = nullptr; // [pred_hidden] fp32 input (decoder out) + ggml_tensor * enc_in = nullptr; // [joint_h] fp32 input (enc_proj frame) ggml_tensor * logits = nullptr; // [joint_n] fp32 output bool ready = false; JointGraph() = default; ~JointGraph() { - if (buf != nullptr) ggml_backend_buffer_free(buf); - if (ctx != nullptr) ggml_free(ctx); - if (backend != nullptr) ggml_backend_free(backend); + if (buf != nullptr) ggml_backend_buffer_free(buf); + if (ctx != nullptr) ggml_free(ctx); + // backend is borrowed from PredGraph — do NOT free here (and PredGraph, + // declared first in the driver, is destroyed after this). } JointGraph(const JointGraph &) = delete; JointGraph & operator=(const JointGraph &) = delete; }; -// Build a per-call compute graph around the model-resident weight j.gw_w/gw_b. -// Returns false (and leaves g.ready == false) if the weight is absent or any -// ggml step fails; the caller then uses the host matmul fallback. n_threads is +// Build the full joint graph on the shared `backend` (PredGraph's pool). Returns +// false (g.ready == false) if the resident weights are absent or any ggml step +// fails; the driver treats a non-ready joint graph as a hard decode error (there +// is no host fallback — the resident weights are the only copy). +bool build_joint_graph(JointGraph & g, const HostJoint & j, ggml_backend_t backend) { + if (backend == nullptr) return false; + if (!j.w_ready || j.gw_w == nullptr || j.gw_b == nullptr || + j.g_pred_w == nullptr || j.g_pred_b == nullptr) return false; + + g.backend = backend; // borrowed + + auto fail = [&]() -> bool { + if (g.buf != nullptr) { ggml_backend_buffer_free(g.buf); g.buf = nullptr; } + if (g.ctx != nullptr) { ggml_free(g.ctx); g.ctx = nullptr; } + g.pred_in = nullptr; g.enc_in = nullptr; g.logits = nullptr; + g.graph = nullptr; g.ready = false; + return false; + }; + + ggml_init_params ip {}; + ip.mem_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + g.ctx = ggml_init(ip); + if (g.ctx == nullptr) return fail(); + + g.pred_in = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, j.pred_hidden); + ggml_set_input(g.pred_in); + g.enc_in = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, j.joint_h); + ggml_set_input(g.enc_in); + + // pred_proj = pred_w @ pred_in + pred_b [joint_h] + ggml_tensor * pred_proj = ggml_add(g.ctx, + ggml_mul_mat(g.ctx, j.g_pred_w, g.pred_in), j.g_pred_b); + // summed = enc_proj + pred_proj [joint_h] + ggml_tensor * summed = ggml_add(g.ctx, g.enc_in, pred_proj); + // activation — loader allow-list guarantees exactly one of relu/sigmoid/tanh. + ggml_tensor * activated; + if (j.activation == "relu") { + activated = ggml_relu(g.ctx, summed); + } else if (j.activation == "sigmoid") { + activated = ggml_sigmoid(g.ctx, summed); + } else { // "tanh" + activated = ggml_tanh(g.ctx, summed); + } + // logits = out_w @ activated + out_b [joint_n] + ggml_tensor * mm = ggml_mul_mat(g.ctx, j.gw_w, activated); + g.logits = ggml_add(g.ctx, mm, j.gw_b); + ggml_set_output(g.logits); + + g.buf = ggml_backend_alloc_ctx_tensors(g.ctx, g.backend); + if (g.buf == nullptr) return fail(); + + g.graph = ggml_new_graph(g.ctx); + ggml_build_forward_expand(g.graph, g.logits); + g.ready = true; + return true; +} + +// CPU-backend threadpool entry points, reached through the registry so the +// library stays DL-safe: under GGML_BACKEND_DL the CPU backend is a loadable +// module and these symbols are NOT linkable directly into libtranscribe (only +// ggml-base symbols like ggml_threadpool_params_default are). Resolve them via +// ggml_backend_reg_get_proc_address, exactly as for ggml_backend_set_n_threads. +typedef ggml_threadpool_t (*pfn_threadpool_new)(ggml_threadpool_params *); +typedef void (*pfn_threadpool_free)(ggml_threadpool_t); +typedef void (*pfn_set_threadpool)(ggml_backend_t, ggml_threadpool_t); + +static void * cpu_backend_proc(ggml_backend_t backend, const char * name) { + ggml_backend_dev_t dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr; + ggml_backend_reg_t reg = dev != nullptr ? ggml_backend_dev_backend_reg(dev) : nullptr; + return reg != nullptr ? ggml_backend_reg_get_proc_address(reg, name) : nullptr; +} + +// Free a CPU-backend threadpool via the registry. Resolving the free fn needs +// the backend alive, so callers must invoke this BEFORE freeing the backend. +static void free_cpu_threadpool(ggml_backend_t backend, ggml_threadpool_t tp) { + if (tp == nullptr || backend == nullptr) return; + if (auto fn = (pfn_threadpool_free) cpu_backend_proc(backend, "ggml_threadpool_free")) { + fn(tp); + } +} + +// --------------------------------------------------------------------------- +// Per-call predictor (LSTM) compute graph +// --------------------------------------------------------------------------- +// +// The MUTABLE half of the predictor: a single per-step LSTM graph built fresh +// per decode call around the model-resident HostPredictor::lstm weights, and +// recomputed in place each step. Inputs are the embedding x and, per layer, the +// previous (h, c); outputs are the new (h, c) per layer. The decode loop feeds +// prev state in and reads new state back into its host LstmState each step, so +// the loop's state machine (swap / predictor_dirty cache / dumps) is unchanged. +// Mirrors JointGraph's reentrancy: every concurrent decode owns its own +// PredGraph (its own backend + threadpool + I/O tensors). +struct PredGraph { + ggml_context * ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_buffer_t buf = nullptr; + ggml_cgraph * graph = nullptr; + ggml_threadpool_t tp = nullptr; + ggml_tensor * x = nullptr; // [H] input embedding + std::vector ph; // [H] prev hidden per layer (input) + std::vector pc; // [H] prev cell per layer (input) + std::vector nh; // [H] new hidden per layer (output) + std::vector nc; // [H] new cell per layer (output) + int H = 0; + int L = 0; + bool ready = false; + + PredGraph() = default; + ~PredGraph() { + if (buf != nullptr) ggml_backend_buffer_free(buf); + if (ctx != nullptr) ggml_free(ctx); + free_cpu_threadpool(backend, tp); // before ggml_backend_free(backend) + if (backend != nullptr) ggml_backend_free(backend); + } + PredGraph(const PredGraph &) = delete; + PredGraph & operator=(const PredGraph &) = delete; +}; + +// Build the per-call LSTM graph around the resident p.lstm[*].g_Wx/g_Wh/g_b. +// Returns false (g.ready == false) if the resident weights are absent or any +// ggml step fails; the driver then returns a hard decode error (no host fallback). n_threads is // the resolved (>0) decode thread count. -bool build_joint_graph(JointGraph & g, const HostJoint & j, int n_threads) { - if (!j.w_ready || j.gw_w == nullptr || j.gw_b == nullptr) return false; +bool build_pred_graph(PredGraph & g, const HostPredictor & p, int n_threads) { + if (!p.lstm_ready) return false; + const int H = p.pred_hidden; + const int L = static_cast(p.lstm.size()); + if (H == 0 || L == 0) return false; auto fail = [&]() -> bool { - if (g.buf != nullptr) { ggml_backend_buffer_free(g.buf); g.buf = nullptr; } - if (g.ctx != nullptr) { ggml_free(g.ctx); g.ctx = nullptr; } - if (g.backend != nullptr) { ggml_backend_free(g.backend); g.backend = nullptr; } - g.act = nullptr; g.logits = nullptr; g.graph = nullptr; g.ready = false; + if (g.buf != nullptr) { ggml_backend_buffer_free(g.buf); g.buf = nullptr; } + if (g.ctx != nullptr) { ggml_free(g.ctx); g.ctx = nullptr; } + free_cpu_threadpool(g.backend, g.tp); g.tp = nullptr; // before freeing backend + if (g.backend != nullptr) { ggml_backend_free(g.backend); g.backend = nullptr; } + g.x = nullptr; g.graph = nullptr; g.ready = false; + g.ph.clear(); g.pc.clear(); g.nh.clear(); g.nc.clear(); return false; }; + g.H = H; g.L = L; + g.backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); if (g.backend == nullptr) return fail(); - // n_threads via proc address: the typed setter is a CPU-module symbol - // under GGML_BACKEND_DL. Absent setter (exotic CPU device) keeps the - // backend's default thread count — correct, just less tuned. if (ggml_backend_dev_t dev = ggml_backend_get_device(g.backend)) { auto set_n_threads = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address( ggml_backend_dev_backend_reg(dev), "ggml_backend_set_n_threads"); - if (set_n_threads != nullptr) { - set_n_threads(g.backend, std::max(1, n_threads)); + if (set_n_threads != nullptr) set_n_threads(g.backend, std::max(1, n_threads)); + } + // Persistent threadpool so the per-step graph_compute reuses workers instead + // of spawning a transient pool each call. ggml's default hybrid-polling + // (poll=50) keeps the workers hot between the sub-millisecond per-step + // dispatches; poll=0 (park-immediately) was measured to regress pred to + // 70-110 ms because parked workers are slow to reschedule. + { + auto tp_new = (pfn_threadpool_new) cpu_backend_proc(g.backend, "ggml_threadpool_new"); + auto tp_set = (pfn_set_threadpool) cpu_backend_proc(g.backend, "ggml_backend_cpu_set_threadpool"); + if (tp_new != nullptr && tp_set != nullptr) { + ggml_threadpool_params tpp = ggml_threadpool_params_default(std::max(1, n_threads)); + g.tp = tp_new(&tpp); + if (g.tp != nullptr) tp_set(g.backend, g.tp); } + // If unresolved (non-CPU exotic backend), g.tp stays null and each + // graph_compute spawns a transient pool — correct, just less tuned. } ggml_init_params ip {}; - ip.mem_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead(); + // ~17 op nodes/layer + (1 + 2L) input tensors; generous headroom. + ip.mem_size = ggml_tensor_overhead() * static_cast(32 * L + 16) + ggml_graph_overhead(); ip.mem_buffer = nullptr; ip.no_alloc = true; g.ctx = ggml_init(ip); if (g.ctx == nullptr) return fail(); - g.act = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, j.joint_h); - ggml_set_input(g.act); - // References the shared resident weight (in j.w_ctx) — ggml stores the - // pointer; the CPU backend reads its bytes directly at compute time. - ggml_tensor * mm = ggml_mul_mat(g.ctx, j.gw_w, g.act); // [joint_n, 1] - g.logits = ggml_add(g.ctx, mm, j.gw_b); // [joint_n, 1] - ggml_set_output(g.logits); + g.x = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, H); + ggml_set_input(g.x); + g.ph.assign(static_cast(L), nullptr); + g.pc.assign(static_cast(L), nullptr); + g.nh.assign(static_cast(L), nullptr); + g.nc.assign(static_cast(L), nullptr); + for (int l = 0; l < L; ++l) { + g.ph[l] = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, H); ggml_set_input(g.ph[l]); + g.pc[l] = ggml_new_tensor_1d(g.ctx, GGML_TYPE_F32, H); ggml_set_input(g.pc[l]); + } + + // gates = Wx@x + Wh@h_prev + b; split [i,f,g,o]; c' = f*c + i*g; h' = o*tanh(c'). + // Gate order [i, f, g, o] (PyTorch standard). + ggml_tensor * in = g.x; + for (int l = 0; l < L; ++l) { + const auto & lh = p.lstm[static_cast(l)]; + ggml_tensor * gates = ggml_add(g.ctx, + ggml_add(g.ctx, + ggml_mul_mat(g.ctx, lh.g_Wx, in), + ggml_mul_mat(g.ctx, lh.g_Wh, g.ph[l])), + lh.g_b); // [4H] + auto part = [&](int k) { + return ggml_view_1d(g.ctx, gates, H, + static_cast(k) * static_cast(H) * ggml_element_size(gates)); + }; + ggml_tensor * i_ = ggml_sigmoid(g.ctx, part(0)); + ggml_tensor * f_ = ggml_sigmoid(g.ctx, part(1)); + ggml_tensor * gg = ggml_tanh (g.ctx, part(2)); + ggml_tensor * o_ = ggml_sigmoid(g.ctx, part(3)); + g.nc[l] = ggml_add(g.ctx, + ggml_mul(g.ctx, f_, g.pc[l]), + ggml_mul(g.ctx, i_, gg)); + ggml_set_output(g.nc[l]); + g.nh[l] = ggml_mul(g.ctx, o_, ggml_tanh(g.ctx, g.nc[l])); + ggml_set_output(g.nh[l]); + in = g.nh[l]; + } - // Allocates only the tensors created in g.ctx (act + op results); the - // shared weight/bias live in another ctx and are already resident. g.buf = ggml_backend_alloc_ctx_tensors(g.ctx, g.backend); if (g.buf == nullptr) return fail(); g.graph = ggml_new_graph(g.ctx); - ggml_build_forward_expand(g.graph, g.logits); + for (int l = 0; l < L; ++l) { + ggml_build_forward_expand(g.graph, g.nh[l]); + ggml_build_forward_expand(g.graph, g.nc[l]); + } g.ready = true; return true; } @@ -309,14 +547,21 @@ HostJoint::~HostJoint() { if (w_backend != nullptr) ggml_backend_free(w_backend); } -// Resolve a decode thread count: n_threads <= 0 means "auto" → min(8, cores), -// matching the encoder default. Threaded through to the joint compute backend -// and linear()'s OpenMP loop per decode call; no process-global state is set, +HostPredictor::~HostPredictor() { + if (lstm_w_buf != nullptr) ggml_backend_buffer_free(lstm_w_buf); + if (lstm_w_ctx != nullptr) ggml_free(lstm_w_ctx); + if (lstm_w_backend != nullptr) ggml_backend_free(lstm_w_backend); +} + +// Resolve a decode thread count: n_threads <= 0 means "auto" → min(8, usable +// cpus) via default_n_threads(), matching the encoder default. Threaded through +// to the joint compute backend +// and the per-call ggml decode graphs; no process-global state is set, // so concurrent decodes on one model do not stomp each other. static int resolve_decode_threads(int n_threads) { return n_threads > 0 ? n_threads - : std::min(8, std::max(1, static_cast(std::thread::hardware_concurrency()))); + : transcribe::default_n_threads(); } transcribe_status build_host_decoder_weights(const ParakeetModel & model, @@ -422,6 +667,15 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, } } + // Make the predictor LSTM weights resident as fp32 ggml tensors for the ggml + // predictor graph. Fatal: there is no host fallback, so a failure here means + // the model cannot decode — fail fast at load rather than at first decode. + if (!build_pred_weights(out.predictor)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet decoder: predictor ggml weight build failed"); + return TRANSCRIBE_ERR_BACKEND; + } + // ----- Joint mirror ----- out.joint.d_enc = hp.enc_d_model; out.joint.pred_hidden = hp.pred_hidden; @@ -429,43 +683,35 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, out.joint.joint_n = hp.joint_n_classes(); out.joint.activation = hp.joint_activation; + // enc/pred/out_b are mirrored to host fp32 (build_joint_weight uploads them + // into resident ggml tensors and frees them). out_w is NOT mirrored here: + // build_joint_weight dequantizes it straight from the model tensor. if (!read_tensor_to_f32(w.joint.enc_w, out.joint.enc_w)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.enc_b, out.joint.enc_b)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.pred_w, out.joint.pred_w)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.pred_b, out.joint.pred_b)) return TRANSCRIBE_ERR_GGUF; - if (!read_tensor_to_f32(w.joint.out_w, out.joint.out_w)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.out_b, out.joint.out_b)) return TRANSCRIBE_ERR_GGUF; - // Make the dominant out_w projection weight resident as a ggml tensor. - // Default to the native-quant weight (faster, WER-validated); fall back to - // the fp32-dequant weight on failure, then to the host matmul (w_ready stays - // false) if even that fails. TRANSCRIBE_JOINT_FP32=1 forces the fp32 path. - if (w.joint.out_w != nullptr) { - const int ne0 = static_cast(w.joint.out_w->ne[0]); - const int ne1 = static_cast(w.joint.out_w->ne[1]); - if (ne0 != out.joint.joint_h || ne1 != out.joint.joint_n) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "parakeet decoder: out_w ne [%d,%d] != [joint_h=%d, joint_n=%d]; " - "using host matmul", - ne0, ne1, out.joint.joint_h, out.joint.joint_n); - } else { - const bool force_fp32 = std::getenv("TRANSCRIBE_JOINT_FP32") != nullptr; - bool ok = false; - if (!force_fp32) { - ok = build_joint_weight(out.joint, w.joint.out_w, /*use_native=*/true); - if (!ok) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "parakeet decoder: native joint weight failed; retrying fp32"); - } - } - if (!ok) { - ok = build_joint_weight(out.joint, w.joint.out_w, /*use_native=*/false); - } - if (!ok) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "parakeet decoder: joint ggml weight build failed; using host matmul"); - } - } + // Make the joint weights resident as fp32 ggml tensors for the joint graph. + // fp32 keeps decode numerically faithful to the historical host path (the + // out_w matmul was always fp32 there); any joint quantization is handled at + // the model/quant level. Fatal: with no host fallback, a failure here means + // the model cannot decode — fail fast at load. + if (w.joint.out_w == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: joint out_w missing"); + return TRANSCRIBE_ERR_GGUF; + } + const int ne0 = static_cast(w.joint.out_w->ne[0]); + const int ne1 = static_cast(w.joint.out_w->ne[1]); + if (ne0 != out.joint.joint_h || ne1 != out.joint.joint_n) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet decoder: out_w ne [%d,%d] != [joint_h=%d, joint_n=%d]", + ne0, ne1, out.joint.joint_h, out.joint.joint_n); + return TRANSCRIBE_ERR_GGUF; + } + if (!build_joint_weight(out.joint, w.joint.out_w)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: joint ggml weight build failed"); + return TRANSCRIBE_ERR_BACKEND; } // ----- TDT params ----- @@ -498,184 +744,26 @@ void LstmState::reset(int n_layers, int pred_hidden) { namespace { -// y = W @ x + b where W is row-major [out, in], x is [in], y is [out]. -// -// "Row-major [out, in]" means W[r, c] = W_flat[r * in + c]. The result -// at row r is the dot product of row r with x. b is added in place; -// pass nullptr for no bias. y is overwritten (not accumulated). -// -// When BLAS is available (Accelerate on Apple, OpenBLAS/MKL/BLIS -// elsewhere), routes through cblas_sgemv for ~10-15x speedup over the -// naive loop. The BLAS reduction order differs from our naive row-wise -// accumulation, so joint logits may drift by ~1e-4 — argmax is robust -// to that. Fallback: scalar loop when no BLAS is linked. -inline void linear(const float * W, - const float * x, - const float * b, - int out_dim, - int in_dim, - float * y, - int n_threads) -{ -#if TRANSCRIBE_HAS_BLAS - (void) n_threads; // BLAS manages its own threads - // y = 1.0 * W @ x + 0.0 * y - cblas_sgemv(CblasRowMajor, CblasNoTrans, - out_dim, in_dim, - 1.0f, W, in_dim, x, 1, - 0.0f, y, 1); - if (b != nullptr) { -#ifdef __APPLE__ - catlas_saxpby(out_dim, 1.0f, b, 1, 1.0f, y, 1); -#else - cblas_saxpy(out_dim, 1.0f, b, 1, y, 1); -#endif - } -#else - // Rows are independent, so larger projections parallelize cleanly across - // cores. Gate on out_dim (>= 2048) to thread the LSTM gate matmuls - // (4*640 = 2560) and the CTC head, while the tiny 640-wide enc/pred - // projections stay serial and don't pay OpenMP fork/join overhead. The - // thread count is passed in per decode (resolved min(8, cores) default) via - // a num_threads clause — no process-global omp_set_num_threads, so - // concurrent decodes don't stomp each other. The joint out_w projection - // runs in its own ggml graph, not here. -#ifdef _OPENMP - #pragma omp parallel for schedule(static) if(out_dim >= 2048) \ - num_threads(n_threads > 0 ? n_threads : 1) -#endif - for (int r = 0; r < out_dim; ++r) { - const float * row = W + static_cast(r) * static_cast(in_dim); - // Eight independent accumulators break the single-acc reduction - // dependency chain so the compiler can auto-vectorize the dot (the - // serial version pins it to one FMA's latency — ~7x off on MSVC, the - // dominant decode cost). Pure fp32, DL-safe, portable: no intrinsics, - // no ggml-cpu symbols, no arch flags. The tree-shaped final sum changes - // float reassociation only (argmax-robust, transcript-identical). - float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; - int c = 0; - const int in8 = in_dim & ~7; - for (; c < in8; c += 8) { - a0 += row[c + 0] * x[c + 0]; a1 += row[c + 1] * x[c + 1]; - a2 += row[c + 2] * x[c + 2]; a3 += row[c + 3] * x[c + 3]; - a4 += row[c + 4] * x[c + 4]; a5 += row[c + 5] * x[c + 5]; - a6 += row[c + 6] * x[c + 6]; a7 += row[c + 7] * x[c + 7]; - } - float acc = ((a0 + a1) + (a2 + a3)) + ((a4 + a5) + (a6 + a7)); - for (; c < in_dim; ++c) { - acc += row[c] * x[c]; - } - y[r] = acc + (b != nullptr ? b[r] : 0.0f); - } -#endif -} - -inline float sigmoidf(float x) { - // The two-branch form keeps the exp argument away from +inf, - // avoiding NaN propagation on extreme logits the model never - // produces in practice. - if (x >= 0.0f) { - const float z = std::exp(-x); - return 1.0f / (1.0f + z); - } - const float z = std::exp(x); - return z / (1.0f + z); -} - -// One LSTM time step. Reads `prev_h` and `prev_c` (each `H` floats), -// applies the input projection, hidden projection, gate equations, -// and writes new state into `new_h` / `new_c`. Caller must ensure -// `new_h` / `new_c` are pre-sized to H. Allocates an inline gate -// buffer of size 4*H on the stack via std::vector with a single -// reusable scratch passed in by the caller. -// -// Gate ordering: [i, f, g, o] (PyTorch standard, matches NeMo -// safetensors layout). -// -// Math (one bias term — NeMo's prednet stores a single concatenated -// bias rather than PyTorch's redundant W_ih + W_hh bias pair): -// -// gates = Wx @ x + Wh @ h_prev + b -// i, f, g, o = split(gates, 4, axis=0) -// i = sigmoid(i) -// f = sigmoid(f) -// g = tanh(g) -// o = sigmoid(o) -// c_new = f * c_prev + i * g -// h_new = o * tanh(c_new) -void lstm_step(const HostLstmLayer & layer, - const float * x, - const float * prev_h, - const float * prev_c, - int H, - std::vector & scratch_gates, - std::vector & scratch_hh, - float * new_h, - float * new_c, - int n_threads) -{ - const int four_H = 4 * H; - if (static_cast(scratch_gates.size()) < four_H) { - scratch_gates.resize(static_cast(four_H)); - } - if (static_cast(scratch_hh.size()) < four_H) { - scratch_hh.resize(static_cast(four_H)); - } - - // gates = Wx @ x + b - linear(layer.Wx.data(), x, layer.b.data(), four_H, H, scratch_gates.data(), n_threads); - // gates += Wh @ prev_h - linear(layer.Wh.data(), prev_h, nullptr, four_H, H, scratch_hh.data(), n_threads); - for (int i = 0; i < four_H; ++i) { - scratch_gates[i] += scratch_hh[i]; - } - - const float * gi = scratch_gates.data() + 0 * H; - const float * gf = scratch_gates.data() + 1 * H; - const float * gg = scratch_gates.data() + 2 * H; - const float * go = scratch_gates.data() + 3 * H; - - for (int j = 0; j < H; ++j) { - const float i_t = sigmoidf(gi[j]); - const float f_t = sigmoidf(gf[j]); - const float g_t = std::tanh(gg[j]); - const float o_t = sigmoidf(go[j]); - const float c_new = f_t * prev_c[j] + i_t * g_t; - new_c[j] = c_new; - new_h[j] = o_t * std::tanh(c_new); - } -} - -// Run the predictor for one decode step. -// -// `last_token` is the token id from the previous decode step, or -1 -// for the start state. The embed lookup goes through the [pred_vocab, -// pred_hidden] table; the start state is an all-zeros embedding -// (matching PredictNetwork's "no previous token" branch). -// -// Reads from `prev_state`, writes new state into `new_state`. The -// caller arranges `new_state` to be the only state mutated, so blank -// emissions can simply discard `new_state` instead of restoring a -// snapshot. -// -// Returns a pointer (borrowed) into `new_state.h.back()` — the last -// LSTM layer's new hidden state — which is the predictor "decoder -// output" the joint network will consume. -const float * predictor_step(const HostPredictor & predictor, - int last_token, - const LstmState & prev_state, - LstmState & new_state, - std::vector & scratch_x, - std::vector & scratch_gates, - std::vector & scratch_hh, - int n_threads) +// ggml-graph variant of predictor_step. Same contract: +// reads `prev_state`, writes the new step's (h, c) into `new_state`, returns a +// borrowed pointer into `new_state.h.back()`. Internally it feeds prev state and +// the embedding into the resident per-call graph, computes, and reads the new +// state back into the host LstmState — so the decode loop's state machine (swap, +// predictor_dirty cache, dumps) is identical to the host path. Caller guarantees +// g.ready and that new_state is sized to (L, H). +const float * predictor_step_ggml(const HostPredictor & predictor, + PredGraph & g, + int last_token, + const LstmState & prev_state, + LstmState & new_state, + std::vector & scratch_x) { const int H = predictor.pred_hidden; if (static_cast(scratch_x.size()) < H) { scratch_x.resize(static_cast(H)); } - // Embed lookup (or start-state zeros). + // Embed lookup (or start-state zeros) — identical to predictor_step. if (last_token < 0) { std::fill_n(scratch_x.data(), H, 0.0f); } else { @@ -686,24 +774,19 @@ const float * predictor_step(const HostPredictor & predictor, static_cast(H) * sizeof(float)); } - // Layer 0 takes the embed; subsequent layers take the previous - // layer's new hidden state. The LSTM step writes new state in - // place into `new_state.h[layer]` and `new_state.c[layer]`. - const float * layer_input = scratch_x.data(); - for (size_t layer = 0; layer < predictor.lstm.size(); ++layer) { - lstm_step(predictor.lstm[layer], - layer_input, - prev_state.h[layer].data(), - prev_state.c[layer].data(), - H, - scratch_gates, - scratch_hh, - new_state.h[layer].data(), - new_state.c[layer].data(), - n_threads); - layer_input = new_state.h[layer].data(); + const size_t hb = static_cast(H) * sizeof(float); + ggml_backend_tensor_set(g.x, scratch_x.data(), 0, hb); + for (int l = 0; l < g.L; ++l) { + ggml_backend_tensor_set(g.ph[l], prev_state.h[static_cast(l)].data(), 0, hb); + ggml_backend_tensor_set(g.pc[l], prev_state.c[static_cast(l)].data(), 0, hb); } + ggml_backend_graph_compute(g.backend, g.graph); + + for (int l = 0; l < g.L; ++l) { + ggml_backend_tensor_get(g.nh[l], new_state.h[static_cast(l)].data(), 0, hb); + ggml_backend_tensor_get(g.nc[l], new_state.c[static_cast(l)].data(), 0, hb); + } return new_state.h.back().data(); } @@ -711,8 +794,9 @@ const float * predictor_step(const HostPredictor & predictor, // // `enc_proj` is the precomputed encoder projection for this frame // (enc_w @ enc_frame + enc_b, `joint_h` floats). The caller batches -// all T_enc projections via sgemm before the decode loop so they -// are computed once rather than redundantly per iteration. +// all T_enc projections via one ggml GEMM (precompute_enc_proj_ggml) +// before the decode loop so they are computed once rather than +// redundantly per iteration. // // pred_proj = pred_w @ pred_state + pred_b [joint_h] // summed = enc_proj + pred_proj [joint_h] @@ -724,50 +808,23 @@ const float * predictor_step(const HostPredictor & predictor, // time. Parakeet 0.6B v2/v3 ship "relu". void joint_step(const HostJoint & j, const JointGraph & g, - int n_threads, const float * enc_proj, const float * pred_state, - std::vector & scratch_pred_proj, - std::vector & scratch_summed, std::vector & out_logits) { - if (static_cast(scratch_pred_proj.size()) < j.joint_h) scratch_pred_proj.resize(static_cast(j.joint_h)); - if (static_cast(scratch_summed.size()) < j.joint_h) scratch_summed.resize(static_cast(j.joint_h)); - if (static_cast(out_logits.size()) < j.joint_n) out_logits.resize(static_cast(j.joint_n)); - - linear(j.pred_w.data(), pred_state, j.pred_b.data(), j.joint_h, j.pred_hidden, - scratch_pred_proj.data(), n_threads); - - for (int i = 0; i < j.joint_h; ++i) { - scratch_summed[i] = enc_proj[i] + scratch_pred_proj[i]; - } - - if (j.activation == "relu") { - for (int i = 0; i < j.joint_h; ++i) { - if (scratch_summed[i] < 0.0f) scratch_summed[i] = 0.0f; - } - } else if (j.activation == "sigmoid") { - for (int i = 0; i < j.joint_h; ++i) { - scratch_summed[i] = sigmoidf(scratch_summed[i]); - } - } else { // "tanh" — loader's allow-list ensures this is the only remaining case. - for (int i = 0; i < j.joint_h; ++i) { - scratch_summed[i] = std::tanh(scratch_summed[i]); - } - } - - // Output projection: per-call ggml graph (native/fp32 weight, threaded - // matmul, bias folded into the graph) when available, else the host matmul. - if (g.ready) { - ggml_backend_tensor_set(g.act, scratch_summed.data(), 0, - static_cast(j.joint_h) * sizeof(float)); - ggml_backend_graph_compute(g.backend, g.graph); - ggml_backend_tensor_get(g.logits, out_logits.data(), 0, - static_cast(j.joint_n) * sizeof(float)); - } else { - linear(j.out_w.data(), scratch_summed.data(), j.out_b.data(), - j.joint_n, j.joint_h, out_logits.data(), n_threads); - } + if (static_cast(out_logits.size()) < j.joint_n) out_logits.resize(static_cast(j.joint_n)); + + // Full joint on the shared decoder pool: pred_w proj + sum + activation + + // out_w proj + bias, one graph, one dispatch. Inputs are the predictor + // output and this frame's enc_proj; the weights (incl. activation) are baked + // into the graph. + ggml_backend_tensor_set(g.pred_in, pred_state, 0, + static_cast(j.pred_hidden) * sizeof(float)); + ggml_backend_tensor_set(g.enc_in, enc_proj, 0, + static_cast(j.joint_h) * sizeof(float)); + ggml_backend_graph_compute(g.backend, g.graph); + ggml_backend_tensor_get(g.logits, out_logits.data(), 0, + static_cast(j.joint_n) * sizeof(float)); // Apply log_softmax over the full joint output (vocab+blank+durations // for TDT, vocab+blank for RNNT) to match NeMo's CPU-inference @@ -813,62 +870,58 @@ void joint_step(const HostJoint & j, } } -// Precompute the joint encoder projection for `T` frames into `out` -// (row-major [T, joint_h]): out[t] = enc_w @ enc_out[t] + enc_b. The -// projection is decode-state-independent, so batching it into one sgemm (or T -// sgemv calls without BLAS) before the greedy loop eliminates redundant -// per-iteration work and gives BLAS a large matrix. Shared by the TDT, RNN-T, -// and streaming RNN-T drivers. -void precompute_enc_proj(const HostJoint & j, - const float * enc_out, - int T, - int d_enc, - std::vector & out, - int n_threads) +// Compute the per-utterance encoder projection out[T, joint_h] = +// enc_out[T, d_enc] @ enc_w^T + enc_b as a single GEMM on `backend` — which +// carries the shared decoder threadpool, so this runs on the same pool as the +// pred/joint graphs (no extra pool). The weight + bias are the model-resident +// j.g_enc_w / j.g_enc_b (no per-call upload); only the activation input and the +// result are allocated here. At n = T the matmul is a real GEMM, so with +// GGML_LLAMAFILE=ON tinyBLAS engages (n >= 2). Returns false on any ggml +// failure; the driver treats it as a hard decode error (no host fallback). +bool precompute_enc_proj_ggml(const HostJoint & j, + ggml_backend_t backend, + const float * enc_out, + int T, + int d_enc, + std::vector & out) { + if (backend == nullptr || j.g_enc_w == nullptr || j.g_enc_b == nullptr) return false; const int joint_h = j.joint_h; out.resize(static_cast(T) * static_cast(joint_h)); -#if TRANSCRIBE_HAS_BLAS - (void) n_threads; - // C[T, joint_h] = enc_out[T, d_enc] @ enc_w^T[d_enc, joint_h] - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, - T, joint_h, d_enc, - 1.0f, enc_out, d_enc, - j.enc_w.data(), d_enc, - 0.0f, out.data(), joint_h); - // Add bias to every row. - for (int t = 0; t < T; ++t) { - float * proj = out.data() + static_cast(t) * static_cast(joint_h); - for (int k = 0; k < joint_h; ++k) { - proj[k] += j.enc_b[static_cast(k)]; - } - } -#else - // No BLAS: this is a [T, d_enc] x [d_enc, joint_h]^T GEMM. The previous code - // ran it as T serial sgemv calls via linear(), and linear() only threads when - // out_dim >= 2048 — joint_h is below that, so the whole projection ran - // single-threaded (~90 ms for T=138 on this model, the dominant decode cost, - // paid on CPU even under the Vulkan backend since the decoder is host code). - // The rows are fully independent, so parallelize over T and fold the bias in. - // The inner dot auto-vectorizes (FMA/AVX-512) under the project's arch flags. - const float * enc_w = j.enc_w.data(); - const float * enc_b = j.enc_b.data(); -#ifdef _OPENMP - #pragma omp parallel for schedule(static) num_threads(n_threads > 0 ? n_threads : 1) -#endif - for (int t = 0; t < T; ++t) { - const float * frame = enc_out + static_cast(t) * static_cast(d_enc); - float * proj = out.data() + static_cast(t) * static_cast(joint_h); - for (int r = 0; r < joint_h; ++r) { - const float * row = enc_w + static_cast(r) * static_cast(d_enc); - float acc = 0.0f; - for (int c = 0; c < d_enc; ++c) { - acc += row[c] * frame[c]; - } - proj[r] = acc + enc_b[static_cast(r)]; - } + + ggml_init_params ip {}; + ip.mem_size = ggml_tensor_overhead() * 6 + ggml_graph_overhead(); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (ctx == nullptr) return false; + + ggml_tensor * in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_enc, T); + ggml_set_input(in); + ggml_tensor * mm = ggml_mul_mat(ctx, j.g_enc_w, in); // [joint_h, T] + ggml_tensor * res = ggml_add(ctx, mm, j.g_enc_b); // + [joint_h] (broadcast over T) + ggml_set_output(res); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (buf == nullptr) { ggml_free(ctx); return false; } + + ggml_backend_tensor_set(in, enc_out, 0, + static_cast(T) * static_cast(d_enc) * sizeof(float)); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, res); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return false; } -#endif + + ggml_backend_tensor_get(res, out.data(), 0, + static_cast(T) * static_cast(joint_h) * sizeof(float)); + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return true; } // Argmax over a contiguous fp32 range. Returns the index of the @@ -963,9 +1016,22 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, // Per-call joint compute graph around the shared resident weight. Owns its // own backend + I/O tensors, so concurrent decodes don't share mutable - // state. Falls back to the host matmul if the graph isn't available. + // state. A graph build failure is a hard decode error (guarded at driver entry). + // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + + // pool; enc_proj and the joint graph borrow it. pg is declared first so it + // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode + // call, reentrant. The CPU backend is always available, so these succeed in + // practice; a build failure is a hard error (the host decode path was + // removed when the migration finalized). + PredGraph pg; JointGraph jg; - build_joint_graph(jg, w.joint, nt); + build_pred_graph(pg, w.predictor, nt); + if (pg.ready) build_joint_graph(jg, w.joint, pg.backend); + if (!pg.ready || !jg.ready) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet decoder: ggml decode graph build failed"); + return TRANSCRIBE_ERR_BACKEND; + } // Two LSTM states, both pre-sized: `state` is the committed // state we read from each iteration; `next_state` is where the @@ -978,19 +1044,18 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, next_state.reset(n_layers, H); // Precompute encoder projections for all T_enc frames (decode-state - // independent; one sgemm before the loop). See precompute_enc_proj. + // independent; one sgemm before the loop). See precompute_enc_proj_ggml. const int joint_h = w.joint.joint_h; std::vector enc_proj_all; const int64_t t_enc_proj_start = ggml_time_us(); - precompute_enc_proj(w.joint, enc_out, T_enc, d_enc, enc_proj_all, nt); + if (!precompute_enc_proj_ggml(w.joint, pg.backend, enc_out, T_enc, d_enc, enc_proj_all)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: enc_proj graph failed"); + return TRANSCRIBE_ERR_BACKEND; + } const int64_t t_enc_proj_us = ggml_time_us() - t_enc_proj_start; // Per-call scratch reused across every decode step. std::vector scratch_x; - std::vector scratch_gates; - std::vector scratch_hh; - std::vector scratch_pred_proj; - std::vector scratch_summed; std::vector scratch_probs; std::vector logits; @@ -1026,9 +1091,8 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, const int64_t t0 = ggml_time_us(); const float * decoder_out; if (predictor_dirty) { - decoder_out = predictor_step( - w.predictor, last_token, state, next_state, - scratch_x, scratch_gates, scratch_hh, nt); + decoder_out = predictor_step_ggml( + w.predictor, pg, last_token, state, next_state, scratch_x); predictor_dirty = false; } else { decoder_out = next_state.h.back().data(); @@ -1038,9 +1102,7 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, // ----- Joint (using precomputed encoder projection) ----- const float * enc_proj = enc_proj_all.data() + static_cast(step) * static_cast(joint_h); - joint_step(w.joint, jg, nt, enc_proj, decoder_out, - scratch_pred_proj, scratch_summed, - logits); + joint_step(w.joint, jg, enc_proj, decoder_out, logits); const int64_t t2 = ggml_time_us(); t_pred_us += t1 - t0; t_joint_us += t2 - t1; @@ -1205,8 +1267,21 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const int blank_id = w.blank_id; // Per-call joint compute graph (see decode_tdt_greedy). + // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + + // pool; enc_proj and the joint graph borrow it. pg is declared first so it + // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode + // call, reentrant. The CPU backend is always available, so these succeed in + // practice; a build failure is a hard error (the host decode path was + // removed when the migration finalized). + PredGraph pg; JointGraph jg; - build_joint_graph(jg, w.joint, nt); + build_pred_graph(pg, w.predictor, nt); + if (pg.ready) build_joint_graph(jg, w.joint, pg.backend); + if (!pg.ready || !jg.ready) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet decoder: ggml decode graph build failed"); + return TRANSCRIBE_ERR_BACKEND; + } LstmState state; LstmState next_state; @@ -1217,14 +1292,13 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const int joint_h = w.joint.joint_h; std::vector enc_proj_all; const int64_t t_enc_proj_start = ggml_time_us(); - precompute_enc_proj(w.joint, enc_out, T_enc, d_enc, enc_proj_all, nt); + if (!precompute_enc_proj_ggml(w.joint, pg.backend, enc_out, T_enc, d_enc, enc_proj_all)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: enc_proj graph failed"); + return TRANSCRIBE_ERR_BACKEND; + } const int64_t t_enc_proj_us = ggml_time_us() - t_enc_proj_start; std::vector scratch_x; - std::vector scratch_gates; - std::vector scratch_hh; - std::vector scratch_pred_proj; - std::vector scratch_summed; std::vector scratch_probs; std::vector logits; @@ -1252,9 +1326,8 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const int64_t t0 = ggml_time_us(); const float * decoder_out; if (predictor_dirty) { - decoder_out = predictor_step( - w.predictor, last_token, state, next_state, - scratch_x, scratch_gates, scratch_hh, nt); + decoder_out = predictor_step_ggml( + w.predictor, pg, last_token, state, next_state, scratch_x); predictor_dirty = false; } else { decoder_out = next_state.h.back().data(); @@ -1263,9 +1336,7 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const float * enc_proj = enc_proj_all.data() + static_cast(step) * static_cast(joint_h); - joint_step(w.joint, jg, nt, enc_proj, decoder_out, - scratch_pred_proj, scratch_summed, - logits); + joint_step(w.joint, jg, enc_proj, decoder_out, logits); const int64_t t2 = ggml_time_us(); t_pred_us += t1 - t0; t_joint_us += t2 - t1; @@ -1396,8 +1467,21 @@ transcribe_status decode_rnnt_greedy_streaming( const int blank_id = w.blank_id; // Per-call joint compute graph (see decode_tdt_greedy). + // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + + // pool; enc_proj and the joint graph borrow it. pg is declared first so it + // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode + // call, reentrant. The CPU backend is always available, so these succeed in + // practice; a build failure is a hard error (the host decode path was + // removed when the migration finalized). + PredGraph pg; JointGraph jg; - build_joint_graph(jg, w.joint, nt); + build_pred_graph(pg, w.predictor, nt); + if (pg.ready) build_joint_graph(jg, w.joint, pg.backend); + if (!pg.ready || !jg.ready) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet decoder: ggml decode graph build failed"); + return TRANSCRIBE_ERR_BACKEND; + } // Validate state_io shape; reset if degenerate (paranoid guard). if (static_cast(state_io.h.size()) != n_layers || @@ -1419,7 +1503,10 @@ transcribe_status decode_rnnt_greedy_streaming( // Precompute encoder projections for this chunk only. const int joint_h = w.joint.joint_h; std::vector enc_proj_all; - precompute_enc_proj(w.joint, enc_out, T_enc_new, d_enc, enc_proj_all, nt); + if (!precompute_enc_proj_ggml(w.joint, pg.backend, enc_out, T_enc_new, d_enc, enc_proj_all)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder (rnnt-stream): enc_proj graph failed"); + return TRANSCRIBE_ERR_BACKEND; + } // Working state. We mutate `state` and `next_state`; on return we // copy `state` (the COMMITTED state after the last non-blank @@ -1430,10 +1517,6 @@ transcribe_status decode_rnnt_greedy_streaming( next_state.reset(n_layers, H); std::vector scratch_x; - std::vector scratch_gates; - std::vector scratch_hh; - std::vector scratch_pred_proj; - std::vector scratch_summed; std::vector scratch_probs; std::vector logits; @@ -1451,9 +1534,8 @@ transcribe_status decode_rnnt_greedy_streaming( const float * decoder_out; if (predictor_dirty) { - decoder_out = predictor_step( - w.predictor, last_token, state, next_state, - scratch_x, scratch_gates, scratch_hh, nt); + decoder_out = predictor_step_ggml( + w.predictor, pg, last_token, state, next_state, scratch_x); predictor_dirty = false; } else { decoder_out = next_state.h.back().data(); @@ -1461,8 +1543,7 @@ transcribe_status decode_rnnt_greedy_streaming( const float * enc_proj = enc_proj_all.data() + static_cast(step) * static_cast(joint_h); - joint_step(w.joint, jg, nt, enc_proj, decoder_out, - scratch_pred_proj, scratch_summed, logits); + joint_step(w.joint, jg, enc_proj, decoder_out, logits); const float * token_logits = logits.data(); const int pred_token = argmax_range(token_logits, n_token_cls); @@ -1537,7 +1618,6 @@ transcribe_status decode_ctc_greedy(const HostDecoderWeights & w, return TRANSCRIBE_ERR_INVALID_ARG; } - const int nt = resolve_decode_threads(n_threads); const int n_classes = w.ctc_head.n_classes; const int blank_id = w.ctc_head.blank_id; @@ -1553,10 +1633,24 @@ transcribe_status decode_ctc_greedy(const HostDecoderWeights & w, w.ctc_head.weight.data(), d_enc, 0.0f, logits_all.data(), n_classes); #else - for (int t = 0; t < T_enc; ++t) { - const float * frame = enc_out + static_cast(t) * static_cast(d_enc); - float * row = logits_all.data() + static_cast(t) * static_cast(n_classes); - linear(w.ctc_head.weight.data(), frame, nullptr, n_classes, d_enc, row, nt); + // No BLAS: project all T frames in parallel over the shared parallel_for_all + // helper (one thread spawn for the utterance; rows within a frame are a + // serial, auto-vectorized dot). The bias is folded in below. (nt is unused + // on BLAS builds — cblas_sgemm owns its own threading — so it lives here.) + { + const int nt = resolve_decode_threads(n_threads); + const float * Wc = w.ctc_head.weight.data(); + transcribe::parallel_for_all(T_enc, nt, [&](int t) { + const float * frame = enc_out + static_cast(t) * static_cast(d_enc); + float * row = logits_all.data() + static_cast(t) * static_cast(n_classes); + for (int r = 0; r < n_classes; ++r) { + const float * wr = Wc + static_cast(r) * static_cast(d_enc); + float acc = 0.0f; + for (int c = 0; c < d_enc; ++c) acc += wr[c] * frame[c]; + row[r] = acc; + } + return true; + }); } #endif for (int t = 0; t < T_enc; ++t) { diff --git a/src/arch/parakeet/decoder.h b/src/arch/parakeet/decoder.h index 14a523e6..04de1cc9 100644 --- a/src/arch/parakeet/decoder.h +++ b/src/arch/parakeet/decoder.h @@ -3,33 +3,30 @@ // Implements the predictor (2-layer LSTM) + joint network forward + // greedy decode driver (TDT, RNN-T, and CTC heads). // -// Execution model. The predictor LSTM, the joint's encoder/predictor -// input projections, the activation, and the greedy search all run on -// host in fp32: each is a small per-step operation (a 640-wide LSTM step -// is ~6.5 MFLOPs; the joint input projections ~2.7 MFLOPs), so keeping -// them on host avoids per-step graph dispatch and keeps the bring-up dump -// points (compare_tensors.py) trivial to wire. +// Execution model. The decoder runs entirely as ggml graphs on ONE shared, +// per-decode CPU threadpool: the predictor LSTM (per-step graph), the +// per-utterance encoder projection (one GEMM), and the joint network (pred_w +// proj + sum + activation + out_w proj, one graph). PredGraph owns the backend +// and threadpool; the enc_proj and joint graphs borrow it, so the whole decode +// shares a single pool with no oversubscription. Every decode call builds its +// own graphs (reentrant); the weights are model-resident fp32 ggml tensors +// (HostPredictor::lstm_w_* and HostJoint::gw_*/g_pred_*), built once at load. // -// The exception is the joint OUTPUT projection (out_w: joint_n × joint_h). -// At small vocab (~1k for the English parakeets) it is also negligible and -// runs on host. At the multilingual vocab (13k) it grows to ~17 MFLOPs per -// step and dominates decode, so it runs through a ggml graph: the weight -// stays resident on the model (build_joint_weight / HostJoint::gw_*) while -// each decode call builds its own stack-local JointGraph (act/logits/graph/ -// backend) around it, keeping concurrent decode reentrant. ggml's threaded, -// SIMD matmul replaces the scalar host loop. The weight keeps its native -// (quantized) dtype -// by default — ggml streams the quantized bytes (~4× less weight bandwidth) and -// quantizes the activation on the fly, the same regime the encoder runs in and -// WER-validated equal to fp32 on English. The caller falls back to an -// fp32-dequant graph (faithful to the host reference: max rel logit diff ~3e-7 -// vs ~3e-3 for the native weight) if the native build fails or when -// TRANSCRIBE_JOINT_FP32=1 forces it; failing both, the host matmul still serves. -// The graph is built once and recomputed in place per step. +// This is a deliberate "one threading system" design: the decoder used to run +// host fp32 matmuls on a bespoke worker pool, which coexisted badly with ggml's +// own pool (teardown / oversubscription / MSVC-vcomp bugs). Routing everything +// through ggml deleted that second system. The CPU backend is always available +// (the decoder runs on host even when the model is on a GPU), so a graph build +// failure is a hard decode error rather than a host fallback. // -// Memory cost: a load-time host mirror of the predictor + joint weights -// (~35 MB on v2, ~73 MB on v3) against the ~2.4 GB encoder footprint. The -// mirror is built once in build_host_decoder_weights via +// All weights are fp32 (the LSTM and joint matmuls are n=1 GEMV per step, where +// quantization buys no usable bandwidth and would only add drift); fp32 keeps +// decode numerically faithful to the historical host reference. Any joint +// quantization is a model/quant-level concern, not done here. +// +// Memory cost: a load-time host + resident-ggml mirror of the predictor + joint +// weights (~35 MB on v2, ~73 MB on v3) against the ~2.4 GB encoder footprint. +// The mirror is built once in build_host_decoder_weights via // ggml_backend_tensor_get (universal across host buffers, Metal unified // memory, and discrete GPUs). // @@ -70,9 +67,19 @@ struct ParakeetHParams; // W_ih + W_hh bias pre-summed at conversion time. struct HostLstmLayer { - std::vector Wx; // [4*pred_hidden, pred_hidden] - std::vector Wh; // [4*pred_hidden, pred_hidden] - std::vector b; // [4*pred_hidden] + // Host fp32 mirrors — LOAD-TIME SCRATCH ONLY: filled at load, uploaded into + // the resident ggml tensors below by build_pred_weights, then freed. + std::vector Wx; // [4*pred_hidden, pred_hidden] (freed after load) + std::vector Wh; // [4*pred_hidden, pred_hidden] (freed after load) + std::vector b; // [4*pred_hidden] (freed after load) + + // Resident fp32 ggml mirrors of the three weights above, consumed by the + // ggml predictor graph. Row-major [4*H, H] host bytes map to ggml ne + // fast-to-slow [H, 4*H] (the mul_mat operand); bias is [4*H]. Borrowed + // pointers into HostPredictor::lstm_w_ctx — owned/freed by ~HostPredictor. + ggml_tensor * g_Wx = nullptr; + ggml_tensor * g_Wh = nullptr; + ggml_tensor * g_b = nullptr; }; struct HostPredictor { @@ -80,6 +87,27 @@ struct HostPredictor { int pred_vocab = 0; // includes the +1 start row std::vector embed_w; // [pred_vocab, pred_hidden] std::vector lstm; // pred_n_layers entries + + // --- resident ggml LSTM weights (immutable, model-owned) --- + // fp32 mirror of the per-layer Wx/Wh/b, made resident once at load so the + // per-decode-call PredGraph reads them without re-uploading. Same + // immutable-half / mutable-per-call split as HostJoint::gw_* + the + // stack-local JointGraph. fp32 (not native quant): the per-step LSTM + // matmuls are n=1 GEMV — tinyBLAS skips them (n<2) and they gain no usable + // weight-bandwidth benefit, so fp32 stays closest to the host reference at + // no speed cost. Owned here; freed by ~HostPredictor. On build failure + // lstm_ready stays false and the decode reports a hard error. + ggml_context * lstm_w_ctx = nullptr; + ggml_backend_t lstm_w_backend = nullptr; // alloc-only; never compute'd + ggml_backend_buffer_t lstm_w_buf = nullptr; + bool lstm_ready = false; + + HostPredictor() = default; + ~HostPredictor(); + HostPredictor(const HostPredictor &) = delete; + HostPredictor & operator=(const HostPredictor &) = delete; + HostPredictor(HostPredictor &&) = delete; + HostPredictor & operator=(HostPredictor &&) = delete; }; struct HostJoint { @@ -88,32 +116,36 @@ struct HostJoint { int joint_h = 0; int joint_n = 0; // total output classes (vocab+blank+durations) std::string activation; // "relu" / "sigmoid" / "tanh" - std::vector enc_w; // [joint_h, d_enc] - std::vector enc_b; // [joint_h] - std::vector pred_w; // [joint_h, pred_hidden] - std::vector pred_b; // [joint_h] - std::vector out_w; // [joint_n, joint_h] (host fallback path) - std::vector out_b; // [joint_n] - // --- resident ggml weight for the out_w projection --- - // The output projection (out_w: joint_n × joint_h) dominates RNN-T - // decode once the vocab is large (13k for the multilingual variant). - // We keep the weight + bias resident as ggml tensors in their native - // (quantized) dtype so the per-step matmul gets ggml's threaded + SIMD - // quantized kernel and ~4× less weight bandwidth than streaming the - // fp32 host mirror every step. - // - // This is the IMMUTABLE, model-owned half: built once at load and only - // ever read after that, so it is safe to share across every context. - // The MUTABLE per-step compute state (activation input, logits output, - // the graph, and a compute backend) lives in a stack-local JointGraph - // built per decode call (see decoder.cpp) — that is what makes - // concurrent decode on contexts sharing this model reentrant. - // Owned here; freed by ~HostJoint. + // Host fp32 weight mirrors — LOAD-TIME SCRATCH ONLY. read_tensor_to_f32 + // fills these at load; build_joint_weight uploads them into the resident + // ggml tensors below and then frees them, so the decode hot path reads only + // the resident tensors. (out_w is not mirrored here — gw_w is built straight + // from the model tensor.) + std::vector enc_w; // [joint_h, d_enc] (freed after load) + std::vector enc_b; // [joint_h] (freed after load) + std::vector pred_w; // [joint_h, pred_hidden] (freed after load) + std::vector pred_b; // [joint_h] (freed after load) + std::vector out_b; // [joint_n] (freed after load) + + // --- resident ggml weights (immutable, model-owned) --- + // The whole joint runs as one ggml graph (build_joint_graph), so every joint + // weight is resident as an fp32 ggml tensor: the encoder projection + // (g_enc_w/g_enc_b), the predictor projection (g_pred_w/g_pred_b), and the + // output projection (gw_w/gw_b). fp32 keeps decode faithful to the host + // reference; any joint quantization is a model/quant-level concern. Built + // once at load and only read after — safe to share across every context; + // the mutable per-decode state lives in a stack-local JointGraph (see + // decoder.cpp), which is what keeps concurrent decode reentrant. Owned here; + // freed by ~HostJoint. ggml_context * w_ctx = nullptr; ggml_backend_t w_backend = nullptr; // alloc-only; never graph_compute'd ggml_backend_buffer_t w_buf = nullptr; - ggml_tensor * gw_w = nullptr; // [joint_h, joint_n] native/fp32 weight + ggml_tensor * g_enc_w = nullptr; // [d_enc, joint_h] fp32 weight + ggml_tensor * g_enc_b = nullptr; // [joint_h] fp32 bias + ggml_tensor * g_pred_w = nullptr; // [pred_hidden, joint_h] fp32 weight + ggml_tensor * g_pred_b = nullptr; // [joint_h] fp32 bias + ggml_tensor * gw_w = nullptr; // [joint_h, joint_n] fp32 weight ggml_tensor * gw_b = nullptr; // [joint_n] fp32 bias bool w_ready = false; diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index 7994d579..f2191e36 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -1497,24 +1497,7 @@ static transcribe_status run_one_shot_inner( // Whisper.cpp pattern: iterate the scheduler's backends and set // n_threads via the registry proc address. GPU backends ignore // this; CPU and BLAS backends use it. - { - int n_threads = pc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(pc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(pc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) { - fn(be, n_threads); - } - } - } + transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); // ----- Compute -------------------------------------------------- const int64_t t_enc_start = ggml_time_us(); @@ -1911,22 +1894,7 @@ static transcribe_status run_batch_encode( 0, mask_buf.size() * sizeof(float)); } - { - int n_threads = pc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(pc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(pc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); const int64_t t_enc_start = ggml_time_us(); if (const ggml_status gs = @@ -2554,22 +2522,7 @@ transcribe_status emit_streaming_chunk( } // Thread count (same recipe as offline run()). - { - int n_threads = pc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(pc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(pc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); if (const ggml_status gs = ggml_backend_sched_graph_compute(pc->sched, eb.graph); @@ -3048,22 +3001,7 @@ transcribe_status emit_buffered_chunk( } // ----- Threading ----- - { - int n_threads = pc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(pc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(pc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); // ----- Compute ----- const int64_t t_enc_start = ggml_time_us(); diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 67ee8647..15c7d400 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -816,22 +816,7 @@ transcribe_status run( } // Thread count. - { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); t_enc_build_us = t_enc_start - t_enc_build_start; @@ -1333,20 +1318,7 @@ namespace { // Apply the session thread count to every backend behind the scheduler. The // setting persists across sched_reset, so callers only need this once. void apply_sched_threads(QwenAsrSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } // Fresh per-utterance graph arena. Frees any prior compute_ctx and inits a @@ -1380,8 +1352,7 @@ transcribe_status encode_all_batched( std::vector mel_nf(n, 0); int n_threads = cc->n_threads; if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); + n_threads = transcribe::default_n_threads(); } const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_threads, [&](int b) { diff --git a/src/arch/sensevoice/model.cpp b/src/arch/sensevoice/model.cpp index f12532bf..eba420d3 100644 --- a/src/arch/sensevoice/model.cpp +++ b/src/arch/sensevoice/model.cpp @@ -249,20 +249,7 @@ int resolve_lid_idx(const SenseVoiceHParams & hp, const char * lang_or_null) { } void apply_thread_policy(SenseVoiceSession * cc) { - int n_threads = cc->n_threads; - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, static_cast( - std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(cc->sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(cc->sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } // Greedy CTC decode + public result-hierarchy build for ONE utterance's CTC diff --git a/src/arch/voxtral/model.cpp b/src/arch/voxtral/model.cpp index 6dcd9f0b..e1dbc267 100644 --- a/src/arch/voxtral/model.cpp +++ b/src/arch/voxtral/model.cpp @@ -522,19 +522,7 @@ transcribe_status init_context( // Apply n_threads to every backend on the scheduler. void set_sched_threads(ggml_backend_sched_t sched, int n_threads) { - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(sched, n_threads); } transcribe_status run( @@ -1066,8 +1054,7 @@ transcribe_status run_batch( int n_mel_threads = cc->n_threads; if (n_mel_threads <= 0) - n_mel_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); + n_mel_threads = transcribe::default_n_threads(); const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_mel_threads, [&](int b) { diff --git a/src/arch/voxtral_realtime/model.cpp b/src/arch/voxtral_realtime/model.cpp index e6c9c980..583d07b3 100644 --- a/src/arch/voxtral_realtime/model.cpp +++ b/src/arch/voxtral_realtime/model.cpp @@ -19,6 +19,7 @@ #include "causal_lm/causal_lm.h" #include "transcribe-arch.h" +#include "transcribe-batch-util.h" #include "transcribe-debug.h" #include "transcribe-flash-policy.h" #include "transcribe-load-common.h" @@ -99,19 +100,7 @@ transcribe_status resolve_specials(const transcribe::Tokenizer & tok, } void apply_threads(ggml_backend_sched_t sched, int n_threads) { - if (n_threads <= 0) { - n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); - } - for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { - ggml_backend_t be = ggml_backend_sched_get_backend(sched, i); - ggml_backend_dev_t dev = ggml_backend_get_device(be); - ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; - if (reg == nullptr) continue; - auto * fn = reinterpret_cast( - ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); - if (fn != nullptr) fn(be, n_threads); - } + transcribe::configure_sched_n_threads(sched, n_threads); } // Read a reference enc.mel.in.f32 ([n_mels, n_frames] mel-major) for numerical diff --git a/src/arch/whisper/model.cpp b/src/arch/whisper/model.cpp index 19404e8f..185dff78 100644 --- a/src/arch/whisper/model.cpp +++ b/src/arch/whisper/model.cpp @@ -757,6 +757,13 @@ transcribe_status run_whisper_encoder_on_window( "whisper run: ggml_backend_sched_new failed"); return TRANSCRIBE_ERR_GGUF; } + + // Apply the caller's CPU thread count to the scheduler's backends once + // at sched creation — it persists across the reused encoder and decoder + // graphs, and cc->n_threads is fixed for the session. GPU backends ignore + // it; CPU/BLAS backends use it. Without this, whisper's graph compute ran + // at ggml's default thread count regardless of params.n_threads. + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } ggml_backend_sched_reset(cc->sched); if (!ggml_backend_sched_alloc_graph(cc->sched, eb.graph)) { @@ -3186,8 +3193,7 @@ transcribe_status whisper_run_batch( valid[b] = 1; } int n_threads = cc->n_threads; - if (n_threads <= 0) n_threads = std::min(8, std::max(1, - static_cast(std::thread::hardware_concurrency()))); + if (n_threads <= 0) n_threads = transcribe::default_n_threads(); int64_t mel_us = 0, enc_us = 0, dec_us = 0; const int64_t t_mel0 = ggml_time_us(); transcribe::parallel_for_all(n, n_threads, [&](int b) { diff --git a/src/transcribe-batch-util.cpp b/src/transcribe-batch-util.cpp index 4d4e6cc8..49a9d7c2 100644 --- a/src/transcribe-batch-util.cpp +++ b/src/transcribe-batch-util.cpp @@ -13,14 +13,76 @@ #include #include +#if defined(__linux__) +#include +#elif defined(_WIN32) +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#ifndef NOMINMAX // else 's min/max macros clobber std::min/std::max +#define NOMINMAX +#endif +#include +#endif + namespace transcribe { +// Number of CPUs the process is actually allowed to run on. Falls back to +// hardware_concurrency() when the platform query is unavailable or fails. +static int usable_cpu_count() { +#if defined(__linux__) + cpu_set_t set; + CPU_ZERO(&set); + if (sched_getaffinity(0, sizeof(set), &set) == 0) { + const int n = CPU_COUNT(&set); + if (n > 0) return n; + } +#elif defined(_WIN32) + DWORD_PTR proc_mask = 0, sys_mask = 0; + if (GetProcessAffinityMask(GetCurrentProcess(), &proc_mask, &sys_mask) && + proc_mask != 0) { + int n = 0; + for (DWORD_PTR m = proc_mask; m != 0; m &= (m - 1)) ++n; // popcount + if (n > 0) return n; + } +#endif + const unsigned hw = std::thread::hardware_concurrency(); + return hw > 0 ? static_cast(hw) : 1; +} + +int default_n_threads(int cap) { + int n = usable_cpu_count(); + if (n < 1) n = 1; + if (cap > 0 && n > cap) n = cap; + return n; +} + + +int configure_sched_n_threads(ggml_backend_sched_t sched, int requested) { + const int n_threads = requested > 0 ? requested : default_n_threads(); + if (sched == nullptr) return n_threads; + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + ggml_backend_t be = ggml_backend_sched_get_backend(sched, i); + ggml_backend_dev_t dev = ggml_backend_get_device(be); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; + if (reg == nullptr) continue; + auto * fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads")); + if (fn != nullptr) { + fn(be, n_threads); + } + } + return n_threads; +} + bool parallel_for_all(int n, int n_threads, const std::function & work) { if (n <= 0) return true; if (n_threads <= 0) { - n_threads = static_cast(std::thread::hardware_concurrency()); + // Affinity-aware, uncapped: use every CPU the process may run on, then + // clamp to the batch size below. (cap <= 0 disables the per-backend cap.) + n_threads = default_n_threads(/*cap=*/0); } n_threads = std::max(1, std::min(n, n_threads)); diff --git a/src/transcribe-batch-util.h b/src/transcribe-batch-util.h index 943c37a0..f3a66d78 100644 --- a/src/transcribe-batch-util.h +++ b/src/transcribe-batch-util.h @@ -37,6 +37,28 @@ namespace transcribe { bool parallel_for_all(int n, int n_threads, const std::function & work); +// Default CPU thread count for the ggml compute backends and the host +// parallel-for, for when the caller passes n_threads <= 0. Counts the CPUs the +// process may actually run on (the affinity mask via sched_getaffinity on Linux +// / GetProcessAffinityMask on Windows), NOT the host's total core count, then +// clamps to [1, cap]. This matters under a constrained scheduler (taskset, a +// cpuset, a CI runner pinned to N vCPUs): std::thread::hardware_concurrency() +// reports the full host, so the old `min(cap, hardware_concurrency())` default +// oversubscribed the usable cores and ggml's spin-wait barriers then livelocked. +// CAVEAT: the affinity mask reflects taskset/cpuset, but NOT a CFS bandwidth +// quota (`docker --cpus=N` without --cpuset) — that case still over-counts. +// A cap <= 0 disables the clamp (use all usable CPUs). +int default_n_threads(int cap = 8); + +// Resolve a CPU thread count and apply it to every CPU/BLAS backend in `sched` +// via the backend registry's ggml_backend_set_n_threads. `requested <= 0` means +// "use default_n_threads()". GPU backends don't expose the setter and are +// skipped; a null sched is a no-op. Returns the resolved thread count so callers +// can reuse it (e.g. for the host parallel-for). This is the ONE place archs +// configure scheduler threads — call it instead of hand-rolling the backend +// loop, so the affinity-aware default can never be forgotten at a new site. +int configure_sched_n_threads(ggml_backend_sched_t sched, int requested); + // --------------------------------------------------------------------------- // Encoder-batch scaffolding (the run_batch() fast-path recipe) // diff --git a/src/transcribe-mel.cpp b/src/transcribe-mel.cpp index 192d6c2d..1933d38a 100644 --- a/src/transcribe-mel.cpp +++ b/src/transcribe-mel.cpp @@ -61,6 +61,8 @@ #include "transcribe-mel.h" +#include "transcribe-batch-util.h" + #ifdef __APPLE__ # include #elif TRANSCRIBE_HAS_BLAS @@ -566,9 +568,7 @@ transcribe_status MelFrontend::compute( int stft_threads = n_threads; if (stft_threads <= 0) { - const unsigned hw = std::thread::hardware_concurrency(); - stft_threads = hw > 0 ? static_cast(hw) : 1; - if (stft_threads > 8) stft_threads = 8; + stft_threads = default_n_threads(); } if (stft_threads > n_frames) { stft_threads = std::max(1, n_frames); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3505d927..64039f46 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -226,6 +226,28 @@ transcribe_apply_warnings(transcribe_log_unit) add_test(NAME transcribe_log_unit COMMAND transcribe_log_unit) +# ----------------------------------------------------------------------------- +# CPU thread-count helpers (affinity-aware default + sched resolution) +# ----------------------------------------------------------------------------- +# +# Pins default_n_threads()/configure_sched_n_threads() in transcribe-batch-util: +# cap+range, the null-sched resolution contract, and (Linux) that the default +# honors the process affinity mask rather than the host core count — the +# property whose absence caused the oversubscription livelock. + +add_executable(transcribe_thread_default_unit + thread_default_unit.cpp) + +target_link_libraries(transcribe_thread_default_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_thread_default_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_thread_default_unit) + +add_test(NAME transcribe_thread_default_unit + COMMAND transcribe_thread_default_unit) + # ----------------------------------------------------------------------------- # Threadpool barrier under CPU oversubscription (regression for the MSVC # non-OpenMP ggml_barrier livelock; see ggml-cpu.c ggml_thread_yield) diff --git a/tests/thread_default_unit.cpp b/tests/thread_default_unit.cpp new file mode 100644 index 00000000..61f2ee45 --- /dev/null +++ b/tests/thread_default_unit.cpp @@ -0,0 +1,169 @@ +// thread_default_unit.cpp - white-box tests for the CPU thread-count helpers +// in transcribe-batch-util (default_n_threads / configure_sched_n_threads). +// +// The bug these pin: the per-arch default used std::thread::hardware_concurrency(), +// which reports the host's total cores and ignores the process CPU affinity +// mask. Under a constrained scheduler (taskset / cpuset / a CI runner pinned to +// N vCPUs) that over-counts, and the resulting oversubscription makes ggml's +// spin-wait barriers livelock. default_n_threads() counts the CPUs the process +// may actually run on instead. So the load-bearing assertion here is the +// affinity one (Linux): pin to K CPUs, expect K. +// +// configure_sched_n_threads() is exercised on the null-scheduler path, which +// isolates its resolution contract (requested>0 passes through; requested<=0 +// falls back to default_n_threads()) without standing up a real backend graph — +// the real per-backend application is covered by the e2e/example runs. + +#include "transcribe-batch-util.h" + +#include + +#if defined(__linux__) +#include +#elif defined(_WIN32) +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#ifndef NOMINMAX // else 's min/max macros clobber std::min/std::max +#define NOMINMAX +#endif +#include +#endif + +namespace { + +int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", \ + __FILE__, __LINE__, #cond); \ + ++g_failures; \ + } \ + } while (0) + +// default_n_threads(): cap and range, independent of platform. +void test_default_range_and_cap() { + const int uncapped = transcribe::default_n_threads(/*cap=*/0); + CHECK(uncapped >= 1); // always at least one thread + + const int capped8 = transcribe::default_n_threads(8); + CHECK(capped8 >= 1); + CHECK(capped8 <= 8); // cap is honored + CHECK(capped8 == (uncapped < 8 ? uncapped : 8)); // cap = min(uncapped, 8) + + CHECK(transcribe::default_n_threads(1) == 1); // cap of 1 pins to 1 + const int cap4 = transcribe::default_n_threads(4); + CHECK(cap4 >= 1 && cap4 <= 4); +} + +// configure_sched_n_threads(nullptr, ...): resolution contract, no backends. +void test_configure_resolution_null_sched() { + // requested > 0 passes through verbatim. + CHECK(transcribe::configure_sched_n_threads(nullptr, 3) == 3); + CHECK(transcribe::configure_sched_n_threads(nullptr, 1) == 1); + // requested <= 0 falls back to default_n_threads(). + const int def = transcribe::default_n_threads(); + CHECK(transcribe::configure_sched_n_threads(nullptr, 0) == def); + CHECK(transcribe::configure_sched_n_threads(nullptr, -5) == def); +} + +#if defined(__linux__) +// The actual bug: default_n_threads() must honor the process affinity mask, +// not the host core count. Pin to K CPUs and expect K. Skips gracefully if the +// environment forbids changing affinity (some sandboxes do). +void test_affinity_honored_linux() { + cpu_set_t original; + CPU_ZERO(&original); + if (sched_getaffinity(0, sizeof(original), &original) != 0) { + std::fprintf(stderr, "SKIP affinity test: sched_getaffinity failed\n"); + return; + } + const int avail = CPU_COUNT(&original); + + // Collect the CPU ids currently allowed, so we pin to real, permitted CPUs. + int cpus[CPU_SETSIZE]; + int n_cpus = 0; + for (int c = 0; c < CPU_SETSIZE && n_cpus < avail; ++c) { + if (CPU_ISSET(c, &original)) cpus[n_cpus++] = c; + } + + // Pin to exactly one CPU -> default_n_threads must be 1. + cpu_set_t one; + CPU_ZERO(&one); + CPU_SET(cpus[0], &one); + if (sched_setaffinity(0, sizeof(one), &one) != 0) { + std::fprintf(stderr, "SKIP affinity test: sched_setaffinity not permitted\n"); + return; + } + CHECK(transcribe::default_n_threads(8) == 1); + CHECK(transcribe::default_n_threads(/*cap=*/0) == 1); // uncapped still sees 1 usable + + // Pin to two CPUs -> expect 2 (only when the box actually has >= 2). + if (n_cpus >= 2) { + cpu_set_t two; + CPU_ZERO(&two); + CPU_SET(cpus[0], &two); + CPU_SET(cpus[1], &two); + if (sched_setaffinity(0, sizeof(two), &two) == 0) { + CHECK(transcribe::default_n_threads(8) == 2); + } + } + + // Restore the process to its original affinity. + CHECK(sched_setaffinity(0, sizeof(original), &original) == 0); +} +#elif defined(_WIN32) +// Windows analog: GetProcessAffinityMask must drive the default, not the host +// core count. Pin the process to K CPUs via SetProcessAffinityMask and expect K. +// Skips gracefully if affinity changes aren't permitted. +void test_affinity_honored_windows() { + const HANDLE proc = GetCurrentProcess(); + DWORD_PTR proc_mask = 0, sys_mask = 0; + if (!GetProcessAffinityMask(proc, &proc_mask, &sys_mask) || proc_mask == 0) { + std::fprintf(stderr, "SKIP affinity test: GetProcessAffinityMask failed\n"); + return; + } + const DWORD_PTR original = proc_mask; + + // Lowest set bit = one permitted CPU. Pin to it -> expect 1. + const DWORD_PTR one = original & (~original + 1); + if (!SetProcessAffinityMask(proc, one)) { + std::fprintf(stderr, "SKIP affinity test: SetProcessAffinityMask not permitted\n"); + return; + } + CHECK(transcribe::default_n_threads(8) == 1); + CHECK(transcribe::default_n_threads(/*cap=*/0) == 1); + + // Pin to two CPUs (lowest two set bits) -> expect 2, when available. + const DWORD_PTR second = (original & ~one) & (~(original & ~one) + 1); + if (second != 0) { + if (SetProcessAffinityMask(proc, one | second)) { + CHECK(transcribe::default_n_threads(8) == 2); + } + } + + // Restore the process to its original affinity. + CHECK(SetProcessAffinityMask(proc, original) != 0); +} +#endif + +} // namespace + +int main() { + test_default_range_and_cap(); + test_configure_resolution_null_sched(); +#if defined(__linux__) + test_affinity_honored_linux(); +#elif defined(_WIN32) + test_affinity_honored_windows(); +#endif + + if (g_failures != 0) { + std::fprintf(stderr, "thread_default_unit: %d failure(s)\n", g_failures); + return 1; + } + std::fprintf(stderr, "thread_default_unit: OK\n"); + return 0; +}