From PyTorch to WebAssembly: Lessons Learned Building and Deploying a Small Language Model at Akamai Function

Where it started: an idea

It began with an interesting idea. I wanted to build a small language model from scratch not fine-tune something, not call an API, but actually construct the model, train it on real data, and understand every moving part. And then I wanted to do something a little unusual with it: infer it using Akamai Functions (Fermyon) running the model as a WebAssembly component at the edge, instead of on a traditional GPU inference server.

That combination is what made this project fun. A tiny model I fully understood, running in an environment almost nobody runs language models in.

This post is the story of that journey: how I started, what I built, the walls I hit, and how I got past them.

The model

The model is deliberately small. It’s a GPT-style transformer, the same architecture family behind modern LLMs, just shrunk down. It’s about 30 million parameters in total, of which the transformer body (the attention and feed-forward layers that do the actual “thinking”) is only around 11 million; most of the rest is the token embedding table. It learns to generate short, coherent children’s-story-style text, trained on the TinyStories dataset (a synthetic corpus of simple stories written with the vocabulary of a young child).

Small is the whole point. A model this size:

• trains in a reasonable time on a single GPU,
• is cheap to run,
• and is small enough to even consider running inside a WebAssembly sandbox.

Training it: a Linode GPU instance

To train the model I used a Linode GPU instance. This kept the whole project grounded in a practical, low-cost setup rather than an expensive cloud ML platform. I prepared the dataset (tokenizing all the stories into numeric token IDs stored efficiently on disk), then ran the training loop with the usual ingredients: mixed-precision compute, gradient accumulation, a learning-rate warmup followed by cosine decay, and checkpointing so the best-performing version of the model was always saved.

After training, I had a single artifact: a checkpoint file containing the learned weights. That checkpoint is the “brain” of the model. Everything after this point was about serving it.

The twist: serving on Akamai Functions (Fermyon)

Here’s where the interesting idea turned into an interesting problem.

Akamai Functions, built on Fermyon Spin, runs code as WebAssembly components. WebAssembly is a portable, sandboxed compilation target fast, secure, and tiny to cold-start. It’s a fantastic place to run edge logic.

But it has one big catch for machine learning: you cannot run PyTorch inside WebAssembly. PyTorch relies on large native libraries that simply don’t exist in a WASM sandbox. So the checkpoint I’d just trained a PyTorch artifact could not be loaded and run directly in the environment I wanted to deploy it to.

This was the first real fork in the road. To serve the model on Fermyon, I needed to move inference out of PyTorch entirely.

Problem 1: Getting the model into WebAssembly at all

The solution was to convert the trained model into ONNX, an open, portable format for describing neural networks. ONNX lets you take a model trained in one framework and run it in a completely different runtime.

On the serving side, I used tract, a pure-Rust ONNX inference engine that compiles cleanly to WebAssembly. Rust plus WASM plus tract meant I could run the model’s math with no Python and no native dependencies. The tokenizer (the piece that turns text into tokens and back) also had to be reimplemented in a WASM-friendly, pure-Rust form.

So the serving pipeline became:

trained checkpoint → convert to ONNX → run with tract inside a Rust WebAssembly component on Fermyon Spin → expose it as a simple HTTP endpoint that takes a prompt and returns generated text.

Getting this working end-to-end was a milestone: a language model I trained from scratch, generating text, running inside a WebAssembly sandbox served by Akamai Functions. No GPU. No PyTorch. Just a small model doing its job at the edge.

Problem 2: It worked… but it was painfully slow

The first working version had a serious flaw. When I asked it to generate a longer piece of text, the request would appear to hang. It wasn’t crashing; it was just crawling.

Measuring it revealed the truth: each generated word (token) was taking around 9 seconds. A hundred-word completion would take roughly fifteen minutes. Completely unusable.

Two things were working against me:

1. No hardware acceleration. The WASM build wasn’t using vectorized CPU instructions, so all the heavy matrix math ran the slow way.
2. Massive redundant computation. More on this below; it was the real killer.

Fix 1: Turning on SIMD (vectorized math)

The first, cheapest win was enabling SIMD “Single Instruction, Multiple Data” in the WebAssembly build. SIMD lets the CPU process many numbers in a single operation, which is exactly what neural-network matrix multiplication needs.

Flipping this on required almost no code change, and the effect was dramatic: per-token time dropped from ~9 seconds to roughly 1 second, about an 8× speedup. Suddenly the model felt like it was actually running rather than struggling.

But 1 second per token still meant a hundred-word answer took a minute and a half. Good, not great. The bigger problem was still lurking.

Fix 2: The KV cache – the real breakthrough

This is where the biggest gain came from, and it’s worth explaining the idea because it’s the single most important optimization in the whole project.

The core inefficiency: language models generate text one token at a time, and each new token depends on all the tokens before it. In the naive approach, to produce token number 50, the model reprocesses all 49 previous tokens from scratch, then adds the new one. To produce token 51, it reprocesses all 50 again. And so on. The work grows and grows, and almost all of it is repeated; you’re recomputing the exact same things over and over.

The idea, a KV cache: inside the attention mechanism (the heart of a transformer), each past token produces two pieces of information, conventionally called the key and the value. The crucial insight is that a token’s key and value never change once computed. So instead of throwing them away and recomputing them every step, you cache them. Each new token only needs to compute its own key and value, then attend over the cached ones.

In plain terms: instead of re-reading the entire story every time it wants to write the next word, the model remembers what it already read and only processes the newest word.

Implementing this meant redesigning how the model is exported and served so that, at each step, it accepts the remembered state, does the small amount of new work, and hands back an updated memory for the next step.

