[555aaditya]

Optimized GPT-style Language Model Engine — Initial Setup → main

Summary

This PR promotes the Initial-Setup branch into main. It delivers a complete,
from-scratch decoder-only transformer language model engine built in PyTorch —
no pretrained weights, no HuggingFace. Every core component (BPE tokenizer, data
engine, GQA + RoPE attention with three interchangeable kernels, training loop,
inference/serving, optimization, and benchmarking) is hand-written and covered by a
test-driven suite.

Scale of change: 72 files, ~11,780 insertions across 41 commits.

---

Highlights

Hand-written stack (the only thing borrowed from PyTorch is tensor math)

• tokenizer/ — byte-level BPE (GPT-2 style rendered-unicode keys) with
  training, greedy encode/decode, special tokens, serialization, and
  pre-tokenization that speeds up BPE encode from ~2.9 KB/s to ~67 MB/s.
• dataset/ — memory-mapped (uint16 mmap) streaming corpora, windowing,
  sharding, content-hashed train/val splits, corpus download, and email/chat
  message document preparation with cleaning & redaction.
• attention/ — grouped-query attention (8 query heads → 2 KV heads, 4x
  smaller KV cache), RoPE, KV cache, and three interchangeable kernels:
  • manual — readable reference / double-precision oracle
  • sdpa — torch fused F.scaled_dot_product_attention
  • flash — a hand-written tiled online-softmax kernel with a recomputing
    backward (O(T) memory in training)
  • All three are asserted to produce identical outputs → correctness oracle.
• training/ — hand-written AdamW, cosine LR + warmup, gradient
  clipping, AMP (bf16 on MPS), gradient accumulation, checkpointing, best-model
  tracking, early stopping, and held-out validation loss / perplexity.
• inference/ — temperature / top-k / top-p sampling, KV cache paged
  generation, and batched generation with left padding that correctly handles
  ragged rows (masked keys match single-stream decoding).
• optimization/ — INT8/INT4 quantization and a KV-cache memory pool.
• benchmark/ — latency (p50/p90/p95/p99), throughput, memory, and evidence
  harness, plus a showcase CLI printing architecture, memory, latency, and live
  completions.

System & Engineering

• Config-driven throughout (configs/default.yaml, configs/tinyshakespeare.yaml).
• Device resolution mps → cuda → cpu for Apple Silicon.
• Free-tier CI: ruff lint + format, mypy typecheck, pytest (3.x matrix +
  coverage), and import-smoke jobs.
• Test-driven: 359 tests passing across unit + integration suites.
• Architecture docs & ADRs in docs/ (TDD, build order) and the global vault.

---

Architecture

raw text → tokenizer → dataset → training → checkpoint → inference
                │          │          │                │
                └──────────┴── model ─┴── optimization ─┘
                                  │
                              benchmark

Module	Responsibility
tokenizer/	Byte-level BPE (train - encode - decode - serialize)
dataset/	mmap + streaming corpora, windowing, sharding, message prep
model/	Transformer blocks (RoPE, RMSNorm, SwiGLU, GQA) -> CausalLM
attention/	manual / sdpa / flash kernels + KV cache
training/	hand-written AdamW, scheduler, AMP, trainer, checkpointing
inference/	sampling, batched generation, KV cache decoding
optimization/	INT8/INT4 quantization, KV memory pool
benchmark/	latency / throughput / memory harness + showcase

---

Why this architecture (engineering rationale)

• Byte-level BPE → zero out-of-vocabulary tokens for arbitrary UTF-8/emoji text.
• Memory-mapped data engine → window count derived from stat() alone; no bulk
  copy into RAM even with multiprocessing workers.
• GQA + RoPE → modern, memory-efficient attention that extrapolates to longer
  contexts and shrinks the KV cache (2048 KiB → 512 KiB per layer at 512 tokens)
  and attention parameters by 37.5%.
• Tiled Flash-style attention with recomputing backward → keeps training memory
  at O(T) instead of O(T²), removing the long-context scalability wall.
• Hand-written AdamW + KV cache + batching → production-style training and
  serving with measurable, gated performance (per TDR-012 benchmarking baselines).

---

Test status

• 359 tests passing (unit + integration), covering tokenizer, dataset, model,
  attention, inference, optimizer, trainer, quantization, KV pool, and benchmarks.
• Integration suite validates tokenizer ↔ dataset ↔ attention ↔ model ↔ trainer
  end to end.
• CI enforced: ruff lint + format, mypy typecheck, pytest on 3.x matrix, and
  import smoke — all must pass before merge.

---

Docs & vault

• docs/TDD.md — full technical design document (v1.1+)
• docs/TDD_PROMPT.md — original TDD prompt
• docs/BUILD_ORDER.md — dependency-aware build plan & module contracts
• Architecture decisions recorded as ADRs/TDRs in the global vault (~/vault/).

---

Notes for reviewers

• main is locked; this PR is the single entry point for the whole feature set.
• After merge, the existing .gitignore keeps caches, checkpoints, datasets,
  secrets, and scratch files out of version control.
• The only runtime file intentionally excluded is the raw arxiv.txt corpus dump
  (regenerable via the dataset pipeline).