inferno

GPT-2 (124M) text generation in pure Rust, CPU-only, zero ML dependencies. Hand-rolled tensor storage, reference kernels, byte-level BPE tokenizer, KV-cached autoregressive runtime — no tch, no candle, no ONNX.

just run "The quick brown fox" --max-new-tokens 50
# The quick brown foxes are a great way to get a little bit of a kick out of your dog.

Layout

src/
  tensor/      f32 storage: shapes, strides, reshape/permute/narrow/concat
  kernels/     CPU reference ops: linear (blocked GEMM), gelu_new,
               softmax, layernorm, embedding, causal + cached attention
  loader/      config.json → Gpt2Config · model.safetensors → Gpt2Weights ·
               tokenizer.json → Gpt2Tokenizer (byte-level BPE)
  model/       KvCache · batch-1 forward pass · greedy/temperature/top-k sampler
  cli.rs       CLI args · server.rs  HTTP layer (planned)
benches/
  gpt2_kernels.rs   per-op Criterion benches at GPT-2 shapes
  generation.rs     end-to-end tok/s on the real weights
gpt2/          config.json · model.safetensors · tokenizer.json (only these three)

Requirements

• Rust 1.85+ (edition 2024), just (cargo install just)
• The three weight files above. To re-fetch from the Hub:

  hf download openai-community/gpt2 \
    --include 'config.json' --include 'model.safetensors' --include 'tokenizer.json' \
    --local-dir gpt2

Usage

just build-release            # optimized binary
just test                     # fast unit tests (52)
just test-heavy               # + real-weight tests (ignored by default)

just run "Hello, world!" --max-new-tokens 30
just sample "Hello, world!"   # temperature 0.8 + top-k 40 preset

cargo run --release -- --help # all flags:
#   [prompt] --model-dir --max-new-tokens --temperature --top-k --seed

Empty prompt = unconditional generation (starts from <|endoftext|>, GPT-2 convention). --temperature 0 (default) is greedy; sampling is seeded and reproducible.

Design notes

• Weights are [in, out]. HF's Conv1D stores c_attn as [768, 2304], so projections are plain x @ w + b — no transpose. There is no LM-head matrix; it is tied to wte (transposed once at startup).
• attn.bias is skipped. That file entry is a causal-mask buffer, not a parameter; kernels build the mask arithmetically.
• Tokenizer = tokenizer.json only. Its vocab/merges are byte-identical to the Hub's vocab.json/merges.txt (verified), but as JSON it avoids merges.txt's #-merge comment ambiguity and carries the EOS id.
• Kernels take &Tensor, return Tensor. Contiguous row-major in/out — the same boundary a future device backend would use.
• Decoding is prefix-aware. One attention kernel serves prefill (offset 0) and cached decode steps (Tq = 1 over the full cache); positions come from absolute wpe slices.

Benchmarks

Machine: AMD Ryzen 7 7435HS (8C/16T), 23 GB RAM, cargo bench (release), single-threaded. The RTX 4050 in this machine is unused — CPU only, by design.

Per-op (just bench-ops — cargo bench --bench gpt2_kernels):

Benchmark	Shape	Time
linear/c_attn_prompt32	32×768 @ 768×2304 + bias	5.63 ms (~20 GFLOP/s)
linear/c_proj_decode1	1×768 @ 768×768 + bias	71.0 µs
linear/mlp_fc_prompt32	32×768 @ 768×3072 + bias	7.80 ms
add/residual_32x768	32×768	2.57 µs
gelu_new/mlp_32x3072	32×3072	568 µs
layernorm/hidden_32x768	32×768, eps 1e-5	36.8 µs
softmax/vocab_50257	1×50257	147 µs
softmax/attn_12x64x64	12 heads × 64×64	149 µs
embedding/gather_32x768	32 ids from 50257×768	2.39 µs (38 GiB/s)
attention/head_64x64	single head, 64 ctx	165 µs

Bottlenecks are the big prompt-width GEMMs and tanh-heavy GELU — both single-threaded scalar code, which is where threading/SIMD would go first.

End-to-end (just bench-gen — cargo bench --bench generation):

Benchmark	Time	Throughput
prefill/64 (64-token prompt)	894 ms	71.5 tok/s in
decode_steps/8 (cached steps)	456 ms	17.5 tok/s steady-state
generate_greedy/16 (full pipeline)	958 ms	16.7 tok/s end-to-end

Matches the CLI: just run "The quick brown fox" --max-new-tokens 20 reports [19 new tokens in 1.2s, 15.7 tok/s].

Validation against llama.cpp

Differential-tested greedy decoding vs llama.cpp (--temp 0, F32 GGUF converted from the same model.safetensors via convert_hf_to_gguf.py):

Prompt	Ours	llama.cpp	Match
The quick brown fox (+20)	[274, 389, …, 13, 198, 198]	same	20/20
Hello, world! (+12)	[198, 198, 40, …, 1654, 644]	same	12/12
The capital of France is (+12)	[262, 3139, …, 286, 262]	same	12/12

Prompt tokenization also matches (Hello, world! → [15496, 11, 995, 0] on both). --show-ids prints our sampled id sequence for such checks.

Roadmap

• Streaming token output, --top-p
• server.rs: tokenize → generate → stream over HTTP
• Multithreaded GEMM (the 16 threads here are idle)