ngxson/fly-llm-onnx

fly-onnx: the fruit fly connectome language model, one step as an ONNX graph

This is a toy model built for fun. It is NOT supported by llama.cpp. This repo is the ONNX export of ngxson/fly-hf, made for running in the browser with onnxruntime-web.

See the fly-hf model card for what the model is: a GPT-style LM whose transformer blocks were replaced by the MaleCNS fruit fly connectome (49,393 central brain neurons, 9M synapses) used as a fixed reservoir, overfitted on TinyStories.

Files

Two variants of the same model:

• v1, self-contained fly_step_q8.onnx (150 MB): the whole step in one graph, including the 9M-edge connectome matvec as Gather + ScatterElements. Simplest for Python.

• v2, for the browser fly_step_v2_dq8.onnx (60 MB; fly_step_v2_q8.onnx is the same graph with DequantizeLinear + MatMul instead of DynamicQuantizeLinear + MatMulInteger, 30x slower on wasm) + edges_offsets.i32, edges_src.u16, edges_val.i16, edges.json (36 MB): the graph takes the recurrent input rec = W @ state as an extra input and you compute the sparse matvec yourself (a CSR loop, about 20 ms in JS). The demo Space ngxson/fly-llm-demo uses this one. Edge layout is described in edges.json: incoming CSR, for target j the edges e in [offsets[j], offsets[j+1]) come from src[e] with weight val[e] * val_scale.

• fly_step_q8.onnx (150 MB, v1): one token step. The readout (49,393 x 1024) is int8 per-column quantized with an in-graph DequantizeLinear, synapse counts are stored as int16 (exact), edge indices are int32. Greedy output is identical to the fp32 PyTorch model on the tested prompts, max logit deviation 0.07.

• tokenizer.json, tokenizer_config.json: byte-level BPE, vocab 1024, <pad>=0, <s>=1, </s>=2.

• fly_step.json: the I/O spec below in machine-readable form.

Graph I/O

The graph computes a single recurrent step. You keep the state and the token window yourself and call it once per token.

name	shape	dtype	meaning
window_ids (input)	[batch, 8]	int64	the last 8 token ids, oldest first, last column is the current token. Pad with 0 at the start.
state (input)	[batch, 49393]	float32	neuron activations, all zeros at the start of a sequence
rec (input, v2 only)	[batch, 49393]	float32	W @ state, the connectome matvec computed outside the graph
logits (output)	[batch, 1024]	float32	next-token logits
new_state (output)	[batch, 49393]	float32	feed back as state for the next token

Ops used: Gather, Mul, ScatterElements (reduction=add), Cast, DequantizeLinear, MatMul, LayerNormalization, Tanh, Add, Slice, Concat. Opset 18.

Example (Python)

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

repo = "ngxson/fly-onnx"
sess = ort.InferenceSession(hf_hub_download(repo, "fly_step_q8.onnx"))
tok = Tokenizer.from_file(hf_hub_download(repo, "tokenizer.json"))

ids = [1] + tok.encode("Once upon a time, there was a").ids  # 1 = <s>
state = np.zeros((1, 49393), np.float32)
window = np.zeros((1, 8), np.int64)
out = list(ids)
for i in range(len(ids) + 40):
    cur = out[i] if i < len(out) else int(logits.argmax())
    if i >= len(out): out.append(cur)
    window = np.concatenate([window[:, 1:], [[cur]]], axis=1)
    logits, state = sess.run(None, {"window_ids": window, "state": state})
print(tok.decode(out[1:]))

Example (browser, onnxruntime-web)

import * as ort from "onnxruntime-web/webgpu";
const session = await ort.InferenceSession.create(url, { executionProviders: ["webgpu", "wasm"] });
let state = new ort.Tensor("float32", new Float32Array(49393), [1, 49393]);
const window = new BigInt64Array(8); // shift left and put the current token id last
const { logits, new_state } = await session.run({ window_ids: new ort.Tensor("int64", window, [1, 8]), state });
state = new_state;

Licence: weights CC BY 4.0 (the connectome is CC BY 4.0, FlyEM / HHMI Janelia et al.). Toy model, no warranty of any kind.