From 1f4a076be01c659a25454862f49c7ea608bc351a Mon Sep 17 00:00:00 2001 From: davidamacey Date: Wed, 19 Aug 2026 20:53:39 -0400 Subject: [PATCH 1/7] Vectorize VBx and thread the AHC pdist: 8x clustering, output-identical At scale (4.7 h meeting: N=21,418 filtered embeddings, K=1,902 AHC seed clusters, D=128) clustering dominates end-to-end wall time: VB-EM 305 s (scalar O(N*K*D) loops x20 iterations, per-speaker penalty recomputed per sample) plus AHC pdist 64.5 s (per-pair scalar loop) = 74% of a 474 s run. - vbx.rs: M-step as one gamma^T*rho matmul; E-step as one rho*alpha^T matmul with the penalty vector computed once per iteration; fused logsumexp/gamma row pass. Same f64 math, same iteration and early-stop semantics. - ahc.rs: condensed distances via blocked Gram matmul (d^2 = |a|^2+|b|^2-2ab) with std::thread::scope over disjoint condensed ranges (lock-free, deterministic ordering). - Cargo.toml: enable ndarray matrixmultiply-threading. Measured (A6000, quiet machine): VB-EM 305.1 -> 36.7 s (8.3x); pdist 64.5 -> 1.2 s (53x); clustering total 348 -> 43.6 s (8x); 4.7 h end-to-end 474 -> 171.6 s. RTTMs bit-identical; clustering tests incl. the Python-parity fixtures (AHC == scipy order; VBx gamma/pi at 1e-4/1e-5) pass; AMI test-16 re-run identical at 13.101% aggregate, 16/16 RTTMs bit-identical. --- Cargo.lock | 28 ++++++++++++++ Cargo.toml | 2 +- src/clustering/ahc.rs | 85 +++++++++++++++++++++++++++++++++-------- src/clustering/vbx.rs | 88 ++++++++++++++++++++----------------------- 4 files changed, 139 insertions(+), 64 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9578705..a87a6e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -888,6 +888,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hf-hub" version = "0.5.0" @@ -1376,7 +1382,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" dependencies = [ "autocfg", + "num_cpus", + "once_cell", "rawpointer", + "thread-tree", ] [[package]] @@ -1540,6 +1549,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "objc2" version = "0.6.4" @@ -2582,6 +2601,15 @@ dependencies = [ "syn", ] +[[package]] +name = "thread-tree" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffbd370cb847953a25954d9f63e14824a36113f8c72eecf6eccef5dc4b45d630" +dependencies = [ + "crossbeam-channel", +] + [[package]] name = "thread_local" version = "1.1.9" diff --git a/Cargo.toml b/Cargo.toml index 46c765f..1ae14f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ hf-hub = { version = "0.5", optional = true, default-features = false } # math kodama = "0.3.0" -ndarray = "0.17.2" +ndarray = { version = "0.17.2", features = ["matrixmultiply-threading"] } ndarray-linalg-mkl = { package = "ndarray-linalg", version = "0.18.1", features = ["intel-mkl-static"], optional = true } ndarray-linalg-static = { package = "ndarray-linalg", version = "0.18.1", features = ["openblas-static"], optional = true } ndarray-linalg-system = { package = "ndarray-linalg", version = "0.18.1", features = ["openblas-system"], optional = true } diff --git a/src/clustering/ahc.rs b/src/clustering/ahc.rs index 007ffb3..77a5f31 100644 --- a/src/clustering/ahc.rs +++ b/src/clustering/ahc.rs @@ -26,30 +26,83 @@ pub fn cluster(embeddings: &ArrayView2, config: AhcConfig) -> Vec { } let normalized = l2_normalize_rows(embeddings); + let t0 = std::time::Instant::now(); let mut condensed = condensed_euclidean(&normalized); + let pdist_ms = t0.elapsed().as_millis(); + let t1 = std::time::Instant::now(); let dendrogram = linkage(&mut condensed, observations, Method::Centroid); - flat_clusters(observations, dendrogram.steps(), config.threshold) + let linkage_ms = t1.elapsed().as_millis(); + let t2 = std::time::Instant::now(); + let labels = flat_clusters(observations, dendrogram.steps(), config.threshold); + tracing::debug!( + observations, + pdist_ms, + linkage_ms, + flat_ms = t2.elapsed().as_millis(), + "AHC stage timing" + ); + labels } fn condensed_euclidean(embeddings: &Array2) -> Vec { + // diar-native patch: blocked Gram-matrix formulation with scoped threads. + // The original per-pair scalar loop cost 64.5 s at N=21k; dist^2 = |a|^2 + |b|^2 - 2ab + // via matmul blocks is ~20x faster and each block writes a disjoint contiguous + // range of the condensed vector, so blocks parallelize without locks. let observations = embeddings.nrows(); - let mut condensed = Vec::with_capacity(observations * (observations - 1) / 2); - for row in 0..observations.saturating_sub(1) { - for col in row + 1..observations { - let lhs = embeddings.row(row); - let rhs = embeddings.row(col); - let distance = lhs - .iter() - .zip(rhs.iter()) - .map(|(left, right)| { - let delta = left - right; - delta * delta - }) - .sum::() - .sqrt(); - condensed.push(distance); + if observations < 2 { + return Vec::new(); + } + let total = observations * (observations - 1) / 2; + let mut condensed = vec![0f32; total]; + let sq_norms: Vec = embeddings + .rows() + .into_iter() + .map(|row| row.dot(&row)) + .collect(); + + const BLOCK: usize = 1024; + // start offset of row i's segment in the condensed vector: + // sum_{r = Vec::new(); + { + let mut rest: &mut [f32] = &mut condensed; + let mut consumed = 0usize; + let mut bi = 0usize; + while bi < observations.saturating_sub(1) { + let bi_end = (bi + BLOCK).min(observations - 1); + let end_offset = seg_start(bi_end); + let (head, tail) = rest.split_at_mut(end_offset - consumed); + blocks.push((bi, bi_end, head)); + consumed = end_offset; + rest = tail; + bi = bi_end; } } + + std::thread::scope(|scope| { + for (bi, bi_end, slice) in blocks { + let emb = &embeddings; + let norms = &sq_norms; + scope.spawn(move || { + let a = emb.slice(ndarray::s![bi..bi_end, ..]); + let b = emb.slice(ndarray::s![bi.., ..]); + let gram = a.dot(&b.t()); // (bi_end-bi) x (observations-bi) + let mut offset = 0usize; + for (local, i) in (bi..bi_end).enumerate() { + for j in (i + 1)..observations { + let dot = gram[[local, j - bi]]; + let d2 = (norms[i] + norms[j] - 2.0 * dot).max(0.0); + slice[offset] = d2.sqrt(); + offset += 1; + } + } + }); + } + }); condensed } diff --git a/src/clustering/vbx.rs b/src/clustering/vbx.rs index 100b53a..1933960 100644 --- a/src/clustering/vbx.rs +++ b/src/clustering/vbx.rs @@ -76,65 +76,59 @@ pub fn vbx( // m-step: compute speaker models // invL[k,d] = 1.0 / (1 + Fa/Fb * N_k * Phi[d]) // alpha[k,d] = Fa/Fb * invL[k,d] * sum_t(gamma[t,k] * rho[t,d]) + // diar-native patch: vectorized — the original per-element loops cost + // O(N*K*D) scalar work per iteration (305 s at N=21k, K=1.9k, D=128). let n_k: Array1 = gamma.sum_axis(Axis(0)); - let mut inv_l = Array2::zeros((n_speakers, dim)); - let mut alpha = Array2::zeros((n_speakers, dim)); - - for speaker_idx in 0..n_speakers { - for dim_idx in 0..dim { - inv_l[[speaker_idx, dim_idx]] = - 1.0 / (1.0 + fa_over_fb * n_k[speaker_idx] * phi_f64[dim_idx]); - } - - // gamma.T @ rho for this speaker - let mut f_k = Array1::::zeros(dim); - for sample_idx in 0..n_samples { - f_k.scaled_add(gamma[[sample_idx, speaker_idx]], &rho.row(sample_idx)); - } - - for dim_idx in 0..dim { - alpha[[speaker_idx, dim_idx]] = - fa_over_fb * inv_l[[speaker_idx, dim_idx]] * f_k[dim_idx]; - } + let mut inv_l = Array2::::zeros((n_speakers, dim)); + for (speaker_idx, mut row) in inv_l.rows_mut().into_iter().enumerate() { + let scale = fa_over_fb * n_k[speaker_idx]; + row.assign(&phi_f64.mapv(|p| 1.0 / (1.0 + scale * p))); } + // f = gamma.T @ rho (K x D), alpha = Fa/Fb * invL ⊙ f + let f = gamma.t().dot(&rho); + let mut alpha = &inv_l * &f; + alpha.mapv_inplace(|v| v * fa_over_fb); + // e-step // log_p_[t,k] = Fa * (rho[t] . alpha[k] - 0.5 * (invL[k] + alpha[k]^2) . Phi + G[t]) - let mut log_p = Array2::::zeros((n_samples, n_speakers)); - for sample_idx in 0..n_samples { - for speaker_idx in 0..n_speakers { - let rho_dot_alpha: f64 = rho.row(sample_idx).dot(&alpha.row(speaker_idx)); - let penalty: f64 = (0..dim) - .map(|dim_idx| { - (inv_l[[speaker_idx, dim_idx]] - + alpha[[speaker_idx, dim_idx]] * alpha[[speaker_idx, dim_idx]]) - * phi_f64[dim_idx] - }) - .sum(); - log_p[[sample_idx, speaker_idx]] = - fa * (rho_dot_alpha - 0.5 * penalty + frame_constants[sample_idx]); - } + // penalty depends only on k — compute once per iteration, not per sample. + let penalty: Array1 = (0..n_speakers) + .map(|speaker_idx| { + inv_l + .row(speaker_idx) + .iter() + .zip(alpha.row(speaker_idx).iter()) + .zip(phi_f64.iter()) + .map(|((&il, &a), &p)| (il + a * a) * p) + .sum() + }) + .collect(); + + let mut log_p = rho.dot(&alpha.t()); // N x K + for (sample_idx, mut row) in log_p.rows_mut().into_iter().enumerate() { + let g = frame_constants[sample_idx]; + row.zip_mut_with(&penalty, |value, &pen| { + *value = fa * (*value - 0.5 * pen + g); + }); } - // GMM-style update with pi priors + // GMM-style update with pi priors (single fused pass per row) let lpi: Array1 = pi.mapv(|p| (p + 1e-8).ln()); - // log_p_x[sample_idx] = logsumexp(log_p[sample_idx] + lpi) let mut log_p_x = Array1::::zeros(n_samples); - for sample_idx in 0..n_samples { - scratch.assign(&log_p.row(sample_idx)); + for ((log_p_row, mut gamma_row), log_p_x_slot) in log_p + .rows() + .into_iter() + .zip(gamma.rows_mut()) + .zip(log_p_x.iter_mut()) + { + scratch.assign(&log_p_row); scratch += &lpi; - log_p_x[sample_idx] = logsumexp_f64(&scratch.view()); - } - - // gamma[sample_idx,speaker_idx] = exp(log_p[sample_idx,speaker_idx] + lpi[speaker_idx] - log_p_x[sample_idx]) - for sample_idx in 0..n_samples { - for speaker_idx in 0..n_speakers { - gamma[[sample_idx, speaker_idx]] = - (log_p[[sample_idx, speaker_idx]] + lpi[speaker_idx] - log_p_x[sample_idx]) - .exp(); - } + let lse = logsumexp_f64(&scratch.view()); + *log_p_x_slot = lse; + gamma_row.zip_mut_with(&scratch, |g, &s| *g = (s - lse).exp()); } // update pi From 90200c1857452d5fa4fe7e3690f0aedc0544f5e0 Mon Sep 17 00:00:00 2001 From: davidamacey Date: Wed, 19 Aug 2026 20:57:02 -0400 Subject: [PATCH 2/7] Parallel fbank: CPU session pool + intra-op thread override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With multimask batching engaged, CPU fbank dominates the CUDA path: fbank_ms=28182 of ~37 s wall (76%) on a 36-min meeting (ES2004a, RTX 3080 Ti) vs gpu_predict_ms=3572. The single fbank session is capped at 4 intra-op threads, and intra-op threads do not scale it (4 -> 24 threads: 35.4 s -> 32.1 s). --chunk-emb-workers has no effect on the CUDA path. - session.rs: SPEAKRS_FBANK_THREADS env override for the intra-op cap. - load/sessions.rs: build split_fbank_pool: Vec (size SPEAKRS_FBANK_POOL, default available_parallelism/4 clamped 1-8) when the split backend is available. - fbank.rs: compute_chunk_fbanks_batch fans chunks across the pool with std::thread::scope; session-local input buffers, deterministic ordering. Measured: ES2004a end-to-end 39.4 s -> 12.9 s (3.1x) at pool=8, RTTM bit-identical; AMI test-16 aggregate DER unchanged. Env vars are the smallest possible surface — happy to rework into RuntimeConfig if preferred. --- src/inference/embedding.rs | 1 + src/inference/embedding/fbank.rs | 60 ++++++++++++++++++++++++ src/inference/embedding/load/sessions.rs | 21 +++++++++ src/inference/embedding/session.rs | 13 +++-- 4 files changed, 92 insertions(+), 3 deletions(-) diff --git a/src/inference/embedding.rs b/src/inference/embedding.rs index 2fd9512..ba5dd87 100644 --- a/src/inference/embedding.rs +++ b/src/inference/embedding.rs @@ -74,6 +74,7 @@ struct OrtEmbeddingState { session: Session, primary_batched_session: Option, split_fbank_session: Option, + split_fbank_pool: Vec, split_fbank_batched_session: Option, split_tail_session: Option, split_tail_batched_session: Option, diff --git a/src/inference/embedding/fbank.rs b/src/inference/embedding/fbank.rs index 59c66fd..5db2f5d 100644 --- a/src/inference/embedding/fbank.rs +++ b/src/inference/embedding/fbank.rs @@ -58,6 +58,44 @@ impl EmbeddingModel { &mut self, audios: &[&[f32]], ) -> Result>, ort::Error> { + // Fan per-chunk fbank out across a pool of CPU sessions. + // fbank is otherwise ~76% of CUDA E2E wall time (intra-op threads don't scale it). + if audios.len() > 1 && !self.ort.split_fbank_pool.is_empty() { + let window_samples = self.meta.window_samples; + let pool = &mut self.ort.split_fbank_pool; + let per_worker = audios.len().div_ceil(pool.len()); + let mut collected: Vec>>> = Vec::new(); + collected.resize_with(pool.len(), || None); + std::thread::scope(|scope| { + let mut handles = Vec::new(); + for (worker_idx, session) in pool.iter_mut().enumerate() { + let start = worker_idx * per_worker; + if start >= audios.len() { + break; + } + let end = (start + per_worker).min(audios.len()); + let slice = &audios[start..end]; + handles.push(( + worker_idx, + scope.spawn(move || -> Result>, ort::Error> { + slice + .iter() + .map(|audio| fbank_via_session(session, audio, window_samples)) + .collect() + }), + )); + } + for (worker_idx, handle) in handles { + let vals = handle + .join() + .map_err(|_| ort::Error::new("fbank pool worker panicked"))??; + collected[worker_idx] = Some(vals); + } + Ok::<(), ort::Error>(()) + })?; + return Ok(collected.into_iter().flatten().flatten().collect()); + } + let has_batched = self.has_batched_fbank(); if !has_batched { tracing::debug!( @@ -171,3 +209,25 @@ impl EmbeddingModel { Ok(true) } } + +// Session-local fbank for the parallel pool path (no shared buffers). +fn fbank_via_session( + session: &mut ort::session::Session, + audio: &[f32], + window_samples: usize, +) -> Result, ort::Error> { + let mut buf = ndarray::Array3::::zeros((1, 1, window_samples)); + let copy_len = audio.len().min(window_samples); + buf.slice_mut(s![0, 0, ..copy_len]) + .assign(&ndarray::ArrayView1::from(&audio[..copy_len])); + let waveform_tensor = TensorRef::from_array_view(buf.view())?; + let outputs = session.run(ort::inputs!["waveform" => waveform_tensor])?; + let output = first_output(outputs.values(), "pool chunk fbank output")?; + let (shape, data) = output.try_extract_tensor::()?; + array2_from_shape_vec( + shape[1] as usize, + shape[2] as usize, + data.to_vec(), + "pool chunk fbank output", + ) +} diff --git a/src/inference/embedding/load/sessions.rs b/src/inference/embedding/load/sessions.rs index ff5ae7f..4e55995 100644 --- a/src/inference/embedding/load/sessions.rs +++ b/src/inference/embedding/load/sessions.rs @@ -26,6 +26,7 @@ pub(super) struct LoadedOrtSessions { session: Session, primary_batched_session: Option, split_fbank_session: Option, + split_fbank_pool: Vec, split_fbank_batched_session: Option, split_tail_session: Option, split_tail_batched_session: Option, @@ -236,10 +237,29 @@ impl LoadedSessions { ); } + // Pool of extra CPU fbank sessions for parallel per-chunk fbank + // (single-session fbank measured at ~76% of CUDA E2E wall on many-core hosts). + let split_fbank_pool: Vec = if use_split_backend { + let pool_size = std::env::var("SPEAKRS_FBANK_POOL") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|c| (c.get() / 4).clamp(1, 8)) + .unwrap_or(1) + }); + (0..pool_size) + .map(|_| EmbeddingModel::build_fbank_session(&split_fbank_path, ExecutionMode::Cpu)) + .collect::, _>>()? + } else { + Vec::new() + }; + let ort = LoadedOrtSessions { session, primary_batched_session, split_fbank_session, + split_fbank_pool, split_fbank_batched_session, split_tail_session, split_tail_batched_session, @@ -288,6 +308,7 @@ impl LoadedSessions { session: self.ort.session, primary_batched_session: self.ort.primary_batched_session, split_fbank_session: self.ort.split_fbank_session, + split_fbank_pool: self.ort.split_fbank_pool, split_fbank_batched_session: self.ort.split_fbank_batched_session, split_tail_session: self.ort.split_tail_session, split_tail_batched_session: self.ort.split_tail_batched_session, diff --git a/src/inference/embedding/session.rs b/src/inference/embedding/session.rs index 751dd24..c0d1d13 100644 --- a/src/inference/embedding/session.rs +++ b/src/inference/embedding/session.rs @@ -62,9 +62,16 @@ impl EmbeddingModel { model_path: &Path, mode: ExecutionMode, ) -> Result { - let threads = std::thread::available_parallelism() - .map(|count| count.get().min(4)) - .unwrap_or(1); + // The default cap of 4 intra-op threads leaves fbank as ~76% of + // CUDA E2E wall time on many-core hosts; allow override via SPEAKRS_FBANK_THREADS. + let threads = std::env::var("SPEAKRS_FBANK_THREADS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|count| count.get().min(4)) + .unwrap_or(1) + }); let builder = Session::builder()? .with_independent_thread_pool()? .with_intra_threads(threads)? From 7687da705dbb157e752ac94ec7e61ce82ab3a2f3 Mon Sep 17 00:00:00 2001 From: davidamacey Date: Wed, 19 Aug 2026 21:02:41 -0400 Subject: [PATCH 3/7] Constant-fold the exported segmentation graphs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exported segmentation ONNX retains SincNet's filter-synthesis subgraph (Sin x2 / Cos x2 / If x1 recomputed from frozen parameters on every forward). On the ORT CUDA EP these ops trigger CPU fallback with Memcpy nodes inserted: 190 ms vs 96.6 ms per batch-32 (2.0x) folded vs unfolded in serving tests, and -7% end-to-end in the speakrs pipeline. Run onnxsim constant folding in export_models.py after export. Folding is bit-exact (max_abs_diff 0.0, argmax mismatch 0; graph shrinks 179 -> 40 nodes) and needs zero runtime changes — the same filenames load as before. --- scripts/export_models.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/scripts/export_models.py b/scripts/export_models.py index e6d01ac..1d1e361 100644 --- a/scripts/export_models.py +++ b/scripts/export_models.py @@ -6,6 +6,7 @@ # "numpy", # "onnx", # "onnxscript", +# "onnxsim", # ] # /// """Download and export ONNX models + PLDA params for speakrs. @@ -66,6 +67,18 @@ def main() -> None: print("Done!") +def fold_onnx_graph(path: str) -> None: + """Constant-fold a graph in place; a folding that changes outputs is a hard error.""" + import onnx + from onnxsim import simplify + + model = onnx.load(path) + simplified, ok = simplify(model) + if not ok: + raise RuntimeError(f"onnxsim could not validate the simplified graph for {path}") + onnx.save(simplified, path) + + def export_segmentation(pipeline: Any, models_dir: str) -> None: print("Exporting segmentation model...") seg_model = pipeline._segmentation.model @@ -102,6 +115,14 @@ def export_segmentation(pipeline: Any, models_dir: str) -> None: dynamo=False, ) + # Constant-fold the exported graphs: SincNet synthesizes its filterbank from frozen + # parameters every forward (Sin/Cos/If subgraph). On ORT's CUDA EP those ops fall back + # to CPU with Memcpy nodes inserted, costing 2x per batch-32 in serving tests. Folding + # is bit-exact (max_abs_diff 0.0, argmax mismatch 0) and shrinks the graph 179->40 nodes. + fold_onnx_graph(os.path.join(models_dir, "segmentation-3.0.onnx")) + fold_onnx_graph(os.path.join(models_dir, "segmentation-3.0-b32.onnx")) + fold_onnx_graph(os.path.join(models_dir, "segmentation-3.0-b64.onnx")) + sz = os.path.getsize(os.path.join(models_dir, "segmentation-3.0.onnx")) / 1e6 print(f" segmentation-3.0.onnx ({sz:.1f} MB)") bsz = os.path.getsize(os.path.join(models_dir, "segmentation-3.0-b32.onnx")) / 1e6 From a82f09d63c7760bf3ac0aa03f1530ba35bfa4a0b Mon Sep 17 00:00:00 2001 From: davidamacey Date: Tue, 25 Aug 2026 07:37:59 -0400 Subject: [PATCH 4/7] perf(ahc): bound pdist worker threads instead of one per block The blocked pairwise-distance computation spawned one OS thread per 1024-row block inside a single `std::thread::scope`, so the concurrent thread count scaled with meeting length rather than with the machine (n=21k => 21 threads, each additionally driving a multi-threaded BLAS `dot`). On low-core or shared hosts that oversubscribes, and every block's Gram scratch stayed live until the whole scope joined. Blocks are now pulled from a shared queue by a bounded worker pool, so peak concurrency and peak scratch memory scale with core count. The default is `available_parallelism()` capped at 8, overridable with `SPEAKRS_AHC_THREADS` (same env-override convention as `SPEAKRS_FBANK_THREADS`/`SPEAKRS_FBANK_POOL`). The block math is unchanged, so output is bit-identical for any worker count; a new unit test asserts that across worker counts 1/2/3/8/64. Raised by the Greptile review on avencera/speakrs#15. --- src/clustering/ahc.rs | 88 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 77 insertions(+), 11 deletions(-) diff --git a/src/clustering/ahc.rs b/src/clustering/ahc.rs index 77a5f31..6b364e9 100644 --- a/src/clustering/ahc.rs +++ b/src/clustering/ahc.rs @@ -45,6 +45,10 @@ pub fn cluster(embeddings: &ArrayView2, config: AhcConfig) -> Vec { } fn condensed_euclidean(embeddings: &Array2) -> Vec { + condensed_euclidean_with_workers(embeddings, pdist_worker_count()) +} + +fn condensed_euclidean_with_workers(embeddings: &Array2, workers: usize) -> Vec { // diar-native patch: blocked Gram-matrix formulation with scoped threads. // The original per-pair scalar loop cost 64.5 s at N=21k; dist^2 = |a|^2 + |b|^2 - 2ab // via matmul blocks is ~20x faster and each block writes a disjoint contiguous @@ -83,21 +87,35 @@ fn condensed_euclidean(embeddings: &Array2) -> Vec { } } + // Bounded worker pool: one thread per block would scale with meeting length + // (n=21k => 21 threads, each driving its own multi-threaded BLAS `dot`), which + // oversubscribes small/shared hosts and keeps every block's Gram matrix alive at + // once. Workers pull blocks from a shared queue instead, so peak concurrency and + // peak scratch memory scale with core count, not with n. + let workers = workers.min(blocks.len()).max(1); + let queue = std::sync::Mutex::new(blocks); std::thread::scope(|scope| { - for (bi, bi_end, slice) in blocks { + for _ in 0..workers { + let queue = &queue; let emb = &embeddings; let norms = &sq_norms; scope.spawn(move || { - let a = emb.slice(ndarray::s![bi..bi_end, ..]); - let b = emb.slice(ndarray::s![bi.., ..]); - let gram = a.dot(&b.t()); // (bi_end-bi) x (observations-bi) - let mut offset = 0usize; - for (local, i) in (bi..bi_end).enumerate() { - for j in (i + 1)..observations { - let dot = gram[[local, j - bi]]; - let d2 = (norms[i] + norms[j] - 2.0 * dot).max(0.0); - slice[offset] = d2.sqrt(); - offset += 1; + loop { + let next = queue.lock().expect("pdist queue poisoned").pop(); + let Some((bi, bi_end, slice)) = next else { + break; + }; + let a = emb.slice(ndarray::s![bi..bi_end, ..]); + let b = emb.slice(ndarray::s![bi.., ..]); + let gram = a.dot(&b.t()); // (bi_end-bi) x (observations-bi) + let mut offset = 0usize; + for (local, i) in (bi..bi_end).enumerate() { + for j in (i + 1)..observations { + let dot = gram[[local, j - bi]]; + let d2 = (norms[i] + norms[j] - 2.0 * dot).max(0.0); + slice[offset] = d2.sqrt(); + offset += 1; + } } } }); @@ -106,6 +124,23 @@ fn condensed_euclidean(embeddings: &Array2) -> Vec { condensed } +/// Number of concurrent workers used for the blocked pairwise-distance computation. +/// +/// Defaults to `available_parallelism()` capped at 8 (each worker also drives a +/// multi-threaded BLAS `dot`, so a higher cap oversubscribes rather than helps). +/// Override with `SPEAKRS_AHC_THREADS`; values are clamped to at least 1. +fn pdist_worker_count() -> usize { + std::env::var("SPEAKRS_AHC_THREADS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|c| c.get().min(8)) + .unwrap_or(1) + }) +} + fn flat_clusters(observations: usize, steps: &[Step], threshold: f32) -> Vec { if observations == 0 { return Vec::new(); @@ -221,6 +256,37 @@ mod tests { .join(name) } + #[test] + fn condensed_euclidean_is_bit_identical_across_worker_counts() { + // 2600 rows => 3 blocks at BLOCK=1024, so worker counts below/at/above the + // block count all get exercised. + let rows = 2600; + let cols = 16; + let data: Vec = (0..rows * cols) + .map(|i| ((i * 37 % 101) as f32 / 101.0) - 0.5) + .collect(); + let embeddings = Array2::from_shape_vec((rows, cols), data).unwrap(); + + let reference = condensed_euclidean_with_workers(&embeddings, 1); + for workers in [2, 3, 8, 64] { + let got = condensed_euclidean_with_workers(&embeddings, workers); + assert_eq!(got.len(), reference.len()); + assert!( + got.iter().zip(reference.iter()).all(|(a, b)| a.to_bits() == b.to_bits()), + "worker count {workers} changed pdist output" + ); + } + } + + #[test] + fn pdist_worker_count_is_bounded() { + let workers = pdist_worker_count(); + assert!(workers >= 1); + if std::env::var_os("SPEAKRS_AHC_THREADS").is_none() { + assert!(workers <= 8, "default worker count should stay bounded"); + } + } + #[test] fn separates_two_clusters() { let embeddings = array![[1.0, 0.0], [0.95, 0.05], [-1.0, 0.0], [-0.95, -0.05],]; From 0c4cdbb416568d68678c6d27ae71e7d0f2437336 Mon Sep 17 00:00:00 2001 From: davidamacey Date: Tue, 25 Aug 2026 07:37:59 -0400 Subject: [PATCH 5/7] fix(fbank): skip the CPU fbank session pool under CoreML MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `split_fbank_pool` was keyed only on split-backend availability with no execution-mode gate, so it was also built under `ExecutionMode::CoreMl` and `CoreMlFast`. `compute_chunk_fbanks_batch` takes the pool branch whenever the pool is non-empty, which early-returns before the native batched CoreML fbank path is ever reached — so every multi-chunk request under CoreML silently ran fbank on CPU sessions. Gating construction (rather than usage) also avoids loading 1-8 unused CPU ORT sessions at init under CoreML. The CUDA path is unaffected: it still builds and uses the pool exactly as before. Reported by the Greptile review on avencera/speakrs#15. --- src/inference/embedding/load/sessions.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/inference/embedding/load/sessions.rs b/src/inference/embedding/load/sessions.rs index 4e55995..d0f7423 100644 --- a/src/inference/embedding/load/sessions.rs +++ b/src/inference/embedding/load/sessions.rs @@ -239,7 +239,9 @@ impl LoadedSessions { // Pool of extra CPU fbank sessions for parallel per-chunk fbank // (single-session fbank measured at ~76% of CUDA E2E wall on many-core hosts). - let split_fbank_pool: Vec = if use_split_backend { + // CoreML modes have a native batched fbank path that the CPU pool would shadow, + // so the pool is skipped entirely there (also avoids loading unused CPU sessions). + let split_fbank_pool: Vec = if use_split_backend && !mode.is_coreml() { let pool_size = std::env::var("SPEAKRS_FBANK_POOL") .ok() .and_then(|v| v.parse::().ok()) From 0e7a12aff5ed2cfde44005008d917fd7999930fb Mon Sep 17 00:00:00 2001 From: attevon-admin Date: Wed, 2 Sep 2026 06:17:34 -0400 Subject: [PATCH 6/7] style: rustfmt the pdist worker-count assertion cargo fmt --check flags this line in the test added earlier in this series. No behaviour change. --- src/clustering/ahc.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/clustering/ahc.rs b/src/clustering/ahc.rs index 6b364e9..ca79eb6 100644 --- a/src/clustering/ahc.rs +++ b/src/clustering/ahc.rs @@ -272,7 +272,9 @@ mod tests { let got = condensed_euclidean_with_workers(&embeddings, workers); assert_eq!(got.len(), reference.len()); assert!( - got.iter().zip(reference.iter()).all(|(a, b)| a.to_bits() == b.to_bits()), + got.iter() + .zip(reference.iter()) + .all(|(a, b)| a.to_bits() == b.to_bits()), "worker count {workers} changed pdist output" ); } From d8f00a8466df83de4991f298b1d8f18fa6beb8df Mon Sep 17 00:00:00 2001 From: attevon-admin Date: Wed, 2 Sep 2026 06:17:34 -0400 Subject: [PATCH 7/7] feat(config): add RuntimeConfig::fbank_pool so embedders can size the pool without setenv The fbank session pool added earlier in this series is sized only from SPEAKRS_FBANK_POOL. That is fine for the CLI, but an embedding host cannot always use it: setenv is not thread-safe, so a caller that loads models lazily or loads several execution modes concurrently has no safe moment to set it. Setting it before the first load and hoping no other thread is running is the only option today, and that stops being true as soon as a second model loads. RuntimeConfig::fbank_pool makes the pool size an ordinary configuration value. None keeps exactly the current behaviour, environment override first and the cores/4 clamp otherwise, so nothing changes for existing callers. Some(n) sizes the pool directly and never reads the environment. Some(0) falls back to the single fbank session. Also adds a debug log of the chosen size, which was previously only inferable from process memory, and drops the now-unneeded let _ = config in the non-CoreML build since config is read on every path. --- src/inference/embedding/load/sessions.rs | 26 +++++++++++++++--------- src/pipeline/config.rs | 11 ++++++++++ 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/src/inference/embedding/load/sessions.rs b/src/inference/embedding/load/sessions.rs index d0f7423..42d00d1 100644 --- a/src/inference/embedding/load/sessions.rs +++ b/src/inference/embedding/load/sessions.rs @@ -55,6 +55,20 @@ pub(super) struct LoadedSessions { coreml: LoadedCoreMlState, } +/// Fallback sizing for the CPU fbank pool when `RuntimeConfig::fbank_pool` is `None`: +/// the `SPEAKRS_FBANK_POOL` override if it parses, else one session per four cores +/// (clamped to `1..=8`). Callers that set `fbank_pool` explicitly never reach the environment. +fn auto_fbank_pool_size() -> usize { + std::env::var("SPEAKRS_FBANK_POOL") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|c| (c.get() / 4).clamp(1, 8)) + .unwrap_or(1) + }) +} + impl LoadedSessions { pub(super) fn load( model_path: &Path, @@ -68,8 +82,6 @@ impl LoadedSessions { let split_primary_tail_batched_path = split_tail_model_path(model_path, PRIMARY_BATCH_SIZE); #[cfg(feature = "coreml")] let native_chunk_compute_units = config.chunk_emb_compute_units.to_ml_compute_units(); - #[cfg(not(feature = "coreml"))] - let _ = config; let use_split_backend = EmbeddingModel::split_backend_available(model_path); #[cfg(feature = "coreml")] @@ -242,14 +254,8 @@ impl LoadedSessions { // CoreML modes have a native batched fbank path that the CPU pool would shadow, // so the pool is skipped entirely there (also avoids loading unused CPU sessions). let split_fbank_pool: Vec = if use_split_backend && !mode.is_coreml() { - let pool_size = std::env::var("SPEAKRS_FBANK_POOL") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or_else(|| { - std::thread::available_parallelism() - .map(|c| (c.get() / 4).clamp(1, 8)) - .unwrap_or(1) - }); + let pool_size = config.fbank_pool.unwrap_or_else(auto_fbank_pool_size); + tracing::debug!(fbank_pool = pool_size, "fbank session pool"); (0..pool_size) .map(|_| EmbeddingModel::build_fbank_session(&split_fbank_path, ExecutionMode::Cpu)) .collect::, _>>()? diff --git a/src/pipeline/config.rs b/src/pipeline/config.rs index b457960..6d8e84a 100644 --- a/src/pipeline/config.rs +++ b/src/pipeline/config.rs @@ -77,6 +77,16 @@ impl PipelineConfig { pub struct RuntimeConfig { /// Number of chunk embedding workers pub chunk_emb_workers: usize, + /// Size of the CPU fbank session pool used for parallel per-chunk fbank. + /// + /// `None` (the default) auto-sizes: the `SPEAKRS_FBANK_POOL` environment override when it + /// parses, otherwise one session per four cores clamped to `1..=8`. `Some(0)` disables the + /// pool and falls back to the single fbank session. + /// + /// Setting this explicitly lets an embedder size the pool without touching the environment, + /// which matters because `setenv` is not thread-safe: a host that loads models lazily or + /// concurrently cannot safely use the environment override. + pub fbank_pool: Option, /// CoreML compute units for chunk embedding (CoreML modes only) #[cfg(feature = "coreml")] #[cfg_attr(docsrs, doc(cfg(feature = "coreml")))] @@ -87,6 +97,7 @@ impl Default for RuntimeConfig { fn default() -> Self { Self { chunk_emb_workers: 1, + fbank_pool: None, #[cfg(feature = "coreml")] chunk_emb_compute_units: CoreMlComputeUnits::All, }