dart_pytorch

A minimal PyTorch-style tensor library for Dart, backed by hand-written CUDA kernels via dart:ffi.

Matmul, elementwise binary ops (+ - * /, scalar and row-broadcast), unary activations (relu, sigmoid, tanh, abs, log, pow), 2D transpose, sum/mean reductions, LayerNorm, row-wise softmax, fused crossEntropy, embedding, scaled dot-product attention, concat, and dropout all run end-to-end on either CPU or GPU. Reverse-mode autograd is wired for the whole op set (with relu/abs backward CPU-only for now). A small nn.Module scaffolding exposes Linear, LayerNorm, Embedding, Dropout, MultiHeadAttention, a pre-LN TransformerBlock, TransformerEncoder (stacked blocks with optional final LN), TransformerLM (token embed + positional encoding

causal encoder + linear head), GPT (weight-tied GPT-2 style causal LM with learned position embeddings, embedding dropout, an EncoderCache KV-cache path for O(N) autoregressive generate(), and greedy / temperature / top-k sampling), and both SinusoidalPositionalEncoding / LearnedPositionalEmbedding as trainable / regularization layers with train() / eval() mode toggling; SGD / Adam optimizers update parameters in place on their native device, clipGradNorm bounds the global gradient L2, and LRSchedulers (StepLR, LinearWarmupCosineDecay) mutate the optimizer's lr on a schedule. Checkpoint (in nn/serialize.dart) persists any Module's parameters to a small binary format on disk and loads them back into a same-shape model. A pure-Dart byte-level BpeTokenizer in core/data/bpe_tokenizer.dart provides train / encode / decode / saveFile / loadFile for tokenizing real text. Runnable char-level demos at bin/lm_demo.dart and bin/gpt_demo.dart overfit a short refrain end-to-end (dart run bin/gpt_demo.dart); bin/gpt_train.dart shows the full pipeline: BPE + GPT + Adam + warmup/cosine schedule + gradient accumulation + checkpoint save/load + sampling.

Layout

lib/
  dart_pytorch.dart              # package entry point (re-exports Tensor)
  core/tensor/
    tensor.dart                  # Tensor class + factories (fromList, fill) + to()
    ops.dart                     # elementwise / activation / reduction ops, CPU + GPU (part)
    mat_mul.dart                 # matmul (CPU loop + GPU tiled kernel) (part)
    layer_norm.dart              # LayerNorm forward + backward, CPU + GPU (part)
    softmax.dart                 # softmax + fused crossEntropy, CPU + GPU (part)
    embedding.dart               # table lookup with scatter-add backward (part)
    attention.dart               # scaled dot-product attention (composition) (part)
    dropout.dart                 # inverted dropout via mask multiply (part)
    concat.dart                  # last-axis 2D concat with slice-back backward (part)
    cuda_engine.dart             # dart:ffi bindings to libmat_mul.so
  core/nn/
    module.dart                  # abstract Module base (parameters, zeroGrad, train/eval)
    linear.dart                  # trainable Linear (y = x @ W.T + b)
    layer_norm.dart              # trainable LayerNorm module wrapper
    embedding.dart               # trainable Embedding module wrapper
    dropout.dart                 # nn.Dropout wrapper with train/eval toggle
    multi_head_attention.dart    # per-head Linear + SDPA + concat + out proj
    transformer.dart             # pre-LN TransformerBlock (MHA + MLP + residuals)
    positional.dart              # sinusoidal + learned positional encodings
    masks.dart                   # causalMask(n) additive attention mask helper
    transformer_encoder.dart     # stacked TransformerBlocks + optional final LN
    transformer_lm.dart          # token embed + posEnc + causal encoder + head
    kv_cache.dart                # MHACache + EncoderCache for GPT.generate fast path
    gpt.dart                     # GPT-2 style: tied weights, learned PE, generate()
    serialize.dart               # Checkpoint.save/loadInto — DPTC binary format
  core/optim/
    optimizer.dart               # abstract Optimizer base (step, zeroGrad, mutable lr)
    sgd.dart                     # SGD with optional momentum + weight decay
    adam.dart                    # Adam with bias correction + decoupled WD
    grad_utils.dart              # clipGradNorm (global L2 clip, in place)
    lr_scheduler.dart            # StepLR + LinearWarmupCosineDecay
  core/data/
    bpe_tokenizer.dart           # byte-level BPE: train / encode / decode / save / load
  native/
    src/
      engine.cu                  # extern "C" DLLEXPORT wrappers (30 symbols)
      kernels/
        common.cuh               # CUDA includes, DLLEXPORT macro, reductions
        matmul.cuh               # tiled matmul_fwd/bwd kernels
        elementwise.cuh          # add/sub/mul/div, scalar/row-bcast, activations, abs/log/pow
        transpose.cuh            # 32x32 tile transpose
        layernorm.cuh            # layernorm_fwd/bwd (block-per-row)
        softmax.cuh              # softmax_fwd/bwd + cross_entropy_fwd/bwd (fused)
        embedding.cuh            # embedding_fwd + scatter-add embedding_bwd
    lib/                         # populated by the nvcc build (gitignored)
