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:

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

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.