skuznetsov/cogni-ml
Crystal ML library: Autograd, Tensors, Neural Networks, Optimizers - skuznetsov/cogni-ml
18 stars · 1 forks · Crystal
README
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. Native Qwen BPE tokenizer with an external llama-tokenize fallback for A/B. Exact prompt-state save/restore, tokenized-prompt reuse, 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 CPU-only GGUF/Qwen metadata smoke on Linux, CUDA hosts, or any environment where Metal is unavailable: crystal build -Dcpu_only bin/qwen35_gguf_info.cr -o build/qwen35_gguf_info ./build/qwen35_gguf_info --model /path/to/Qwen3.5-9B-Q4_K_M.gguf ./build/qwen35_gguf_info --model /path/to/Qwen3.5-0.8B-Q8_0.gguf --load-weights This entrypoint intentionally does not run inference. It verifies GGUF parsing, Qwen 3.5/3.6 hparams, tensor inventory, and the structured Qwen35Weights loader without pulling the Metal bridge into a Linux build. Build the minimal Crystal CUDA Driver API smoke on NVIDIA/Linux hosts: crystal build bin/cuda_driver_smoke.cr -o build/cuda_driver_smoke ./build/cuda_driver_smoke 4096 The CUDA smoke is a backend boundary probe only: it links libcuda, loads embedded PTX, launches a vector-add kernel, and checks the result. CUDA probe code uses src/ml/cuda/driver.cr for reusable CUDA context, module/function, launch, copy, synchronize, and device-buffer ownership. This is intentionally small: it owns the raw CUDA Driver API lifecycle and calls, while higher-level layer execution is still probe-local until the CUDA backend split is promoted. It also provides ML::CUDA::ResidentSequenceRunner, a thin lifecycle facade for resident sequence probes with explicit upload_weights, reset_sequence, run_sequence, and read_outputs phases. src/ml/cuda/qwen_recurrent_layer_runner.cr is the first Qwen-specific runner extraction: it owns one recurrent layer's CUDA modules, device buffers, kernel parameters, weight upload, sequence reset, token launch graph, and output readback. QwenRecurrentLayerRunner::Weights.load owns GGUF tensor lookup, tensor-shape/type validation, and raw weight reads for the runner, including recurrent-layer ffn_down tensors stored as either Q4_K or Q6_K. CPU-reference comparison intentionally remains in the probe. Build the first quantized CUDA correctness probe on NVIDIA/Linux hosts: crystal build -Dcpu_only bin/cuda_q8_gemv_probe.cr -o build/cuda_q8_gemv_probe ./build/cuda_q8_gemv_probe \ --model /path/to/Qwen3.5-0.8B-Q8_0.gguf \ --tensor blk.0.ffn_up.weight \ --kernel warp4 \ --reps 100 \ --warmup 10 cuda_q8_gemv_probe loads a real GGUF Q8_0 tensor, launches a Crystal-driven CUDA Driver API GEMV kernel over the raw GGUF block layout, and compares against the existing CPU QuantMatmul reference. --kernel scalar keeps the first one-thread-per-output-row correctness kernel; the default --kernel warp4 maps four output rows to four warps per thread block and is the current faster probe shape. This is still a standalone backend-boundary probe, not an optimized Qwen CUDA inference path yet. The current full qwen35_generate CLI remains Metal-first. Build the first Q4_K CUDA correctness probe for Qwen 9B/27B-style target tensors: crystal build -Dcpu_only bin/cuda_q4k_gemv_probe.cr -o build/cuda_q4k_gemv_probe ./build/cuda_q4k_gemv_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --tensor blk.0.attn_gate.weight \ --kernel warp4 \ --reps 20 \ --warmup 3 cuda_q4k_gemv_probe uses the raw GGUF Q4_K block layout (d, dmin, 12-byte packed scales/mins, 128-byte packed nibbles) and checks the CUDA output against the CPU QuantMatmul Q4_K reference. --kernel scalar keeps the first correctness kernel; the default --kernel warp4 maps four output rows to four warps per block and is the current faster probe shape. Build the Q6_K CUDA correctness/speed probe for Q4_K_M tensors that remain in Q6_K: crystal build -Dcpu_only bin/cuda_q6k_gemv_probe.cr -o build/cuda_q6k_gemv_probe ./build/cuda_q6k_gemv_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --tensor blk.0.ffn_down.weight \ --kernel warp4 \ --reps 10 \ --warmup 2 cuda_q6k_gemv_probe covers the GGUF Q6_K block layout (ql, qh, signed scales, d) used by output/value/down projections in mixed-quant target models. Like the Q4_K/Q8_0 probes, it is a standalone backend primitive check; full CUDA Qwen execution is still a separate backend split. Build the first GPU-resident FFN sequence probe: crystal build -Dcpu_only bin/cuda_ffn_sequence_probe.cr -o build/cuda_ffn_sequence_probe ./build/cuda_ffn_sequence_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 0 \ --reps 10 \ --warmup 2 cuda_ffn_sequence_probe composes the checked CUDA primitives as Q4_K ffn_gate + Q4_K ffn_up -> SwiGLU -> Q6_K ffn_down while keeping the input, intermediate activations, and output projection input GPU-resident. Only the final hidden vector is copied back for comparison against the CPU QuantMatmul FFN reference. Build the full-attention input projection bundle probe: crystal build -Dcpu_only bin/cuda_attn_projection_probe.cr -o build/cuda_attn_projection_probe ./build/cuda_attn_projection_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 3 \ --reps 10 \ --warmup 2 cuda_attn_projection_probe runs Q4_K attn_q + Q4_K attn_k + Q6_K attn_v from one GPU-resident hidden vector and copies Q/K/V back only after all projections complete. It targets full-attention layers such as blk.3 in Qwen3.5 9B. The probe now routes through ML::CUDA::QwenFullAttnProjectionRunner, supports --tokens N, and keeps Q/K/V outputs GPU-resident until the final correctness readback. It is the reusable input-projection boundary for future full-attention/KV CUDA work, not a complete full-attention layer runner yet. Build the full-attention Q/K normalization + RoPE + KV-cache boundary probe: crystal build -Dcpu_only bin/cuda_full_attn_kv_probe.cr -o build/cuda_full_attn_kv_probe ./build/cuda_full_attn_kv_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 3 \ --tokens 4 \ --start-pos 2 \ --max-seq 12 cuda_full_attn_kv_probe now routes through ML::CUDA::QwenFullAttnLayerRunner, a residual-hidden-to-final-hidden wrapper around the projection runner and ML::CUDA::QwenFullAttnKVRunner. The projection runner can apply the initial attn_norm on CUDA from residual hidden states before Q/K/V projection; Q is split into normalized/RoPE'd Q and gate, K is RMSNormed/RoPE'd, K/V rows are appended to a CUDA-resident cache at start_pos, a correctness-first serial CUDA kernel computes GQA scores, softmax, value reduction, and Q-gate multiplication, the resident gated attention output is projected through attn_output.weight, and the layer tail runs residual add, post-attention RMSNorm, FFN gate/up/SwiGLU/down, and final residual. It checks Q, gate, K, gated attention output, projected attention output, final hidden, K-cache, and V-cache against the CPU Qwen reference. This is now a clean one-layer semantics probe with device input/output hooks for mixed-stack composition, but it is not yet an end-to-end Linux decode path: full/recurrent stack scheduling, logits/top1, tokenizer/sampling, restored nonzero prefix KV, and a faster attention kernel remain separate gates. Build the Q5_K CUDA recurrent-QKV probe: crystal build -Dcpu_only bin/cuda_q5k_gemv_probe.cr -o build/cuda_q5k_gemv_probe ./build/cuda_q5k_gemv_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --tensor blk.0.attn_qkv.weight \ --reps 10 \ --warmup 2 cuda_q5k_gemv_probe covers the GGUF Q5_K block layout used by recurrent-layer combined attn_qkv.weight tensors in the current Qwen3.5 9B Q4_K_M file. Build the recurrent-layer projection bundle probe: crystal build -Dcpu_only bin/cuda_recurrent_projection_probe.cr -o build/cuda_recurrent_projection_probe ./build/cuda_recurrent_projection_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 0 \ --reps 10 \ --warmup 2 cuda_recurrent_projection_probe runs Q5_K attn_qkv + Q4_K attn_gate + Q4_K ssm_alpha + Q4_K ssm_beta from one GPU-resident hidden vector and copies the four outputs back only after all kernels complete. It is the first CUDA recurrent projection-bundle proof; DeltaNet recurrence, convolution, state updates, and ssm_out remain separate work. Build the synthetic DeltaNet output slice probe: crystal build -Dcpu_only bin/cuda_deltanet_output_probe.cr -o build/cuda_deltanet_output_probe ./build/cuda_deltanet_output_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 0 \ --reps 10 \ --warmup 2 cuda_deltanet_output_probe runs a synthetic CUDA DeltaNet state update, applies post RMSNorm/SiLU gating on GPU, and feeds the result directly into the real Q4_K ssm_out.weight projection. It is a stateful boundary probe, not a full recurrent layer: recurrent conv prep, alpha/beta transforms, residuals, and FFN remain separate work. Build the recurrent prep/output slice probe: crystal build -Dcpu_only bin/cuda_recurrent_prep_output_probe.cr -o build/cuda_recurrent_prep_output_probe ./build/cuda_recurrent_prep_output_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layer 0 \ --tokens 4 \ --reps 10 \ --warmup 2 cuda_recurrent_prep_output_probe now composes one full recurrent layer token slice: input RMSNorm, real recurrent projection bundle (attn_qkv, attn_gate, ssm_alpha, ssm_beta), recurrent conv prep, alpha/beta transforms, DeltaNet, post RMSNorm/SiLU, Q4_K ssm_out, residual add, post-attention RMSNorm, Q4_K FFN gate/up, SwiGLU, Q6_K FFN down, and final residual. --tokens N runs a GPU-resident sequence through persistent conv/SSM state and compares all token outputs plus final recurrent states against the CPU reference. The probe now separates one-time weight upload from per-sequence input/state reset and prints weight_upload_ms; timed cuda_ms_per_token excludes the persistent weight upload. GGUF recurrent-layer tensor loading is routed through QwenRecurrentLayerRunner::Weights.load, so the probe no longer manually passes every raw tensor into the runner constructor. It is still a standalone one-layer probe, not an end-to-end Linux decoder. Build the recurrent multi-layer stack scaffold: crystal build -Dcpu_only bin/cuda_recurrent_stack_probe.cr -o build/cuda_recurrent_stack_probe ./build/cuda_recurrent_stack_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layers 0,2,4 \ --tokens 2 cuda_recurrent_stack_probe chains multiple QwenRecurrentLayerRunner instances and compares the final hidden sequence plus each layer's recurrent conv/SSM state against the CPU reference. The default path hands each recurrent layer's CUDA output buffer directly to the next layer's CUDA input; --host-handoff keeps the older host-copy route as a debug oracle. This is still a recurrent-only scaffold, not an end-to-end Linux decoder. Build the mixed recurrent/full-attention CUDA stack probe: crystal build -Dcpu_only bin/cuda_mixed_stack_probe.cr -o build/cuda_mixed_stack_probe ./build/cuda_mixed_stack_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --layers 0,1,2,3,4 \ --tokens 2 \ --start-pos 2 \ --max-seq 12 cuda_mixed_stack_probe composes QwenRecurrentLayerRunner and QwenFullAttnLayerRunner in model layer order with device-resident hidden handoff across recurrent/full-attention boundaries, then runs QwenOutputHeadRunner for output RMSNorm, quantized lm-head projection, and resident top1. The layer/head loop is now owned by ML::CUDA::QwenMixedStackRunner, which is the first model-slice decode-state object for CUDA. By default the probe copies back only the CUDA top1 id/value plus hidden/state debug outputs; pass --read-logits to also copy full logits for attribution, and --profile-phases to insert per-layer/head synchronizations and print attribution lines. It compares the final hidden sequence, top1, recurrent conv/SSM states, and full-attention KV cache rows against the CPU reference. This is the first mixed-stack CUDA correctness scaffold through resident top1; it still stops before tokenizer/sampling, repeated full-model decode ownership, and an optimized topK/sampling kernel. The current resident top1 is a simple two-phase partial-scan/reduce kernel: correct, but not yet promoted as a speed-optimized head. Experimental CUDA exact cache-replay fast lane: crystal build --release --no-debug -Dcpu_only \ bin/cuda_mixed_stack_probe.cr \ -o build/cuda_mixed_stack_probe
./build/cuda_mixed_stack_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --all-layers \ --max-seq 256 \ --greedy-loop-tokens 64 \ --greedy-loop-probe-chunk-gamma 16 \ --greedy-loop-probe-chunk-active-verify \ --greedy-loop-probe-ngram \ --greedy-loop-probe-ngram-source-history "$SOURCE_TOKEN_IDS" \ --greedy-loop-probe-ngram-replay-start 1 \ --greedy-loop-probe-ngram-cursor-only \ --greedy-loop-probe-ngram-trusted-source \ --greedy-loop-probe-ngram-schedule 64 \ --skip-debug-readback This path is exact verification of a validated source/cache cursor, not a general-purpose sampler. It first checks that the live prefix matches the source history at the replay cursor. If the prefix gate fails, the active verifier path is disabled before any proposals are trusted and the probe falls back to plain greedy target decode. If the gate passes, proposal chunks are verified through the same resident CUDA stack and rejected chunks restore the exact target state. Use smaller schedules such as 4,4,8,16 for weaker proposal sources where early reject economics matter; use bulk schedules such as 64 only for artifact-level trusted replay where the source history was produced by the same model/cache contract and every proposed token is still target-verified before commit. Trusted artifact restore lower-bound: ./build/cuda_mixed_stack_probe \ --model /path/to/Qwen3.5-9B-Q4_K_M.gguf \ --all-layers \ --max-seq 256 \ --known-replay-history "$SOURCE_TOKEN_IDS" \ --known-replay-start 1 \ --known-replay-tokens 64 \ --known-replay-trusted-artifact-restore \ --skip-debug-readback For a conservative host-backed artifact simulation, add: --known-replay-trusted-artifact-host-restore To snapshot only KV rows up to the restored cursor instead of the full max_seq cache, use: --known-replay-trusted-artifact-live-kv To include diagnostic artifact file IO and hash timing, add: --known-replay-trusted-artifact-io-probe \ --known-replay-trusted-artifact-io-path /path/to/artifact.bin This is a different contract from verified replay. It simulates a cache artifact that already contains the exact token span and the post-span decode state for the same model/config/source hash. The timed region restores that state and emits cached tokens; it does not recompute the verifier body. Use it only for session/cache artifacts whose model hash, tokenizer hash, prefix hash, token span, and state artifact hash have been validated. The default trusted-artifact probe snapshots/restores state inside device memory and is a pure lower bound. The host-backed variant copies the complete decode state to host memory, poisons the runner, then times host-to-device restore. On the RTX 5060 Ti Qwen3.5-9B snapshot with max_seq=128, that state is 61,079,552 bytes; restore costs 12.259ms for a 64-token artifact and 12.331ms for a non-zero start9/24-token artifact. This is still not full production IO: durable lookup, disk reads, hashing, and live-length KV trimming remain outside the measured region. The live-KV variant preserves exact output while reducing only the full-attention cache portion. On the same host/model, start1/gen64 stores 56,885,248 bytes and restores in 11.358ms; start9/gen24 stores 54,788,096 bytes and restores in 10.984ms. The modest delta is a useful finding: recurrent DeltaNet state, not KV capacity, dominates this short-context artifact. On the persistent reefy storage path, the live-KV start1/gen64 artifact breaks down into 52,690,944 recurrent-state bytes and 4,194,304 KV bytes. The diagnostic file path measured write 17.695ms, read 42.046ms, SHA-256 verification 59.556ms, and H2D restore 13.003ms. For start9/gen24, the artifact is 52,690,944 recurrent bytes plus 2,097,152 KV bytes; write/read/ hash/restore measured 18.249/42.994/53.814/12.825ms. The SHA-256 timing uses the host sha256sum/shasum tool to avoid adding OpenSSL linkage to the CUDA probe, so treat it as a product-boundary diagnostic rather than a kernel metric. An experimental contiguous-read reconstruction reduced start1/gen64 read time from 42.046ms to 36.310ms while preserving 64/64 exact output, but start9/ gen24 remained noisy (43.618ms read). This points to some avoidable allocation/ fragmentation cost, but not enough to change the main conclusion: cache artifacts need async prefetch or a resident artifact service to stay off the decode critical path. A first scalar CPU recurrent-state block-INT8 codec diagnostic compresses the 52,690,944 recurrent bytes to 13.18-14.00MB (25.0-26.6%) depending on block size. On start1/gen64, block sizes 64/256/1024/4096 measured relative RMSE 0.008798/0.012300/0.015476/0.019087; encode/decode stayed around 457-469ms / 121-125ms. A follow-up restore gate now decodes the block-INT8 recurrent buffers back into a host snapshot, restores that state, and runs one continuation token against an exact uncompressed restored state. On the RTX 5060 Ti host, block sizes 64/256/1024/4096 preserved the same next top1 on start9/gen24, and block 256/4096 also preserved the same next top1 on a later start33/gen8 cursor. The m