Serve a model, then make it cheaper

# server.py — uvicorn server:app --port 8000
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
state = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load once at startup. Loading is seconds; generating is milliseconds.
    state["tok"] = AutoTokenizer.from_pretrained(MODEL)
    state["model"] = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.float16, device_map="auto"
    ).eval()
    yield
    state.clear()

app = FastAPI(lifespan=lifespan)

class Req(BaseModel):
    prompt: str
    max_new_tokens: int = 128

# A plain def, not async def: FastAPI runs it in a threadpool, so the blocking
# forward pass does not stall the event loop and every other request with it.
@app.post("/generate")
def generate(req: Req):
    tok, model = state["tok"], state["model"]
    enc = tok(req.prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        out = model.generate(**enc, max_new_tokens=req.max_new_tokens, do_sample=False)
    new = out[0][enc["input_ids"].shape[1]:]
    return {
        "text": tok.decode(new, skip_special_tokens=True),
        "prompt_tokens": int(enc["input_ids"].shape[1]),
        "output_tokens": int(new.shape[0]),
    }

# load.py
import asyncio, math, time
import httpx

URL = "http://localhost:8000/stream"
PROMPT = "Summarise the following in two sentences: " + "the quick brown fox. " * 40

async def one(client):
    t0 = time.perf_counter()
    ttft, n = None, 0
    async with client.stream("POST", URL, json={"prompt": PROMPT,
                                                "max_new_tokens": 128}) as r:
        r.raise_for_status()
        async for _ in r.aiter_text():
            if ttft is None:
                ttft = time.perf_counter() - t0
            n += 1
    total = time.perf_counter() - t0
    # Per-token decode cost excludes the first token, which was prefill.
    tpot = (total - ttft) / max(1, n - 1)
    return {"e2e": total, "ttft": ttft, "tpot": tpot, "tokens": n}

async def run(concurrency: int, requests: int):
    sem = asyncio.Semaphore(concurrency)
    async with httpx.AsyncClient(timeout=120) as client:
        async def guarded():
            async with sem:
                return await one(client)
        t0 = time.perf_counter()
        rows = await asyncio.gather(*(guarded() for _ in range(requests)))
        wall = time.perf_counter() - t0
    return rows, wall

def pct(values, p):
    """Nearest-rank percentile: the number you print is a request that
    actually happened, rather than an interpolation between two."""
    s = sorted(values)
    return s[max(1, math.ceil(p / 100 * len(s))) - 1]

if __name__ == "__main__":
    for c in (1, 4, 16):
        rows, wall = asyncio.run(run(c, 64))
        tokens = sum(r["tokens"] for r in rows)
        print(f"concurrency={c:>2}  "
              f"e2e p50={pct([r['e2e'] for r in rows], 50):.2f}s "
              f"p95={pct([r['e2e'] for r in rows], 95):.2f}s "
              f"p99={pct([r['e2e'] for r in rows], 99):.2f}s  "
              f"ttft p95={pct([r['ttft'] for r in rows], 95)*1000:.0f}ms  "
              f"tput={tokens / wall:.1f} tok/s")

# Decode is a bandwidth problem. The floor on one step is: read the weights.
params      = 8e9        # an 8B model
bytes_each  = 2          # fp16
bandwidth   = 2.0e12     # A100 80GB: about 2 TB/s of HBM bandwidth

step_seconds = params * bytes_each / bandwidth
print(f"{step_seconds*1000:.1f} ms per decode step")   # 8.0 ms
print(f"{1/step_seconds:.0f} tokens/s ceiling, batch 1")  # 125

# The same read serves every sequence in the batch, so batching multiplies
# throughput almost for free until something else binds.
for batch in (1, 8, 32):
    print(batch, f"{batch/step_seconds:.0f} tok/s aggregate")

def kv_bytes_per_token(n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    return 2 * n_layers * n_kv_heads * head_dim * dtype_bytes   # K and V

# Llama-3.1-8B: 32 layers, 8 KV heads (grouped-query attention), head_dim 128
per_token = kv_bytes_per_token(32, 8, 128, 2)
print(per_token / 1024, "KiB per token")                    # 128.0
print(per_token * 4096 / 2**30, "GiB per 4096-token seq")   # 0.5

# On an 80 GB card: 16 GB of fp16 weights, leave ~4 GB for activations.
free = (80 - 16 - 4) * 2**30
print(int(free / (per_token * 4096)), "concurrent 4k sequences")   # 120

import asyncio
from dataclasses import dataclass, field

BATCH_WINDOW = 0.01   # seconds of latency you agree to pay to fill a batch
MAX_BATCH = 16

@dataclass
class Job:
    prompt: str
    max_new_tokens: int
    future: asyncio.Future = field(default_factory=asyncio.get_event_loop().create_future)

queue: asyncio.Queue[Job] = asyncio.Queue()

async def batcher():
    loop = asyncio.get_running_loop()
    while True:
        batch = [await queue.get()]                 # block until there is work
        deadline = loop.time() + BATCH_WINDOW
        while len(batch) < MAX_BATCH:
            remaining = deadline - loop.time()
            if remaining <= 0:
                break
            try:
                batch.append(await asyncio.wait_for(queue.get(), remaining))
            except asyncio.TimeoutError:
                break
        # to_thread keeps the blocking forward pass off the event loop, so the
        # next batch is already filling while this one runs.
        try:
            texts = await asyncio.to_thread(generate_batch, batch)
            for job, text in zip(batch, texts):
                job.future.set_result(text)
        except Exception as exc:                    # one bad batch, not a dead server
            for job in batch:
                if not job.future.done():
                    job.future.set_exception(exc)

def generate_batch(batch):
    tok, model = state["tok"], state["model"]
    enc = tok([j.prompt for j in batch], return_tensors="pt",
              padding=True, padding_side="left").to(model.device)
    with torch.inference_mode():
        out = model.generate(**enc, max_new_tokens=max(j.max_new_tokens for j in batch),
                             do_sample=False)
    return tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)

from collections import OrderedDict
from hashlib import blake2b

class LRU:
    def __init__(self, capacity=1024):
        self.capacity, self.store = capacity, OrderedDict()
        self.hits = self.misses = 0

    def get(self, key):
        if key in self.store:
            self.store.move_to_end(key)
            self.hits += 1
            return self.store[key]
        self.misses += 1
        return None

    def put(self, key, value):
        self.store[key] = value
        self.store.move_to_end(key)
        if len(self.store) > self.capacity:
            self.store.popitem(last=False)

    @property
    def hit_rate(self):
        total = self.hits + self.misses
        return self.hits / total if total else 0.0

cache = LRU()

def cache_key(req) -> str:
    # Every parameter that changes the output belongs in the key. Miss one
    # and you will serve a 128-token answer to a request that asked for 512.
    return blake2b(
        f"{MODEL}|{req.prompt}|{req.max_new_tokens}|greedy".encode(), digest_size=16
    ).hexdigest()

from transformers import BitsAndBytesConfig

four_bit = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",             # normal-float 4, not plain int4
    bnb_4bit_compute_dtype=torch.float16,  # weights stored 4-bit, maths in fp16
    bnb_4bit_use_double_quant=True,        # quantize the quantization constants too
)

