Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3829,6 +3829,7 @@ pub fn bound_process_to_decode_budget() {

match rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.thread_name(prefill_worker_name)
.build_global()
{
Ok(()) => eprintln!(
Expand All @@ -3843,6 +3844,25 @@ pub fn bound_process_to_decode_budget() {
}
}

/// Name for a worker in the process-wide prefill/MLAS Rayon pool.
///
/// Rayon does not name `build_global` workers by default, and an unnamed
/// thread's `comm` defaults to the *process* name — so an unnamed pool of
/// `N` workers reads, in `ps`/`top`/`/proc/<pid>/task`, as `N` extra copies
/// of the host binary rather than as a pool. That made this pool the single
/// largest unattributed block of threads in a budgeted process (an explicit
/// budget of `N` more than doubles the thread count, and until now `N` of
/// those were anonymous).
///
/// Deliberately short: Linux truncates `comm` to 15 bytes, so the longer
/// `onnx-genai-`-prefixed convention used elsewhere in this crate collapses
/// to `onnx-genai-deco`/`onnx-genai-spmd` and loses the index. This form
/// survives truncation intact for every width this pool is ever built at,
/// which is the whole point of naming it.
fn prefill_worker_name(index: usize) -> String {
format!("nxgn-prefill-{index}")
}

/// Default persistent-pool worker count for `available` logical CPUs: half of
/// them (at least one), derived purely from topology (Rule 2).
///
Expand Down Expand Up @@ -13861,6 +13881,34 @@ mod tests {
assert_eq!(resolve_rayon_global_threads(Some(8), None, 0), None);
}

#[test]
fn prefill_worker_names_survive_linux_comm_truncation() {
// Linux stores `comm` in 15 bytes plus a NUL. A name longer than that
// is silently truncated, which is how the `onnx-genai-`-prefixed
// convention elsewhere in this crate loses its worker index. These
// names must stay legible for every width the pool is built at, so
// that a thread census can attribute them without guesswork.
const COMM_LEN: usize = 15;
for index in [0, 1, 9, 15, 31, 63, 99] {
let name = prefill_worker_name(index);
assert!(
name.len() <= COMM_LEN,
"`{name}` is {} bytes and would be truncated in /proc/<pid>/task/*/comm",
name.len()
);
assert!(
name.starts_with("nxgn-prefill-"),
"unexpected prefix: {name}"
);
assert!(
name.ends_with(&index.to_string()),
"`{name}` lost its worker index"
);
}
// Distinct workers must stay distinguishable.
assert_ne!(prefill_worker_name(0), prefill_worker_name(1));
}

#[test]
fn explicit_budget_precedes_env_for_every_decode_pool() {
assert_eq!(
Expand Down
Loading