Repository navigation
feat(qmoe): expose per-expert region candidates for the expert bank (#82) - #1779
Merged
Merged
Conversation
) First implementation slice for issue #82: introduce a typed per-expert resource identity/region candidate seam for the QMoE expert bank so a future pluggable residency policy has real per-expert units, without changing residency behavior, kernel ABI/math, or adding a policy trait. - Add onnx_runtime_loader::qmoe_expert_tensor_layout(bits, block_size, blocks_per_row, dims) as the single source of truth for deriving a QMoE tensor's ExpertTensorLayout from node attributes + tensor dims; reused by both the existing advisory placement pass (engine::placement, refactored to call it instead of constructing ExpertTensorLayout inline) and the new executor build-time seam, so the two call sites cannot drift onto different byte-range formulas. - Add Executor::expert_region_candidates: HashMap<ValueId, WeightRegionCatalog>, populated additively alongside the existing lazy-weight handle map (build_lazy_weight_handles_and_candidates) for com.microsoft::QMoE expert-bank initializers. Existing WeightHandle/LazyWeight construction is unchanged in shape and count; nothing yet reads this new map outside tests. - Every QMoE-boundary initializer that is a lazy-weight candidate now gets an entry: a derivable expert-major layout is classified as Pageable with its per-expert byte ranges, and a value whose layout cannot be derived or fails validation is classified NonPageable with an explicit NonPageableReason (e.g. NotExpertMajor) rather than silently omitted from the map. - Tests: unit tests for qmoe_expert_tensor_layout's rank-3 requirement and field population, a WeightRegionCatalog pageable-partition test proving expert regions are contiguous/non-overlapping and cover the full tensor, an inline-tensor rejection-reason test, and two executor-level tests proving (a) region candidates partition the expert-major bank while weight_handles/regions are byte-identical to pre-change behavior, and (b) a non-rank-3 layout is recorded with an explicit NonPageableReason instead of vanishing from the map. No kernel, CUDA weight_paging.rs, or ABI changes; whole-bank-resident execution and existing binder/allocation counts are preserved by construction and by the added tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1779 +/- ##
==========================================
- Coverage 80.70% 80.63% -0.07%
==========================================
Files 411 411
Lines 199384 199648 +264
Branches 199384 199648 +264
==========================================
+ Hits 160906 160985 +79
- Misses 33024 33200 +176
- Partials 5454 5463 +9
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
justinchuby
added a commit
that referenced
this pull request
Aug 22, 2026
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> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
pushed a commit
that referenced
this pull request
Aug 22, 2026
…orkflow Architecture A, completed. A bare-decoder package keeps `model.io` as its sole serialized answer; the loader compiles that ABI into a canonical WorkflowSpec in memory and the interpreter's canonical decode loop executes it. There is no longer a direct decode path to fall back to. The split that makes this real rather than a wrapper is in `decode_loop.rs`: * `forward_step` is the `onnx-genai.autoregressive-decode` primitive -- one decoder forward pass, taking the device greedy-argmax or device-sampling fast path when it applies. KV stays the executor's business, exactly as the lowered workflow declares (`management: runtime`), so device residency is untouched. * `select_and_commit_step` is the `onnx-genai.token-policy` primitive -- the logit-processor chain, sampler, logprobs, KV commit, and stop/EOS detection. Both are the *only* implementations. `step_decode_loop` is now written on them, so a caller cannot reach a second copy of the policy. `pipeline::canonical_decode::run_canonical_decode` is the one token loop. It owns iteration count, the context-limit stop, the stop decision, and the emitted token stream; the body is forward-then-policy. `run_decode_loop` -- the second token loop only the direct `Engine` path could drive -- is deleted, and its unit tests re-point at the loop that remains. The ORT session path, both native decode sessions, and the workflow path now share it. Fail-closed, not mode-switched. `install_canonical_workflow` runs at load for every decoder package and asserts the lowered workflow declares the two contracts this runtime implements; a package that cannot be lowered fails to load. `generate_with_callbacks` then refuses outright if no canonical form is present, naming the fact that the old path no longer exists. Two hand-built test doubles tripped that guard on the first run -- they skip the loader -- and now carry the same lowering a real minimal decoder gets, rather than the guard being weakened. Parity, on real models: greedy goldens for phi35-mini int4, phi4-mini CUDA, qwen2.5-0.5b CUDA, qwen3-0.6b, and gpt-oss-20b are byte-identical to the pre-change baseline captured at c58eb5b. The tokens did not move; only who produced them did. `canonical_execution_parity` (8) pins what goldens cannot: that a bare decoder is lowered at load and reports `provenance: lowered` while still serializing only `model.io`; that prefill, cached decode on a reused session, deterministic greedy generation, EOS stopping, seeded sampling, and batching (where the decode path supports it) all run through the canonical loop; that an authored workflow is executed as authored and never lowered beside itself; and -- the case the whole change rests on -- that the legacy direct path cannot be selected: no constructor skips lowering, and a runtime with its canonical form removed refuses to generate instead of falling back. Verified: onnx-genai-engine 72 targets green under `native-backend`; server / CLI / C ABI / metadata / bench 30 green; H200 native-cuda green (canonical_execution_parity 8, native_workflow_parity 12, corpus 5, one_runtime_e2e 5, fixture floor 3); `cargo fmt --all --check`; clippy `-D warnings` clean apart from a `too_many_arguments` lint on `engine/placement.rs`, which arrived with main (#1779) and this branch does not touch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
…paging authorities (#82) (#1789) Part of #82 (DeepSeek-V4/GLM-5.2 QMoE vertical slice). 4th slice of the pluggable residency-policy line: #1779 (region candidates) → #1784 (policy seam) → #1787 (migrate existing decision logic into `CudaHotSetResidencyPolicy`) → this PR (elastic resize). ## Goal Introduce FreeToken-style elastic expert-residency **resizing** as a policy-issued resize *intent*, executed transactionally by the existing Resource Governor + PMM/VMM/weight-paging authorities. The policy never allocates/frees/copies/relocates a virtual address — it only asks; existing infrastructure executes. ## New contract (`onnx-runtime-ep-api::weight`) - `ResidencyResizeRequest { direction: Grow | Shrink, target_bytes, priority }` - `ResizeSafePoint { capturing, pending_deferred_releases, admission_in_flight, multi_device }` with `is_safe()` / `blocking_reason()` - `ResizeRejection` (`NotSafePoint` / `WouldExposeColdExpert` / `ExecutionFailed` / `NoOp`) - `ResidencyResizePlan::{Accepted(request), Rejected { request, reason }}` produced by the pure `plan_resize(request, safe_point)` — no allocation, copy, or VA touched during planning. - `ResidencyResizeOutcome` (before/after bytes, `accepted_bytes`, `rollback_count`, `safe_point`, `is_success()`). ## Executor side (`onnx-runtime-ep-cuda::weight_paging`) - `CudaWeightResidency::resize_safe_point(device_count)` builds the snapshot from **existing** signals only: `CudaRuntime::is_capturing()`, the provider deferred-release queue's pending count, and the existing in-flight-fill quarantine. No new state. - `CudaWeightResidency::execute_resize(plan, device_count)` is the *only* place a plan becomes bytes actually moved: - **Grow** reuses the existing `MemoryLease::grow` path — the same one `admit_committed_span`'s over-budget handling already uses. - **Shrink** reuses the existing `ReclaimableMappedHolder::reclaim_mapped` (governor-driven LRU eviction + deferred-release-queue-idle wait) when a mapped allowance is installed, or an equivalent plain-lease evict-then-shrink through the same `remove_page` helper otherwise (so process-global resident-byte/eviction telemetry stays in lockstep, not duplicated). - Re-validates the safe point immediately before committing (planning and execution can be separated in time). - Refuses to shrink an *ungoverned* budget (no lease/allowance) — mirrors `set_ungoverned_budget`'s existing refusal to shrink once nothing governs the budget. ## Safe-point invariant A resize may only commit when: no CUDA graph slot in use is capturing/replaying; the provider's deferred-release queue has nothing pending; no page admission is currently in flight; and the call is **not** running under multi-device/TP. No existing barrier/authority in this codebase coordinates a resize across devices, so multi-device fails closed with an explicit reason rather than inventing distributed synchronization (confirmed via search — no TP/multi-device coordination primitive exists yet). ## What can actually be reclaimed this slice Honest scope: only the existing dense/MatMul weight-paging cache's ordinary LRU-evictable, non-pinned resident bytes — the same population `reclaim_mapped` already touches. **No QMoE per-expert reclaim is enabled.** Whole-bank QMoE residency is unchanged; slice 2/3's "no cold expert may reach QMoE" invariant is preserved because nothing here changes per-expert eviction. Per-expert-aware eviction (requiring a routed-residency guarantee) is the next dependency gate. ## Default behavior Unchanged: `budget()`/`adopt_governed_budget()`/`set_ungoverned_budget()` are untouched, and every existing `weight_paging` test still passes without this seam ever being invoked. ## Tests - `crates/onnx-runtime-ep-api/src/weight.rs` — `plan_resize` acceptance/no-op/capture/pending-release/in-flight/multi-device rejection, deterministic blocking-reason priority, outcome success reporting. - `crates/onnx-runtime-ep-cuda/src/weight_paging.rs` — ungoverned cache refuses grow/shrink leaving budget untouched; governed grow moves exactly the requested bytes / fails over-capacity with budget unchanged; zero-byte resize is a rejected no-op; multi-device fails closed; repeated grow/shrink oscillation keeps accounting consistent; rejected-plan outcome preserves the original request's direction. ## Local validation (A100) - `cargo fmt` / `cargo clippy -D warnings` on both touched crates — clean. - `cargo test -p onnx-runtime-ep-cuda --features cuda --lib` — 524 passed, 0 failed, 31 ignored (GPU-only helpers not applicable), including a single-threaded rerun to rule out cross-test device contention. - `cargo test -p onnx-runtime-ep-api --lib` — 79 passed (+ new resize tests, 88 total counting weight.rs additions). - `cargo test -p onnx-runtime-session --lib --features cuda` — 205 passed. ## Independent review A code-review agent inspected the diff before this commit and found two substantive issues, both fixed prior to committing: 1. The plain-lease shrink fallback updated only per-instance accounting and skipped the process-global resident-byte/eviction gauges every other eviction call site updates — now routed through the existing `remove_page` helper so telemetry stays consistent. 2. A plan-time-rejected outcome always reported `direction: Grow` regardless of the original request — `ResidencyResizePlan::Rejected` now carries the original request so telemetry reports the correct direction/bytes even on rejection. ## Explicitly out of scope q* CPU/GPU split, predictor-based prefetch, a new allocator/cache, and enabling per-expert QMoE eviction/reclaim (next dependency gate, once a routed-residency guarantee exists). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
First implementation slice for #82, per Roy's architecture verdict: introduce a typed per-expert resource identity/region-candidate seam for the QMoE expert bank so a future pluggable residency policy has real per-expert units — without changing residency behavior, adding a
ResidencyPolicytrait, enabling paging, or touching kernel ABI/math. This slice is measurement-neutral: whole-bank-resident execution is unchanged.Architecture invariant preserved
PMM/VMM allocation/accounting, stable VA ownership, synchronization, and correctness stay non-pluggable single authorities. No second cache/allocator was created.
crates/onnx-runtime-ep-cuda/src/weight_paging.rsandkernels/qmoe.rsare not touched — this slice lives entirely in CPU-side/EP-agnostic code (onnx-runtime-loader,onnx-genai-engine,onnx-runtime-session), so CUDA/kernel behavior is unchanged by construction, not just by test.Exact new type/contract
onnx_runtime_loader::qmoe_expert_tensor_layout(bits, block_size, blocks_per_row, dims) -> Option<ExpertTensorLayout>— the single source of truth for deriving a QMoE tensor'sExpertTensorLayoutfrom node attributes + tensor dims. Reused by both the existing advisory placement pass (onnx-genai-engine::engine::placement, refactored to call it instead of constructingExpertTensorLayoutinline) and the new executor build-time seam, so the two call sites cannot silently drift onto different byte-range formulas.Executor::expert_region_candidates: HashMap<ValueId, WeightRegionCatalog>— populated additively inbuild_lazy_weight_handles_and_candidatesfor everycom.microsoft::QMoEexpert-bank initializer that is a lazy-weight candidate. Nothing production-side reads this map yet (by design); a#[cfg(test)]accessor exists so tests can assert against it.Supported / rejected layouts
WeightRegionCatalog::classify's validation is recorded asNonPageable(NonPageableReason::{NotExpertMajor, ShapeMismatch, InvalidQuantization, InlineTensor, ...})rather than silently omitted from the map.Tests (added in this commit)
onnx-runtime-loader::weights::tests:qmoe_expert_tensor_layoutrank-3 requirement + field population; a full pageable-partition test proving expert regions are contiguous/non-overlapping and exactly cover the tensor; an inline-tensor rejection-reason test.onnx-runtime-session::executor::tests:qmoe_expert_region_candidates_partition_expert_major_bank_without_changing_handle_output(region ranges partition the bank;weight_handles/LazyWeight.regionscounts are byte-identical to pre-change behavior) andqmoe_expert_region_candidates_record_reason_for_non_rank3_layout(an unsupported layout is recorded with an explicitNonPageableReasoninstead of vanishing).Results
cargo test -p onnx-runtime-loader --lib: 19 passedcargo test -p onnx-genai-engine --lib: 433 passed (placement.rs refactor unaffected)cargo test -p onnx-runtime-session --lib: 198 passedcargo clippy -p onnx-runtime-loader -p onnx-genai-engine -p onnx-runtime-session --tests: cleancargo fmt: cleanweight_paging.rs/kernels/qmoe.rsuntouched), so CUDA/paging behavior is unchanged by construction rather than by GPU measurement. A100s are available in this environment for the next slice, which does touch the CUDA paging path.Reviewer verdict
An independent code-review pass flagged one substantive issue: a value whose layout could not even be derived (e.g. non-rank-3 dims) was silently absent from
expert_region_candidatesrather than recorded with aNonPageableReason, contradicting the field's own doc comment. Fixed:qmoe_push_expert_layoutnow always records an entry for any QMoE-boundary initializer, falling back to anInterleaved-ordered layout (whichclassifyturns intoNonPageableReason::NotExpertMajor) when a layout cannot be derived. Re-ran all targeted tests after the fix — all green.Next dependency gate
Per Roy's verdict: the next slice introduces a
ResidencyPolicytrait as a byte-identical refactor that readsexpert_region_candidatesfor telemetry/decision-making, still without enabling paging or changing kernel behavior — this PR's map is exactly the seam that slice consumes.Closes #82 (partially — first slice only; residency/paging enforcement is out of scope here).