From e6494add09b840ef1dc20efd8506a9f3aeee836f Mon Sep 17 00:00:00 2001 From: Paul Sorensen Date: Fri, 11 Sep 2026 20:57:27 -0700 Subject: [PATCH 1/2] feat: expose ONNX Runtime session config entries through InitOptions fastembed builds the ONNX Runtime SessionBuilder internally. Callers had no way to set session-only configuration entries such as mlas.disable_kleidiai, which controls a per-session GEMM buffer switch on Apple SME hardware. Add a session_config field to every init-options struct. Add a with_session_config(key, value) builder method that appends one entry. init_session_builder applies each entry with SessionBuilder::with_config_entry after setting intra-op threads and before the DirectML block. Update all 10 call sites that build a session from init options. All From impls between default and user-defined option structs now carry session_config across. --- README.md | 1 + src/bgem3_embedding/impl.rs | 9 ++++++--- src/common.rs | 16 ++++++++++++++++ src/image_embedding/impl.rs | 6 ++++-- src/image_embedding/init.rs | 13 +++++++++++++ src/init.rs | 26 ++++++++++++++++++++++++++ src/reranking/impl.rs | 7 +++++-- src/reranking/init.rs | 31 +++++++++++++++++++++++++++++++ src/sparse_text_embedding/impl.rs | 3 ++- src/text_embedding/impl.rs | 7 +++++-- src/text_embedding/init.rs | 31 +++++++++++++++++++++++++++++++ 11 files changed, 140 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 9ad6096..c759d91 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ let mut model = TextEmbedding::try_new(Default::default())?; let mut model = TextEmbedding::try_new( TextInitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true).with_intra_threads(4), )?; +// Also available: `.with_session_config(key, value)` to pass raw ONNX Runtime session config entries. let documents = vec![ "passage: Hello, World!", diff --git a/src/bgem3_embedding/impl.rs b/src/bgem3_embedding/impl.rs index e0e7f69..f934be0 100644 --- a/src/bgem3_embedding/impl.rs +++ b/src/bgem3_embedding/impl.rs @@ -35,6 +35,7 @@ impl Bgem3Embedding { show_download_progress, execution_providers, intra_threads, + session_config, } = options; let model_repo = Bgem3Embedding::retrieve_model( @@ -63,7 +64,7 @@ impl Bgem3Embedding { } } - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_file_reference)?; let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?; @@ -79,10 +80,11 @@ impl Bgem3Embedding { execution_providers, max_length, intra_threads, + session_config, .. } = options; - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_memory(&model.onnx_file)?; let tokenizer = load_tokenizer(model.tokenizer_files, max_length)?; @@ -100,10 +102,11 @@ impl Bgem3Embedding { execution_providers, max_length, intra_threads, + session_config, .. } = options; - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_path.as_ref().join("model.onnx"))?; let tokenizer = load_tokenizer(tokenizer_files, max_length)?; diff --git a/src/common.rs b/src/common.rs index 2cad637..026c859 100644 --- a/src/common.rs +++ b/src/common.rs @@ -266,6 +266,7 @@ pub fn pull_from_hf( pub(crate) fn init_session_builder( execution_providers: Vec, intra_threads: Option, + session_config: Vec<(String, String)>, ) -> Result { let threads = match intra_threads { Some(n) => n, @@ -289,6 +290,12 @@ pub(crate) fn init_session_builder( .with_intra_threads(threads) .map_err(builder_error)?; + for (key, value) in session_config { + builder = builder + .with_config_entry(&key, &value) + .map_err(builder_error)?; + } + if has_directml { builder = builder .with_memory_pattern(false) @@ -362,4 +369,13 @@ mod tests { "error message was: {err}" ); } + #[test] + fn init_session_builder_applies_config_entry() { + let builder = init_session_builder( + vec![], + Some(1), + vec![("session.disable_prepacking".into(), "1".into())], + ); + assert!(builder.is_ok()); + } } diff --git a/src/image_embedding/impl.rs b/src/image_embedding/impl.rs index 6c0008d..29d2fbe 100644 --- a/src/image_embedding/impl.rs +++ b/src/image_embedding/impl.rs @@ -35,6 +35,7 @@ impl ImageEmbedding { cache_dir, show_download_progress, intra_threads, + session_config, } = options; let model_repo = ImageEmbedding::retrieve_model( @@ -61,7 +62,7 @@ impl ImageEmbedding { source: Box::new(e), })?; - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_file_reference)?; Ok(Self::new(preprocessor, session)) @@ -77,11 +78,12 @@ impl ImageEmbedding { let ImageInitOptionsUserDefined { execution_providers, intra_threads, + session_config, } = options; let preprocessor = Compose::from_bytes(model.preprocessor_file)?; - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_memory(&model.onnx_file)?; Ok(Self::new(preprocessor, session)) diff --git a/src/image_embedding/init.rs b/src/image_embedding/init.rs index b68828a..401e4a1 100644 --- a/src/image_embedding/init.rs +++ b/src/image_embedding/init.rs @@ -16,6 +16,10 @@ pub struct ImageInitOptionsUserDefined { /// every available CPU core via `std::thread::available_parallelism`. /// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput. pub intra_threads: Option, + /// ONNX Runtime session configuration entries, applied with + /// `SessionBuilder::with_config_entry`. Use this for settings that have + /// no dedicated builder method, such as `mlas.disable_kleidiai`. + pub session_config: Vec<(String, String)>, } impl ImageInitOptionsUserDefined { @@ -38,6 +42,14 @@ impl ImageInitOptionsUserDefined { self.intra_threads = Some(intra_threads); self } + + /// Add an ONNX Runtime session configuration entry, applied with + /// `SessionBuilder::with_config_entry`. Call it once per entry. + /// Example: `.with_session_config("mlas.disable_kleidiai", "1")`. + pub fn with_session_config(mut self, key: impl Into, value: impl Into) -> Self { + self.session_config.push((key.into(), value.into())); + self + } } /// Convert ImageInitOptions to ImageInitOptionsUserDefined @@ -48,6 +60,7 @@ impl From for ImageInitOptionsUserDefined { ImageInitOptionsUserDefined { execution_providers: options.execution_providers, intra_threads: options.intra_threads, + session_config: options.session_config, } } } diff --git a/src/init.rs b/src/init.rs index 5f445f4..6f7572e 100644 --- a/src/init.rs +++ b/src/init.rs @@ -18,6 +18,10 @@ pub struct InitOptionsWithLength { /// every available CPU core via `std::thread::available_parallelism`. /// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput. pub intra_threads: Option, + /// ONNX Runtime session configuration entries, applied with + /// `SessionBuilder::with_config_entry`. Use this for settings that have + /// no dedicated builder method, such as `mlas.disable_kleidiai`. + pub session_config: Vec<(String, String)>, } #[derive(Debug, Clone)] @@ -31,6 +35,10 @@ pub struct InitOptions { /// every available CPU core via `std::thread::available_parallelism`. /// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput. pub intra_threads: Option, + /// ONNX Runtime session configuration entries, applied with + /// `SessionBuilder::with_config_entry`. Use this for settings that have + /// no dedicated builder method, such as `mlas.disable_kleidiai`. + pub session_config: Vec<(String, String)>, } impl Default for InitOptionsWithLength { @@ -42,6 +50,7 @@ impl Default for InitOptionsWithLength { show_download_progress: true, max_length: M::MAX_LENGTH, intra_threads: None, + session_config: Vec::new(), } } } @@ -54,6 +63,7 @@ impl Default for InitOptions { cache_dir: get_cache_dir().into(), show_download_progress: true, intra_threads: None, + session_config: Vec::new(), } } } @@ -96,6 +106,14 @@ impl InitOptionsWithLength { self } + /// Add an ONNX Runtime session configuration entry, applied with + /// `SessionBuilder::with_config_entry`. Call it once per entry. + /// Example: `.with_session_config("mlas.disable_kleidiai", "1")`. + pub fn with_session_config(mut self, key: impl Into, value: impl Into) -> Self { + self.session_config.push((key.into(), value.into())); + self + } + /// Set whether to show download progress pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self { self.show_download_progress = show_download_progress; @@ -135,6 +153,14 @@ impl InitOptions { self } + /// Add an ONNX Runtime session configuration entry, applied with + /// `SessionBuilder::with_config_entry`. Call it once per entry. + /// Example: `.with_session_config("mlas.disable_kleidiai", "1")`. + pub fn with_session_config(mut self, key: impl Into, value: impl Into) -> Self { + self.session_config.push((key.into(), value.into())); + self + } + /// Set whether to show download progress pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self { self.show_download_progress = show_download_progress; diff --git a/src/reranking/impl.rs b/src/reranking/impl.rs index abdd81f..7a5dd8c 100644 --- a/src/reranking/impl.rs +++ b/src/reranking/impl.rs @@ -52,6 +52,7 @@ impl TextRerank { cache_dir, show_download_progress, intra_threads, + session_config, } = options; let model_repo = pull_from_hf(model_name.to_string(), cache_dir, show_download_progress)?; @@ -75,7 +76,7 @@ impl TextRerank { })?; } - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_file_reference)?; let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?; @@ -95,9 +96,11 @@ impl TextRerank { intra_threads, disable_cpu_fallback, dimension_overrides, + session_config, } = options; - let mut session_builder = init_session_builder(execution_providers, intra_threads)?; + let mut session_builder = + init_session_builder(execution_providers, intra_threads, session_config)?; let builder_error = |err: ort::Error| { Error::OrtBuilder(err.to_string()) }; diff --git a/src/reranking/init.rs b/src/reranking/init.rs index f6d7d9c..1685bda 100644 --- a/src/reranking/init.rs +++ b/src/reranking/init.rs @@ -39,6 +39,10 @@ pub struct RerankInitOptionsUserDefined { /// Override named free dimensions before ORT optimizes and places the /// user-defined reranker graph. pub dimension_overrides: Vec<(String, i64)>, + /// ONNX Runtime session configuration entries, applied with + /// `SessionBuilder::with_config_entry`. Use this for settings that have + /// no dedicated builder method, such as `mlas.disable_kleidiai`. + pub session_config: Vec<(String, String)>, } impl Default for RerankInitOptionsUserDefined { @@ -49,6 +53,7 @@ impl Default for RerankInitOptionsUserDefined { intra_threads: None, disable_cpu_fallback: false, dimension_overrides: Vec::new(), + session_config: Vec::new(), } } } @@ -90,6 +95,14 @@ impl RerankInitOptionsUserDefined { self.dimension_overrides.push((name.into(), size)); self } + + /// Add an ONNX Runtime session configuration entry, applied with + /// `SessionBuilder::with_config_entry`. Call it once per entry. + /// Example: `.with_session_config("mlas.disable_kleidiai", "1")`. + pub fn with_session_config(mut self, key: impl Into, value: impl Into) -> Self { + self.session_config.push((key.into(), value.into())); + self + } } /// Convert RerankInitOptions to RerankInitOptionsUserDefined @@ -103,6 +116,7 @@ impl From for RerankInitOptionsUserDefined { intra_threads: options.intra_threads, disable_cpu_fallback: false, dimension_overrides: Vec::new(), + session_config: options.session_config, } } } @@ -171,4 +185,21 @@ mod tests { assert!(o.disable_cpu_fallback); assert_eq!(o.dimension_overrides, vec![("sequence_length".into(), 128)]); } + + #[test] + fn session_config_is_collected_and_carried_by_from() { + let opts = RerankInitOptions::new(RerankerModel::default()) + .with_session_config("a", "1") + .with_session_config("b", "2"); + assert_eq!( + opts.session_config, + vec![("a".into(), "1".into()), ("b".into(), "2".into())] + ); + + let user_defined = RerankInitOptionsUserDefined::from(opts); + assert_eq!( + user_defined.session_config, + vec![("a".into(), "1".into()), ("b".into(), "2".into())] + ); + } } diff --git a/src/sparse_text_embedding/impl.rs b/src/sparse_text_embedding/impl.rs index 000c615..f560c66 100644 --- a/src/sparse_text_embedding/impl.rs +++ b/src/sparse_text_embedding/impl.rs @@ -36,6 +36,7 @@ impl SparseTextEmbedding { show_download_progress, execution_providers, intra_threads, + session_config, } = options; let model_repo = SparseTextEmbedding::retrieve_model( @@ -64,7 +65,7 @@ impl SparseTextEmbedding { } } - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_file_reference)?; let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?; diff --git a/src/text_embedding/impl.rs b/src/text_embedding/impl.rs index 1e40323..e9a5f80 100644 --- a/src/text_embedding/impl.rs +++ b/src/text_embedding/impl.rs @@ -38,6 +38,7 @@ impl TextEmbedding { cache_dir, show_download_progress, intra_threads, + session_config, } = options; let model_repo = TextEmbedding::retrieve_model( @@ -68,7 +69,7 @@ impl TextEmbedding { // prioritise loading pooling config if available, if not (thanks qdrant!), look for it in hardcoded let post_processing = TextEmbedding::get_default_pooling_method(&model_name); - let session = init_session_builder(execution_providers, intra_threads)? + let session = init_session_builder(execution_providers, intra_threads, session_config)? .commit_from_file(model_file_reference)?; let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?; @@ -94,13 +95,15 @@ impl TextEmbedding { intra_threads, disable_cpu_fallback, dimension_overrides, + session_config, } = options; let session = { let builder_error = |err: ort::Error| { Error::OrtBuilder(err.to_string()) }; - let mut session_builder = init_session_builder(execution_providers, intra_threads)?; + let mut session_builder = + init_session_builder(execution_providers, intra_threads, session_config)?; if disable_cpu_fallback { session_builder = session_builder diff --git a/src/text_embedding/init.rs b/src/text_embedding/init.rs index 7d599a4..0ad313d 100644 --- a/src/text_embedding/init.rs +++ b/src/text_embedding/init.rs @@ -39,6 +39,10 @@ pub struct InitOptionsUserDefined { /// graph. Static-shape execution providers such as QNN can use this with /// one user-defined model session per admitted shape. pub dimension_overrides: Vec<(String, i64)>, + /// ONNX Runtime session configuration entries, applied with + /// `SessionBuilder::with_config_entry`. Use this for settings that have + /// no dedicated builder method, such as `mlas.disable_kleidiai`. + pub session_config: Vec<(String, String)>, } impl InitOptionsUserDefined { @@ -78,6 +82,14 @@ impl InitOptionsUserDefined { self.dimension_overrides.push((name.into(), size)); self } + + /// Add an ONNX Runtime session configuration entry, applied with + /// `SessionBuilder::with_config_entry`. Call it once per entry. + /// Example: `.with_session_config("mlas.disable_kleidiai", "1")`. + pub fn with_session_config(mut self, key: impl Into, value: impl Into) -> Self { + self.session_config.push((key.into(), value.into())); + self + } } impl Default for InitOptionsUserDefined { @@ -88,6 +100,7 @@ impl Default for InitOptionsUserDefined { intra_threads: None, disable_cpu_fallback: false, dimension_overrides: Vec::new(), + session_config: Vec::new(), } } } @@ -103,6 +116,7 @@ impl From for InitOptionsUserDefined { intra_threads: options.intra_threads, disable_cpu_fallback: false, dimension_overrides: Vec::new(), + session_config: options.session_config, } } } @@ -192,4 +206,21 @@ mod tests { ] ); } + + #[test] + fn session_config_is_collected_and_carried_by_from() { + let opts = TextInitOptions::new(EmbeddingModel::AllMiniLML6V2) + .with_session_config("a", "1") + .with_session_config("b", "2"); + assert_eq!( + opts.session_config, + vec![("a".into(), "1".into()), ("b".into(), "2".into())] + ); + + let user_defined = InitOptionsUserDefined::from(opts); + assert_eq!( + user_defined.session_config, + vec![("a".into(), "1".into()), ("b".into(), "2".into())] + ); + } } From 764920b9b2e8bc55345e3c2715e9bf1ae18a1431 Mon Sep 17 00:00:00 2001 From: Anush Date: Sat, 12 Sep 2026 11:20:28 +0530 Subject: [PATCH 2/2] Update README.md --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index c759d91..9ad6096 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,6 @@ let mut model = TextEmbedding::try_new(Default::default())?; let mut model = TextEmbedding::try_new( TextInitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true).with_intra_threads(4), )?; -// Also available: `.with_session_config(key, value)` to pass raw ONNX Runtime session config entries. let documents = vec![ "passage: Hello, World!",