feat: use jemallocator instead of default malloc

Cargo.toml

@@ -106,6 +106,10 @@ ts-rs = { git = "https://github.com/arilotter/ts-rs.git", rev = "92ce1752227fec9
 ] }
 rayon = "1.10.0"
 
+tikv-jemallocator = { version = "0.6.1", features = [
+  "unprefixed_malloc_on_supported_platforms",
+] }
+
 # dev only
 test-log = { version = "0.2.16", features = ["trace"] }
 pretty_assertions = "1.4.1"

architectures/centralized/client/Cargo.toml

@@ -26,6 +26,7 @@ bytemuck.workspace = true
 clap-markdown.workspace = true
 hex = "0.4.3"
 psyche-python-extension-impl = { workspace = true, optional = true }
+tikv-jemallocator.workspace = true
 
 [features]
 parallelism = ["psyche-client/parallelism"]

architectures/centralized/client/src/main.rs

@@ -16,6 +16,9 @@ use tracing::info;
 
 mod app;
 
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[derive(Parser, Debug)]
 struct Args {
     #[command(subcommand)]

architectures/centralized/server/Cargo.toml

@@ -31,6 +31,7 @@ bytemuck.workspace = true
 toml.workspace = true
 clap-markdown.workspace = true
 psyche-python-extension-impl = { workspace = true, optional = true }
+tikv-jemallocator.workspace = true
 
 [features]
 python = ["psyche-python-extension-impl"]

architectures/centralized/server/src/main.rs

@@ -16,6 +16,9 @@ use std::{
 };
 use tracing::{error, info};
 
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[derive(Parser, Debug)]
 struct Args {
     #[command(subcommand)]

architectures/decentralized/solana-client/Cargo.toml

@@ -30,6 +30,7 @@ tokio.workspace = true
 tokio-util.workspace = true
 tracing.workspace = true
 psyche-python-extension-impl = { workspace = true, optional = true }
+tikv-jemallocator.workspace = true
 
 [features]
 parallelism = ["psyche-client/parallelism"]

architectures/decentralized/solana-client/src/main.rs

@@ -29,6 +29,9 @@ use tracing::info;
 mod app;
 mod network_identity;
 
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[derive(Parser, Debug)]
 struct CliArgs {
     #[command(subcommand)]

architectures/inference-only/inference-node/Cargo.toml

@@ -33,3 +33,5 @@ iroh-blobs.workspace = true
 iroh-gossip.workspace = true
 
 serde_json.workspace = true
+
+tikv-jemallocator.workspace = true

architectures/inference-only/inference-node/src/main.rs

@@ -18,6 +18,9 @@ use std::sync::Arc;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, warn};
 
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[derive(Parser, Debug)]
 #[command(name = "psyche-inference-node")]
 struct Cli {