native/lib/libmat_mul.so         # actual load path used by cuda_engine.dart
doc/device-placement.md         # per-op CPU vs GPU decisions + implementation status
test/dart_pytorch_test.dart      # matmul + CPU-op correctness tests

Note: cuda_engine.dart loads ${cwd}/native/lib/libmat_mul.so, so the .so lives at the repo root, not under lib/native/lib/.

Requirements

Dart SDK ^3.9.3
NVIDIA GPU + CUDA toolkit (tested with CUDA 12.0, driver 576.x)
Linux / WSL2

Build the native library

nvcc --shared -Xcompiler -fPIC \
     -o native/lib/libmat_mul.so \
     lib/native/src/engine.cu

No local GPU? See doc/colab.md for a copy/paste recipe that runs on a free Google Colab, Kaggle, Paperspace, or Lightning AI GPU. scripts/setup_colab.sh handles the whole install + nvcc build in one command.

Have friends with GPUs, or want to pool donated compute? See doc/coop-training.md for three cooperative training modes (single-process replicas, HTTP coordinator + workers, peer-to-peer gossip). All three use DiLoCo-style periodic parameter averaging so one full model shipment every K local steps is enough bandwidth for training over slow / free internet links.

Run the tests

dart pub get
dart test

Expected output: +137: All tests passed! (matmul CPU + GPU paths, mixed-device rejection, device round-trip, every CPU op, every GPU op, a CPU/GPU parity chain, 22 autograd tests, 9 LayerNorm tests, 16 softmax / cross-entropy / embedding tests, 13 optimizer tests, 9 attention / Linear tests, 14 regularization tests, 17 concat / MultiHeadAttention / TransformerBlock tests, plus 15 positional-encoding tests including PE + TransformerBlock trained end-to-end with Adam).

Usage

import 'package:dart_pytorch/dart_pytorch.dart';

void main() {
  // Small tensors default to CPU (below Tensor.autoDeviceThreshold = 4096).
  final a = Tensor.fromList([2, 3], [1, 2, 3, 4, 5, 6]);
  final b = Tensor.fromList([3, 2], [7, 8, 9, 10, 11, 12]);
  final c = a.matmul(b);            // CPU path, shape [2, 2]
  print(c.toList());                // [58, 64, 139, 154]

  // CPU elementwise + activations.
  final x = Tensor.fromList([4], [-1.0, 0.0, 1.0, 2.0]);
  print((x + 1).toList());          // [0, 1, 2, 3]
  print(x.relu().toList());         // [0, 0, 1, 2]

  // Autograd.
  final w = Tensor.fromList([1, 2], [0.5, -0.3], requiresGrad: true);
  final xVec = Tensor.fromList([2, 1], [1.0, 2.0]);
  final loss = (w.matmul(xVec) - 4.0).pow(2).sum();
  loss.backward();
  print(w.grad!.toList()); // gradient wrt w

  // Opt in to GPU for large workloads.
  final big = Tensor.fromList([64, 64],
      List<double>.generate(4096, (i) => i.toDouble()),
      device: Device.GPU);
  final bigResult = big.matmul(big); // tiled 32x32 kernel
  big.dispose();
  bigResult.dispose();

  // Or transfer explicitly.
  final gpuVersion = a.to(Device.GPU);
  gpuVersion.dispose();
}

