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.

tiny-llm: the 100M model is trained, and RAG is wired in

The 100M parameter run is done.
976,562 steps on a 2B-token corpus, and chat.py now pulls retrieved passages before generating instead of relying on memorized weights.

The problem, and the fix

Started out wanting this model trained purely on security content (HackTricks, PayloadsAllTheThings, SecLists, GTFOBins, CheatSheetSeries, wstg).
All of that combined comes to about 6.73M tokens. Compute-optimal for a 100M model is around 2B tokens, real security wikis just don’t reach that scale, and no amount of repo-hunting was going to fix that 300x gap.

The solution i found was blend the 6.73M security tokens into a 2B-token corpus, topped up with FineWeb-Edu for general English coverage. The model gets real security-domain exposure plus enough general text to actually hold together grammatically, not by inventing security content that doesn’t exist, but by being upfront about what’s actually available.

The training run

d_model=768, n_heads=12, n_layer=12 4x the previous 24M config. Config changes were straightforward, getting a resumable 50-hour run right was the actual work:

• Checkpoint pruning — early on, checkpoints/ silently grew to 561GB before a torch.save call failed with a disk-full error mid-write. Fixed by keeping only the 5 most recent checkpoints on every save, removing older ones automatically.
• torch.compile ordering — resuming from a checkpoint failed with a state_dict key mismatch the first time, because model = torch.compile(model) was running before the resume logic loaded the weights. Compiled models get every parameter name prefixed with _orig_mod., so loading a plain checkpoint into a compiled model (or vice versa) breaks. Fixed by moving the resume block to before the torch.compile call — load into the “bare” model first, compile after.
• Verified resume actually works — killed the process mid-run, restarted, confirmed it picked up exactly where it left off (same step, same loss) before trusting it for the full 50 hours unattended.

Final loss landed around 3.5–3.8 (train and val tracking closely) — expected for a corpus this heterogeneous (security docs + general web text) compared to the tighter 1.4–1.8 range on TinyStories.

Wiring in retrieval

A 100M model isn’t going to reliably “know” security facts from 6.73M tokens seen once inside a 2B-token corpus — so instead of hoping the weights memorized it, chat.py now retrieves first:

1. User’s prompt goes to a local SQLite FTS5 search index built over the same repos/ content used in training
2. Top-k matching passages come back with source, heading, and BM25 relevance score
3. Those passages get prepended to the prompt as context before the model generates

Straightforward in concept, but ran into block_size=256 limits fast, retrieved passages plus the question can easily blow past what the model was trained to handle positionally. Fixed by truncating the encoded prompt to the last 256 tokens before generation, keeping the question and “Answer:” cue intact even if some of the retrieved context gets cut from the front.

What actually happens now

Asked “What is SQL injection?” — with retrieval wired in, the model opens with something closer to on-topic (“one of the most powerful tools available to attackers”) instead of pure noise, but drifts into plausible-sounding but fabricated details a few sentences in. That’s the honest result at this scale: retrieval gives it real information to start from, but staying faithful to that context for 100 generated tokens is past what 100M parameters can reliably do. Better than without retrieval, not a fix for the underlying scale limit.

tiny-llm: 24M model trained, and it actually generates coherent stories

Picked up from the last devlog — train.py was working (fp16, grad clipping, val loss, checkpointing), but only tested on 2000 steps. This one covers finishing the training setup, the real 4-hour run, and getting the model talking through a CLI.

Cosine LR with warmup

Added a proper learning rate schedule before committing to a long run:

• Warmup — LR ramps up linearly from 0 to the target (3e-4) over the first 2000 steps, instead of hitting the model with full LR while the weights are still random.
• Cosine decay — after warmup, LR follows a cosine curve down toward 0 over the rest of training, so the model settles into finer adjustments by the end instead of taking large steps the whole way through.

Nothing exotic here — PyTorch doesn’t need a library for this, just a function that computes the right LR for the current step and gets applied to the optimizer’s param groups each iteration.

The real run

474M tokens (the full TinyStories corpus), which comes out to ~231,445 steps at batch_size=8, block_size=256. Ran for about 4 hours on the RTX 2080 Super:
step 0: train loss 11.01 val loss 10.62
step 231000: train loss 1.48 val loss 1.49
step 231400: train loss 1.82 val loss 1.69

Train and val loss tracked closely the entire run — no overfitting. Loss in the 1.4–1.8 range by the end, a solid result for a from-scratch 24M model on this corpus.

sample.py: watching it actually generate

Wrote a script to load a checkpoint and generate text autoregressively — feed in a prompt, predict the next token, append it, repeat. Had to fix two small but important bugs along the way:

• torch.arange and torch.triu inside the model default to creating tensors on CPU even when everything else is on GPU — same device-mismatch class of bug as during training, fixed the same way (device=x.device).
• Loading the checkpoint failed at first because torch.compile prefixes every parameter name with _orig_mod. — stripped it with a quick dict comprehension before calling load_state_dict.

First real generation, prompted with “once upon a time”:

once upon a time, there was a fish named Fin. Fin was very brave and always had a big challenge to do. One day, Fin asked his friend, Sally, “Will you marry me?” Sally said, “Yes, Fin! Let’s go!” Fin and Sally had a big adventure in the ocean…

