connectome-kernels

Fused CUDA kernels for training a recurrent network whose wiring is a fixed sparse directed graph, such as a connectome. One scalar state per node, one trainable scalar per edge, any graph you can express as (src, dst) index lists. Built for FlyGPT, which trains the fruit-fly connectome (QuixiAI/MaleCNS) on Shakespeare, but nothing in it is fly-specific.

proposal_i = tanh( Σ_j W_ij h_j + drive_i + bias_i )
h_i       ← (1 − leak_i) h_i + leak_i · proposal_i          repeated `microsteps` times per input step

W is sparse with a fixed pattern; drive enters at a designated set of input nodes. The whole window of T input steps × microsteps runs as one autograd Function: one fused kernel per microstep forward, three backward (elementwise grads with row reductions; the transposed sparse product for the state gradient; a per-edge dot product for the edge gradient). Everything is fp32.

Why

The natural PyTorch implementation (torch.sparse.mm inside a Python time loop) launches ~14k kernels per training step, spends its backward on sparse-sparse additions, and is CPU-launch-bound. On the FlyGPT launch shape, one step of training:

N nodes	E edges	batch	T × microsteps	GPU	torch.sparse COO	connectome-kernels
5,000	524,324	32	64 × 2	B200, torch 2.14+cu132	207 ms, 2.43 GB	15.5 ms, 0.55 GB

Same loss to 6 decimals, edge gradients equal to 3e-11 (bench/bench.py).

Install

Needs a CUDA GPU, PyTorch ≥ 2.4 with CUDA, and an nvcc whose major version matches torch's CUDA build (torch.version.cuda). Pre-built for sm_80 (A100), sm_90 (H100), sm_100 (B200); set CK_CUDA_ARCHS to change.

pip install --no-build-isolation git+https://github.com/QuixiAI/connectome-kernels
# or, from a checkout:
CK_CUDA_ARCHS="100" pip install --no-build-isolation -e .

If the compiled module is absent, the package JIT-builds itself on first use (also needs nvcc and ninja).

Use

import torch
from connectome_kernels import SparseGraph, sparse_recurrence

# any directed graph: src[e] -> dst[e], int64, on the GPU
graph = SparseGraph(src, dst, num_nodes=N, input_nodes=input_nodes)

edge_values = torch.nn.Parameter(torch.randn(E, device="cuda"))     # one trainable weight per edge, YOUR edge order
raw_leak    = torch.nn.Parameter(torch.zeros(N, device="cuda"))
bias        = torch.nn.Parameter(torch.zeros(N, device="cuda"))

# drives: [T, n_in, B] external input to the input nodes at each step; state0: [B, N]
out = sparse_recurrence(edge_values * scale, torch.sigmoid(raw_leak), bias, drives, state0, graph, microsteps=2)
# out: [T, B, N] node states after each input step; fully differentiable wrt every argument

Degree normalization, sign constraints, or initialization from synapse counts are the caller's business: pass whatever per-edge values you like (edge_values * scale above). Edge parameters stay in your edge order; the internal CSR permutation is invisible.

connectome_kernels.dense_reference is a slow dense implementation of the same recurrence; the tests check outputs and all five gradients (edge values, leak, bias, drives, initial state) against it for batch sizes that exercise partial warps and multi-chunk rows, and for 1 to 3 microsteps.

Design notes

• States are stored [N, B] so a node's batch row is one contiguous 128-byte line; one warp handles 32 batch columns of one row, and rows with many edges are split across 16 warps with a shared-memory reduction.
• The backward never forms a sparse tensor: the state gradient uses a precomputed transposed CSR, the edge gradient is a warp-per-edge dot product over the batch, and per-node reductions (bias, leak) ride along in the elementwise kernel.
• Saved for backward: the state and pre-activation after every microstep, 2 × T × M × N × B floats.

Citation

@misc{hartford2026connectomekernels,
  title        = {connectome-kernels: fused CUDA kernels for training recurrent networks on fixed sparse graphs},
  author       = {Hartford, Eric},
  year         = {2026},
  publisher    = {GitHub},
  howpublished = {\url{https://github.com/QuixiAI/connectome-kernels}}
}

MIT license.