stylus-nanoGPT

Imported from OffchainLabs/stylus-nanoGPT (AGENTS.md). Install upstream with npx skills add OffchainLabs/stylus-nanoGPT. Copyright stays with the author.

AGENTS.md — reproduction guide & gotchas

Notes for agents/humans who need to reproduce, re-measure, or extend this project. It encodes the traps that cost real time the first time around — read the Gotchas section before touching the Solana build or any gas/CU measurement.

What this repo is

A 20,304-param, 1-layer char-level GPT (nanoGPT, trained on tiny-Shakespeare, out-20k, val loss 1.98) run on-chain three ways, all in integer fixed-point (no floats), all reading the same 24,215-byte weights blob, all producing byte-identical output:

• Stylus (Rust→WASM) — the main project, repo root (src/).
• EVM (Solidity) — benchmark-compare/evm/.
• Solana (Rust→SBF) — benchmark-compare/solana/.

The point is a cross-runtime cost benchmark (benchmark-compare/README.md). It's a tech demo, not a useful generator.

Layout

src/                       Stylus contract (lib.rs reads shared blob; gpt.rs = math; layout.rs = offsets)
weights.bin                20,304 int8 weights (packed into the blob)
export/                    export_weights.py → pack_blob.py → blob; gen_rust.py is LEGACY/reference; render_readme.py → HTML
benchmark-compare/
  evm/                     foundry: src/NanoGptEVM.sol + blob.bin + blob_initcode.hex
  solana/                  SBF program (src/) + measure/ (litesvm CU harness)
  README.md                THE comparison doc (numbers live here)

Toolchains (all already installed on this machine)

Canonical correctness check (the invariant)

All three runtimes MUST agree, byte-for-byte:

generateCached("To be", 8, 42, 80)  => "To be her bee"
generateCached("To be", 59, 42, 80) => "To be her beer,\nFor or and thiou by comme ision lest,\nWhecher, t"

• Stylus: cargo test --release (5 tests: Rust == validated numpy reference, and cache==recompute).
• EVM / Solana: produce the same strings (see measurement commands below).

If you change gpt.rs (any runtime), re-run all three and confirm the strings still match.

Reproduce the artifacts (only if the model changes)

VENV=/Users/jasonwan1/Work/code/nanoGPT/.venv/bin/python
$VENV export/export_weights.py    # ckpt.pt -> weights.bin + export/weights_meta.json
$VENV export/ref_fixedpoint.py    # numpy reference (validation, optional)
$VENV export/pack_blob.py         # -> benchmark-compare/evm/blob.bin (the shared blob, 24,215 B)
# Then hand-update offsets/scales in src/layout.rs AND benchmark-compare/evm/src/NanoGptEVM.sol.
# export/gen_rust.py prints those consts for reference (writes export/weights_ref.rs); it is NOT
# the compiled source — the contract reads the blob and uses layout.rs.

Reproduce the numbers

Stylus gas — measure on Arbitrum Sepolia (a REAL chain), not a dev node

SEP=https://sepolia-rollup.arbitrum.io/rpc          # or an Alchemy arb-sepolia URL
A=0x6b232ce37d4b790769c0f83f9e030f985b56dd0e         # deployed contract (generate + generateCached)
cast call     $A 'generateCached(string,uint8,uint64,uint8)(string)' 'To be' 8 42 80 --rpc-url $SEP
cast estimate $A 'generateCached(string,uint8,uint64,uint8)'  'To be' 8  42 80 --rpc-url $SEP   # ~765K
cast estimate $A 'generateCached(string,uint8,uint64,uint8)'  'To be' 59 42 80 --rpc-url $SEP  # ~4.14M

Deployed addresses (Arb Sepolia): blob 0x9a5d7a51b52cc87933f2df7495a069d5fecea4d6, Stylus logic 0x3204819276f1598d4822e5a26d19bd11fd4a5af1, Solidity logic 0x8a38d81dAFfe295A5895Bdb272B90fdA5011E671.

EVM gas — measure on local anvil (uncapped; anvil EVM gas == Arbitrum EVM gas)

anvil --gas-limit 20000000000 --silent &            # needs a HUGE gas limit (calls are 50M–900M)
RPC=http://localhost:8545
KEY=0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80  # anvil dev acct #0
cd benchmark-compare/evm && forge build
# deploy the blob (NOTE arg order — see Gotchas):
BLOB=$(cast send --rpc-url $RPC --private-key $KEY --create "$(cat blob_initcode.hex)" --json | python3 -c 'import sys,json;print(json.load(sys.stdin)["contractAddress"])')
# deploy the contract (NOTE --constructor-args must be LAST — see Gotchas):
NANO=$(forge create src/NanoGptEVM.sol:NanoGptEVM --rpc-url $RPC --private-key $KEY --broadcast --constructor-args $BLOB | grep -i 'Deployed to:' | awk '{print $NF}')
cast estimate $NANO 'generateCached(string,uint8,uint64,uint8)' 'To be' 8  42 80 --rpc-url $RPC  # ~162.7M
cast estimate $NANO 'generateCached(string,uint8,uint64,uint8)' 'To be' 59 42 80 --rpc-url $RPC # ~904M

