From d3b678420f816395da33ada5377550ef1c4cf061 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 22 Aug 2026 21:50:06 +0000 Subject: [PATCH] feat(qmoe): expose per-expert region candidates for the expert bank (#82) 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, 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> --- .../onnx-genai-engine/src/engine/placement.rs | 62 ++---- crates/onnx-runtime-loader/src/lib.rs | 2 +- crates/onnx-runtime-loader/src/weights.rs | 130 ++++++++++++ .../src/executor/build.rs | 200 +++++++++++++++++- .../src/executor/state.rs | 17 ++ .../src/executor/tests.rs | 163 ++++++++++++++ 6 files changed, 519 insertions(+), 55 deletions(-) diff --git a/crates/onnx-genai-engine/src/engine/placement.rs b/crates/onnx-genai-engine/src/engine/placement.rs index 3119a4007a..ab7908280b 100644 --- a/crates/onnx-genai-engine/src/engine/placement.rs +++ b/crates/onnx-genai-engine/src/engine/placement.rs @@ -1,9 +1,7 @@ use super::*; use onnx_runtime_ep_cpu::{LayerWeightRegions, PlacementPlan, plan_placement}; use onnx_runtime_ir::{Graph, Node, WeightRef}; -use onnx_runtime_loader::{ - ExpertQuantization, ExpertStorageOrder, ExpertTensorLayout, WeightRegionCatalog, -}; +use onnx_runtime_loader::{WeightRegionCatalog, qmoe_expert_tensor_layout}; use onnx_runtime_session::InferenceSession; impl From for WeightPlacementReport { @@ -92,17 +90,12 @@ fn qmoe_layer( if block_size == 0 { return Ok(None); } - let quantization = |blocks_per_row| ExpertQuantization { - bits: bits as usize, - block_size, - blocks_per_row, - }; - + let bits = bits as usize; let mut regions = Vec::new(); - push_qmoe_pair(graph, node, 2, 3, 11, quantization, &mut regions)?; - push_qmoe_pair(graph, node, 5, 6, 12, quantization, &mut regions)?; + push_qmoe_pair(graph, node, 2, 3, 11, bits, block_size, &mut regions)?; + push_qmoe_pair(graph, node, 5, 6, 12, bits, block_size, &mut regions)?; if node.inputs.get(8).and_then(|slot| *slot).is_some() { - push_qmoe_pair(graph, node, 8, 9, 13, quantization, &mut regions)?; + push_qmoe_pair(graph, node, 8, 9, 13, bits, block_size, &mut regions)?; } if regions.is_empty() { return Ok(None); @@ -125,7 +118,8 @@ fn push_qmoe_pair( packed_index: usize, scales_index: usize, zero_points_index: usize, - quantization: impl Fn(usize) -> ExpertQuantization, + bits: usize, + block_size: usize, regions: &mut Vec, ) -> anyhow::Result<()> { let packed = required_initializer(graph, node, packed_index)?; @@ -133,41 +127,17 @@ fn push_qmoe_pair( let packed_dims = qmoe_dims(packed, packed_index)?; let scale_dims = qmoe_dims(scales, scales_index)?; let blocks_per_row = scale_dims[2]; - regions.push(WeightRegionCatalog::classify( - packed, - ExpertTensorLayout { - version: 1, - experts: packed_dims[0], - rows_per_expert: packed_dims[1], - storage_elements_per_row: packed_dims[2], - order: ExpertStorageOrder::ExpertMajor, - quantization: Some(quantization(blocks_per_row)), - }, - )); - regions.push(WeightRegionCatalog::classify( - scales, - ExpertTensorLayout { - version: 1, - experts: scale_dims[0], - rows_per_expert: scale_dims[1], - storage_elements_per_row: scale_dims[2], - order: ExpertStorageOrder::ExpertMajor, - quantization: Some(quantization(blocks_per_row)), - }, - )); + let packed_layout = qmoe_expert_tensor_layout(bits, block_size, blocks_per_row, packed_dims) + .with_context(|| format!("QMoE input {packed_index} must be rank-3"))?; + regions.push(WeightRegionCatalog::classify(packed, packed_layout)); + let scales_layout = qmoe_expert_tensor_layout(bits, block_size, blocks_per_row, scale_dims) + .with_context(|| format!("QMoE input {scales_index} must be rank-3"))?; + regions.push(WeightRegionCatalog::classify(scales, scales_layout)); if let Some(zero_points) = optional_initializer(graph, node, zero_points_index)? { let zero_dims = qmoe_dims(zero_points, zero_points_index)?; - regions.push(WeightRegionCatalog::classify( - zero_points, - ExpertTensorLayout { - version: 1, - experts: zero_dims[0], - rows_per_expert: zero_dims[1], - storage_elements_per_row: zero_dims[2], - order: ExpertStorageOrder::ExpertMajor, - quantization: Some(quantization(blocks_per_row)), - }, - )); + let zero_layout = qmoe_expert_tensor_layout(bits, block_size, blocks_per_row, zero_dims) + .with_context(|| format!("QMoE input {zero_points_index} must be rank-3"))?; + regions.push(WeightRegionCatalog::classify(zero_points, zero_layout)); } Ok(()) } diff --git a/crates/onnx-runtime-loader/src/lib.rs b/crates/onnx-runtime-loader/src/lib.rs index 5bdfcdb057..e1b18178ee 100644 --- a/crates/onnx-runtime-loader/src/lib.rs +++ b/crates/onnx-runtime-loader/src/lib.rs @@ -66,7 +66,7 @@ pub use epcontext::{ pub use error::LoaderError; pub use weights::{ ExpertQuantization, ExpertStorageOrder, ExpertTensorLayout, ExpertWeightRegion, - NonPageableReason, Pageability, WeightRegionCatalog, WeightStore, + NonPageableReason, Pageability, WeightRegionCatalog, WeightStore, qmoe_expert_tensor_layout, }; pub use writer::{EpContextDumpConfig, EpContextPartition, dump_ep_context}; diff --git a/crates/onnx-runtime-loader/src/weights.rs b/crates/onnx-runtime-loader/src/weights.rs index 9f771cbc23..4c83bbf17b 100644 --- a/crates/onnx-runtime-loader/src/weights.rs +++ b/crates/onnx-runtime-loader/src/weights.rs @@ -498,6 +498,40 @@ impl WeightRegionCatalog { } } +/// Derive the expert-major [`ExpertTensorLayout`] for one `com.microsoft::QMoE` +/// input tensor (a packed weight, scale, or zero-point tensor), given the +/// node-level quantization attributes and this tensor's own dims. +/// +/// This is the single source of truth for "which axis is which" on a QMoE +/// expert-bank tensor, shared by advisory host/device placement +/// (`onnx-genai-engine::engine::placement`) and the lazy-weight paging seam +/// (`onnx-runtime-session::executor::build`) so the two call sites cannot +/// silently drift onto different byte-range formulas. Returns `None` when +/// `dims` is not rank-3 (i.e. not an expert-major tensor at all); callers +/// treat that as "not applicable", not as a paging failure. +pub fn qmoe_expert_tensor_layout( + bits: usize, + block_size: usize, + blocks_per_row: usize, + dims: &[usize], +) -> Option { + if dims.len() != 3 { + return None; + } + Some(ExpertTensorLayout { + version: 1, + experts: dims[0], + rows_per_expert: dims[1], + storage_elements_per_row: dims[2], + order: ExpertStorageOrder::ExpertMajor, + quantization: Some(ExpertQuantization { + bits, + block_size, + blocks_per_row, + }), + }) +} + /// Overflow error shared by loader range catalogs and paging-aware kernels. #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] #[error("{0}")] @@ -1296,4 +1330,100 @@ mod tests { if message.contains("endpoint") )); } + + #[test] + fn qmoe_expert_tensor_layout_rejects_non_rank3_dims() { + assert_eq!(qmoe_expert_tensor_layout(4, 32, 2, &[8, 16]), None); + assert_eq!(qmoe_expert_tensor_layout(4, 32, 2, &[8, 16, 4, 2]), None); + assert_eq!(qmoe_expert_tensor_layout(4, 32, 2, &[]), None); + } + + #[test] + fn qmoe_expert_tensor_layout_populates_expert_major_fields_for_rank3_dims() { + let layout = qmoe_expert_tensor_layout(4, 32, 3, &[8, 16, 24]) + .expect("rank-3 dims must derive a layout"); + assert_eq!(layout.version, 1); + assert_eq!(layout.experts, 8); + assert_eq!(layout.rows_per_expert, 16); + assert_eq!(layout.storage_elements_per_row, 24); + assert_eq!(layout.order, ExpertStorageOrder::ExpertMajor); + assert_eq!( + layout.quantization, + Some(ExpertQuantization { + bits: 4, + block_size: 32, + blocks_per_row: 3, + }) + ); + } + + #[test] + fn qmoe_expert_tensor_layout_classifies_as_pageable_and_partitions_the_bank() { + let stamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock") + .as_nanos(); + let dir = std::env::current_dir().expect("cwd").join("target"); + std::fs::create_dir_all(&dir).expect("create target"); + let path = dir.join(format!( + "qmoe-expert-region-catalog-{}-{stamp}.bin", + std::process::id() + )); + // 4 experts * 16 rows * 24 storage elements (Uint8) per row. + let tensor_len = 4 * 16 * 24; + std::fs::write(&path, vec![0u8; tensor_len]).expect("write external data"); + + let weight = WeightRef::External { + path, + offset: 0, + length: tensor_len, + dtype: DataType::Uint8, + dims: vec![4, 16, 24], + }; + let layout = qmoe_expert_tensor_layout(4, 32, 3, weight.dims()) + .expect("rank-3 dims must derive a layout"); + let catalog = WeightRegionCatalog::classify(&weight, layout); + assert!(catalog.is_pageable()); + + // Regions exactly partition the tensor: contiguous, non-overlapping, + // and covering exactly `tensor_len` bytes across the 4 experts. + let mut expected_offset = 0usize; + let per_expert_len = 16 * 24; + for expert in 0..4 { + let range = catalog + .relative_range(expert) + .unwrap_or_else(|| panic!("expert {expert} must have a region")); + assert_eq!(range.start, expected_offset); + assert_eq!(range.end, expected_offset + per_expert_len); + expected_offset = range.end; + } + assert_eq!(expected_offset, tensor_len); + assert!(catalog.region(4).is_none()); + + std::fs::remove_file(weight_path_for_test(&catalog)).ok(); + } + + /// Test-only helper: re-derive the external file path a catalog built in + /// this module's own tests was classified from, so the test can clean up + /// after itself without threading the path through every assertion. + fn weight_path_for_test(catalog: &WeightRegionCatalog) -> PathBuf { + catalog.path.clone().unwrap_or_default() + } + + #[test] + fn qmoe_expert_tensor_layout_rejects_inline_tensor_with_reason() { + let inline = WeightRef::Inline(TensorData::from_raw( + DataType::Uint8, + vec![2, 4, 8], + vec![0u8; 2 * 4 * 8], + )); + let layout = qmoe_expert_tensor_layout(4, 32, 1, inline.dims()) + .expect("rank-3 dims must derive a layout"); + let catalog = WeightRegionCatalog::classify(&inline, layout); + assert!(!catalog.is_pageable()); + assert_eq!( + catalog.pageability(), + &Pageability::NonPageable(NonPageableReason::InlineTensor) + ); + } } diff --git a/crates/onnx-runtime-session/src/executor/build.rs b/crates/onnx-runtime-session/src/executor/build.rs index 43e198c3f2..854902aeb8 100644 --- a/crates/onnx-runtime-session/src/executor/build.rs +++ b/crates/onnx-runtime-session/src/executor/build.rs @@ -364,21 +364,44 @@ pub(super) fn format_node_identity(scope: &str, node_id: NodeId, node: &Node) -> } } -pub(super) fn build_lazy_weight_handles( +/// Build lazy-weight handles exactly as before, plus (additively) per-expert +/// region candidates for `com.microsoft::QMoE` expert-bank initializers. +/// +/// The candidate map is measurement-neutral bookkeeping only (issue #82's +/// first slice): nothing here changes which initializers become lazy, how +/// many [`ExternalMmapRegion`]s a [`LazyWeight`] carries, or how many +/// allocations/H2D copies the CUDA paging path performs. A QMoE initializer +/// whose layout cannot be proven pageable is recorded with its +/// [`onnx_runtime_loader::NonPageableReason`] rather than silently omitted. +pub(super) fn build_lazy_weight_handles_and_candidates( graph: &Graph, weights: &Arc, ep: &dyn ExecutionProvider, -) -> Result> { +) -> Result<( + HashMap, + HashMap, +)> { let capabilities = ep.capabilities(); if !capabilities.advertises(onnx_runtime_ep_api::NXRT_WEIGHT_PAGING_CAPABILITY) { - return Ok(HashMap::new()); + return Ok((HashMap::new(), HashMap::new())); } + let qmoe_layouts = qmoe_expert_layouts_by_value(graph); + let mut handles = HashMap::new(); + let mut candidates = HashMap::new(); for candidate in lazy_weight_candidates(graph) { let value = candidate.value; let boundary = candidate.boundary; let weight = &graph.initializers[&value]; + if boundary == LazyWeightBoundary::QMoe + && let Some(layout) = qmoe_layouts.get(&value) + { + candidates.insert( + value, + onnx_runtime_loader::WeightRegionCatalog::classify(weight, layout.clone()), + ); + } let Some((mapping_id, offset, len)) = weights.external_mmap_provenance(weight) else { continue; }; @@ -407,10 +430,168 @@ pub(super) fn build_lazy_weight_handles( })?; handles.insert(value, WeightHandle::Lazy(lazy)); } - Ok(handles) + Ok((handles, candidates)) +} + +/// Derive an [`onnx_runtime_loader::ExpertTensorLayout`] for every QMoE +/// expert-bank initializer value in `graph`, keyed by the initializer's +/// [`ValueId`]. +/// +/// Mirrors `onnx-genai-engine::engine::placement::qmoe_layers`'s input-index +/// convention (fc1 packed/scales/zero-points at 2/3/11, fc2 at 5/6/12, and an +/// optional fc3 triple at 8/9/13) so both call sites derive layouts from the +/// same node attributes/initializer shapes and cannot silently drift onto +/// different byte-range formulas. A value absent from the returned map either +/// isn't a QMoE expert-bank tensor, or its node/attributes/shape could not +/// support deriving a layout at all (as opposed to deriving one that then +/// fails [`onnx_runtime_loader::WeightRegionCatalog::classify`]'s validation, +/// which is instead recorded as a `NonPageable` catalog by the caller). +fn qmoe_expert_layouts_by_value( + graph: &Graph, +) -> HashMap { + let mut layouts = HashMap::new(); + for (_, node) in graph.nodes.iter() { + if !(node.op_type == "QMoE" && (node.domain.is_empty() || node.domain == "com.microsoft")) { + continue; + } + let Some(bits) = qmoe_int_attr(node, "expert_weight_bits", 4) + .ok() + .and_then(|value| usize::try_from(value).ok()) + .filter(|bits| matches!(bits, 1 | 2 | 4 | 8)) + else { + continue; + }; + let Some(block_size) = qmoe_int_attr(node, "block_size", 0) + .ok() + .and_then(|value| usize::try_from(value).ok()) + .filter(|block_size| *block_size > 0) + else { + continue; + }; + for &(packed_idx, scales_idx, zero_idx) in &[(2, 3, 11), (5, 6, 12), (8, 9, 13)] { + if packed_idx == 8 && node.inputs.get(8).and_then(|slot| *slot).is_none() { + continue; + } + qmoe_push_expert_layout( + graph, + node, + scales_idx, + packed_idx, + bits, + block_size, + &mut layouts, + ); + qmoe_push_expert_layout( + graph, + node, + scales_idx, + scales_idx, + bits, + block_size, + &mut layouts, + ); + if node.inputs.get(zero_idx).and_then(|slot| *slot).is_some() { + qmoe_push_expert_layout( + graph, + node, + scales_idx, + zero_idx, + bits, + block_size, + &mut layouts, + ); + } + } + } + layouts +} + +/// Derive and record the layout for one QMoE input tensor at `input_index`, +/// using the sibling scales tensor at `scales_index` to read `blocks_per_row` +/// (the scales tensor's own storage-elements-per-row axis). +/// +/// Whenever `value` resolves to a graph initializer, an entry is always +/// recorded (never silently omitted): a derivable expert-major layout is +/// recorded as-is, and a value whose shape/sibling scales cannot support +/// deriving one (e.g. non-rank-3 dims) is instead recorded with +/// [`onnx_runtime_loader::ExpertStorageOrder::Interleaved`], which +/// [`onnx_runtime_loader::WeightRegionCatalog::classify`] turns into an +/// explicit `NonPageableReason::NotExpertMajor` rather than the value +/// vanishing from the candidate map entirely. +fn qmoe_push_expert_layout( + graph: &Graph, + node: &Node, + scales_index: usize, + input_index: usize, + bits: usize, + block_size: usize, + layouts: &mut HashMap, +) -> Option<()> { + let value = node.inputs.get(input_index).and_then(|slot| *slot)?; + let weight = qmoe_initializer(graph, value)?; + let derived = (|| { + let scales_value = node.inputs.get(scales_index).and_then(|slot| *slot)?; + let scales_weight = qmoe_initializer(graph, scales_value)?; + let scales_dims = scales_weight.dims(); + if scales_dims.len() != 3 { + return None; + } + let blocks_per_row = scales_dims[2]; + onnx_runtime_loader::qmoe_expert_tensor_layout( + bits, + block_size, + blocks_per_row, + weight.dims(), + ) + })(); + let layout = derived.unwrap_or(onnx_runtime_loader::ExpertTensorLayout { + version: 1, + experts: 0, + rows_per_expert: 0, + storage_elements_per_row: 0, + order: onnx_runtime_loader::ExpertStorageOrder::Interleaved, + quantization: None, + }); + layouts.insert(value, layout); + Some(()) +} + +/// Resolve `value` to the [`onnx_runtime_ir::WeightRef`] it initializes, +/// following one `Cast` producer (the same convention +/// `onnx-genai-engine::engine::placement` uses for QMoE inputs that a graph +/// rewrite has cast to a wider compute dtype). +fn qmoe_initializer(graph: &Graph, value: ValueId) -> Option<&onnx_runtime_ir::WeightRef> { + if let Some(weight) = graph.initializers.get(&value) { + return Some(weight); + } + let producer = graph.values.get(value)?.producer?; + let node = graph.nodes.get(producer)?; + if node.domain.is_empty() && node.op_type == "Cast" { + let source = node.inputs.first().and_then(|slot| *slot)?; + return graph.initializers.get(&source); + } + None +} + +fn qmoe_int_attr(node: &Node, name: &'static str, default: i64) -> std::result::Result { + match node.attr(name) { + Some(attr) => attr.as_int().ok_or(()), + None => Ok(default), + } } impl Executor { + /// Per-expert region candidates recorded for `com.microsoft::QMoE` lazy + /// weights, keyed by initializer [`ValueId`]. See + /// [`Executor::expert_region_candidates`](state) for the additive, + /// measurement-neutral contract this exposes. + #[cfg(test)] + pub(super) fn expert_region_candidates( + &self, + ) -> &HashMap { + &self.expert_region_candidates + } + /// Compile a graph + weights into a runnable executor on the CPU EP. pub(crate) fn build( graph: Graph, @@ -435,17 +616,19 @@ impl Executor { Self::place_graph(&mut graph, &weights, &mut ep, require_cuda)?; // Topological order up front: also validates the selected graph is a DAG. let order = graph.topological_order()?; - let weight_handles = { + let (weight_handles, expert_region_candidates) = { let mut span = trace_span("session.lazy_weight_handles", "session"); - let handles = build_lazy_weight_handles(&graph, &weights, ep.as_ref())?; + let (handles, candidates) = + build_lazy_weight_handles_and_candidates(&graph, &weights, ep.as_ref())?; if let Some(span) = span.as_mut() { span.set_args( Args::new() .with("handles", handles.len() as u64) - .with("initializers", graph.initializers.len() as u64), + .with("initializers", graph.initializers.len() as u64) + .with("expert_region_candidates", candidates.len() as u64), ); } - handles + (handles, candidates) }; let (mut value_shapes, mut value_dtypes, buffers, buffer_shapes) = @@ -482,6 +665,7 @@ impl Executor { ep, graph_slot: DeviceGraphSlot::Primary, weight_handles, + expert_region_candidates, 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 52814c5cd7..936fda24b7 100644 --- a/crates/onnx-runtime-session/src/executor/state.rs +++ b/crates/onnx-runtime-session/src/executor/state.rs @@ -88,6 +88,23 @@ pub(crate) struct Executor { /// Lazy external initializers available only at the nxrt fused-MoE boundary. /// Stock EPs ignore this map and keep receiving the resident buffers below. pub(super) weight_handles: HashMap, + /// Per-expert region candidates for `com.microsoft::QMoE` expert-bank + /// initializers, keyed by the same [`ValueId`] as [`Self::weight_handles`]. + /// + /// This is purely additive bookkeeping for a future pluggable residency + /// policy (see issue #82): it records *where each expert's bytes would + /// live* inside an already-fully-resident lazy weight, but nothing reads it + /// to change binding, allocation, or paging behavior today. A value present + /// here is classified [`onnx_runtime_loader::Pageability::Pageable`]; a + /// value classified `NonPageable` records its rejection reason for + /// diagnostics instead, so no entry silently vanishes — either the + /// candidate is a validated partition of the expert bank, or the exact + /// reason it isn't is recorded. + /// + /// Read only under `#[cfg(test)]` via [`Executor::expert_region_candidates`] + /// today; production consumers land in a follow-up slice per issue #82. + #[allow(dead_code)] + pub(super) expert_region_candidates: HashMap, 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 5e4bb20f74..d30210b339 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -3286,6 +3286,7 @@ impl ExecutionProvider for WeightDeliveryEp { } if LazyWeightBoundary::BlockQuantizedMoe.matches(&op.domain, &op.op_type) || LazyWeightBoundary::MatMulNBits.matches(&op.domain, &op.op_type) + || LazyWeightBoundary::QMoe.matches(&op.domain, &op.op_type) || (op.is_default_domain() && op.op_type == "Identity") { KernelMatch::Supported { @@ -3881,6 +3882,168 @@ fn two_node_weight_delivery_fixture() -> (Graph, Arc, std::path::Pa (graph, Arc::new(store), path) } +/// A `com.microsoft::QMoE` node with one expert-major fc1 packed/scales pair +/// (2 experts, 3 rows/expert, 4 storage elements/row) as external +/// initializers, plus an activation input and output so it type-checks as a +/// minimal QMoE graph. Used to exercise issue #82's per-expert region +/// candidate bookkeeping without touching kernel execution. +fn qmoe_expert_region_fixture() -> (Graph, Arc, std::path::PathBuf) { + static NEXT_FILE: AtomicU64 = AtomicU64::new(0); + let root = std::env::var_os("CARGO_TARGET_DIR") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| std::env::current_dir().unwrap().join("target")) + .join("weight-handle-tests"); + std::fs::create_dir_all(&root).unwrap(); + let id = NEXT_FILE.fetch_add(1, Ordering::Relaxed); + let path = root.join(format!( + "qmoe-expert-region-{}-{id}.bin", + std::process::id() + )); + // fc1 packed: 2 experts * 3 rows * 4 storage elements (Uint8) = 24 bytes. + // fc1 scales: 2 experts * 3 rows * 1 block = 6 bytes (Float32 dtype but we + // only need declared length/dims to match for this fixture). + std::fs::write(&path, [0u8; 24 + 6]).unwrap(); + + let mut graph = Graph::new(); + graph.opset_imports.insert("com.microsoft".into(), 1); + let activation = + graph.create_named_value("activation", DataType::Float32, static_shape([1, 4])); + graph.add_input(activation); + let packed = graph.create_named_value("fc1_packed", DataType::Uint8, static_shape([2, 3, 4])); + graph.set_initializer( + packed, + WeightRef::External { + path: path.clone(), + offset: 0, + length: 24, + dtype: DataType::Uint8, + dims: vec![2, 3, 4], + }, + ); + let scales = graph.create_named_value("fc1_scales", DataType::Uint8, static_shape([2, 3, 1])); + graph.set_initializer( + scales, + WeightRef::External { + path: path.clone(), + offset: 24, + length: 6, + dtype: DataType::Uint8, + dims: vec![2, 3, 1], + }, + ); + let output = graph.create_named_value("output", DataType::Float32, static_shape([1, 4])); + let mut node = Node::new( + NodeId(0), + "QMoE", + vec![ + Some(activation), + None, + Some(packed), + Some(scales), + None, + Some(packed), + Some(scales), + ], + vec![output], + ); + node.domain = "com.microsoft".to_owned(); + node.attributes + .insert("expert_weight_bits".to_owned(), Attribute::Int(4)); + node.attributes + .insert("block_size".to_owned(), Attribute::Int(32)); + graph.insert_node(node); + graph.add_output(output); + + let mut store = WeightStore::new(); + store.map_external(&path).unwrap(); + (graph, Arc::new(store), path) +} + +#[test] +fn qmoe_expert_region_candidates_partition_expert_major_bank_without_changing_handle_output() { + 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(); + + // Existing behavior unchanged: still exactly one lazy weight handle per + // initializer (fc1 packed + fc1 scales), each with exactly one region. + assert_eq!(executor.weight_handles.len(), 2); + for handle in executor.weight_handles.values() { + let WeightHandle::Lazy(lazy) = handle else { + panic!("expected a lazy weight handle"); + }; + assert_eq!(lazy.regions.len(), 1); + } + + // Additive: both QMoE initializers got a per-expert region candidate. + let candidates = executor.expert_region_candidates(); + assert_eq!(candidates.len(), 2); + for catalog in candidates.values() { + assert!(catalog.is_pageable(), "{:?}", catalog.pageability()); + // Regions exactly partition the tensor: contiguous, non-overlapping. + let mut expected_start = 0usize; + for expert in 0..2 { + let range = catalog + .relative_range(expert) + .unwrap_or_else(|| panic!("expert {expert} must have a region")); + assert_eq!(range.start, expected_start); + expected_start = range.end; + } + assert!(catalog.region(2).is_none()); + } +} + +#[test] +fn qmoe_expert_region_candidates_record_reason_for_non_rank3_layout() { + let (mut graph, weights, path) = qmoe_expert_region_fixture(); + // Corrupt the fc1 packed initializer to rank-2 so it can no longer derive + // an expert-major layout; the candidate map must still record a reason + // via `NonPageableReason`, not silently omit or panic. + 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(); + + // Rank-2 dims cannot derive an expert-major layout, but the value must + // still be recorded (not silently dropped) with an explicit + // `NonPageableReason`; the untouched rank-3 scales tensor is still + // pageable. + let candidates = executor.expert_region_candidates(); + assert_eq!(candidates.len(), 2); + let packed_catalog = candidates + .get(&packed_value) + .expect("rank-2 fc1_packed must still be recorded, not silently omitted"); + assert!(!packed_catalog.is_pageable()); + assert_eq!( + packed_catalog.pageability(), + &onnx_runtime_loader::Pageability::NonPageable( + onnx_runtime_loader::NonPageableReason::NotExpertMajor + ) + ); + + // The weight handle map is unaffected either way: still exactly one lazy + // handle per initializer, proving the invalid layout does not change + // existing binder/allocation behavior. + assert_eq!(executor.weight_handles.len(), 2); +} + #[test] fn claim_time_optional_input_dtype_is_undefined_not_silently_float32() { let mut graph = Graph::new();