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

Merged
RM
3d ago
+268-36

User description

Summary

Implements RM-1462: MoE list_experts / extract_expert now work for memory-mapped GGUF layouts without loading the full file into a Vec<u8>.

  • Add GgufTensorAccess (src/gguf/access.rs) so GgufLayout and GgufLayoutMmap share path, tensor directory, and payload-byte access.
  • Refactor MoE extraction to list_experts_from / extract_expert_from over that trait; existing list_experts / extract_expert (&GgufLayout) remain unchanged for API compatibility.
  • With the mmap feature, list_experts_mmap / extract_expert_mmap are thin aliases for mmap layouts.
  • Copy policy: returned RawTensor::bytes are always owned Vec<u8> (mmap sources are borrowed only during the call).
  • Integration test mmap_moe_extraction_matches_owned_stacked_and_per_expert asserts owned vs mmap parity for stacked and per-expert fixtures.

directory_matches

Reviewed GgufLayoutMmap::directory_matches against the issue description and existing tests (mmap_directory_matches_rejects_different_tensor_names, etc.). It already compares alignment, tensor-data offset, architecture, tensor count, and per-tensor name/dims/dtype/length/offsets. No code change was required.

Testing

  • cargo fmt
  • cargo clippy --all-features --all-targets --locked -- -D warnings
  • cargo test --all-features --locked

Unresolved / follow-up

  • Expert payloads are still cloned on return; zero-copy expert views would be a separate API decision.
  • Aikido SAST scan was not run (MCP sign-in required in this environment).
Open in Web Open in Cursor 

Summary by cubic

Adds MoE expert listing and extraction for mmap-backed GGUF layouts, so downstream consumers no longer need to copy the whole file into a Vec<u8> to read expert weights.

  • Introduces GgufTensorAccess so GgufLayout and GgufLayoutMmap share tensor-directory, path, and payload access.
  • Refactors MoE extraction into list_experts_from / extract_expert_from over that trait; the existing list_experts / extract_expert (&GgufLayout) and new *_mmap aliases (with the mmap feature) reuse the same parsing.
  • Returned expert payloads are always owned Vec<u8>; mmap memory is borrowed only during the call.
  • directory_matches already compares tensor names, dims, dtypes, lengths, and offsets, so no change was needed there.
  • Adds an integration test asserting owned/mmap parity on stacked and per-expert fixtures.

CodeAnt-AI Description

Enable MoE expert extraction from memory-mapped GGUF files

What Changed

  • List and extract MoE experts directly from memory-mapped GGUF checkpoints without loading the entire file into a byte buffer.
  • Support stacked and per-expert tensor layouts, with results matching extraction from owned checkpoints.
  • Return expert payloads as owned buffers that can be used independently of the mapped file.

Impact

✅ MoE analysis on memory-mapped GGUF files
✅ No full-file byte-buffer copy for MoE analysis
✅ Consistent expert results across owned and mapped checkpoints

💡 Usage Guide

Checking Your Pull Request

Every time you make a pull request, our system automatically looks through it. We check for security issues, mistakes in how you're setting up your infrastructure, and common code problems. We do this to make sure your changes are solid and won't cause any trouble later.

Talking to CodeAnt AI

Got a question or need a hand with something in your pull request? You can easily get in touch with CodeAnt AI right here. Just type the following in a comment on your pull request, and replace "Your question here" with whatever you want to ask:

@codeant-ai ask: Your question here

This lets you have a chat with CodeAnt AI about your pull request, making it easier to understand and improve your code.

Example

@codeant-ai ask: Can you suggest a safer alternative to storing this secret?

Preserve Org Learnings with CodeAnt

You can record team preferences so CodeAnt AI applies them in future reviews. Reply directly to the specific CodeAnt AI suggestion (in the same thread) and replace "Your feedback here" with your input:

@codeant-ai: Your feedback here

This helps CodeAnt AI learn and adapt to your team's coding style and standards.

Example

@codeant-ai: Do not flag unused imports.

Retrigger review

Ask CodeAnt AI to review the PR again, by typing:

@codeant-ai: review

Check Your Repository Health

To analyze the health of your code repository, visit our dashboard at https://app.codeant.ai. This tool helps you identify potential issues and areas for improvement in your codebase, ensuring your repository maintains high standards of code health.

Summary by CodeRabbit

  • New Features
    • MoE expert discovery and weight extraction now work with both standard and memory-mapped GGUF checkpoints.
    • Extracted weights are returned in owned buffers, including when read from memory-mapped files.
    • Added support for shared expert discovery and extraction APIs across stacked and per-expert tensor layouts.
  • Bug Fixes
    • Verified that expert listings and extracted gate weights match between standard and memory-mapped readers.