UPSTREAM PR #17052: # Add Megrez-MoE Architecture Support ggml-org#16724

[DajanaV]

Mirrored from ggml-org/llama.cpp#17052

Summary

This PR adds full support for the Megrez-MoE (Mixture of Experts) architecture to llama.cpp, enabling inference on Megrez2-3x7B models and similar MoE variants.

Architecture Details

Megrez-MoE is a Mixture of Experts architecture with:

• 64 experts with top-6 selection per layer
• 4 shared experts across all tokens
• Sigmoid + bias gating mechanism (different from standard softmax gating)
• 30 MoE layers with 2048 embedding dimension
• Context length up to 163,840 tokens

Changes Made

1. Architecture Registration

• Added LLM_ARCH_MEGREZ_MOE architecture enum
• Registered MoE-specific hyperparameters (expert counts, FFN dimensions)
• Added tensor mapping for 64 experts × 30 layers

2. MoE FFN Implementation

Implemented build_mergez_moe_ffn() with:

• Sigmoid gating with bias (unique to Megrez-MoE)
• Top-K expert selection using ggml_top_k()
• Shared experts processing for all tokens
• Per-expert feed-forward computation

3. Model Loading

• Added llm_build_megrez_moe class
• Implemented hyperparameter loading (expert_count, expert_used_count, etc.)
• Implemented tensor loading for all expert weights

4. Graph Memory Fix

• Problem: Warmup crashed with "not enough space in context's memory pool"
• Root Cause: MoE FFN creates ~35 intermediate tensors per layer (sigmoid, reshape, top_k, etc.)
• Solution: Added 4096 node overhead to graph_max_nodes() for Megrez-MoE (30 layers × 35 tensors ≈ 1050 nodes, doubled for safety)

Testing

• All 39 existing tests pass
• No regression in other architectures (verified with Gemma-2)
• Warmup works without crashes
• Output generation verified (up to 200 tokens)
• Performance: ~17 tokens/second on CPU

Comparison

# Without this PR:
$ ./build/bin/llama-cli -m Megrez2-3x7B.gguf -p "Test"
error: unknown model architecture: 'megrez-moe'

# With this PR:
$ ./build/bin/llama-cli -m Megrez2-3x7B.gguf -p "Test"
Works perfectly with warmup enabled

[loci-agentic-ai]

Access the complete analysis in the LOCI Dashboard

Performance Analysis Summary: Megrez-MoE Architecture Support

Overview

This analysis examines the performance impact of adding Megrez-MoE (Mixture of Experts) architecture support to llama.cpp. The changes introduce comprehensive MoE functionality including 64 experts with top-6 selection, shared experts, and sigmoid gating mechanisms.

Key Findings

Performance Impact

• Highest degradation identified: _M_const_cast function shows 209% Response Time increase (85 ns → 262 ns) and 277% Throughput degradation (64 ns → 241 ns)
• Core function impact: The affected function is part of C++ Standard Library red-black tree operations used for architecture lookups, not directly impacting inference functions like llama_decode, llama_encode, or llama_tokenize
• Inference performance: No direct impact on tokens per second expected since core inference functions remain unchanged

Power Consumption Analysis

• Primary impact: build.bin.libllama.so shows minimal +0.427% power consumption increase
• Secondary impact: build.bin.llama-cvector-generator shows negligible increase
• Overall assessment: Power impact is contained and acceptable for the new functionality added

Technical Analysis

• Flame graph insights: Performance regression concentrated in the function's own logic (92% of execution time), with additional overhead from stack protection checks and red-black tree iterator operations
• CFG comparison: Code reorganization separated stack canary setup into distinct basic blocks, adding one extra branch instruction and changing memory access patterns
• Root cause: Architecture map expansion increases iterator traversal overhead during model initialization

Code Review Insights

• Memory allocation: Added 4096 graph nodes overhead for Megrez-MoE models to prevent warmup crashes
• Architecture complexity: Introduced 25 new tensor mappings and sophisticated layer-specific expert allocation patterns
• Implementation quality: Follows established MoE patterns while accommodating unique Megrez-MoE characteristics

Conclusion

The implementation successfully adds Megrez-MoE support with localized performance impact limited to architecture lookup operations. The observed regression represents an acceptable trade-off for enabling an entirely new model architecture family without affecting core inference performance.