Add pluggable ResidencyPolicy seam for QMoE expert bank (#82)

[justinchuby]

Working as Batty (CUDA & Performance Engineer).

Goal

Next #82 slice per Roy's architecture verdict: introduce the smallest real, pluggable expert-residency policy boundary that consumes the per-expert region candidates from #1779, while default behavior stays exactly whole-bank-resident. This is a policy-boundary refactor, not a residency behavior change.

Exact new contract (crates/onnx-runtime-ep-api/src/weight.rs)

• ResidencyDecision — WholeBankResident { reason: Option<ResidencyDegradationReason> } | PerExpertCandidate { experts: Vec<usize> }.
• ResidencyDegradationReason — NonPageableCatalog(NonPageableReason) | PolicyDeclinedSplit | RejectedByValidation(String).
• ResidencyPlan — pure data: HashMap<ValueId, ResidencyDecision> + policy_name. No allocator/stream/pointer fields; cannot allocate/copy/synchronize by construction. policy_name(), decision(), ordered_values() (deterministic ascending order), len()/is_empty(), resident_count()/degraded_count() (disjoint).
• ResidencyPolicyInput<'a> — boundary, catalog: &WeightRegionCatalog, budget_bytes: Option<u64>.
• ResidencyPolicy trait — name(), decide(&self, &ResidencyPolicyInput) -> ResidencyDecision. Immutable borrow only in, owned value out — no way to mutate external state through the signature (the "non-owning" proof).
• WholeBankResidentPolicy — the only shipped policy; always resident, surfacing the catalog's own NonPageableReason (if any) rather than inventing one.
• plan_residency(candidates, policy, budget_bytes) -> ResidencyPlan — the single plan-creation call site.
• validate_decision(catalog, decision) — fail-closed: any invalid PerExpertCandidate (out-of-range index, duplicate index, proposed over a non-pageable catalog, or an index with no real region) degrades that single entry to WholeBankResident { reason: Some(RejectedByValidation(..)) }. Never panics, never fails the whole plan.

Call sites (requirement #2)

• Plan creation: Executor::build_with_cuda_requirement (build.rs) calls plan_default_residency() immediately after computing expert_region_candidates, using WholeBankResidentPolicy. Result stored as Executor.residency_plan (new field, state.rs), read-only under #[cfg(test)] today.
• Execution/validation: plan_residency() itself validates every decision before it's accepted into the plan (see above) — this is "existing infrastructure" in the sense that no new dispatch-time consumer was introduced; no production path acts on PerExpertCandidate yet, matching the explicit scope (no kernel ABI changes, no CUDA/weight_paging.rs touch).
• Telemetry: a new "session.residency_plan" trace span records policy, resident, degraded, total, using the existing trace_span/Args infra — no duplicate counter/state.

Non-owning proof

ResidencyPlan/ResidencyPolicy/ResidencyPolicyInput hold no allocator, stream, or pointer/handle fields — only value-level data (ValueId, usize indices, &'static str, catalog references). The trait's decide() takes &self and an immutable input reference and returns an owned enum — it has no path to allocate, copy, or synchronize.

Substitutability proof

A test-only AlwaysSplitPolicy/AlwaysSplit policy (in both weight.rs and executor/tests.rs) proposes PerExpertCandidate for every pageable catalog and is accepted by plan_residency()/wired through plan_residency_with() (a test/measurement-only seam alongside the production plan_default_residency()), proving the trait boundary is real, not an inert marker — without shipping LRU/static-hot/q*/prefetch-execution behavior.

Tests (same commit)

• weight.rs: default-policy equivalence for pageable and non-pageable catalogs, alternate-policy substitutability, three invalid-decision rejections (out-of-range index, duplicate index, per-expert-over-non-pageable), deterministic ordering.
• executor/tests.rs: default plan matches whole-bank-resident with unchanged weight_handles output; non-pageable reason surfaced without changing handles; seam substitutability at the executor's own candidate wiring, with the default Executor.residency_plan unaffected by the substitution.

Local evidence

• cargo test -p onnx-runtime-ep-api --lib: 71 passed.
• cargo test -p onnx-runtime-session --lib: 201 passed (+3 new vs. pre-slice 198).
• cargo test -p onnx-genai-engine --lib: 445 passed, 1 ignored (unaffected, run as a broader regression check).
• cargo clippy -p onnx-runtime-ep-api -p onnx-runtime-session --lib --tests: clean.
• cargo fmt on touched crates: clean.
• No CUDA/weight_paging.rs/kernel-ABI files touched in this diff — a CUDA/VMM comparison run is not applicable to this slice (verified via git diff --stat, only onnx-runtime-ep-api/onnx-runtime-session files changed); 5x A100 confirmed available for future slices that do touch CUDA paths.

Reviewer verdict

Independent code-review agent found one substantive issue: resident_count()'s doc comment said "no degradation reason" but the implementation counted all WholeBankResident entries including degraded ones, so resident_count() + degraded_count() weren't disjoint as documented (misleading telemetry). Fixed: resident_count() now filters on reason: None to match its documented semantics and stay disjoint from degraded_count(); corresponding tests updated. Re-ran all targeted tests after the fix — all green.

Explicitly out of scope (per Roy's staged plan, .squad/decisions/inbox/roy-residency-policy-gated-on-82.md)

No static-hot algorithm, LRU mutation, plan_double_buffer/prefetch execution wiring, q*, elastic resize, or kernel ABI changes. weight_paging.rs's CudaWeightResidency/eviction_for_boundary remain untouched and the single non-pluggable authority for allocation/accounting/VA/sync.

Next policy-implementation slice

Per Roy: refactor eviction_for_boundary/static-pin-threshold logic in weight_paging.rs into a real ResidencyPolicy implementation (byte-identical behavior first), then wire the already-implemented plan_double_buffer (executor/prefetch.rs) through the plan for actual prefetch execution.

Part of #82 — this is one slice of a staged plan; issue should remain open for the follow-ups above.

Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com

[codecov]

Codecov Report

❌ Patch coverage is 94.81481% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.28%. Comparing base (fbccba5) to head (09c126c).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-runtime-ep-api/src/weight.rs	94.81%	13 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1784      +/-   ##
==========================================
- Coverage   80.64%   80.28%   -0.37%     
==========================================
  Files         411      411              
  Lines      199652   199922     +270     
  Branches   199652   199922     +270     
==========================================
- Hits       161019   160501     -518     
- Misses      33175    33962     +787     
- Partials     5458     5459       +1     

Flag	Coverage Δ
cli-ort-linux	72.47% <ø> (ø)
cli-ort-windows	72.06% <ø> (+0.09%)	⬆️
mlas	85.20% <ø> (-0.14%)	⬇️
offline	80.41% <94.81%> (-0.39%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
crates/onnx-runtime-session/src/executor/state.rs	83.33% <ø> (ø)
crates/onnx-runtime-ep-api/src/weight.rs	83.82% <94.81%> (+10.05%)	⬆️

... and 11 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.