Not perfect logic (fish swimming through an ocean “full of fish to eat” is a little off), but grammatically solid, consistent character names throughout, and a real narrative arc. Exactly what you’d expect from a well-trained 24M model on TinyStories.

chat.py: a quick interactive loop

Wrapped the generation code in a while True loop so the model loads once and takes prompts repeatedly instead of restarting for every single generation. Worth noting for anyone trying this themselves: this is a base model, not an instruction-tuned chat model — it doesn’t “answer” prompts, it just continues them as text. Say “hello” and it treats that as the opening of a story, not a greeting to respond to. That’s expected behavior at this stage, not a bug — instruction-following is a separate fine-tuning step for later.

What’s next

The full pipeline — prepare.py → train.py → sample.py/chat.py — is confirmed working end-to-end, from raw dataset to a model you can actually talk to (in the “continues your text” sense). Next up: scaling to the 100M-param run on a 2B-token corpus, this time blending real security-focused text (HackTricks, PayloadsAllTheThings, and others) with general web text, plus a retrieval index so the model can pull in relevant passages instead of relying purely on memorized weights.

tiny-llm: train.py works, and the model is actually learning

model.py and prepare.py were done, but nothing had actually been trained yet. This devlog covers train.py and the first real smoke test on the 24M-parameters config.

Writing train.py

Wrote the training loop the same way as the model piece by piece, understanding each part before moving on:

• get_batch — pulls random windows from the memmapped .bin files. x is a block of tokens, y is the same block shifted by one position (next-token prediction). Understanding why y is just x shifted by 1 was one of those small moments where next-token prediction actually clicked.
• Mixed precision (fp16 + GradScaler) — my GPU (RTX 2080 Super, Turing) doesn’t have real bf16 tensor cores, so fp16 with a GradScaler is the right call. The scaler exists because fp16 has a much narrower range than fp32 — small gradients can underflow to zero and silently vanish. The scaler multiplies the loss up before backward() and unscales before step(), so nothing gets lost in the process.
• Gradient clipping — caps the gradient norm at 1.0 before the optimizer step, to stop occasional exploding gradients from wrecking a good run.
• Validation loss — evaluates on a held-out batch from val.bin every 100 steps, wrapped in model.eval() + @torch.no_grad(), so I can actually tell if the model is generalizing instead of just memorizing.
• Checkpointing — saves model + optimizer state every 1000 steps. Training a real run for tens of hours without this would be a great way to lose everything to one crash.

Debugging session

Hit two device-mismatch errors early on — torch.arange and torch.triu both default to creating tensors on CPU, even when the rest of the model is on the GPU. Fixed by explicitly passing device=x.device wherever a new tensor gets created inside the model. Small bug, but a good reminder that PyTorch doesn’t infer device placement for you.

First real run

Ran the 24M-param smoke test (n_layer=6, n_head=6, n_embd=384, block_size=256) on TinyStories, ~470M tokens, for 2000 steps on my RTX 2080 Super:
step 0: train loss 11.01 val loss 10.62
step 1000: train loss 3.34 val loss 3.47
step 1900: train loss 2.89 val loss 3.10

Loss starts right around ln(50257) ≈ 10.8, which is exactly what you’d expect from a model guessing uniformly over the vocabulary before it’s learned anything. Train and val loss track each other closely the whole way down — no sign of overfitting, which is what I wanted to confirm before committing to a much longer run.

What’s next

The full pipeline is confirmed working end-to-end: prepare.py → get_batch → model.py → train.py, all on my own hardware. Next step is adding a cosine learning rate schedule with warmup, then moving on to the actual 100M-param run.

PS: train.py had some AI usage

tiny-llm: getting the data pipeline and model architecture in place

Before writing any model code, I needed a way to actually get training data into a usable format. prepare.py handles that: it streams a dataset from HuggingFace (starting with TinyStories, ~470M tokens), tokenizes it with the GPT-2 BPE tokenizer via tiktoken, and writes it out as flat uint16 arrays split into train.bin and val.bin.

A few details mattered here even though I didn’t write this part myself:

• Tokens are stored as uint16 instead of the more obvious int32/int64 — GPT-2’s vocab is 50,257 tokens, which fits comfortably under 65,536, so uint16 cuts file size in half for free.
• The split between train/val happens per document, not per token. Each story gets an end-of-text token appended, and documents are randomly assigned to train or val as a whole — so a story never gets cut in half across the train/val boundary and leaks information.
• Files are written so they can be read later with np.memmap, meaning the corpus never has to fit in RAM — training just maps the file and reads slices from disk on demand.

With the data pipeline sorted, I moved on to the part I actually wanted to understand deeply: the model itself. That’s model.py, CausalSelfAttention, MLP, Block, and GPT, all written and shape-tested from scratch.

Having both pieces done means the two “boring but necessary” and “the actual point of the project” halves are finally connected: real tokenized text on one end, a transformer architecture ready to consume it on the other.

Next: train.py — wiring prepare.py’s output into model.py via memmap batching, cross-entropy loss, AdamW, and a training loop.