Skip to content

perf: use available cores for embedding ONNX session intra-op threads - #6

Open
ryoma0421 wants to merge 1 commit into
avencera:masterfrom
ryoma0421:perf/embedding-intra-threads
Open

perf: use available cores for embedding ONNX session intra-op threads#6
ryoma0421 wants to merge 1 commit into
avencera:masterfrom
ryoma0421:perf/embedding-intra-threads

Conversation

@ryoma0421

@ryoma0421 ryoma0421 commented Aug 10, 2026

Copy link
Copy Markdown

What

The embedding ONNX session is created with with_intra_threads(1), which bottlenecks CPU-mode execution: on Apple Silicon the full pipeline runs at only ~1-2x realtime (measured with a 5.7-minute meeting recording; embedding inference dominates wall time). This PR lets ORT use up to 6 cores (available_parallelism().min(6), same cap as the segmentation session).

Measurements (M-series, CPU mode, 5.7 min Japanese meeting audio, 4 speakers)

config speed diarization output
intra_threads(1), minimal models 1.03x RT baseline
intra_threads(1), + batched models 2x RT identical
this patch, + batched models 8-9x RT identical

Speaker assignment agreement against a labeled ground-truth set is unchanged (96.1% in all three configs, same speaker count).

Happy to adjust the cap or gate it behind a config option if you prefer.

Summary by CodeRabbit

  • Performance Improvements
    • Improved embedding processing performance by automatically using available CPU parallelism, capped at six threads.
    • Retained a conservative fallback for environments where parallelism information is unavailable.

With intra_threads(1) the CPU pipeline is bottlenecked by embedding
inference and runs at ~1-2x realtime on Apple Silicon. Allowing ORT to
use up to 6 cores brings the full pipeline to ~8-9x realtime (measured
on M-series with a 5.7-minute meeting recording, batched ONNX models
present) with identical diarization output.
@coderabbitai

coderabbitai Bot commented Aug 10, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: cdb397bb-994d-40b7-8422-1b9d144adb4a

📥 Commits

Reviewing files that changed from the base of the PR and between b0756b1 and e70e3ea.

📒 Files selected for processing (1)
  • src/inference/embedding/session.rs

📝 Walkthrough

Walkthrough

build_session_with_graph now configures ONNX Runtime intra-op execution with available parallelism capped at six threads. It uses one thread when parallelism detection is unavailable. Inter-op execution remains single-threaded.

Changes

Embedding session threading

Layer / File(s) Summary
Configure session thread counts
src/inference/embedding/session.rs
build_session_with_graph derives the intra-op thread count from available parallelism, caps it at six, and uses one as the fallback. Inter-op threading remains set to one.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main performance change: using available CPU cores for ONNX session intra-operation threads.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@greptile-apps

greptile-apps Bot commented Aug 10, 2026

Copy link
Copy Markdown

Greptile Summary

This PR removes the single-thread bottleneck from embedding inference by sizing the ONNX Runtime intra-op thread pool from available CPU parallelism, capped at six threads with a one-thread fallback.

  • Uses std::thread::available_parallelism() to select the embedding session’s intra-op thread count.
  • Retains the existing independent thread pool, single inter-op thread, and memory-pattern configuration.
  • Aligns embedding inference with the bounded threading strategy already used by segmentation inference.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness, security, or resource-management failure identified.

The new thread count is bounded between one and six, preserves a safe fallback when parallelism cannot be determined, and does not alter inference outputs or session selection behavior.

Important Files Changed

Filename Overview
src/inference/embedding/session.rs Replaces the fixed one-thread embedding session configuration with a bounded available-core count; no actionable correctness issue was identified.

Reviews (1): Last reviewed commit: "perf: use available cores for embedding ..." | Re-trigger Greptile

@praveenperera
praveenperera force-pushed the master branch 2 times, most recently from 34993d5 to b0756b1 Compare August 20, 2026 23:33
attevon-admin pushed a commit to attevon-llc/diar-native that referenced this pull request Aug 25, 2026
The embedding ONNX session builder (build_session_with_graph, used for the
tail/multimask/primary embedding models) hardcoded .with_intra_threads(1),
leaving the embedding tail single-threaded. Under ExecutionMode::Cpu -- the
mode our CPU-only image tier (docker/Dockerfile.server-cpu) runs in -- that
tail dominates wall time. CUDA/CoreML were unaffected because the heavy ops
are off-CPU there.

Replace the constant with available_parallelism().min(6), overridable via
SPEAKRS_INTRA_THREADS, matching the existing SPEAKRS_FBANK_THREADS pattern in
the same file. The cap of 6 is the same one the already-shipped segmentation
session builder uses, so both model families now scale identically.

Oversubscription: several embedding sessions are built per pipeline and each
has an independent thread pool, but only one executes at a time within a
request, so concurrent thread demand is bounded by inflight-requests x 6, not
sessions x 6 -- the same bound segmentation already imposes.

Measured on AMI EN2002c, first 360s, 16 kHz mono, in diar-bench-builder
(48-core host, no GPU visible to the container), median of 3 alternating runs:

  CPU mode  intra=1: 218.9s / 217.2s / 220.1s  (1.64x realtime)
  CPU mode  intra=6:  57.8s /  59.2s /  67.7s  (6.1x realtime)  -> 3.7x faster

CUDA regression check (RTX A6000, steady state after warmup):

  CUDA mode intra=1: 4.83s / 4.86s
  CUDA mode intra=6: 4.88s / 4.85s   -> no regression

RTTM output is bit-identical across thread counts in both modes (single md5
across all legs), confirming this is a scheduling change only. speakrs suite:
96 passed, 0 failed.

Approach and the original CPU-mode measurement come from upstream PR
avencera/speakrs#6 by @ryoma0421; adopted here into our fork rather than
waiting on upstream merge.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant