Cogni-ML

Crystal machine learning library with native Apple Silicon GPU acceleration.

Cogni-ML is currently two things:

• A general Crystal ML toolkit: tensors, autograd, NN layers, optimizers, GGUF readers, and llama.cpp bindings.
• A native Metal inference lab for GGUF models, with production-oriented work on nomic-embed-text-v2-moe embeddings and Qwen 3.5 text generation.

Highlights

• Native Metal embedding pipeline for nomic-embed-text-v2-moe.
• Native Qwen 3.5 9B GGUF inference path for Apple Silicon Metal.
• Q4_K/Q5_K/Q6_K/Q8_0 quantized matmul kernels.
• Chunked Qwen 3.5 prefill, decode wave scheduling, prompt-state cache restore, and exact speculative decode harnesses.
• ComputeGraph wave scheduling with offset-aware barrier optimization.
• Crystal autograd engine, NN layers, and Adam/AdamW optimizers.
• llama.cpp FFI bindings for general GGUF model access.

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, Qwen35, NomicBertMoE
  metal/        Device, ComputeEncoder, ComputeGraph, GraphEncoder

Qwen 3.5 Native Metal

The native Qwen path targets Qwen3.5-9B-Q4_K_M.gguf on Apple Silicon. The code supports:

• Qwen 3.5 GGUF metadata and tokenizer loading.
• Q4_K, Q5_K, Q6_K, and Q8_0 quantized projections.
• Full-attention layers with GQA, partial RoPE, KV cache writes, and fused output projection.
• DeltaNet/recurrent layers with GPU-resident recurrent state and chunked prefill scan.
• Chunked prefill with final-token top1 shortcut.
• Decode wave scheduling to reduce command-buffer boundaries.
• Exact prompt-state save/restore and longest-prefix prompt cache.
• Exact speculative decode harnesses:
  • neural draft with Qwen 3.5 0.8B Q8_0,
  • n-gram/cache draft for repeated/generated-template text,
  • target-verifier chunks with row-batched top1 for larger accepted chunks.

The 9B Q4_K_M path is the primary verified target. Qwen 3.6 27B is a scale-up target, but it should be treated as experimental until local correctness and performance runs are completed.

Model Layout

The developer CLIs default to local LM Studio / llama.cpp-style paths:

~/.cache/lm-studio/models/lmstudio-community/Qwen3.5-9B-GGUF/Qwen3.5-9B-Q4_K_M.gguf
~/.cache/lm-studio/models/lmstudio-community/Qwen3.5-0.8B-GGUF/Qwen3.5-0.8B-Q8_0.gguf
~/SrcArchives/AI/llama.cpp/build/bin/llama-tokenize
~/SrcArchives/AI/llama.cpp/build/bin/llama-bench

Most benchmark/probe CLIs also accept --model, --target, --draft, --tokenizer-bin, or environment overrides. bin/qwen35_generate.cr is intentionally a small demo and currently uses its constants at the top of the file.

Build Qwen CLIs

Build the Metal bridge once:

make build/bridge.o

Build the practical generation demo:

crystal build --release --no-debug \
  --link-flags="$(pwd)/build/bridge.o -framework Metal -framework Foundation -lc++" \
  bin/qwen35_generate.cr \
  -o build/qwen35_generate

Run greedy generation:

./build/qwen35_generate "The capital of France is" 64

Enable exact n-gram speculative decode for repeated text:

QWEN35_NGRAM_DECODE=1 ./build/qwen35_generate "The capital of France is" 64

Use the conservative automatic decode policy:

QWEN35_DECODE_POLICY=auto ./build/qwen35_generate "The capital of France is" 64

Enable exact neural speculative decode with the Qwen 3.5 0.8B draft:

QWEN35_SPECULATIVE_DECODE=1 \
QWEN35_HEAD_FULL_ROWS_GUARDED=1 \
./build/qwen35_generate "The capital of France is" 64

Enable exact prompt cache:

QWEN35_PROMPT_CACHE=1 \
QWEN35_SESSION_ID=demo \
./build/qwen35_generate "The capital of France is" 64

Useful Qwen environment switches:

Library Integration

The current Qwen API is low-level and intended for native inference experiments:

require "ml/gguf/qwen35_cpu"
require "ml/gguf/qwen35_weights"

model = "/path/to/Qwen3.5-9B-Q4_K_M.gguf"
weights = ML::GGUF::Qwen35Weights.from_gguf(model)
state = ML::GGUF::Qwen35CPU::State.new(weights.hparams, max_seq: 1024)

prompt_ids = [760_i32, 6511_i32, 314_i32, 9338_i32, 13_i32]
next_id, next_logit = ML::GGUF::Qwen35CPU.prefill_tokens_top1(weights, prompt_ids, 0, state)

64.times do |i|
  puts next_id
  next_id, next_logit = ML::GGUF::Qwen35CPU.forward_top1(weights, next_id, prompt_ids.size + i, state)
end

When linking an executable that uses Metal, include the bridge object and Apple frameworks:

crystal build your_app.cr \
  --link-flags="$(pwd)/build/bridge.o -framework Metal -framework Foundation -lc++"

For CPU-only builds:

crystal build -Dcpu_only your_app.cr

Benchmark Against llama.cpp

Build the matched benchmark:

crystal build --release --no-debug \
  --link-flags="$(pwd)/build/bridge.o -framework Metal -framework Foundation -lc++" \
  bin/benchmark_qwen_vs_llama.cr \
  -o build/benchmark_qwen_vs_llama

Run a normal first-run prefill/decode comparison:

./build/benchmark_qwen_vs_llama \
  --model ~/.cache/lm-studio/models/lmstudio-community/Qwen3.5-9B-GGUF/Qwen3.5-9B-Q4_K_M.gguf \
  --llama-bench ~/SrcArchives/AI/llama.cpp/build/bin/llama-bench \
  --prompt=64 \
  --gen=64 \
  --reps=5 \
  --warmup=2

For publishable measurements, wait for a quiet host:

./build/benchmark_qwen_vs_llama \
  --prompt=64 \
  --gen=64 \
  --reps=5 \
  --warmup=2 \
  --wait-quiet-ms=60000 \
  --require-quiet

Additional benchmark modes:

# Fresh State per repetition, but Metal state buffers are prepared before
# the timed prefill. This measures prompt ingest without first-touch buffer
# allocation/zeroing in the timed region.
./build/benchmark_qwen_vs_llama --native-prefill-prepare-state

# State buffers allocated once, then reset between reps.
./build/benchmark_qwen_vs_llama --native-prefill-prealloc

# Exact prompt-cache restore after one seeded native prefill.
./build/benchmark_qwen_vs_llama --native-prefill-cache

Fresh local M2 Max 64GB relaxed-load snapshot after the shared-H16 recurrent projection cleanup, Qwen 3.5 9B Q4_K_M, llama.cpp llama-bench, prompt=64, gen=64, reps=3, warmup=1, flash-attention off:

Notes:

• The table is a local engineering snapshot, not a lab-clean public benchmark.
• First-run prefill is still behind llama.cpp on this machine. The native wins currently come from state reuse, prompt-cache restore, and exact speculative decode.
• --native-prefill-prepare-state uses a fresh State per repetition but calls Qwen35CPU.prepare_state_metal! before timing. This is useful for server-style latency where a session object can be prepared before the prompt arrives.
• --native-prefill-cache measures exact restore of a previously computed prompt state; it is not a first-run prefill replacement.
• Short decode runs are noisy on a desktop system. The two plain decode rows above are intentionally both shown: treat plain decode as parity-to-faster, not as a stable public margin without a quiet rerun.

Speculative Decode Harnesses

