Cogni-ML

Crystal machine learning library with native Apple Silicon GPU acceleration.

Highlights:

• Native Metal GPU embedding pipeline — 43ms for 260 tokens on M2 Max (2.2x faster than baseline)
• GGUF model loading with Q5_K/Q6_K quantization support
• simdgroup_matrix_multiply_accumulate GEMM kernels
• Compute graph with automatic wave-based barrier optimization
• Autograd engine, NN layers, Adam optimizer
• llama.cpp bindings for any GGUF model

Architecture

src/ml/
  core/         Tensor, Shape, MetalBuffer
  autograd/     Variable, GradFn (backward pass)
  nn/           Linear, LayerNorm, MultiHeadAttention, ViT
  optim/        Adam/AdamW
  llm/          llama.cpp FFI bindings
  gguf/         GGUF reader, tokenizer, dequantization, NomicBertMoE
  metal/        Device, ComputeEncoder, ComputeGraph, GraphEncoder

GPU Embedding Pipeline

The crown jewel: a fully native Metal compute pipeline for nomic-embed-text-v2-moe BERT embeddings.

require "ml"
require "ml/gguf/nomic_bert"
require "ml/gguf/metal_backend"
require "ml/metal/compute_graph"

ML::Metal::Device.init!
model = ML::GGUF::NomicBertMoE.from_gguf("path/to/model.gguf", ML::GGUF::MetalBackend.new)

embedding = model.embed("Your text here")  # → Array(Float32), dim=768

Performance (Apple M2 Max, 38 GPU cores)

| Tokens | Latency | |--------|---------| | 20 | 14ms | | 94 | 16ms | | 196 | 33ms | | 433 | 70ms |

What's inside

• simdgroup_matrix GEMM — hardware-accelerated 8x8 matrix tiles for Q5_K/Q6_K dequant+multiply
• Batched expert GEMM — all 8 MoE experts in 1 dispatch (LTP Diamond surgery)
• ComputeGraph — automatic wave scheduling with offset-aware + Block Integrity dependency analysis
• GraphEncoder — drop-in ComputeEncoder replacement that builds the compute graph
• Fused kernels — QKV split+RoPE, gate+softmax+topk, atomic scatter, f32 norm2
• Indirect dispatch — GPU-driven threadgroup counts, zero CPU-GPU sync for MoE routing

Supported models

| Model | Format | Status | |-------|--------|--------| | nomic-embed-text-v2-moe | GGUF Q5_K_M | Full native Metal pipeline | | Any BERT-like encoder | GGUF | Via NomicBertMoE (if architecture matches) | | Llama, Qwen, Mistral, etc. | GGUF | Via llama.cpp bindings |

Installation

# shard.yml
dependencies:
  cogni-ml:
    github: anthropics/cogni-ml  # or local path
    version: ~> 0.10.0

Build with Metal GPU

make build    # Compiles bridge.mm + links Metal frameworks
make spec     # Run tests with GPU

CPU-only build

crystal build -Dcpu_only your_app.cr

Quick Start

Tensor + Autograd (CPU)

require "ml"

x = ML::Autograd::Variable.rand(2, 3, requires_grad: true, device: ML::Tensor::Device::CPU)
layer = ML::NN::Linear.new(3, 4, device: ML::Tensor::Device::CPU)

out = layer.forward(x)
loss = out.mean
loss.backward

opt = ML::Optim::Adam.new(layer.parameters)
opt.step
opt.zero_grad

LLM Inference (llama.cpp)

require "ml/llm/llama"

ML::LLM.init
model = ML::LLM::Model.new("path/to/model.gguf")
gen = ML::LLM::Generator.new(model)
puts gen.ask("What is Crystal?", max_tokens: 100)
ML::LLM.cleanup

GGUF Embeddings (Metal GPU)

require "ml"
require "ml/gguf/nomic_bert"
require "ml/gguf/metal_backend"
require "ml/metal/compute_graph"

ML::Metal::Device.init!
model = ML::GGUF::NomicBertMoE.from_gguf(
  "nomic-embed-text-v2-moe.Q5_K_M.gguf",
  ML::GGUF::MetalBackend.new
)

# Single embedding
vec = model.embed("Crystal programming language")
puts "dim=#{vec.size}"  # 768

# Batch embedding
vecs = model.embed_batch(["Hello", "World", "Crystal"])

Metal Kernels

11 Metal shader files implementing:

| Kernel | Purpose | |--------|---------| | gemm_mm.metal | simdgroup_matrix GEMM for Q5_K/Q6_K + batched expert variants | | gemm_simd.metal | Scalar SIMD GEMM (small batch fallback) | | attention_matmul.metal | Flash attention with simdgroup_matrix Q*K^T | | bert_fp16.metal | Fused ops: QKV split+RoPE, gate+softmax+topk, norms, scatter, routing | | gemm_mm_f16.metal | FP16 GEMM (experimental) | | nn.metal | General NN ops (linear, layernorm, GELU) |

Platform Support

| Platform | GPU | CPU | Status | |----------|-----|-----|--------| | macOS (Apple Silicon) | Metal | Yes | Primary target | | macOS (Intel) | Metal | Yes | Supported | | Linux | - | Yes | -Dcpu_only |

Build Flags

| Flag | Effect | |------|--------| | -Dcpu_only | Disable Metal, pure CPU | | -Duse_gguf | Enable GGUF model loading (requires llama.cpp for LLM, standalone for embeddings) |

License

MIT