Currently supported

See doc/device-placement.md for the full policy and per-op status. Short version:

Area	CPU	GPU	Notes
`Tensor.fromList` / `Tensor.fill`	yes	yes	Device chosen by size (threshold 4096) or explicit
`to(Device)`, `toList()`, `dispose()`	yes	yes
`matmul` (forward)	yes	yes	CPU naive loop; GPU tiled 32x32
`matmul` (backward)	—	—	GPU kernels compiled, not wired to Dart
Elementwise `+ - * /` (same shape)	yes	yes
Scalar broadcast (`t + num` or `t + [1]`)	yes	yes	`num` on GPU uploads a 1-elem tensor + disposes
Row broadcast `[1,N]` into `[M,N]`	yes (all)	`+` only	Sub/mul/div row-bcast on GPU throws — use `.to(Device.CPU)`
`relu`, `sigmoid`, `tanh`, `abs`, `log`, `pow`	yes	yes
`transpose` (2D)	yes	yes	CPU strided copy; GPU 32x32 tile kernel
`sum`, `mean`	yes	yes	GPU: `atomicAdd`-based reduction into `[1,1]`
`layerNorm`	yes	yes	2D `[R,C]` over `C`; also `nn.LayerNorm(dim)`
`softmax`	yes	yes	Row-wise on 2D `[R,C]`, numerically stable
`crossEntropy(targets)`	yes	yes	Fused softmax + NLL; returns `[R,1]` per-sample
`embedding(indices)`	yes	yes	`[V,D]` table gathered by `[N]`; also `nn.Embedding`
Autograd graph	yes	yes	Dart-side tape; `relu`/`abs` backward CPU-only
`SGD` / `Adam` optimizers	yes	yes	State (velocity, m/v) lives on parameter's device; `Tensor.assign` swaps updates in-place
`scaledDotProductAttention`	yes	yes	2D single-head SDPA; composed from matmul/transpose/softmax/matmul with optional additive mask
`nn.Linear(in, out)`	yes	yes	`y = x @ W.T + b`, Kaiming-uniform init; `bias: false` supported
`dropout(p)` / `nn.Dropout`	yes	yes	Inverted dropout via `mask` multiply; `train()`/`eval()` toggles it
`TensorConcat.concat(list, axis=1)`	yes	yes	Last-axis 2D concat; GPU round-trips through host
`nn.MultiHeadAttention(D, H)`	yes	yes	Per-head Linear + SDPA + concat + out proj; optional attn Dropout
`nn.TransformerBlock(D, H, {ffnDim})`	yes	yes	Pre-LN encoder block: MHA + residual + MLP + residual, all submodules toggled by `train()`/`eval()`
`nn.SinusoidalPositionalEncoding(D)`	yes	yes	Fixed sin/cos PE, no params, recomputed per forward for the exact seqLen
`nn.LearnedPositionalEmbedding(maxLen, D)`	yes	yes	Trainable position table gathered via `Embedding`; scatter-add backward comes free
`clipGradNorm(params, maxNorm)`	yes	yes	Global L2 grad clip, in-place via `Tensor.assign`

Placement rule: ops respect the input tensor's device. Mixed-device inputs to a binary op raise ArgumentError — call .to(...) yourself so the copy cost stays visible.

FFI surface (C entry points)

Defined in lib/native/src/engine.cu (30 symbols):

Category	Symbols
Lifecycle	`create_tensor`, `destroy_tensor`, `get_tensor_data`
Matmul	`matmul_tensors`
Elementwise	`add_tensors`, `sub_tensors`, `mul_tensors`, `div_tensors`
Scalar bcast	`add_tensor_scalar`, `sub_tensor_scalar`, `mul_tensor_scalar`, `div_tensor_scalar`
Row bcast	`add_tensor_row_broadcast`
Unary math	`abs_tensor`, `log_tensor`, `pow_tensor` (`float exp`)
Activations	`relu_tensor`, `sigmoid_tensor`, `tanh_tensor`
Rearrangement	`transpose_tensor`
Reductions	`sum_tensor`, `mean_tensor` (both return a `[1,1]` handle)
LayerNorm	`layernorm_forward`, `layernorm_backward`
Softmax / CE	`softmax_forward`, `softmax_backward`, `cross_entropy_forward`, `cross_entropy_backward`
Embedding	`embedding_forward`, `embedding_backward` (scatter-add into `gTable`)