Neural draft harness with Qwen 3.5 0.8B Q8_0:

crystal build --release --no-debug \
  --link-flags="$(pwd)/build/bridge.o -framework Metal -framework Foundation -lc++" \
  bin/qwen35_speculative_accept.cr \
  -o build/qwen35_speculative_accept

./build/qwen35_speculative_accept \
  --target ~/.cache/lm-studio/models/lmstudio-community/Qwen3.5-9B-GGUF/Qwen3.5-9B-Q4_K_M.gguf \
  --draft ~/.cache/lm-studio/models/lmstudio-community/Qwen3.5-0.8B-GGUF/Qwen3.5-0.8B-Q8_0.gguf \
  --tokens 64 \
  --ngram \
  "The capital of France is"

Target-only n-gram speculative harness:

crystal build --release --no-debug \
  --link-flags="$(pwd)/build/bridge.o -framework Metal -framework Foundation -lc++" \
  bin/qwen35_ngram_speculative.cr \
  -o build/qwen35_ngram_speculative

./build/qwen35_ngram_speculative \
  --tokens 64 \
  --gamma 32 \
  --min-ngram 6 \
  "The capital of France is"

Both harnesses replay/check exact greedy target output by default unless their CLI explicitly says otherwise.

Fresh local speculative smoke, same M2 Max 64GB and Qwen 3.5 9B target:

Speculative decode caveats:

• The speculative paths are exact greedy verification paths, not approximate sampling shortcuts.
• Neural speculative speed depends on draft acceptance. High-accept prompts are faster; rejection-heavy prompts quickly fall back to plain target decode.
• In qwen35_generate, neural speculative decode is useful for longer high-accept generations. In a local 64-token smoke, The capital of France is measured 20.40 ms/tok greedy, 16.61 ms/tok neural speculative, and 15.10 ms/tok neural speculative with guarded full-row verification. A 32-token smoke was slower due fixed draft/verifier overhead.
• N-gram speculation is a workload-specialized path for repeated/generated-template text. It is intentionally fail-closed after a rejected n-gram chunk by default.
• N-gram verifier chunks temporarily disable guarded full-row verification even if QWEN35_HEAD_FULL_ROWS_GUARDED=1, because partial n-gram rejection exposed a close-row guard failure during adversarial CLI testing.
• QWEN35_HEAD_FULL_ROWS_GUARDED=1 is still an experimental research switch. The harness checks final output against plain greedy target output, but the route is not broad-defaulted because it relies on a full-row F16 top1 margin guard.
• These numbers are effective decode throughput after prompt prefill; they do not make first-run prefill faster.

Native Metal Embeddings

The embedding path targets nomic-embed-text-v2-moe with a fully native Metal compute pipeline.

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")

Embedding Performance

Apple M2 Max, 38 GPU cores:

Embedding Pipeline Internals

• simdgroup-matrix GEMM for Q5_K/Q6_K dequant+multiply.
• Batched expert GEMM for MoE experts.
• ComputeGraph wave scheduling with offset-aware dependency analysis.
• Fused QKV split/RoPE, gate/softmax/top-k, scatter, and norm kernels.
• GPU-driven dispatch where useful.

Supported Models

Installation

# shard.yml
dependencies:
  cogni-ml:
    github: skuznetsov/cogni-ml
    version: ~> 0.40.0

Build And Test

make build
make spec

CPU-only:

make build_cpu
make spec_cpu

llama.cpp helper targets:

make llama
make llama_env

The Makefile searches common local, Homebrew, and system library locations for libllama. Override with LLAMA_DIR, LLAMA_BUILD, or LLAMA_LIB_DIR if needed.

Quick Start

Tensor + Autograd

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 Through 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

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
)

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

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

Metal Kernels

Platform Support

NVIDIA/CUDA support is not implemented. The Qwen native path is Metal-first.

Build Flags

License

MIT