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.

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.