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,30 +26,83 @@ 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> {
+    // 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;
         }
     }
+
+    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
 }
 

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

src/inference/embedding.rs

@@ -74,6 +74,7 @@ struct OrtEmbeddingState {
     session: Session,
     primary_batched_session: Option<Session>,
     split_fbank_session: Option<Session>,
+    split_fbank_pool: Vec<Session>,
     split_fbank_batched_session: Option<Session>,
     split_tail_session: Option<Session>,
     split_tail_batched_session: Option<Session>,

src/inference/embedding/fbank.rs

@@ -58,6 +58,44 @@ impl EmbeddingModel {
         &mut self,
         audios: &[&[f32]],
     ) -> Result<Vec<Array2<f32>>, 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<Option<Vec<Array2<f32>>>> = 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<Vec<Array2<f32>>, 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<Array2<f32>, ort::Error> {
+    let mut buf = ndarray::Array3::<f32>::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::<f32>()?;
+    array2_from_shape_vec(
+        shape[1] as usize,
+        shape[2] as usize,
+        data.to_vec(),
+        "pool chunk fbank output",
+    )
+}