The payoff was enormous. Per-token time fell from ~1 second to around 0.065 seconds, roughly a further 15× on top of the SIMD gain. A hundred-token completion that once took fifteen minutes now finished in about six seconds.

A subtle but important lesson along the way

There was one non-obvious trap. When I first switched to the KV-cache design, it actually got slower, not faster. The reason: the more flexible, “remembering” model had variable-sized inputs, and the inference engine couldn’t pre-plan its computations as aggressively.

The fix was to explicitly ask the runtime to fully optimize the model graph even in this flexible form. That one change was another order-of-magnitude difference. The lesson stuck with me: an optimization is only real once you’ve measured it; intuition about performance is frequently wrong, and the profiler is the only honest judge.

Problem 3: It ran fast, but it wouldn’t deploy

By now the model was quick. But when I actually went to push it up to Akamai Functions, I hit a wall that had nothing to do with speed: size.

Akamai Functions caps how large a deployed application can be at 50 MB. My model, stored in full 32-bit precision, weighed in at around 120 MB. It simply would not fit through the door. The thing ran beautifully on my own machine and then flatly refused to deploy. Frustrating in a very specific way: the hard engineering was done, and I was blocked by a packaging limit.

Fix 3: Quantization – a smaller model, and a new engine to run it

The fix was quantization. The idea is simple: a model’s weights are just numbers, and by default each one is stored in high precision as a 32-bit floating-point value, four bytes each. But that much precision usually isn’t needed to get good output. Quantization stores each weight in a smaller, coarser form here: 8-bit integers, one byte each for a roughly 4× reduction in size at a very small accuracy cost.

I used a scheme called Q8_0 (8-bit quantization with a shared scale per small block of weights). It took the model from ~120 MB down to about 32 MB comfortably under the 50 MB limit, small enough to bundle the model right into the deployed component with room to spare.

There was one catch, and it forced a second decision. ONNX and tract couldn’t run the quantized model; tract wouldn’t load 8-bit weights on this flexible, KV-cached graph. Quantization and my existing inference engine were mutually exclusive.

So I swapped the engine. I moved from tract to candle, Hugging Face’s pure-Rust tensor library, which, like tract, compiles cleanly to WebAssembly, but unlike tract supports Q8 quantization natively in Rust. I rebuilt the serving component around candle: the same model, the same KV cache, the same pure-Rust tokenizer, now reading the compact 8-bit weights directly. To be sure I hadn’t changed the model’s behaviour, I checked the quantized version against the original PyTorch model token by token; the tiny rounding from int8 never once changed the word it chose.

Then two good things happened at once. First, the deployment went straight through no size wall, no chunking tricks, just a clean push to Akamai Functions. Second, it was faster. Eight-bit weights are a quarter the size of the originals, so they move through the CPU far more quickly, and the numbers showed it immediately: in about one second, the model delivered a full 50-token completion. That works out to roughly 20 milliseconds per token, about 50 tokens every second.

Put the whole arc next to each other and it’s almost hard to believe it’s the same model on the same kind of hardware:

~9 seconds per token → ~1 second (SIMD) → ~0.065 seconds (KV cache) → ~0.02 seconds (Q8 + candle).

The very first working version needed about fifteen minutes to write a short piece of text. The deployed version writes the same thing in a second or two.

Where it landed

Stacking these techniques together a genuinely small model, converting it to a portable format, running it in pure Rust WebAssembly, enabling SIMD, implementing a KV cache, and finally quantizing to 8-bit and serving it with candle took inference from ~15 minutes for a short piece of text down to about a second: 50 tokens delivered in roughly one second, all inside a WebAssembly sandbox deployed live on Akamai Functions, with no GPU at serving time.

That’s the headline: a language model trained from scratch on a modest Linode GPU instance, quantized down to 32 MB, served at the edge as a lightweight WASM component, and answering in about a second.

How I could push tokens/sec even further

I stopped at “fast enough,” but there’s clear headroom left. A few directions, roughly in order of effort-to-reward:

• Going further with quantization (int4): I already moved to 8-bit (Q8), which shrank the model ~4× and sped things up. Dropping to 4-bit would roughly halve the model again (to ~16 MB) and cut cold-starts further, for a larger accuracy trade-off worth trying if size or memory ever gets tight.
• Multi-threading: the current inference runs on a single core; spreading the matrix math across cores can give a near-linear speedup where the runtime supports it.
• Running natively instead of in WASM: the same model run natively on a small VM (no sandbox overhead, easy multi-threading) is usually several times faster; the trade-off is giving up edge deployment for raw speed.
• A leaner model: distilling into fewer layers or a smaller vocabulary cuts work per token directly, at the cost of retraining.

The bigger takeaway

The most valuable realization from this project isn’t any single trick. It’s the overall shape of what becomes possible.

You don’t always need a giant, general-purpose model running on expensive accelerators. For many real-world tasks, you can build a domain-specific small language model, one trained narrowly on the kind of text you actually care about. A focused small model can be:

• Stronger on its specific domain than a huge generalist model,
• Far cheaper to run,
• and small enough that you don’t need heavyweight infrastructure to serve it.

You could serve such a model on low-cost virtual machines; even a small Linode instance would comfortably handle it or, as this project showed, right at the edge on Akamai Functions. No GPU fleet, no massive bills, just a tidy, purpose-built model doing exactly what it was made to do.

That’s the future I find genuinely exciting: not only ever-bigger models, but small, sharp, affordable ones that anyone can train and deploy starting from nothing more than an idea and a single GPU instance.