model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=four_bit,
                                             device_map="auto")
print(model.get_memory_footprint() / 2**30, "GiB")

# Predicted decode step time from the footprint alone:
gib = model.get_memory_footprint() / 2**30
print(f"{gib * 2**30 / 2.0e12 * 1000:.1f} ms per step at 2 TB/s")

import math

@torch.inference_mode()
def perplexity(model, tok, text, window=1024):
    """Perplexity over held-out text. Cheap, sensitive, and not sufficient
    on its own — see the note below."""
    ids = tok(text, return_tensors="pt").input_ids.to(model.device)
    total_nll, total_tokens = 0.0, 0
    for i in range(0, ids.shape[1], window):
        chunk = ids[:, i : i + window]
        if chunk.shape[1] < 2:
            break
        # transformers shifts labels internally, so loss is mean NLL per token.
        loss = model(chunk, labels=chunk).loss.item()
        total_nll += loss * (chunk.shape[1] - 1)
        total_tokens += chunk.shape[1] - 1
    return math.exp(total_nll / total_tokens)

held_out = open("held_out.txt").read()
print("fp16  ", perplexity(fp16_model, tok, held_out))
print("4-bit ", perplexity(model, tok, held_out))

FROM python:3.11-slim

# Dependencies before code: this layer is cached across every code change,
# and it is the layer that takes four minutes.
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY server.py .

# Weights are not baked in. A 16 GB image is slow to pull and impossible to
# patch; mount a volume or download at startup behind the health check.
ENV HF_HOME=/cache
VOLUME /cache

RUN useradd --create-home app && chown -R app /cache
USER app

EXPOSE 8000
HEALTHCHECK --start-period=120s --interval=15s \
  CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]