Solana CU — in-process via litesvm (NO validator, NO network)

export PATH="$HOME/.local/share/solana/install/active_release/bin:$PATH"
cd benchmark-compare/solana && ./reproduce.sh         # build-sbf (fails fast) THEN measure -> CU tables
# or manually: cargo build-sbf ; cd measure && cargo run --release
# (the harness refuses a .so older than the program source, so you can't measure a stale binary)

The harness loads the .so, injects the weights account directly, and reads compute_units_consumed. It lifts litesvm's CU cap so the true CU is visible even above Solana's real 1.4M/tx limit. The "txs needed" column is ceil(CU / 1.4M) — a lower bound, not a built multi-tx workflow (the program has no cross-tx state machine).

Gotchas (read before measuring or touching Solana)

1. Dev-node Stylus gas is INFLATED — measure Stylus on a real chain

A fresh nitro --dev charges a much higher Stylus per-call init gas (~2.2M for n=1) than a real chain (Sepolia ~292K). This is the dev chain's init-gas configuration, not a missing cache: Sepolia's program is codehashIsCached = false (uncached) yet still ~292K. So:

• Do NOT use dev-node numbers for Stylus gas. Use Arbitrum Sepolia (or One).
• The dev node is fine for EVM gas (EVM gas is identical everywhere) — but EVM gas is also identical on anvil, so just use anvil.
• A prior write-up wrongly framed this as a CacheManager/"cold init" issue; it isn't. (ArbWasmCache codehashIsCached keys on the module hash, and the program is uncached on Sepolia regardless.)

2. Solana SBF dependency-edition hell — DON'T bump these pins

The 1.16.24 SBF rustc is ~1.68; modern crates.io needs ≥1.85 (edition2024) or ≥1.71 (hash_one). benchmark-compare/solana/Cargo.lock is hand-pinned and lock version must stay 3 (old cargo rejects v4): blake3 = 1.3.3, ahash = 0.8.6 (0.8.7+ uses BuildHasher::hash_one, unstable on the SBF rustc), jobserver = 0.1.32. If you regenerate the lock, re-apply these or the build breaks.

3. Solana heap is a hardcoded 32 KB — needs a custom allocator

solana-program 1.16's default allocator has HEAP_LENGTH = 32 KB (a compile-time constant); request_heap_frame is ignored by it, and 32 KB can't even hold the 36 KB KV-cache. Fix (already in the code): enable the custom-heap cargo feature of this crate (NOT of solana-program — that errors "does not have that feature") so entrypoint! skips its allocator, then install a 256 KB BumpAllocator. The allocator never frees, so gpt.rs also does a per-token heap-frame reset (snapshot/restore the bump pointer) to reach the full 64-token context.

4. solana-test-validator won't start on this macOS

Darwin 25.x + the 2023 solana 1.16.24 build → dies right after genesis with a generic "blockstore error" (native-lib incompat). Use the litesvm harness instead (in-process, no RocksDB/network). litesvm enforces the real 256 KB heap_size cap even in simulation (you can't lift it); you CAN lift the CU limit via with_compute_budget.

5. foundry CLI arg-order traps

• cast send --create is a subcommand: put --rpc-url/--private-key before --create.
• forge create --constructor-args is variadic: put it LAST, or it eats --private-key /--rpc-url and you get "Error accessing local wallet".
• anvil needs --gas-limit 20000000000 so eth_estimateGas can binary-search the 50M–900M calls.
• macOS has no timeout command.

6. nitro dev account nonce ≠ 0 at start

nitro --dev starts the funded account at nonce ~26, so a CREATE won't land on a hardcoded address. The Stylus contract hardcodes BLOB = 0x9a5d…; on a dev node the blob lands elsewhere → you'd have to update the const + rebuild + redeploy. On Sepolia it's already at 0x9a5d…. (Another reason to measure Stylus on Sepolia, not a dev node.) A fresh dev chain also has no CacheManager (ArbWasmCache.allCacheManagers() == []).

7. Naming / limits

• Stylus maps snake_case Rust → camelCase ABI: call generateCached, not generate_cached.
• The 30 KB compressed Stylus program needs ArbOS ≥ 60 (96 KB limit). Sepolia/One and recent nitro builds (this one is ArbOS 115) are fine.
• The EVM version needs >32M gas even for 1 token, so it can't be sent as a tx on Arbitrum — only eth_call/eth_estimateGas.
• Model context is 64 tokens (learned absolute wpe); with the 5-char prompt "To be" that's n ≤ 59. Asking for more silently stops at 59.

8. Numbers drift slightly on rebuild

Rebuilding the .so/contract can shift CU/gas by <0.001% (codegen relayout). If you rebuild, re-run the measurement and update the README tables so a re-runner matches byte-for-byte.

Where the README numbers come from

benchmark-compare/README.md tables: Stylus = Sepolia cast estimate; EVM = anvil cast estimate; Solana = the litesvm harness. Headline: ~200× cheaper than EVM (188–223×, recompute & KV-cache, all n); 1 Stylus tx vs ≥26–30 Solana tx (Solana 1.4M-CU/256 KB-heap caps). After any code change, re-measure and re-render: python export/render_readme.py (writes the 3 README.html).