Skip to content

Repository files navigation

cortex-tensor

Pure-Rust tensor, transformer, and Mixture-of-Experts building blocks. No CUDA, no Julia FFI, no framework dependencies — just Vec<f32> and honest math.

Rust License codecov

Overview

cortex-tensor is a minimal, framework-free foundation for building transformer-based language models and MoE routers in Rust. It was surgically extracted from a larger hybrid codebase (corinth-canal) and then stripped of every GPU / CUDA / Julia / SNN-specific concern so it can stand alone as a reusable, open-source numerical kernel.

Design goals:

  • Zero GPU coupling. No cust, no libc pinned-host registration, no #[cfg(feature = "gpu")] branches.
  • Zero framework dependency. No candle, no tch, no ort. The tensor type is a row-major Vec<f32> with explicit shape + strides.
  • Small, auditable dependency set. serde, serde_json, thiserror, rand, rayon, memmap2, half — nothing else.
  • Inference-ready MoE. A GGUF checkpoint bridge with family-aware adapter resolution for OLMoE, Qwen3-MoE, Gemma-4, DeepSeek-2, and Llama-MoE.

Architecture

src/
├── lib.rs            # re-exports Tensor, CortexError, HybridError, Result
├── error.rs          # CortexError + HybridError alias
├── types.rs          # EMBEDDING_DIM, ModelFamily, RoutingMode
├── tensor/
│   ├── mod.rs        # row-major Tensor { data, shape, strides }
│   └── ops.rs        # matmul, batched_matmul, causal_mask, softmax, ...
├── transformer/
│   ├── attention.rs  # MultiHeadAttention (scaled dot-product, causal mask)
│   ├── block.rs      # TransformerBlock (attn + MLP + LayerNorm)
│   ├── model.rs      # TransformerConfig + TransformerLM (decoder-only)
│   └── mod.rs
└── moe/
    ├── 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
    ├── gguf.rs       # GGUF magic/version + GGML type constants
    ├── routing.rs    # softmax, top-k, L2 normalize, embedding resample
    └── tests.rs      # router + checkpoint unit tests

Modules

tensor

Item Purpose
Tensor Row-major f32 tensor (data: Vec<f32>, shape, strides), Serialize/Deserialize.
ops::matmul / batched_matmul Cache-friendly tiled CPU matmul.
ops::causal_mask Additive mask for auto-regressive attention.
ops::softmax / layer_norm Standard building blocks.

transformer

Item Purpose
MultiHeadAttention Multi-head scaled dot-product attention with causal masking. Weights are dense Tensors; no external framework needed.
TransformerBlock Attention → residual → MLP → residual, with pre-LayerNorm.
TransformerLM / TransformerConfig Decoder-only transformer LM: token + positional embedding → N × block → LayerNorm → LM head.

moe

Item Purpose
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 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 crate. The in-crate reader is frozen for enhancements while that extraction lands — see GGUF parser boundary.

Future formats (planning, see #9): Safetensors support will arrive via a dedicated reusable safetensors-parser crate (header inspection + deterministic manifest + MoE candidate discovery), extracted as a one-way copy of reference logic from rmems/corinth-canal (see corinth-canal#116, engram-parser#10, and cortex #7/#8 for the GGUF precedent with engram-parser). No implementation or dependency is present yet — this keeps the reusable parser boundary clean. Cross-links and notes are maintained for alignment.

GGUF adapter + synapse source + SAAQ flow (code paths)

  • MoeRouter::load / load_with_family_and_modeprobe_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).
  • extract_named_token_embedding_from_checkpoint (checkpoint.rs) supports dequant for Q8_0/Q5_K (and F32/F16); IQ3_S errors for embeddings.
  • Public metadata exposes preferred_gpu_synapse_tensor_name, real_gpu_synapse_tensor_name, synapse_source for SAAQ experiment / Surrogate_Viz consumers to choose dequant vs. synthetic path and load the right tensor (f16 path or f32 routing path).
  • SAAQ artifacts (external): calibration runs consume the router to emit artifacts for viz; see labels on related issues for campaign.

Scope / Boundaries

This crate owns:

  • Row-major f32 Tensor and CPU tensor ops (matmul, batched matmul, causal mask, softmax, layer norm).
  • Decoder-only transformer building blocks (attention, block, TransformerLM).
  • MoE routing math — gate scores, softmax, top-k selection, L2 normalization, embedding resampling — and the simulation routing modes.
  • Model-family adapters and tensor selection (Olmoe, Qwen3Moe, Gemma4, DeepSeek2, LlamaMoe), including synapse-source resolution.
  • Dequantization of supported GGUF quants to f32 (Q8_0, Q5_K, F16).
  • The consumer-side GGUF bridge it needs today: mmap'd tensor access and token-embedding extraction.

