perf(kv-router): serve the extension's Rust heap from a linked-in jemalloc

components/src/dynamo/frontend/__main__.py

@@ -51,6 +51,12 @@ def _maybe_preload_jemalloc() -> None:
 if __name__ == "__main__":
     _maybe_preload_jemalloc()
 
+    # dynamo._core keeps its mimalloc heap off transparent huge pages, which bloat
+    # many-threaded workers. The frontend's dense heap serves more requests at saturation
+    # on huge pages, so opt back in unless the operator set it. mimalloc reads this once,
+    # when the extension loads.
+    os.environ.setdefault("MIMALLOC_ALLOW_THP", "1")
+
     # Import the runtime only after configuring the allocator.
     from dynamo.frontend.main import main
 

lib/bench/Cargo.toml

@@ -60,7 +60,9 @@ multiturn = [
   "dep:rand",
 ]
 offline-replay = ["dep:blake3", "dep:dynamo-mocker", "dep:libc"]
-router-bench = ["dynamo-kv-router/bench"]
+# Router benches link mimalloc as the global allocator, matching the Python
+# extension that runs the router in production (`lib/bindings/python`).
+router-bench = ["dynamo-kv-router/bench", "dep:mimalloc", "dep:libmimalloc-sys"]
 router-metrics = ["dynamo-kv-router/metrics"]
 router-test-support = [
   "router-bench",
@@ -162,6 +164,8 @@ flate2 = { version = "1", optional = true }
 futures-util = { workspace = true, optional = true }
 hf-hub = { workspace = true, features = ["ureq"], optional = true }
 indicatif = { version = "0.18", optional = true }
+mimalloc = { version = "0.1.52", default-features = false, features = ["local_dynamic_tls"], optional = true }
+libmimalloc-sys = { version = "0.1.49", default-features = false, features = ["extended"], optional = true }
 rand = { workspace = true, optional = true }
 reqwest = { workspace = true, optional = true }
 rustc-hash = { version = "2", optional = true }

lib/bench/kv_router/INDEXER_BENCH.md

@@ -28,7 +28,7 @@ Preparation consumes and releases the parsed trace, generated replay artifacts,
 worker timelines, and intermediate payload owners. The timed process retains
 only the prepared schedule, one flattened query-hash slab, and backend state.
 On Linux, the benchmark returns free preparation pages with `malloc_trim` and
-waits a fixed five seconds before constructing backend workers. It then
+`mi_collect` and waits a fixed five seconds before constructing backend workers. It then
 page-touches the schedule once and performs the fixed lookup warm-up before
 timing. This explicit quiescence prevents the parallel event-generation phase
 and released allocator arenas from affecting the timed backend drain; JSON
@@ -89,6 +89,16 @@ the engine re-chunks synthesized prompt tokens at that size.
 Linux uses absolute `CLOCK_MONOTONIC` sleeps followed by the configured spin.
 macOS uses a portable sleep-plus-spin timer for correctness tests only.
 
+The router benches (`mooncake_bench`, `active_sequences_bench`,
+`approximate_lru_bench`) use mimalloc as the global allocator, matching the
+Python extension that runs the router in production (`lib/bindings/python`).
+They keep mimalloc's default transparent-huge-page advice, as `dynamo.frontend`
+does; other processes that load the extension turn it off.
+The allocator is part of the measured system. Under glibc malloc, CRTC event
+workers serialize on arena locks: write capacity stops rising after about 16
+event workers and falls beyond that. Compare allocators or event-worker counts
+only between builds and runs that differ in that one variable.
+
 ### Single CRTC trial
 
 Select one logical CPU per physical core from the process's actual affinity

lib/bench/kv_router/active_sequences_bench.rs

@@ -1,6 +1,11 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+// Match the production allocator: the router runs inside the Python extension,
+// which routes Rust allocations through mimalloc (`lib/bindings/python`).
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 #[path = "active_sequences_open_loop.rs"]
 mod active_sequences_open_loop;
 #[path = "active_sequences_shared.rs"]

lib/bench/kv_router/approximate_lru_bench.rs

@@ -1,6 +1,11 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+// Match the production allocator: the router runs inside the Python extension,
+// which routes Rust allocations through mimalloc (`lib/bindings/python`).
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 use std::collections::{HashMap, HashSet};
 use std::hint::black_box;
 use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
@@ -1149,6 +1154,7 @@ async fn quiesce(milliseconds: u64) {
     #[cfg(target_os = "linux")]
     unsafe {
         libc::malloc_trim(0);
+        libmimalloc_sys::mi_collect(true);
     }
     if milliseconds > 0 {
         tokio::time::sleep(Duration::from_millis(milliseconds)).await;

lib/bench/kv_router/mooncake_bench.rs

@@ -1,6 +1,11 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+// Match the production allocator: the router runs inside the Python extension,
+// which routes Rust allocations through mimalloc (`lib/bindings/python`).
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 #[path = "mooncake_open_loop.rs"]
 mod mooncake_open_loop;
 #[path = "mooncake_shared.rs"]
@@ -305,6 +310,7 @@ fn quiesce_prepared_heap() {
     // Return their free pages and let reclamation settle before backend workers start.
     unsafe {
         libc::malloc_trim(0);
+        libmimalloc_sys::mi_collect(true);
     }
     std::thread::sleep(std::time::Duration::from_millis(PRE_RUN_QUIESCENCE_MS));
 }

lib/bindings/python/Cargo.toml

@@ -22,7 +22,13 @@ name = "_core"
 crate-type = ["cdylib", "rlib"]
 
 [features]
-default = ["custom-policy"]
+default = ["custom-policy", "mimalloc"]
+# Route the extension's Rust allocations through mimalloc. glibc's arena locks
+# serialize the KV router's multi-threaded event workers; Python and C/C++ libraries
+# in the process keep their own allocators. Local-dynamic TLS keeps mimalloc safe
+# to load through `dlopen` without consuming the static TLS reserve. mimalloc 0.1.52
+# builds the mimalloc v3 series (its `v2` feature selects 2.x).
+mimalloc = ["dep:mimalloc", "dep:libmimalloc-sys"]
 # Links the router plugins Dynamo ships plus the replaceable catalog selected by the
 # dependency alias below. Enabled by default so a deployment can select a shipped policy through
 # router-policy YAML without rebuilding. A private image replaces only the catalog dependency, so
@@ -75,6 +81,8 @@ async-trait = { version = "0.1" }
 bytes = { version = "1" }
 dashmap = { version = "6.1" }
 futures = { version = "0.3" }
+mimalloc = { version = "0.1.52", default-features = false, features = ["local_dynamic_tls"], optional = true }
+libmimalloc-sys = { version = "0.1.49", default-features = false, features = ["extended"], optional = true }
 rmp = { version = "0.8" }
 rmpv = { version = "1", features = ["with-serde"] }
 once_cell = { version = "1.20.3" }

lib/bindings/python/rust/lib.rs

@@ -1,6 +1,34 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+#[cfg(feature = "mimalloc")]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+/// Keeps mimalloc's arenas out of transparent huge pages unless `MIMALLOC_ALLOW_THP` says
+/// otherwise; glibc, which held these allocations before, never asked for them.
+///
+/// mimalloc advises each arena it reserves for huge pages. In a many-threaded process each
+/// thread's sparse pages then fault in whole 2 MiB pages, and khugepaged refills partly freed
+/// ones while idle. `dynamo.frontend` opts back in for its dense heap.
+///
+/// Must run before the extension's first Rust allocation, which reserves the first arena.
+#[cfg(feature = "mimalloc")]
+fn configure_allocator() {
+    // libmimalloc-sys does not export this option; pin the mimalloc 3.3 enum it indexes.
+    const _: () = assert!(libmimalloc_sys::_mi_option_last == 47);
+    const MI_OPTION_ALLOW_THP: libmimalloc_sys::mi_option_t = 43;
+    // SAFETY: both calls only touch mimalloc's option state, and module init runs before
+    // this extension starts any thread. Initializing first keeps the environment
+    // authoritative and means mimalloc saw THP allowed at startup, so it skips the
+    // process-wide PR_SET_THP_DISABLE and leaves other libraries' memory to the host's
+    // THP policy.
+    unsafe {
+        libmimalloc_sys::mi_process_init();
+        libmimalloc_sys::mi_option_set_default(MI_OPTION_ALLOW_THP, 0);
+    }
+}
+
 use dynamo_llm::local_model::{
     LocalModel, register_model_card, update_model_taints as update_model_taints_rs,
 };
@@ -285,6 +313,14 @@ fn wait_for_bridge_tasks_at_exit(py: Python<'_>) {
 }
 
 fn register_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // Model teardown trims glibc arenas, which no longer hold this extension's Rust
+    // allocations; also return mimalloc's freed memory.
+    #[cfg(feature = "mimalloc")]
+    dynamo_llm::discovery::register_allocator_trim_hook(|| {
+        // SAFETY: mi_collect only releases memory mimalloc already considers free.
+        unsafe { libmimalloc_sys::mi_collect(true) }
+    });
+
     // OTLP export no longer requires a pre-existing runtime, so initialize at import.
     if std::env::var_os(SKIP_PYTHON_LOG_INIT_ENV).is_none() {
         rs::logging::init();
@@ -468,13 +504,17 @@ fn register_core_with_router_plugins(m: &Bound<'_, PyModule>) -> PyResult<()> {
 #[cfg(feature = "custom-policy")]
 #[pymodule]
 fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    #[cfg(feature = "mimalloc")]
+    configure_allocator();
     register_core_with_router_plugins(m)
 }
 
 /// The stock extension-module entrypoint.
 #[cfg(not(feature = "custom-policy"))]
 #[pymodule]
 fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    #[cfg(feature = "mimalloc")]
+    configure_allocator();
     register_core(m)
 }