All wrappers return void* (Tensor handle) except lifecycle helpers.

Provenance

The CUDA kernels and FFI patterns were lifted from ../dart_cuda, keeping only forward-mode wrappers: kernels/common.cuh, kernels/matmul.cuh, kernels/elementwise.cuh, kernels/transpose.cuh, and the forward-only slice of engine.cu.

Libraries

core/coop/lm_factory: Small factory that produces either a GPT (attention-based) or an AFTLanguageModel (attention-free) language model behind a uniform interface, so the coop_* training scripts can flip architectures with a single --arch=gpt|aft flag without duplicating the training loop.
core/coop/param_avg: Parameter averaging for DiLoCo-style cooperative training.
core/coop/shard: Shared data helpers for the coop_* demos:
core/data/bpe_tokenizer: Byte-level Byte Pair Encoding tokenizer.
core/data/char_tokenizer: Character-level tokenizer.
core/data/csv_dataset: Tabular CSV dataset for regression and classification.
core/data/dataset: PyTorch-style Dataset + DataLoader abstractions.
core/data/hf_bpe_tokenizer: Loads a HuggingFace tokenizer.json (byte-level BPE, e.g. GPT-2, GPT-NeoX/Pythia, Llama-style) and provides encode/decode.
core/data/image_folder_dataset: Image-folder classification / triplet dataset.
core/data/text_token_dataset: Language-model sliding-window dataset.
core/data/wordpiece_tokenizer: BERT WordPiece tokenizer (uncased BERT / MiniLM / all-MiniLM-L6-v2 style). Faithful enough to reproduce HuggingFace BertTokenizer output on the vast majority of English text.
core/nn/aft_transformer: Attention-Free Transformer (AFT) block + language-model stack.
core/nn/attention/aft_attention: Attention-Free Transformer (AFT) attention module.
core/nn/attention/multi_head_attention: Multi-head self / cross attention (2D single-sequence).
core/nn/attention/multi_head_cross_attention: Multi-head cross-attention.
core/nn/bert: HuggingFace BERT-style encoder. Post-LN, learned absolute positions, WordPiece vocab, token-type embeddings folded into a single bias since sentence-transformer inference always uses type id 0. Matches the layout of bert-base-*, sentence-transformers/ all-MiniLM-L6-v2, etc.
core/nn/bert_hf_loader: Loads HuggingFace BERT-style safetensors (bert-base, MiniLM, sentence-transformers/all-MiniLM-, all-mpnet-) into a BertModel. Pooler tensors (pooler.dense.*), position_ids and buffers are ignored — sentence-transformer models don't use the CLS pooler.
core/nn/clip_hf_loader: Loads HuggingFace CLIP vision weights (CLIPVisionModel / CLIPModel safetensors) into a CLIPVisionModel.
core/nn/conv2d: 2D convolution — inference-only, im2col + matmul.
core/nn/dropout: nn.Dropout — stochastic zeroing of activations during training.
core/nn/embedding: Embedding module — trainable table of shape [V, D].
core/nn/encoder_decoder: Full encoder-decoder Transformer for seq2seq tasks (translation, summarization, etc.).
core/nn/ffn/swiglu: SwiGLU feed-forward block — the gated FFN used by Llama, Mistral, Qwen, Phi-3 and DeepSeek. Formally:
core/nn/gpt: A minimal GPT-style causal language model.
core/nn/gpt2_hf_loader: Loads HuggingFace gpt2 / gpt2-medium / gpt2-large / gpt2-xl weights into a GPT module.
core/nn/gptj: GPT-J causal language model (EleutherAI, 2021).
core/nn/gptj_hf_loader: Loads HuggingFace GPT-J (EleutherAI/gpt-j-6B) safetensors weights into a GPTJModel.
core/nn/kv_cache: Per-block KV cache for autoregressive attention.
core/nn/layer_norm: LayerNorm module — a trainable wrapper around Tensor.layerNorm.
core/nn/lc0: LC0 classical CNN (with optional SE units) — GPU inference.
core/nn/lc0_input: Chess FEN -> 112-plane LC0 input tensor.
core/nn/lc0_proto: Minimal LC0 weights.pb.gz reader.
core/nn/linear: Linear (affine) layer — y = x @ W.T + b.
core/nn/llama: Llama-style causal language model.
core/nn/llama_hf_loader: Loads HuggingFace Llama (llama / LlamaForCausalLM architecture) weights from a safetensors file into a Llama.
core/nn/llama_vision: Vision-conditioned Llama — from-scratch multi-modal wrapper.
core/nn/masks: Attention masks — helpers that build the additive [N, N] masks consumed by scaledDotProductAttention.
core/nn/modalities/audio_transformer: Audio Transformer — sequence encoder for pre-extracted audio features (e.g. MFCCs / mel-spectrogram frames).
core/nn/modalities/multi_modal_classifier: Multimodal classifier — fuses per-modality features via mean pooling + concatenation + a small MLP head.
core/nn/modalities/multi_modal_encoder: Multimodal encoder — fuses per-modality sequences into a joint sequence context, suitable for feeding into a text decoder for tasks like captioning or spoken-Q&A.
core/nn/modalities/multi_modal_generator: Multimodal generative transformer — the "spits out text" analog of Gemini / VLM architectures. Encodes any subset of {image, audio, video, text} into a joint memory sequence, then autoregressively decodes a target text stream via causal self- attention plus cross-attention over the memory.
core/nn/modalities/multi_modal_lm: Decoder-only multimodal language model — the "packed causal stream" architecture used by Gemini / GPT-4V / Fuyu.
core/nn/modalities/text_transformer: Text Transformer — token-index encoder.
core/nn/modalities/video_transformer: Video Transformer — sequence encoder over per-frame feature embeddings (e.g. CNN-extracted features per video frame).
core/nn/module: Base class for stateful neural-network modules.
core/nn/moe: Mixture-of-Experts feed-forward block.
core/nn/moe_transformer: Transformer stack with the dense FFN sublayer replaced by a Mixture-of-Experts block.
core/nn/muzero_lm: MuZero-style latent-dynamics network with a language-model representation function.
core/nn/positional: Positional encodings for sequence models.
core/nn/pythia: GPT-NeoX / Pythia causal language model.
core/nn/pythia_hf_loader: Loads HuggingFace Pythia (gpt_neox architecture) weights from a safetensors file into a PythiaModel.
core/nn/rms_norm: RMSNorm module — a trainable wrapper around Tensor.rmsNorm.
core/nn/rotary: Rotary Positional Embeddings (RoPE) — see https://arxiv.org/abs/2104.09864.
core/nn/safetensors: Minimal reader for the safetensors on-disk format.
core/nn/sentence/cross_encoder: Cross-encoder reranker.
core/nn/sentence/losses: Training losses for sentence-transformers style bi-encoders.
core/nn/sentence/sentence_encoder: Sentence-Transformers style embedder.
core/nn/sentence/similarity: Similarity + semantic-search helpers for sentence embeddings.
core/nn/serialize: Simple parameter checkpointing for Modules.
core/nn/transformer: Pre-LayerNorm Transformer encoder block.
core/nn/transformer_decoder: Stack of TransformerDecoderBlocks with an optional final LayerNorm.
core/nn/transformer_decoder_block: Pre-LayerNorm Transformer decoder block (seq2seq style).
core/nn/transformer_encoder: Stack of TransformerBlocks with an optional final LayerNorm.
core/nn/transformer_lm: A minimal Transformer language model.
core/nn/vision/clip_vision_model: OpenAI CLIP vision transformer — implementation compatible with HuggingFace openai/clip-vit-base-patch{16,32} and openai/clip-vit-large-patch14 checkpoints.
core/nn/vision/vision_encoder: Small interface shared by vision backbones that plug into LlamaVision (and, more generally, anything that treats an image as a sequence of tokens).
core/nn/vision/vit_backbone: Vision Transformer (ViT) backbone.
core/nn/vision/vit_classifier: Image classifier on top of a Vision Transformer backbone.
core/nn/vision/vit_face_embedding: Face-embedding head on top of a Vision Transformer backbone.
core/nn/vision/vit_object_detector: DETR-style multi-object detector on top of a Vision Transformer.
core/nn/vision_projector: Small MLP that projects vision-encoder features into a language model's embedding space.
core/optim/adam: Adam optimizer (Kingma & Ba, 2015) with bias correction.
core/optim/grad_utils: Gradient regularization utilities.
core/optim/lr_scheduler: Learning-rate schedulers.
core/optim/optimizer: Base class for parameter optimizers.
core/optim/sgd: Stochastic gradient descent with optional heavy-ball momentum.
core/tensor/cuda_engine: dart:ffi bindings for the CUDA shared library libmat_mul.so.
core/tensor/dtype: Tensor storage dtypes and codec helpers.
core/tensor/native_lib_download: Auto-fetch the prebuilt native CUDA library (libmat_mul.so / mat_mul.dll / libmat_mul.dylib).
core/tensor/tensor: Device-aware Tensor with a CPU (Float32List) or GPU (FFI handle) backing, plus a Dart-side reverse-mode autograd tape.
core/utils/hungarian_algorithm: Hungarian (Kuhn–Munkres) algorithm for minimum-cost assignment.
core/vector_store/auto_tune: Auto-tuning for approximate index parameters.
core/vector_store/bench: Reusable recall / latency benchmark harness for vector indexes.
core/vector_store/faiss_io: FAISS binary-format interop.
core/vector_store/gpu_index_flat: GPU-backed brute-force flat index.
core/vector_store/index: Dart port of a subset of FAISS' vector-index toolkit.
core/vector_store/index_binary: Binary-vector companion to Index.
core/vector_store/index_binary_flat: Brute-force Hamming-distance index over fixed-length bit strings.
core/vector_store/index_binary_ivf: Inverted-list search over binary vectors.
core/vector_store/index_convert: Cross-index conversion helpers.
core/vector_store/index_factory: indexFactory — build a composed Index from a compact FAISS- style description string.
core/vector_store/index_flat: Flat (brute-force) indexes — IndexFlatL2 and IndexFlatIP.
core/vector_store/index_hnsw: IndexHNSW — Hierarchical Navigable Small World graph index.
core/vector_store/index_id_map: IndexIDMap — wraps any Index to support custom int64 ids.
core/vector_store/index_io: Persistence layer for the FAISS-in-Dart index toolkit.
core/vector_store/index_ivf_flat: Cell-probe IVF index with a flat encoding (IndexIVFFlat).
core/vector_store/index_ivf_pq: IndexIVFPQ — a.k.a. IVFADC. FAISS' workhorse for billion-scale search.
core/vector_store/index_lsh: IndexLSH — random-projection locality-sensitive hashing.
core/vector_store/index_pq: Product Quantizer + IndexPQ.
core/vector_store/index_pre_transform: IndexPreTransform — wraps an inner Index with an ordered chain of VectorTransforms. All inputs to train / add / search / rangeSearch flow through the chain first; the inner index sees only transformed vectors.
core/vector_store/index_refine_flat: IndexRefineFlat — pair a fast approximate index with an exact IndexFlat for post-verification.
core/vector_store/index_replicas: IndexReplicas — hold N functionally-identical copies of an index.
core/vector_store/index_scalar_quantizer: IndexScalarQuantizer (SQ8) — per-dimension 8-bit scalar quantization.
core/vector_store/index_shards: IndexShards — partition the corpus across a fixed set of inner indexes and merge search results at query time.
core/vector_store/kmeans: Lloyd's k-means clustering used by IVF-family indexes.
core/vector_store/l2_norm_transform: L2NormTransform — divides each vector by its L2 norm so that all outputs lie on the unit hypersphere. Data-independent; no training.
core/vector_store/pca_transform: PCATransform — principal component analysis pre-transform.
core/vector_store/random_rotation_transform: RandomRotationTransform — multiplies each vector by a fixed random orthogonal matrix R ∈ ℝ^{d×d}. Sampled once from Random(seed) via QR of a Gaussian matrix, so distances and inner products are preserved exactly (up to floating-point rounding).
core/vector_store/vector_transform: Vector pre-transforms — reversible or non-reversible mappings applied to database and query vectors before they reach an underlying Index. Mirrors FAISS' VectorTransform hierarchy.
dart_pytorch: Public entry point for the dart_pytorch package.