feat(moe): MoE extraction on mmap-backed GGUF layouts (RM-1462)

CHANGELOG.md

@@ -6,6 +6,14 @@ All notable changes to this project are documented in this file.
 
 ### Added
 
+- **MoE extraction on mmap-backed GGUF ([RM-1462](https://linear.app/rpd-34/issue/RM-1462)):**
+  [`GgufTensorAccess`] shares tensor-directory and payload access between
+  [`GgufLayout`] and [`GgufLayoutMmap`]. [`list_experts_from`] /
+  [`extract_expert_from`] (and `list_experts_mmap` / `extract_expert_mmap`
+  when the `mmap` feature is enabled) reuse the same stacked and per-expert
+  parsing as the existing owned-layout helpers. Expert payloads are still
+  copied into owned `Vec<u8>` buffers on return. Integration tests assert
+  owned/mmap parity on synthetic fixtures.
 - **Packed dequant for Q4_0 / Q4_1 / Q4_K / Q3_K / Q2_K (#125):**
   `dequantize_q4_0`, `dequantize_q4_1`, `dequantize_q2_k`, `dequantize_q3_k`,
   `dequantize_q4_k`, wired into `dequantize_packed` and `packed_row_size`

src/analysis/moe/extract.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-//! Per-expert weight extraction from a parsed [`GgufLayout`].
+//! Per-expert weight extraction from a parsed GGUF layout (owned or mmap).
 //!
 //! Two on-disk conventions are supported:
 //!
@@ -12,29 +12,50 @@
 //!    `blk.{b}.ffn_{role}.{e}.weight`. Whole-tensor payload is cloned.
 //!
 //! No arithmetic is performed — only byte-level slicing plus a `Vec`
-//! clone so the returned [`MoeExpertWeights`] is self-owning.
+//! clone so the returned [`MoeExpertWeights`] is self-owning. Payloads
+//! are always copied into `RawTensor::bytes` even when the layout borrows
+//! from an mmap mapping (callers get an owned buffer; see copy policy on
+//! [`extract_expert_from`]).
 
 use std::collections::BTreeSet;
 
 use super::expert::{MoeExpertWeights, RawTensor};
 use crate::error::{ParserError, Result};
-use crate::gguf::{GgufLayout, Tensor};
+use crate::gguf::{GgufLayout, GgufTensorAccess, Tensor};
+
+/// Enumerate every `(block, expert)` pair in an owned [`GgufLayout`].
+///
+/// Convenience wrapper around [`list_experts_from`].
+pub fn list_experts(layout: &GgufLayout) -> Vec<(usize, usize)> {
+    list_experts_from(layout)
+}
 
 /// Enumerate every `(block, expert)` pair present in the checkpoint.
 ///
 /// A pair is reported if *any* of the three projections (`gate`, `up`,
 /// `down`) is discoverable for it under either the stacked or the
 /// per-expert naming scheme.
-pub fn list_experts(layout: &GgufLayout) -> Vec<(usize, usize)> {
+///
+/// Works for any [`GgufTensorAccess`] implementation (owned buffer or
+/// `mmap` mapping).
+pub fn list_experts_from<L: GgufTensorAccess>(layout: &L) -> Vec<(usize, usize)> {
     let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
 
-    for tensor in layout.tensors.values() {
+    for tensor in layout.gguf_tensors().values() {
         record_expert_pairs(tensor, &mut pairs);
     }
 
     pairs.into_iter().collect()
 }
 
+#[cfg(feature = "mmap")]
+/// Enumerate MoE `(block, expert)` pairs from a memory-mapped layout.
+///
+/// Identical to [`list_experts_from`] on [`crate::GgufLayoutMmap`].
+pub fn list_experts_mmap(layout: &crate::GgufLayoutMmap) -> Vec<(usize, usize)> {
+    list_experts_from(layout)
+}
+
 fn record_expert_pairs(tensor: &Tensor, pairs: &mut BTreeSet<(usize, usize)>) {
     if let Some((block, _role)) = parse_stacked_name(&tensor.name) {
         if let Some(n_experts) = stacked_expert_count(tensor) {
@@ -50,14 +71,31 @@ fn record_expert_pairs(tensor: &Tensor, pairs: &mut BTreeSet<(usize, usize)>) {
     }
 }
 
+/// Extract the raw weights for a single `(block, expert)` pair from an
+/// owned [`GgufLayout`].
+///
+/// Convenience wrapper around [`extract_expert_from`].
+pub fn extract_expert(
+    layout: &GgufLayout,
+    block: usize,
+    expert: usize,
+) -> Result<MoeExpertWeights> {
+    extract_expert_from(layout, block, expert)
+}
+
 /// Extract the raw weights for a single `(block, expert)` pair.
 ///
 /// Returns [`MoeExpertWeights`] whose `gate / up / down` fields are
 /// `None` when that projection was not found in the checkpoint. If
 /// *none* of the three projections can be located, returns
 /// [`ParserError::MissingTensor`].
-pub fn extract_expert(
-    layout: &GgufLayout,
+///
+/// **Copy policy:** tensor payloads are always cloned into owned
+/// `Vec<u8>` in the returned [`RawTensor`] values, including stacked
+/// slices and mmap-backed sources. The layout is only borrowed for
+/// directory lookup and byte views during the call.
+pub fn extract_expert_from<L: GgufTensorAccess>(
+    layout: &L,
     block: usize,
     expert: usize,
 ) -> Result<MoeExpertWeights> {
@@ -68,7 +106,7 @@ pub fn extract_expert(
     if gate.is_none() && up.is_none() && down.is_none() {
         return Err(ParserError::MissingTensor {
             name: format!("blk.{block}.ffn_(gate|up|down)[_exps].weight"),
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
         });
     }
 
@@ -81,15 +119,27 @@ pub fn extract_expert(
     })
 }
 
-fn extract_role(
-    layout: &GgufLayout,
+#[cfg(feature = "mmap")]
+/// Extract one MoE expert from a memory-mapped layout.
+///
+/// Identical to [`extract_expert_from`] on [`crate::GgufLayoutMmap`].
+pub fn extract_expert_mmap(
+    layout: &crate::GgufLayoutMmap,
+    block: usize,
+    expert: usize,
+) -> Result<MoeExpertWeights> {
+    extract_expert_from(layout, block, expert)
+}
+
+fn extract_role<L: GgufTensorAccess>(
+    layout: &L,
     block: usize,
     expert: usize,
     role: &str,
 ) -> Result<Option<RawTensor>> {
     // Prefer the stacked convention first.
     let stacked_name = format!("blk.{block}.ffn_{role}_exps.weight");
-    if let Some(tensor) = layout.tensors.get(&stacked_name) {
+    if let Some(tensor) = layout.gguf_tensors().get(&stacked_name) {
         return Ok(Some(slice_stacked_expert(layout, block, tensor, expert)?));
     }
 
@@ -100,7 +150,7 @@ fn extract_role(
         format!("blk.{block}.ffn_{role}.{expert}.weight"),
         format!("blk.{block}.ffn_{role}_{expert}.weight"),
     ] {
-        if let Some(tensor) = layout.tensors.get(&candidate) {
+        if let Some(tensor) = layout.gguf_tensors().get(&candidate) {
             let bytes = layout.tensor_bytes(tensor)?.to_vec();
             return Ok(Some(RawTensor {
                 source_name: tensor.name.clone(),
@@ -123,8 +173,8 @@ fn extract_role(
 /// dimension in the `dims` vector. The per-expert chunk consists of
 /// the first `n_elements / n_experts` elements per expert, laid out
 /// contiguously in the tensor buffer.
-fn slice_stacked_expert(
-    layout: &GgufLayout,
+fn slice_stacked_expert<L: GgufTensorAccess>(
+    layout: &L,
     block: usize,
     tensor: &Tensor,
     expert: usize,
@@ -144,14 +194,14 @@ fn slice_stacked_expert(
     })
 }
 
-fn validate_stacked_expert_count(
-    layout: &GgufLayout,
+fn validate_stacked_expert_count<L: GgufTensorAccess>(
+    layout: &L,
     block: usize,
     tensor: &Tensor,
     expert: usize,
 ) -> Result<usize> {
     let n_experts = stacked_expert_count(tensor).ok_or_else(|| ParserError::InvalidLayout {
-        path: layout.path.clone(),
+        path: layout.gguf_path().to_owned(),
         reason: format!(
             "stacked tensor '{}' has rank {} (dims={:?}); expected rank >= 2",
             tensor.name,
@@ -162,7 +212,7 @@ fn validate_stacked_expert_count(
 
     if n_experts == 0 {
         return Err(ParserError::InvalidLayout {
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
             reason: format!("stacked tensor '{}' has zero experts", tensor.name),
         });
     }
@@ -178,16 +228,16 @@ fn validate_stacked_expert_count(
     Ok(n_experts)
 }
 
-fn stacked_slice_range(
-    layout: &GgufLayout,
+fn stacked_slice_range<L: GgufTensorAccess>(
+    layout: &L,
     tensor: &Tensor,
     expert: usize,
     n_experts: usize,
     buffer_len: usize,
 ) -> Result<(usize, usize)> {
     if !tensor.byte_len.is_multiple_of(n_experts) {
         return Err(ParserError::InvalidLayout {
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
             reason: format!(
                 "stacked tensor '{}' byte length {} is not divisible by n_experts={n_experts}",
                 tensor.name, tensor.byte_len,
@@ -199,18 +249,18 @@ fn stacked_slice_range(
     let start = expert
         .checked_mul(stride)
         .ok_or_else(|| ParserError::InvalidLayout {
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
             reason: format!("stacked stride overflow for tensor '{}'", tensor.name),
         })?;
     let end = start
         .checked_add(stride)
         .ok_or_else(|| ParserError::InvalidLayout {
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
             reason: format!("stacked end overflow for tensor '{}'", tensor.name),
         })?;
     if end > buffer_len {
         return Err(ParserError::InvalidLayout {
-            path: layout.path.clone(),
+            path: layout.gguf_path().to_owned(),
             reason: format!(
                 "stacked slice range [{start}..{end}] for tensor '{}' exceeds buffer len {buffer_len}",
                 tensor.name,

src/analysis/moe/mod.rs

@@ -2,16 +2,24 @@
 
 //! Mixture-of-Experts weight extraction (GGUF-specific).
 //!
-//! Locates MoE expert tensors inside a parsed [`GgufLayout`](crate::gguf::GgufLayout)
-//! and returns raw per-expert byte buffers. Performs no neural-network
-//! math — callers receive `Vec<u8>` + shape/dtype metadata and are
-//! responsible for any downstream compute.
+//! Locates MoE expert tensors inside a parsed GGUF layout and returns raw
+//! per-expert byte buffers. Performs no neural-network math — callers receive
+//! `Vec<u8>` + shape/dtype metadata and are responsible for any downstream
+//! compute.
+//!
+//! Owned checkpoints use [`list_experts`] / [`extract_expert`] on
+//! [`GgufLayout`](crate::gguf::GgufLayout). Mmap-backed checkpoints use
+//! the same helpers via [`list_experts_from`] / [`extract_expert_from`]
+//! (or the `*_mmap` aliases when the `mmap` feature is enabled) on
+//! [`GgufLayoutMmap`](crate::gguf::GgufLayoutMmap). Both paths share
+//! [`GgufTensorAccess`](crate::gguf::GgufTensorAccess) so parsing cannot
+//! drift between layout types.
 //!
 //! # Scope
 //!
 //! This is a model-analysis specialization, not part of the core
-//! checkpoint contract: [`extract_expert`] and [`list_experts`] take a
-//! concrete `&GgufLayout`, and [`RawTensor`] carries the GGUF-specific
+//! checkpoint contract: entry points take a concrete GGUF layout type or
+//! [`GgufTensorAccess`], and [`RawTensor`] carries the GGUF-specific
 //! [`DType`](crate::gguf::DType) / `ggml_type` fields of its source
 //! tensor. It does not accept [`Checkpoint`](crate::checkpoint::Checkpoint)
 //! — a format-independent MoE/topology layer is planned v0.4 work.
@@ -32,4 +40,6 @@ mod expert;
 mod extract;
 
 pub use expert::{MoeExpertWeights, RawTensor};
-pub use extract::{extract_expert, list_experts};
+pub use extract::{extract_expert, extract_expert_from, list_experts, list_experts_from};
+#[cfg(feature = "mmap")]
+pub use extract::{extract_expert_mmap, list_experts_mmap};

src/checkpoint/open.rs

@@ -56,7 +56,10 @@ impl AnyCheckpoint {
         }
     }
 
-    /// Memory-mapped GGUF layout.
+    /// Memory-mapped GGUF layout for
+    /// [`crate::analysis::moe::list_experts_from`] /
+    /// [`crate::analysis::moe::extract_expert_from`] (or the `*_mmap`
+    /// aliases).
     #[cfg(feature = "mmap")]
     pub fn as_gguf_mmap(&self) -> Option<&crate::GgufLayoutMmap> {
         match self {

src/gguf/access.rs

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+//! Shared tensor-directory + payload access for owned and mmap GGUF layouts.
+
+use std::collections::HashMap;
+
+use super::layout::GgufLayout;
+use super::tensor::Tensor;
+use crate::error::Result;
+
+/// Tensor directory and payload access shared by [`GgufLayout`] and
+/// [`GgufLayoutMmap`](super::map::GgufLayoutMmap).
+///
+/// MoE helpers ([`crate::analysis::moe::list_experts_from`],
+/// [`crate::analysis::moe::extract_expert_from`]) and other analysis code
+/// use this trait so owned and mmap-backed checkpoints cannot drift.
+pub trait GgufTensorAccess {
+    /// Source path label (filesystem path or synthetic label).
+    fn gguf_path(&self) -> &str;
+
+    /// Parsed tensor directory (`name` -> [`Tensor`] metadata).
+    fn gguf_tensors(&self) -> &HashMap<String, Tensor>;
+
+    /// Borrowed raw payload bytes for `tensor`.
+    fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]>;
+}
+
+impl GgufTensorAccess for GgufLayout {
+    fn gguf_path(&self) -> &str {
+        &self.path
+    }
+
+    fn gguf_tensors(&self) -> &HashMap<String, Tensor> {
+        &self.tensors
+    }
+
+    fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]> {
+        GgufLayout::tensor_bytes(self, tensor)
+    }
+}
+
+#[cfg(feature = "mmap")]
+impl GgufTensorAccess for super::map::GgufLayoutMmap {
+    fn gguf_path(&self) -> &str {
+        &self.path
+    }
+
+    fn gguf_tensors(&self) -> &HashMap<String, Tensor> {
+        &self.tensors
+    }
+
+    fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]> {
+        super::map::GgufLayoutMmap::tensor_bytes(self, tensor)
+    }
+}

src/gguf/map.rs

@@ -22,7 +22,10 @@ use crate::error::{ParserError, Result};
 /// GGUF layout backed by a read-only file mapping instead of an owned `Vec<u8>`.
 ///
 /// Header, KV metadata, and the tensor directory are identical to
-/// [`GgufLayout`]. Tensor payloads are borrowed from the mapping.
+/// [`GgufLayout`]. Tensor payloads are borrowed from the mapping. MoE
+/// expert listing and extraction use [`super::access::GgufTensorAccess`]
+/// via [`crate::analysis::moe::list_experts_from`] /
+/// [`crate::analysis::moe::extract_expert_from`].
 pub struct GgufLayoutMmap {
     /// Path the checkpoint was mapped from (kept for error messages).
     pub path: String,

src/gguf/mod.rs

@@ -7,6 +7,7 @@
 //! byte buffer. Packed K-quant dequant helpers are re-exported from this module.
 //! Optional mmap-backed loading is behind the `mmap` feature.
 
+mod access;
 mod cursor;
 mod dequant;
 mod layout;
@@ -18,6 +19,7 @@ mod tensor;
 use std::fs;
 use std::path::Path;
 
+pub use access::GgufTensorAccess;
 pub use dequant::{
     dequantize_iq3_m, dequantize_packed, dequantize_q2_k, dequantize_q3_k, dequantize_q4_0,
     dequantize_q4_1, dequantize_q4_k, dequantize_q5_k, dequantize_q6_k, dequantize_q8_0,