feat(qmoe): expose per-expert region candidates for the expert bank (#82)

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<PlacementPlan> 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,49 +118,26 @@ 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<WeightRegionCatalog>,
 ) -> anyhow::Result<()> {
     let packed = required_initializer(graph, node, packed_index)?;
     let scales = required_initializer(graph, node, scales_index)?;
     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(())
 }

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};
 

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<ExpertTensorLayout> {
+    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)
+        );
+    }
 }