CUDA pipeline performance series: vectorized VBx + threaded pdist, fbank session pool, folded segmentation export

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 }

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

src/clustering/ahc.rs

@@ -26,33 +26,121 @@ pub fn cluster(embeddings: &ArrayView2<f32>, config: AhcConfig) -> Vec<usize> {
     }
 
     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<f32>) -> Vec<f32> {
+    condensed_euclidean_with_workers(embeddings, pdist_worker_count())
+}
+
+fn condensed_euclidean_with_workers(embeddings: &Array2<f32>, workers: usize) -> Vec<f32> {
+    // 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::<f32>()
-                .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<f32> = 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<i}(n-1-r) = i*(n-1) - i*(i-1)/2
+    let seg_start = |i: usize| i * (observations - 1) - i * i.saturating_sub(1) / 2;
+
+    // hand each block its contiguous slice
+    let mut blocks: Vec<(usize, usize, &mut [f32])> = 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;
         }
     }
+
+    // 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 _ in 0..workers {
+            let queue = &queue;
+            let emb = &embeddings;
+            let norms = &sq_norms;
+            scope.spawn(move || {
+                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;
+                        }
+                    }
+                }
+            });
+        }
+    });
     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::<usize>().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<f32>], threshold: f32) -> Vec<usize> {
     if observations == 0 {
         return Vec::new();
@@ -168,6 +256,39 @@ 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<f32> = (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],];

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<f64> = 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::<f64>::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::<f64>::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::<f64>::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<f64> = (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<f64> = pi.mapv(|p| (p + 1e-8).ln());
 
-        // log_p_x[sample_idx] = logsumexp(log_p[sample_idx] + lpi)
         let mut log_p_x = Array1::<f64>::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