This crate does not own:

  • Canonical GGUF v3 deserialization (header, KV metadata, tensor directory) and per-expert raw weight extraction — see engram-parser and the parser-boundary note below.
  • Safetensors header inspection, deterministic manifests, and MoE candidate discovery — planned for a dedicated safetensors-parser crate (see #9).
  • CUDA / GPU / SIMD execution, and any GPU host registration.
  • SNN neuron dynamics (neuromod) and ANN→SNN orchestration (hybrid-fusion).
  • Tokenization and automatic differentiation (see Non-goals).

Allowed dependencies: the current small set — serde, serde_json, thiserror, rand, rayon, memmap2, half, plus optional sentry — and, in future, the zero-dependency Limen-Neural parser crates.

Forbidden dependencies: GPU backends (cust), inference frameworks (candle, tch, ort), domain/SNN orchestration crates, and any dependency on rmems/corinth-canal. Extraction from corinth-canal is a one-way copy; that repo keeps an unmodified reference copy per its PROMOTION_RULES.md.

Crate Role
engram-parser GGUF parse + per-expert raw weight extraction
cortex-tensor (this crate) Tensor math + MoE routing on extracted weights
hybrid-fusion ANN→SNN orchestration
neuromod SNN neuron dynamics (downstream consumer)

See LIM-9 for the full Rust runtime/deployment boundary matrix, and issues #5 (boundary doc), #8 (GGUF parser coordination), and #9 (Safetensors coordination) for this repo's tracking.

GGUF parser boundary (see #8)

engram-parser is the parser-layer provider for this ecosystem: it is the canonical, zero-dependency home for GGUF v3 layout parsing and MoE per-expert raw weight extraction, being extracted from the experimental rmems/corinth-canal reference implementation (see engram-parser#7 and corinth-canal#115). cortex-tensor stays the consumer: f32 math, Tensor ops, routing, and model adapters on top of parsed layout / extracted weights.

Layer Canonical owner Where it lives in this crate today
GGUF magic + v3 header, KV metadata, tensor directory engram-parser src/moe/checkpoint.rs (parse_checkpoint_layout)
GGML type + GGUF value-type constants engram-parser src/moe/gguf.rs
Per-expert raw weight extraction engram-parser not implemented here
mmap'd tensor access for the router cortex-tensor src/moe/checkpoint.rs (probe_and_map_checkpoint)
Dequantization to f32 cortex-tensor src/moe/dequant.rs
Routing math, top-k, family adapters cortex-tensor src/moe/routing.rs, src/moe/adapter.rs

Freeze while the extraction lands: no new parser code and no dtype/GGUF format enhancements in src/moe/checkpoint.rs, src/moe/gguf.rs, or src/moe/dequant.rs until engram-parser#7 lands (or an explicit sub-issue is opened under it). Known gaps versus the corinth-canal reference — additional dtypes (BF16, Q6_K, IQ3_*), a ggml_type_label helper, and the "GGUF wire type 31 is Q4_0_4_4, not IQ3_M" discipline — are deliberately parked on engram-parser#7 rather than duplicated here. Cross-repo planning is tracked in Linear LIM-88 (under LIM-9).

Install

[dependencies]
cortex-tensor = { git = "https://github.com/Limen-Neural/cortex-tensor", branch = "main" }

Quick start

use cortex_tensor::tensor::ops::matmul;
use cortex_tensor::Tensor;

let a = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = Tensor::from_vec(vec![1.0, 0.0, 0.0, 1.0], &[2, 2]);
let c = matmul(&a, &b);
assert_eq!(c.shape(), &[2, 2]);

Building a transformer block:

use cortex_tensor::transformer::{MultiHeadAttention, TransformerBlock};

let attn = MultiHeadAttention::new(/* dim */ 512, /* num_heads */ 8);
let block = TransformerBlock::new(/* dim */ 512, /* num_heads */ 8, /* mlp_dim */ 2048);

Loading a family-aware MoE GGUF and running the router:

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

Enable with the sentry feature (uses sentry-rust 0.48):

[dependencies]
cortex-tensor = { git = "...", features = ["sentry"] }

Init guard example (call early in main or lib init; keep guard alive for duration of process). Note: the consuming crate/binary must explicitly enable the sentry feature on its cortex-tensor dependency (transitive features do not auto-activate). Then use the re-exported path (or add sentry as direct dep):

#[cfg(feature = "sentry")]
let _sentry_guard = cortex_tensor::sentry::init((
    "https://<key>@sentry.io/<project>",
    cortex_tensor::sentry::ClientOptions {
        release: cortex_tensor::sentry::release_name!(),
        environment: Some("production".into()),
        ..Default::default()
    },
));

// Your app code; errors auto captured when panics or cortex_tensor::sentry::capture_message etc used.
// (The re-export brings the full sentry crate API under the feature gate.)

When the feature is off the re-export is not present (guarded).

Non-goals

  • No GPU backend. Ever. If you need CUDA, consume this crate's Tensor into your own kernels.
  • No automatic differentiation. This is an inference and forward-pass library.
  • No tokenizer. Pair it with tokenizers or llama.cpp's tokenizer of choice.
  • No SNN / neuromorphic logic. Those live in upstream projects.

Status

Extracted and compiling cleanly under Rust edition 2024. Public API is subject to change until 0.1.0 is tagged. Tests and benchmarks are incoming.

License

This project is licensed under either of

at your option.

About

Pure-Rust matrix multiplication and Transformer execution. Originally SpikeLMo that ended up being configured into corinth-canal.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages