From 962a839bd17e90782f1ffa0f9ef50c91b90cc1b8 Mon Sep 17 00:00:00 2001 From: Raul Montoya Cardenas Date: Wed, 29 Jul 2026 06:06:29 -0500 Subject: [PATCH] refactor(moe): rename OlmoeRouter to MoeRouter; finish Codecov #14 Replace family-branded OlmoeRouter/OlmoeOutput with generic MoeRouter/MoeOutput so the public API matches multi-family MoE support. ModelFamily::Olmoe and GGUF olmoe.* metadata keys are unchanged. Complete remaining issue #14 acceptance items: Codecov README badge, top-level ignore paths in codecov.yml (target/**, examples/**), and local lcov.info gitignore. README MoE example uses a Result-returning main so `?` is valid. Implemented by Grok Build: Grok 4.5 Closes #14 --- .gitignore | 1 + README.md | 34 +++++++++++++++++++++------------- codecov.yml | 4 ++++ src/moe/mod.rs | 22 +++++++++++----------- src/moe/tests.rs | 11 +++++------ 5 files changed, 42 insertions(+), 30 deletions(-) diff --git a/.gitignore b/.gitignore index 5bddbaf..573aebc 100644 --- a/.gitignore +++ b/.gitignore @@ -42,3 +42,4 @@ Thumbs.db *.tmp tmp/ cache/ +lcov.info diff --git a/README.md b/README.md index 84d42e4..b97d72f 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,7 @@ Pure-Rust tensor, transformer, and Mixture-of-Experts building blocks. No CUDA, [![Rust](https://img.shields.io/badge/rust-edition%202024-orange)](https://www.rust-lang.org/) [![License](https://img.shields.io/badge/license-MIT%2FApache-blue)](./LICENSE-APACHE-2.0) +[![codecov](https://codecov.io/gh/Limen-Neural/cortex-tensor/branch/main/graph/badge.svg)](https://codecov.io/gh/Limen-Neural/cortex-tensor) ## Overview @@ -32,7 +33,7 @@ src/ │ ├── model.rs # TransformerConfig + TransformerLM (decoder-only) │ └── mod.rs └── moe/ - ├── mod.rs # OlmoeRouter public API, RoutingMode + ├── mod.rs # MoeRouter public API, RoutingMode ├── adapter.rs # model-family detection + tensor selection ├── checkpoint.rs # GGUF parser, mmap'd F32/F16/Q8_0/Q5_K access ├── dequant.rs # Q8_0 / Q5_K row dequant, f16→f32, row sizing @@ -64,11 +65,11 @@ src/ | Item | Purpose | |---|---| -| `OlmoeRouter` | Family-aware MoE router. Loads a GGUF checkpoint, detects model family, and produces top-k expert selections. | +| `MoeRouter` | Family-aware MoE router. Loads a GGUF checkpoint, detects model family, and produces top-k expert selections. | | `RoutingMode` | `StubUniform`, `DenseSim`, `SpikingSim` (simulation-only; no GPU dispatch). | | `ModelFamily` | `Olmoe`, `Qwen3Moe`, `Gemma4`, `DeepSeek2`, `LlamaMoe`. | -Supported GGUF tensor types: `F32`, `F16`, `Q8_0`, `Q5_K`. `IQ3_S` is detected and rejected (for token embeddings) with a clear error so callers can fall back to `llama.cpp` prompt embeddings. For the preferred GPU synapse tensor (e.g. attn_q on qwen3_moe_iq3_m), unsupported quants now correctly route to a checkpoint-backed `routing-f32` source (using the F32 routing tensor) instead of synthetic fallback. See `synapse_source()`, `real_gpu_synapse_tensor_name()`, and `OlmoeRouter` metadata. +Supported GGUF tensor types: `F32`, `F16`, `Q8_0`, `Q5_K`. `IQ3_S` is detected and rejected (for token embeddings) with a clear error so callers can fall back to `llama.cpp` prompt embeddings. For the preferred GPU synapse tensor (e.g. attn_q on qwen3_moe_iq3_m), unsupported quants now correctly route to a checkpoint-backed `routing-f32` source (using the F32 routing tensor) instead of synthetic fallback. See `synapse_source()`, `real_gpu_synapse_tensor_name()`, and `MoeRouter` metadata. **Parser layer (planning, see #8):** the canonical home for GGUF v3 deserialization and per-expert raw weight extraction is `engram-parser`, not this @@ -79,7 +80,7 @@ crate. The in-crate reader is frozen for enhancements while that extraction land ### GGUF adapter + synapse source + SAAQ flow (code paths) -- `OlmoeRouter::load` / `load_with_family_and_mode` → `probe_and_map` calls `resolve_adapter` (adapter.rs). +- `MoeRouter::load` / `load_with_family_and_mode` → `probe_and_map` calls `resolve_adapter` (adapter.rs). - `resolve_adapter` infers family from arch, validates routing tensor (must be F32 rank-2), selects token_embd or tok_embeddings, sets `preferred_gpu_synapse_tensor` to `blk.0.attn_q.weight` when present. - Synapse source selection (updated for qwen3 IQ3_S): if attn_q is F16 rank-2 containing hidden_size (relaxed from strict square to support GQA) → `real`; elif attn_q present → `routing-f32` (real name = routing tensor name); else `synthetic-fallback`. - Routing always uses `routing_tensor` via `checkpoint_gate_scores` (routing.rs) when checkpoint loaded (never synthetic for real loads). @@ -194,17 +195,24 @@ let attn = MultiHeadAttention::new(/* dim */ 512, /* num_heads */ 8); let block = TransformerBlock::new(/* dim */ 512, /* num_heads */ 8, /* mlp_dim */ 2048); ``` -Loading an OLMoE-family GGUF and running the router: +Loading a family-aware MoE GGUF and running the router: ```rust -use cortex_tensor::moe::{OlmoeRouter, RoutingMode}; - -let mut router = OlmoeRouter::load( - "path/to/olmoe.gguf", - RoutingMode::DenseSim, - /* top_k */ 2, -)?; -let (experts, weights) = router.route_for_token(/* token_id */ 42)?; +use cortex_tensor::moe::{MoeRouter, RoutingMode}; + +fn main() -> cortex_tensor::Result<()> { + let mut router = MoeRouter::load_with_mode( + "path/to/model.gguf", + /* num_experts */ 0, // 0 → take count from checkpoint metadata + /* top_k */ 2, + RoutingMode::DenseSim, + )?; + let embedding = vec![0.0f32; cortex_tensor::types::EMBEDDING_DIM]; // or extract_token_embedding + let out = router.forward(&embedding)?; + // out.selected_experts, out.expert_weights, out.hidden + let _ = out; + Ok(()) +} ``` ## Optional Sentry monitoring diff --git a/codecov.yml b/codecov.yml index 6debacc..5f140d2 100644 --- a/codecov.yml +++ b/codecov.yml @@ -9,6 +9,10 @@ coverage: default: informational: true +ignore: + - "target/**" + - "examples/**" + comment: layout: "reach,diff,flags,files" behavior: default diff --git a/src/moe/mod.rs b/src/moe/mod.rs index f381490..1e0a728 100644 --- a/src/moe/mod.rs +++ b/src/moe/mod.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -// NOTE: This module implements the main OlmoeRouter. High LOC is due to the +// NOTE: This module implements the main MoeRouter. High LOC is due to the // full public API, loading logic, multiple routing modes, and extensive tests. // Further modularization planned. @@ -64,7 +64,7 @@ pub(crate) use self::gguf::{ GGUF_VALUE_TYPE_UINT64, GGUF_VERSION, }; -pub struct OlmoeRouter { +pub struct MoeRouter { model_path: String, num_experts: usize, top_k: usize, @@ -94,13 +94,13 @@ pub struct RouterMetadata { } #[derive(Debug, Clone)] -pub struct OlmoeOutput { +pub struct MoeOutput { pub expert_weights: Vec, pub selected_experts: Vec, pub hidden: Vec, } -impl OlmoeRouter { +impl MoeRouter { pub fn load(model_path: &str, num_experts: usize, top_k: usize) -> Result { Self::load_with_family_and_mode( model_path, @@ -212,7 +212,7 @@ impl OlmoeRouter { Ok(metadata) } - pub fn forward(&mut self, embedding: &[f32]) -> Result { + pub fn forward(&mut self, embedding: &[f32]) -> Result { if embedding.len() != EMBEDDING_DIM { return Err(HybridError::InputLengthMismatch { expected: EMBEDDING_DIM, @@ -292,7 +292,7 @@ impl OlmoeRouter { Ok((metadata, checkpoint)) } - fn simulate_moe_routing(&self, embedding: &[f32]) -> Result { + fn simulate_moe_routing(&self, embedding: &[f32]) -> Result { let gate_scores = self.compute_gate_scores(embedding)?; let expert_weights = softmax(&gate_scores); let selected_experts = top_k_indices(&expert_weights, self.top_k); @@ -302,14 +302,14 @@ impl OlmoeRouter { .sum(); let hidden: Vec = embedding.iter().map(|&v| v * selected_mass).collect(); - Ok(OlmoeOutput { + Ok(MoeOutput { expert_weights, selected_experts, hidden, }) } - fn spiking_moe_routing(&mut self, embedding: &[f32]) -> Result { + fn spiking_moe_routing(&mut self, embedding: &[f32]) -> Result { let gate_scores = self.compute_gate_scores(embedding)?; let n = self.num_experts; let mut membrane_scores = Vec::with_capacity(n); @@ -356,7 +356,7 @@ impl OlmoeRouter { *value = spike * 0.3; } - Ok(OlmoeOutput { + Ok(MoeOutput { expert_weights, selected_experts, hidden, @@ -379,9 +379,9 @@ impl OlmoeRouter { Ok(synthetic_gate_scores(self.num_experts, embedding)) } - fn stub_output(&self) -> OlmoeOutput { + fn stub_output(&self) -> MoeOutput { let n = self.num_experts.max(1); - OlmoeOutput { + MoeOutput { expert_weights: vec![1.0 / n as f32; n], selected_experts: (0..self.top_k.min(n)).collect(), hidden: vec![0.0; EMBEDDING_DIM], diff --git a/src/moe/tests.rs b/src/moe/tests.rs index 1d62950..a0794a1 100644 --- a/src/moe/tests.rs +++ b/src/moe/tests.rs @@ -4,9 +4,8 @@ use super::test_fixtures::*; use super::*; use std::fs::remove_file; -fn stub() -> OlmoeRouter { - OlmoeRouter::load_with_mode("", 8, 1, RoutingMode::StubUniform) - .expect("stub load should succeed") +fn stub() -> MoeRouter { + MoeRouter::load_with_mode("", 8, 1, RoutingMode::StubUniform).expect("stub load should succeed") } #[test] @@ -35,7 +34,7 @@ fn test_dense_sim_uses_real_gate_weights() { let path = write_temp_file(&build_real_size_checkpoint(gate_bytes), "dense-real"); let mut model = - OlmoeRouter::load_with_mode(path.to_str().unwrap(), 8, 2, RoutingMode::DenseSim).unwrap(); + MoeRouter::load_with_mode(path.to_str().unwrap(), 8, 2, RoutingMode::DenseSim).unwrap(); let mut embedding = vec![0.0f32; EMBEDDING_DIM]; embedding[0] = 1.0; let out = model.forward(&embedding).unwrap(); @@ -48,7 +47,7 @@ fn test_dense_sim_uses_real_gate_weights() { #[test] fn test_spiking_sim_state_can_reset() { - let mut model = OlmoeRouter::load_with_mode("", 8, 2, RoutingMode::SpikingSim).unwrap(); + let mut model = MoeRouter::load_with_mode("", 8, 2, RoutingMode::SpikingSim).unwrap(); let _ = model.forward(&vec![1.0; EMBEDDING_DIM]).unwrap(); assert!(model.has_state_activity()); model.reset_state(); @@ -61,7 +60,7 @@ fn test_real_checkpoint_probe_via_env() { return; }; - let metadata = OlmoeRouter::probe_model(&path, None).unwrap(); + let metadata = MoeRouter::probe_model(&path, None).unwrap(); assert!(!metadata.architecture.is_empty()); assert!(metadata.hidden_size > 0); assert!(metadata.num_experts > 0);