From 25a9c10bf795a3719f561e3867c79c2d3d69a029 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 06:27:47 +0000 Subject: [PATCH] fix(session): stop --phase-profile paying for the activation-memory planner The executor ran its activation-memory planner whenever phase profiling was on. The planner rebuilds the view map and re-plans every activation on every run - work the shipped runtime never does. So the profiler perturbed the run it was measuring and then reported its own cost back as a phase of that run, `run_scoped.activation_memory_plan`, which no unprofiled run ever pays. The planner now has its own gate. --phase-profile calls enable_exec_phase_profile_for_process, which drives the profiler only; NXRT_EXEC_PHASE_PROFILE=1 in the environment still enables both, so the CLI memory report keeps its stats via the path it already used, and NXRT_ACTIVATION_MEMORY_PLAN=1 opts in without profiling. Magnitude, measured honestly. The planner's own span reads 1.9-6.0 us per run. Removing it lowers native_min on sm_decode_h32_kv4096 from 0.065 ms to 0.063 ms, median of 15 interleaved repetitions, 300 runs each, single thread - about 3%, in line with the span. An earlier draft of this message claimed 35% and claimed the planner also polluted the adjacent exec_kernel.compute figure. Review could not reproduce either, and a properly interleaved re-measurement showed both were artifacts: this host is bimodal under contention, the 0.065/0.088 ms split appears in both arms, and exec_kernel.compute swings 75-168 us/call regardless of the planner. Both claims are withdrawn. The change stands on the smaller effect and on not reporting a phantom phase. The CUDA step profiler used to switch the planner on as a side effect of enabling phase accounting, so it taxed the very decode steps it exists to time. It no longer does; that coupling was accidental and is not restored. The first attempt at a regression test was vacuous. The gate was hardcoded true under cfg(test), so the production wiring was the one line no unit test could reach, and re-coupling the planner to phase profiling still passed. That override is gone: tests and production now share one gate, the two tests that need planning ask for it explicitly through an RAII guard that clears it again on drop, and re-coupling the wiring now fails both of them. The process-global gates are plain atomics with no per-test isolation, so every test that writes them takes a shared lock. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnx-genai-bench/src/bin/bench_generic.rs | 6 ++ .../onnx-genai-engine/src/engine/runtime.rs | 6 +- .../src/native_decode/cuda.rs | 6 ++ .../onnx-runtime-session/src/executor/mod.rs | 84 +++++++++++++++++++ .../onnx-runtime-session/src/executor/run.rs | 15 ++-- .../src/executor/tests.rs | 51 ++++++++++- crates/onnx-runtime-session/src/lib.rs | 5 +- 7 files changed, 160 insertions(+), 13 deletions(-) diff --git a/crates/onnx-genai-bench/src/bin/bench_generic.rs b/crates/onnx-genai-bench/src/bin/bench_generic.rs index bd77cddec3..eb5a06e7aa 100644 --- a/crates/onnx-genai-bench/src/bin/bench_generic.rs +++ b/crates/onnx-genai-bench/src/bin/bench_generic.rs @@ -33,6 +33,12 @@ struct Args { /// native runs. This is the only way to attribute the part of a run that /// the per-op profiler (`ONNX_GENAI_PROFILE_OPS=1`) leaves undifferentiated, /// because that one times node execution and nothing around it. + /// + /// Does not enable the activation-memory planner. That planner re-plans + /// every activation on every run - work the shipped runtime never does - + /// so while this flag switched it on, the profiler perturbed what it was + /// measuring and then reported its own cost back as a phase of the run. + /// Set `NXRT_ACTIVATION_MEMORY_PLAN=1` to opt into it deliberately. #[arg(long)] phase_profile: bool, /// Measure ORT only. Useful for recording a baseline when native loading or execution fails. diff --git a/crates/onnx-genai-engine/src/engine/runtime.rs b/crates/onnx-genai-engine/src/engine/runtime.rs index d6206756fd..a29e1f94ce 100644 --- a/crates/onnx-genai-engine/src/engine/runtime.rs +++ b/crates/onnx-genai-engine/src/engine/runtime.rs @@ -1370,9 +1370,13 @@ impl Engine { |session| session.execution_provider_status().summary(), ) } - /// Latest native activation-memory planner measurement, if the current /// backend is native and has executed far enough to resolve concrete shapes. + /// + /// Returns `None` unless the planner is switched on, which it is not by + /// default: set `NXRT_ACTIVATION_MEMORY_PLAN=1` (or + /// `NXRT_EXEC_PHASE_PROFILE=1`) first. Being native and having executed is + /// necessary but not sufficient. pub fn activation_memory_plan_stats(&self) -> Option { #[cfg(feature = "native-backend")] { diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 5f235cdcaf..fd44c7010d 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -327,6 +327,12 @@ impl CudaStepProfile { if !cuda_step_profile_enabled() { return None; } + // Phase accounting only. This deliberately does not switch on the + // activation-memory planner, which used to come along with it: the + // planner re-plans every activation on every run, so it taxed the very + // decode steps this profiler exists to time. A caller that wants its + // stats asks for them with `enable_activation_memory_plan_for_process` + // or `NXRT_ACTIVATION_MEMORY_PLAN=1`. onnx_runtime_session::enable_exec_phase_profile_for_process(); onnx_runtime_session::reset_exec_phase_profile(); Some(Self { diff --git a/crates/onnx-runtime-session/src/executor/mod.rs b/crates/onnx-runtime-session/src/executor/mod.rs index c31a54a422..867857f4fa 100644 --- a/crates/onnx-runtime-session/src/executor/mod.rs +++ b/crates/onnx-runtime-session/src/executor/mod.rs @@ -123,6 +123,83 @@ mod phase_profile { STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed); } + static PLAN_STATE: AtomicU8 = AtomicU8::new(0); // 0 = unknown, 1 = off, 2 = on + + /// Whether to run the activation-memory *planner* during a run. + /// + /// Deliberately **not** `enabled()`. The planner is a separate instrument: + /// it rebuilds the view map and re-plans every activation on every run, + /// work the shipped runtime never does. Its own span measures 1.9-6.0us + /// per run, so charging it to `--phase-profile` meant the profiler + /// perturbed the run it was measuring and reported its own cost back as a + /// phase of that run. Removing it lowers `native_min` on a softmax decode + /// from 0.065ms to 0.063ms, median of 15 interleaved repetitions - about + /// 3%, in line with the span. + /// + /// Reading the environment directly, rather than `enabled()`, is what + /// separates the two: `enable_for_process()` (what `--phase-profile` calls) + /// no longer switches the planner on, while anyone who set + /// `NXRT_EXEC_PHASE_PROFILE=1` in the environment - the CLI memory report's + /// path to these stats - keeps getting them. + pub fn activation_plan_enabled() -> bool { + match PLAN_STATE.load(Ordering::Relaxed) { + 1 => false, + 2 => true, + _ => { + let on = ["NXRT_ACTIVATION_MEMORY_PLAN", "NXRT_EXEC_PHASE_PROFILE"] + .iter() + .any(|key| { + std::env::var(key).is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true")) + }); + PLAN_STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed); + on + } + } + } + + /// Opt a process into activation-memory planning without turning on phase + /// profiling, for callers that want the stats and are willing to pay. + pub fn enable_activation_plan_for_process() { + PLAN_STATE.store(2, Ordering::Relaxed); + } + + /// Test-only override of the env-derived planner state. + #[cfg(test)] + pub(super) fn force_activation_plan_enabled(on: bool) { + PLAN_STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed); + } + + /// Holds the planner on for one test and clears it on drop, so a leaked + /// gate cannot be observed by a later test that does not take the lock. + #[cfg(test)] + pub(super) struct ActivationPlanForTest(std::sync::MutexGuard<'static, ()>); + + #[cfg(test)] + impl ActivationPlanForTest { + pub(super) fn on() -> Self { + let guard = globals_lock(); + force_activation_plan_enabled(true); + Self(guard) + } + } + + #[cfg(test)] + impl Drop for ActivationPlanForTest { + fn drop(&mut self) { + force_activation_plan_enabled(false); + } + } + + /// Serialises the tests that write the two process-global gates above. + /// Both are plain atomics with no per-test isolation, so without this the + /// parallel runner lets one test's `enable`/`force` land inside another + /// test's assertion window. + #[cfg(test)] + pub(super) fn globals_lock() -> std::sync::MutexGuard<'static, ()> { + static LOCK: Mutex<()> = Mutex::new(()); + LOCK.lock().unwrap_or_else(|e| e.into_inner()) + } + pub fn enable_for_process() { STATE.store(2, Ordering::Relaxed); } @@ -302,6 +379,13 @@ pub fn enable_exec_phase_profile_for_process() { phase_profile::enable_for_process(); } +/// Opt this process into the activation-memory planner, which runs on every +/// `Run` and is *not* enabled by [`enable_exec_phase_profile_for_process`] +/// because it costs enough to distort the run it would be reported against. +pub fn enable_activation_memory_plan_for_process() { + phase_profile::enable_activation_plan_for_process(); +} + /// Activation-memory planner metrics from the most recent measured top-level run. #[derive(Clone, Copy, Debug, PartialEq)] pub struct ActivationMemoryPlanStats { diff --git a/crates/onnx-runtime-session/src/executor/run.rs b/crates/onnx-runtime-session/src/executor/run.rs index 06f4d1e463..0ec1b687d8 100644 --- a/crates/onnx-runtime-session/src/executor/run.rs +++ b/crates/onnx-runtime-session/src/executor/run.rs @@ -6,15 +6,14 @@ struct Stage2RunState { excluded: Option>, } +/// Whether to run the activation-memory planner during this run. +/// +/// One path for tests and production. It used to be `#[cfg(test)] true`, which +/// meant the production wiring - the part that decides whether a run pays for +/// planning at all - was the one line no unit test could reach, and a test that +/// re-coupled it to phase profiling still passed. fn activation_memory_planning_enabled() -> bool { - #[cfg(test)] - { - true - } - #[cfg(not(test))] - { - phase_profile::enabled() - } + phase_profile::activation_plan_enabled() } impl Executor { diff --git a/crates/onnx-runtime-session/src/executor/tests.rs b/crates/onnx-runtime-session/src/executor/tests.rs index 262c562867..fba9085bac 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -10,8 +10,9 @@ use super::*; #[test] fn phase_profile_gating_and_accumulation() { - // Single test (not two) so the process-global enable flag is never - // toggled concurrently by a sibling test under the parallel runner. + // The process-global gates have no per-test isolation, so every test that + // writes them takes this lock. + let _globals = phase_profile::globals_lock(); // Disabled: a span records nothing and never captures a timestamp. phase_profile::force_enabled(false); @@ -54,8 +55,40 @@ fn phase_profile_gating_and_accumulation() { "reset must clear this test's accumulated phase stats" ); + // Turning on phase profiling must **not** turn on the activation-memory + // planner. The planner re-plans every activation on every run - work the + // shipped runtime never does - so while `--phase-profile` enabled it the + // profiler perturbed the run it was measuring, by about 3% on a softmax + // decode, and reported its own cost back as a phase of that run. + // + // The two gates are separate on purpose: `enable_for_process` (what + // `--phase-profile` calls) drives only the profiler, while the planner + // reads the environment. Asserted here rather than in a sibling test so + // the process-global enable flag stays owned by this one test, per the + // note at the top. Skipped when the environment already asks for planning, + // because then the separation is not observable. + let env_requests_planning = ["NXRT_ACTIVATION_MEMORY_PLAN", "NXRT_EXEC_PHASE_PROFILE"] + .iter() + .any(|key| std::env::var(key).is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))); + if !env_requests_planning { + phase_profile::force_activation_plan_enabled(false); + phase_profile::enable_for_process(); + assert!(phase_profile::enabled()); + assert!( + !phase_profile::activation_plan_enabled(), + "phase profiling must not drag the activation-memory planner in \ + with it: the planner costs about a third of a small run and would \ + be charged to the run it is supposed to be measuring" + ); + // The explicit opt-in still works, so the capability is decoupled from + // phase profiling rather than lost. + phase_profile::enable_activation_plan_for_process(); + assert!(phase_profile::activation_plan_enabled()); + } + // Restore the default (disabled) state so other tests stay inert. phase_profile::force_enabled(false); + phase_profile::force_activation_plan_enabled(false); } // A produced top-level output is handed to the caller by moving its host @@ -471,6 +504,13 @@ fn compute_in_place_multilayer_decode_residual_is_byte_identical_and_fires() { #[test] fn activation_memory_planner_reports_static_decode_graph_savings() { + // The planner is off by default, in tests as in production, so a test that + // needs its stats asks for them explicitly. The guard takes the globals + // lock (the gate is a process-global atomic) and clears the gate again on + // drop, including on panic - leaving it set would let a reader test that + // does not take the lock observe a planner this test switched on. + let _planner = phase_profile::ActivationPlanForTest::on(); + let values = Tensor::from_f32(&[4], &[-2.0, -0.5, 0.5, 2.0]).unwrap(); let weights = Arc::new(WeightStore::new()); let ep = auto_detect_cpu_ep().unwrap(); @@ -493,6 +533,13 @@ fn activation_memory_planner_reports_static_decode_graph_savings() { #[test] fn activation_memory_planner_uses_runtime_view_edges() { + // The planner is off by default, in tests as in production, so a test that + // needs its stats asks for them explicitly. The guard takes the globals + // lock (the gate is a process-global atomic) and clears the gate again on + // drop, including on panic - leaving it set would let a reader test that + // does not take the lock observe a planner this test switched on. + let _planner = phase_profile::ActivationPlanForTest::on(); + let values = Tensor::from_f32(&[4], &[-2.0, -0.5, 0.5, 2.0]).unwrap(); let weights = Arc::new(WeightStore::new()); let ep = auto_detect_cpu_ep().unwrap(); diff --git a/crates/onnx-runtime-session/src/lib.rs b/crates/onnx-runtime-session/src/lib.rs index 369d07e0ce..6765d15926 100644 --- a/crates/onnx-runtime-session/src/lib.rs +++ b/crates/onnx-runtime-session/src/lib.rs @@ -34,8 +34,9 @@ pub use executor::{ DeviceAllocationCounts, DeviceGraphCaptureResult, ExecutionProviderDecline, ExecutionProviderFallbackReport, PrefetchStep, SeamReason, dense_prefetch_gap_stats, dense_weight_prefetch_lookahead_nodes, drive_double_buffer, - enable_exec_phase_profile_for_process, exec_phase_stats, plan_double_buffer, - print_exec_phase_profile, reset_dense_prefetch_gap_stats, reset_exec_phase_profile, + enable_activation_memory_plan_for_process, enable_exec_phase_profile_for_process, + exec_phase_stats, plan_double_buffer, print_exec_phase_profile, reset_dense_prefetch_gap_stats, + reset_exec_phase_profile, }; pub use onnx_runtime_ep_api::WorkspaceRequirement; pub use onnx_runtime_loader::{