diff --git a/crates/onnx-runtime-ep-api/Cargo.toml b/crates/onnx-runtime-ep-api/Cargo.toml index 07586db770..4be3ef0670 100644 --- a/crates/onnx-runtime-ep-api/Cargo.toml +++ b/crates/onnx-runtime-ep-api/Cargo.toml @@ -16,12 +16,12 @@ memmap2 = { workspace = true } libloading = { workspace = true } onnx-genai-ort-sys = { workspace = true } onnx-runtime-ir = { workspace = true } +onnx-runtime-loader = { workspace = true } onnx-runtime-optimizer = { workspace = true } onnx-runtime-tracer = { workspace = true } thiserror = { workspace = true } [dev-dependencies] -onnx-runtime-loader = { workspace = true } [lints] workspace = true diff --git a/crates/onnx-runtime-ep-api/src/lib.rs b/crates/onnx-runtime-ep-api/src/lib.rs index edea762b6c..b44c72fb9e 100644 --- a/crates/onnx-runtime-ep-api/src/lib.rs +++ b/crates/onnx-runtime-ep-api/src/lib.rs @@ -66,8 +66,10 @@ pub use tensor::{ pub use weight::{ ExecutionProviderCapabilities, LazyDeviceWeightBinder, LazyWeight, LazyWeightBoundary, LazyWeightCandidate, MmapRegionSource, NXRT_WEIGHT_PAGING_CAPABILITY, NegotiatedWeight, - PagedWeight, Phase3aHostOnlyBinder, ResidentWeight, ResidentWeightMaterializer, WeightHandle, - WeightHandleError, lazy_weight_candidates, + PagedWeight, Phase3aHostOnlyBinder, ResidencyDecision, ResidencyDegradationReason, + ResidencyPlan, ResidencyPolicy, ResidencyPolicyInput, ResidentWeight, + ResidentWeightMaterializer, WeightHandle, WeightHandleError, WholeBankResidentPolicy, + lazy_weight_candidates, plan_residency, }; // Re-export the device vocabulary from the IR so EP authors have one import. diff --git a/crates/onnx-runtime-ep-api/src/weight.rs b/crates/onnx-runtime-ep-api/src/weight.rs index eca76c42cf..dd4205bf57 100644 --- a/crates/onnx-runtime-ep-api/src/weight.rs +++ b/crates/onnx-runtime-ep-api/src/weight.rs @@ -465,6 +465,256 @@ pub enum WeightHandleError { DeviceBinding(String), } +/// Why a candidate expert region did not enter a [`ResidencyPlan`] as +/// resident. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ResidencyDegradationReason { + /// The catalog itself was non-pageable (see + /// [`onnx_runtime_loader::NonPageableReason`]) — the plan cannot promise + /// per-expert placement for this value at all. + NonPageableCatalog(onnx_runtime_loader::NonPageableReason), + /// The policy chose to keep this value fully resident rather than emit + /// per-expert placement (the default/only shipped behavior today). + PolicyDeclinedSplit, + /// The policy asked for a plan this build-time seam cannot safely honor + /// (e.g. it referenced an expert index out of range, or overlapping + /// ranges) — the executor falls back to whole-bank residency and records + /// why. + RejectedByValidation(String), +} + +/// One entry in a [`ResidencyPlan`]: what an execution boundary should do for +/// one lazy-weight value. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ResidencyDecision { + /// Keep the whole initializer resident, exactly like today's shipped + /// behavior. `reason` is `None` for the default policy's decisions and + /// `Some(..)` when a catalog/validation problem forced this fallback. + WholeBankResident { + reason: Option, + }, + /// Placement/admission is unconstrained per expert (informational only — + /// this slice does not execute prefetch/eviction from this decision). + /// `experts` lists the candidate expert indices in ascending order (the + /// deterministic ordering requirement). + PerExpertCandidate { experts: Vec }, +} + +/// A validated, typed residency decision set for one build's lazy-weight +/// candidates. +/// +/// A `ResidencyPlan` is pure data: it names *which* value gets *which* +/// decision. It never allocates, copies bytes, opens a CUDA stream, or owns a +/// VA range — those stay the exclusive responsibility of +/// `onnx-runtime-ep-cuda::weight_paging`'s `CudaWeightResidency` (or the +/// executor's existing whole-bank materialization path on non-CUDA EPs). +/// Building a plan and *acting* on it are two different lifecycle steps by +/// construction: nothing in this type can mutate memory. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ResidencyPlan { + decisions: std::collections::HashMap, + /// Name of the policy that produced this plan, for telemetry. + policy_name: &'static str, +} + +impl ResidencyPlan { + pub fn policy_name(&self) -> &'static str { + self.policy_name + } + + pub fn decision(&self, value: ValueId) -> Option<&ResidencyDecision> { + self.decisions.get(&value) + } + + /// Values in deterministic (ascending [`ValueId`]) order, for stable + /// telemetry/log output and reproducible tests. + pub fn ordered_values(&self) -> impl Iterator + '_ { + let mut values: Vec = self.decisions.keys().copied().collect(); + values.sort_unstable_by_key(|value| value.0); + values.into_iter() + } + + pub fn len(&self) -> usize { + self.decisions.len() + } + + pub fn is_empty(&self) -> bool { + self.decisions.is_empty() + } + + /// Count of values with a resident decision that carries no degradation + /// reason (i.e. the policy's ordinary output, not a validation fallback). + /// Disjoint from [`Self::degraded_count`]: every entry is counted by + /// exactly one of the two. + pub fn resident_count(&self) -> usize { + self.decisions + .values() + .filter(|decision| { + matches!( + decision, + ResidencyDecision::WholeBankResident { reason: None } + ) + }) + .count() + } + + /// Count of values whose decision carries an explicit degradation reason + /// (nonpageable catalog, policy decline, or validation rejection). + pub fn degraded_count(&self) -> usize { + self.decisions + .values() + .filter(|decision| { + matches!( + decision, + ResidencyDecision::WholeBankResident { reason: Some(_) } + ) + }) + .count() + } +} + +/// Structural inputs a [`ResidencyPolicy`] may use to decide placement. +/// +/// Deliberately excludes model names, op allowlists, and quantization-format +/// branches: a policy may only look at [`LazyWeightBoundary`], the +/// [`onnx_runtime_loader::WeightRegionCatalog`] (byte layout/pageability), and +/// the caller-supplied budget/capability context. +pub struct ResidencyPolicyInput<'a> { + pub boundary: LazyWeightBoundary, + pub catalog: &'a onnx_runtime_loader::WeightRegionCatalog, + /// Advisory device byte budget available for expert-bank residency, if + /// the caller has one. `None` means "unconstrained" (today's default). + pub budget_bytes: Option, +} + +/// A pluggable placement/admission/eviction/prefetch/resize *decision* +/// surface. A policy answers "which experts (if any) does this value split +/// into for planning purposes", nothing else. +/// +/// Implementations must not allocate, copy, synchronize, or hold VA/pointer +/// state — see the module-level `ResidencyPlan` doc for why that split is +/// enforced by type rather than convention. +pub trait ResidencyPolicy: Send + Sync { + /// Stable identifier for telemetry (e.g. `"whole_bank_resident"`). + fn name(&self) -> &'static str; + + /// Decide one value's residency for planning purposes only. + fn decide(&self, input: &ResidencyPolicyInput<'_>) -> ResidencyDecision; +} + +/// The only shipped policy today: always keep every candidate value fully +/// resident, exactly matching pre-#82 behavior byte-for-byte/handle-for-handle. +#[derive(Clone, Copy, Debug, Default)] +pub struct WholeBankResidentPolicy; + +impl ResidencyPolicy for WholeBankResidentPolicy { + fn name(&self) -> &'static str { + "whole_bank_resident" + } + + fn decide(&self, input: &ResidencyPolicyInput<'_>) -> ResidencyDecision { + let reason = if input.catalog.is_pageable() { + None + } else { + match input.catalog.pageability() { + onnx_runtime_loader::Pageability::NonPageable(reason) => Some( + ResidencyDegradationReason::NonPageableCatalog(reason.clone()), + ), + onnx_runtime_loader::Pageability::Pageable => None, + } + }; + ResidencyDecision::WholeBankResident { reason } + } +} + +/// Build a validated [`ResidencyPlan`] from the per-expert region candidates +/// gathered at build time, running `policy` over each and validating its +/// output before accepting it. +/// +/// This is the one call site where a plan is *created*; nothing here +/// allocates, copies, or synchronizes. Validation failures (an expert index +/// referencing an expert count the catalog does not have, or a duplicate +/// expert index) degrade that single value to `WholeBankResident` with +/// [`ResidencyDegradationReason::RejectedByValidation`] rather than failing +/// the whole plan or the model load — correctness always wins over honoring a +/// policy's request. +pub fn plan_residency<'a>( + candidates: impl IntoIterator< + Item = ( + ValueId, + LazyWeightBoundary, + &'a onnx_runtime_loader::WeightRegionCatalog, + ), + >, + policy: &dyn ResidencyPolicy, + budget_bytes: Option, +) -> ResidencyPlan { + let mut decisions = std::collections::HashMap::new(); + for (value, boundary, catalog) in candidates { + let input = ResidencyPolicyInput { + boundary, + catalog, + budget_bytes, + }; + let decision = policy.decide(&input); + let validated = validate_decision(catalog, decision); + decisions.insert(value, validated); + } + ResidencyPlan { + decisions, + policy_name: policy.name(), + } +} + +/// Validate a policy's decision against the catalog it was derived from. +/// Never panics: any inconsistency degrades to `WholeBankResident` with an +/// explicit `RejectedByValidation` reason. +fn validate_decision( + catalog: &onnx_runtime_loader::WeightRegionCatalog, + decision: ResidencyDecision, +) -> ResidencyDecision { + match decision { + ResidencyDecision::PerExpertCandidate { experts } => { + if !catalog.is_pageable() { + return ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation( + "policy proposed per-expert placement over a nonpageable catalog".into(), + )), + }; + } + let expert_count = catalog.layout().experts; + let mut seen = std::collections::HashSet::with_capacity(experts.len()); + for &expert in &experts { + if expert >= expert_count { + return ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation(format!( + "expert index {expert} out of range for a {expert_count}-expert bank" + ))), + }; + } + if !seen.insert(expert) { + return ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation(format!( + "expert index {expert} appears more than once in one plan entry" + ))), + }; + } + if catalog.region(expert).is_none() { + return ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation(format!( + "expert {expert} has no validated byte range in its catalog" + ))), + }; + } + } + let mut ordered = experts; + ordered.sort_unstable(); + ResidencyDecision::PerExpertCandidate { experts: ordered } + } + whole_bank @ ResidencyDecision::WholeBankResident { .. } => whole_bank, + } +} + #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; @@ -572,4 +822,234 @@ mod tests { &[1, 2, 3, 4] ); } + + // -- ResidencyPlan / ResidencyPolicy tests ----------------------------- + + use onnx_runtime_ir::WeightRef; + use onnx_runtime_loader::{ + ExpertQuantization, ExpertStorageOrder, ExpertTensorLayout, WeightRegionCatalog, + }; + + fn expert_layout() -> ExpertTensorLayout { + ExpertTensorLayout { + version: 1, + experts: 3, + rows_per_expert: 2, + storage_elements_per_row: 4, + order: ExpertStorageOrder::ExpertMajor, + quantization: Some(ExpertQuantization { + bits: 4, + block_size: 16, + blocks_per_row: 1, + }), + } + } + + fn pageable_catalog() -> WeightRegionCatalog { + let layout = expert_layout(); + let weight = WeightRef::External { + path: std::path::PathBuf::from("/nonexistent/weights.bin"), + offset: 16, + length: layout.experts * layout.rows_per_expert * layout.storage_elements_per_row, + dtype: DataType::Uint8, + dims: vec![ + layout.experts, + layout.rows_per_expert, + layout.storage_elements_per_row, + ], + }; + WeightRegionCatalog::classify(&weight, layout) + } + + fn non_pageable_catalog() -> WeightRegionCatalog { + let mut layout = expert_layout(); + layout.order = ExpertStorageOrder::Interleaved; + let weight = WeightRef::External { + path: std::path::PathBuf::from("/nonexistent/weights.bin"), + offset: 16, + length: layout.experts * layout.rows_per_expert * layout.storage_elements_per_row, + dtype: DataType::Uint8, + dims: vec![ + layout.experts, + layout.rows_per_expert, + layout.storage_elements_per_row, + ], + }; + WeightRegionCatalog::classify(&weight, layout) + } + + /// Test-only alternate policy proving the trait boundary is + /// substitutable: it proposes splitting every pageable catalog into + /// per-expert candidates. This is never shipped as a default and must + /// not ship LRU/static-hot/q* behavior. + struct AlwaysSplitPolicy; + + impl ResidencyPolicy for AlwaysSplitPolicy { + fn name(&self) -> &'static str { + "test_always_split" + } + + fn decide(&self, input: &ResidencyPolicyInput<'_>) -> ResidencyDecision { + if input.catalog.is_pageable() { + ResidencyDecision::PerExpertCandidate { + experts: (0..input.catalog.layout().experts).collect(), + } + } else { + ResidencyDecision::WholeBankResident { reason: None } + } + } + } + + struct OutOfRangePolicy; + + impl ResidencyPolicy for OutOfRangePolicy { + fn name(&self) -> &'static str { + "test_out_of_range" + } + + fn decide(&self, _input: &ResidencyPolicyInput<'_>) -> ResidencyDecision { + ResidencyDecision::PerExpertCandidate { + experts: vec![9999], + } + } + } + + struct DuplicatePolicy; + + impl ResidencyPolicy for DuplicatePolicy { + fn name(&self) -> &'static str { + "test_duplicate" + } + + fn decide(&self, _input: &ResidencyPolicyInput<'_>) -> ResidencyDecision { + ResidencyDecision::PerExpertCandidate { + experts: vec![0, 0], + } + } + } + + fn value(id: u32) -> ValueId { + ValueId(id) + } + + #[test] + fn whole_bank_resident_policy_matches_default_behavior_for_pageable_catalog() { + let catalog = pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &WholeBankResidentPolicy, + None, + ); + assert_eq!(plan.policy_name(), "whole_bank_resident"); + assert_eq!(plan.len(), 1); + assert_eq!(plan.resident_count(), 1); + assert_eq!(plan.degraded_count(), 0); + assert_eq!( + plan.decision(value(1)), + Some(&ResidencyDecision::WholeBankResident { reason: None }) + ); + } + + #[test] + fn whole_bank_resident_policy_surfaces_non_pageable_reason() { + let catalog = non_pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &WholeBankResidentPolicy, + None, + ); + assert_eq!( + plan.decision(value(1)), + Some(&ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::NonPageableCatalog( + onnx_runtime_loader::NonPageableReason::NotExpertMajor + )) + }) + ); + } + + #[test] + fn alternate_policy_is_substitutable_and_produces_per_expert_candidates() { + let catalog = pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &AlwaysSplitPolicy, + None, + ); + assert_eq!(plan.policy_name(), "test_always_split"); + assert_eq!( + plan.decision(value(1)), + Some(&ResidencyDecision::PerExpertCandidate { + experts: vec![0, 1, 2] + }) + ); + } + + #[test] + fn out_of_range_expert_index_degrades_to_whole_bank_with_reason() { + let catalog = pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &OutOfRangePolicy, + None, + ); + assert_eq!(plan.resident_count(), 0); + match plan.decision(value(1)) { + Some(ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation(_)), + }) => {} + other => panic!("expected rejected validation, got {other:?}"), + } + assert_eq!(plan.degraded_count(), 1); + } + + #[test] + fn duplicate_expert_index_degrades_to_whole_bank_with_reason() { + let catalog = pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &DuplicatePolicy, + None, + ); + match plan.decision(value(1)) { + Some(ResidencyDecision::WholeBankResident { + reason: Some(ResidencyDegradationReason::RejectedByValidation(_)), + }) => {} + other => panic!("expected rejected validation, got {other:?}"), + } + } + + #[test] + fn per_expert_candidate_over_non_pageable_catalog_is_rejected() { + let catalog = non_pageable_catalog(); + let plan = plan_residency( + [(value(1), LazyWeightBoundary::QMoe, &catalog)], + &AlwaysSplitPolicy, + None, + ); + // AlwaysSplitPolicy itself checks is_pageable(), so this exercises + // the policy's own fallback rather than validation -- assert it is + // still whole-bank resident, never a candidate over a non-pageable + // catalog. + assert!(matches!( + plan.decision(value(1)), + Some(ResidencyDecision::WholeBankResident { .. }) + )); + } + + #[test] + fn plan_residency_orders_values_deterministically() { + let catalog = pageable_catalog(); + let plan = plan_residency( + [ + (value(5), LazyWeightBoundary::QMoe, &catalog), + (value(1), LazyWeightBoundary::QMoe, &catalog), + (value(3), LazyWeightBoundary::QMoe, &catalog), + ], + &WholeBankResidentPolicy, + None, + ); + let ordered: Vec = plan.ordered_values().map(|v| v.0).collect(); + assert_eq!(ordered, vec![1, 3, 5]); + } } diff --git a/crates/onnx-runtime-session/src/executor/build.rs b/crates/onnx-runtime-session/src/executor/build.rs index 854902aeb8..18ece16b68 100644 --- a/crates/onnx-runtime-session/src/executor/build.rs +++ b/crates/onnx-runtime-session/src/executor/build.rs @@ -433,6 +433,46 @@ pub(super) fn build_lazy_weight_handles_and_candidates( Ok((handles, candidates)) } +/// Build the default [`onnx_runtime_ep_api::ResidencyPlan`] for this build's +/// per-expert region candidates. +/// +/// This is the one call site where a plan is *created*. It always uses +/// [`onnx_runtime_ep_api::WholeBankResidentPolicy`] today — the only shipped +/// policy — so the produced plan is handle-for-handle/byte-for-byte +/// equivalent to pre-#82 behavior: every candidate value resolves to +/// `WholeBankResident`. A future slice may take the policy as a parameter; +/// this slice intentionally does not expose that knob outside tests. +pub(super) fn plan_default_residency( + graph: &Graph, + candidates: &HashMap, +) -> onnx_runtime_ep_api::ResidencyPlan { + plan_residency_with( + graph, + candidates, + &onnx_runtime_ep_api::WholeBankResidentPolicy, + ) +} + +/// Test/measurement seam: build a plan with an arbitrary +/// [`onnx_runtime_ep_api::ResidencyPolicy`], to prove the boundary is +/// substitutable without wiring a second production call site. +pub(super) fn plan_residency_with( + graph: &Graph, + candidates: &HashMap, + policy: &dyn onnx_runtime_ep_api::ResidencyPolicy, +) -> onnx_runtime_ep_api::ResidencyPlan { + let boundary_by_value: HashMap = lazy_weight_candidates(graph) + .into_iter() + .map(|candidate| (candidate.value, candidate.boundary)) + .collect(); + let inputs = candidates.iter().filter_map(|(value, catalog)| { + boundary_by_value + .get(value) + .map(|&boundary| (*value, boundary, catalog)) + }); + onnx_runtime_ep_api::plan_residency(inputs, policy, None) +} + /// Derive an [`onnx_runtime_loader::ExpertTensorLayout`] for every QMoE /// expert-bank initializer value in `graph`, keyed by the initializer's /// [`ValueId`]. @@ -592,6 +632,15 @@ impl Executor { &self.expert_region_candidates } + /// The default [`onnx_runtime_ep_api::ResidencyPlan`] built for this + /// executor's expert region candidates. Read-only outside tests today — + /// no dispatch-time consumer exists yet; this exposes the plan purely + /// for telemetry/measurement and future policy wiring. + #[cfg(test)] + pub(super) fn residency_plan(&self) -> &onnx_runtime_ep_api::ResidencyPlan { + &self.residency_plan + } + /// Compile a graph + weights into a runnable executor on the CPU EP. pub(crate) fn build( graph: Graph, @@ -631,6 +680,21 @@ impl Executor { (handles, candidates) }; + let residency_plan = { + let mut span = trace_span("session.residency_plan", "session"); + let plan = plan_default_residency(&graph, &expert_region_candidates); + if let Some(span) = span.as_mut() { + span.set_args( + Args::new() + .with("policy", plan.policy_name().to_string()) + .with("resident", plan.resident_count() as u64) + .with("degraded", plan.degraded_count() as u64) + .with("total", plan.len() as u64), + ); + } + plan + }; + let (mut value_shapes, mut value_dtypes, buffers, buffer_shapes) = Self::materialize_initializers(&graph, &weights, ep.as_ref(), &weight_handles)?; @@ -666,6 +730,7 @@ impl Executor { graph_slot: DeviceGraphSlot::Primary, weight_handles, expert_region_candidates, + residency_plan, prefetch_issue_nodes: std::sync::Mutex::new(HashMap::new()), prefetch_lookahead_nodes: dense_weight_prefetch_lookahead_nodes(), buffers, diff --git a/crates/onnx-runtime-session/src/executor/state.rs b/crates/onnx-runtime-session/src/executor/state.rs index 936fda24b7..3e482571b3 100644 --- a/crates/onnx-runtime-session/src/executor/state.rs +++ b/crates/onnx-runtime-session/src/executor/state.rs @@ -105,6 +105,14 @@ pub(crate) struct Executor { /// today; production consumers land in a follow-up slice per issue #82. #[allow(dead_code)] pub(super) expert_region_candidates: HashMap, + /// The validated [`onnx_runtime_ep_api::ResidencyPlan`] derived from + /// `expert_region_candidates` using [`onnx_runtime_ep_api::WholeBankResidentPolicy`] + /// (the only shipped policy today). Pure bookkeeping: it does not own + /// allocation, copying, or synchronization, and no production dispatch + /// path reads it yet. Read only under `#[cfg(test)]` via + /// [`Executor::residency_plan`]. + #[allow(dead_code)] + pub(super) residency_plan: onnx_runtime_ep_api::ResidencyPlan, pub(super) prefetch_issue_nodes: std::sync::Mutex>, pub(super) prefetch_lookahead_nodes: usize, /// One device buffer per backed value. Static values are allocated once at diff --git a/crates/onnx-runtime-session/src/executor/tests.rs b/crates/onnx-runtime-session/src/executor/tests.rs index d30210b339..0da7481d10 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -8352,3 +8352,125 @@ fn weight_derived_caches_are_cleared_before_their_buffers_are_freed() { ); } } + +#[test] +fn qmoe_residency_plan_default_policy_matches_whole_bank_resident_for_pageable_candidates() { + let (graph, weights, _path) = qmoe_expert_region_fixture(); + let deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let ep = WeightDeliveryEp::new(true, deliveries); + let executor = Executor::build(graph, weights, Arc::new(ep)).unwrap(); + + let plan = executor.residency_plan(); + let candidates = executor.expert_region_candidates(); + assert_eq!(plan.policy_name(), "whole_bank_resident"); + // Same cardinality as the expert region candidates: the plan covers + // exactly the values that have a QMoE catalog, no more, no fewer. + assert_eq!(plan.len(), candidates.len()); + assert_eq!(plan.resident_count(), plan.len()); + assert_eq!(plan.degraded_count(), 0); + for value in candidates.keys() { + assert_eq!( + plan.decision(*value), + Some(&onnx_runtime_ep_api::ResidencyDecision::WholeBankResident { reason: None }) + ); + } + + // Existing behavior fully unchanged: still one lazy handle per + // initializer, each with exactly one region. + assert_eq!(executor.weight_handles.len(), 2); +} + +#[test] +fn qmoe_residency_plan_surfaces_non_pageable_reason_without_changing_handles() { + let (mut graph, weights, path) = qmoe_expert_region_fixture(); + let packed_value = graph + .initializers + .keys() + .copied() + .find(|value| graph.value(*value).name.as_deref() == Some("fc1_packed")) + .expect("fc1_packed initializer must exist"); + graph.set_initializer( + packed_value, + WeightRef::External { + path: path.clone(), + offset: 0, + length: 24, + dtype: DataType::Uint8, + dims: vec![6, 4], + }, + ); + + let deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let ep = WeightDeliveryEp::new(true, deliveries); + let executor = Executor::build(graph, weights, Arc::new(ep)).unwrap(); + + let plan = executor.residency_plan(); + match plan.decision(packed_value) { + Some(onnx_runtime_ep_api::ResidencyDecision::WholeBankResident { + reason: + Some(onnx_runtime_ep_api::ResidencyDegradationReason::NonPageableCatalog(reason)), + }) => { + assert_eq!( + reason, + &onnx_runtime_loader::NonPageableReason::NotExpertMajor + ); + } + other => panic!("expected non-pageable whole-bank reason, got {other:?}"), + } + assert_eq!( + plan.degraded_count(), + 1, + "non-pageable reason counts as degraded" + ); + assert_eq!(executor.weight_handles.len(), 2); +} + +/// Prove the residency-policy seam is substitutable at the executor's own +/// candidate/boundary wiring, without any production dispatch path consuming +/// the result -- this is the "test-only alternate policy" required to show +/// the trait is not an inert marker. +#[test] +fn qmoe_residency_plan_seam_is_substitutable_with_an_alternate_policy() { + use onnx_runtime_ep_api::{ResidencyDecision, ResidencyPolicy, ResidencyPolicyInput}; + + struct AlwaysSplit; + impl ResidencyPolicy for AlwaysSplit { + fn name(&self) -> &'static str { + "test_always_split" + } + fn decide(&self, input: &ResidencyPolicyInput<'_>) -> ResidencyDecision { + if input.catalog.is_pageable() { + ResidencyDecision::PerExpertCandidate { + experts: (0..input.catalog.layout().experts).collect(), + } + } else { + ResidencyDecision::WholeBankResident { reason: None } + } + } + } + + let (graph, weights, _path) = qmoe_expert_region_fixture(); + let deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let ep = WeightDeliveryEp::new(true, deliveries); + let executor = Executor::build(graph, weights, Arc::new(ep)).unwrap(); + + let candidates = executor.expert_region_candidates(); + let plan = + crate::executor::build::plan_residency_with(&executor.graph, candidates, &AlwaysSplit); + assert_eq!(plan.policy_name(), "test_always_split"); + assert!(!plan.is_empty()); + for value in candidates.keys() { + assert!(matches!( + plan.decision(*value), + Some(ResidencyDecision::PerExpertCandidate { .. }) + )); + } + + // Substituting the policy must not touch the default plan stored on the + // executor, nor the weight-handle/binding output. + assert_eq!( + executor.residency_plan().policy_name(), + "whole_bank_resident" + ); + assert_eq!(executor.weight_handles.len(), 2); +}