tiny-llm: KV-cache, a public API, and shipping the whole thing
Last devlog left off with the 100M model trained and RAG wired into chat.py, generating in ~50s per reply. This one covers making that fast, and getting the whole pipeline live on the internet.
Understanding KV-cache
Without a cache, every new token recalculates K/V for all tokens already processed, even though those values haven’t changed. The KV cache keeps them around and only computes K/V for the new token.
Three real changes to model.py:
CausalSelfAttention.forward now takes an optional kv_cache, concatenates new K/V onto it along the sequence dimension, and returns the updated cache.
The causal mask only applies when there’s no cache. With a cache, the new token’s query attends over the entire K/V (past + new), and since generation only ever feeds one new token at a time once a cache exists, there’s no future position to accidentally see. “No mask” is the correct causal mask in that case, not an approximation.
Position embeddings need to be absolute, not relative to the current call. With a cache, x is only the new tokens, so position 0 here is really past_length in the full sequence, pulled from the cached K’s sequence-length dimension.
Before relying on the cache, I added a correctness check: run the whole sequence normally, then one token at a time through the cache, compare logits with torch.allclose. They matched. Wanted to confirm correctness before worrying about performance.
Wiring it into generation
chat.py and backend.py now do a prefill (whole prompt through the model once, building the initial cache) followed by decode steps that only pass the newest token plus the accumulated cache, instead of re-encoding a sliding window every step.
Deploying: Nest, Vercel, and bugs that only show up live
Backend/model: FastAPI on Hack Club Nest (Stardance’s rules don’t allow Hugging Face for hosting), running as a systemd service, bound to 127.0.0.1 behind Caddy for TLS.
Frontend: static HTML/CSS/JS on Vercel, auto-deployed via GitHub Actions.
Real bugs surfaced by going live:
addMessage() never returned the div it created, so every successful reply crashed setting .textContent on undefined, landing in the catch block and showing an error even when the backend answered correctly. Live on the public demo before it got caught.
The rate limiter keyed on a spoofable query parameter instead of the real client IP. Fixed by reading X-Forwarded-For’s last hop, which only Caddy can set here.
No protection against concurrent requests on a 2-vCPU, 1.5GB box with a documented prior OOM kill. Queuing behind a lock still pushed memory to the ceiling, so it became a 503 (fail fast) instead.
A missing message field crashed inside search(). Now returns 400 first.
The model’s sanity-check block ran on every import, not just direct execution. Gated behind if __name__ == "__main__":.
The result
/generate went from ~50s to ~4.0s, a ~12.5x speedup, same model, same box, real public HTTPS endpoint. Verified live, not assumed.
What’s running now
A GPT trained from scratch, retrieval-augmented over real security docs, served with KV-cache inference behind a rate-limited public API. Every piece went through a review pass that caught real, live bugs before calling it done. The model’s output quality at 100M params is the physical ceiling of this scale, not a bug to fix, that was always the outcome of training something this size from nothing. Getting the whole stack this far, correctly, is the milestone.