"""What a model server's three main settings change, on the GPU, model and step-time model of lesson 24 (how-an-ai-server-answers-many-people-at-once). One A100 80GB holds three things: the model's weights, a working reserve, and the cache of every chat in progress. The server's settings do not make memory; they decide how the same bytes are shared: memory share how much of the card the server may use at all (vLLM: gpu_memory_utilization) longest chat the most tokens one chat may reach, prompt and reply together (vLLM: max_model_len) most chats the most chats the server runs in one step (vLLM: max_num_seqs) cache pool = memory share * card - weights - reserve, cut into blocks of 16 tokens (lesson 22). `serve` runs lesson 24's continuous-batching server with that pool: a waiting chat joins only if the pool has room for it, and when the running chats need a block that is not there, the newest chat is paused, its cache dropped, and it is rebuilt later (vLLM V1's default, "recompute"). The chip and the model are lesson 24's, so the times are best cases, not measurements. The reserve and the chats are invented. Python 3 standard library only. The widget (../widget.html) repeats `serve` in JavaScript; test_widget.py runs both and compares them. """ from __future__ import annotations import importlib.util import math import random import sys from collections import deque from pathlib import Path sys.dont_write_bytecode = True # loading lesson 24's file below must not leave a __pycache__ in its folder LESSON_24 = Path(__file__).resolve().parents[2] / "how-an-ai-server-answers-many-people-at-once" / "code" / "batching_sim.py" def _load_lesson_24(): spec = importlib.util.spec_from_file_location("batching_sim_lesson_24", LESSON_24) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module # dataclasses looks the module up by name spec.loader.exec_module(module) return module L24 = _load_lesson_24() Request = L24.Request # --- the card and the model: lesson 24's ---------------------------------------------------------------------------------------------------- GPU_BYTES = int(L24.GPU_MEMORY_BYTES) # the data sheet says "80GB"; read there, and here, as 80 billion bytes GPU_BYTES_GIB = 80 * 2**30 # the other reading, as lessons 18 and 26 take it: 80 GiB. Only `memory_parts(share, GPU_BYTES_GIB)` uses it WEIGHT_BYTES = L24.FP16_BYTES * L24.PARAMS # 16-bit weights KV_BYTES = L24.KV_BYTES_PER_TOKEN # 131,072 bytes of cache for each token of a chat BLOCK_TOKENS = 16 # lesson 22 / the vLLM paper's default block size BLOCK_BYTES = BLOCK_TOKENS * KV_BYTES # 2 MiB RESERVE_BYTES = 2_000_000_000 # invented: working memory a pass needs plus the server's own overhead. A real server measures it at start-up. TYPICAL_CHAT = 2000 # tokens in a typical chat (the invented chats below average about this) # --- the settings the widget offers --------------------------------------------------------------------------------------------------------- SHARES = list(range(30, 99)) # memory share, whole percent (vLLM's documented default is 92) LONGEST = [4000, 8000, 16000, 32000, 64000, 128000] # longest chat, tokens CAPS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] # most chats at once OPENING = {"share": 90, "longest": 8000, "cap": 64, "reserved": False} KNEE_GAIN = 1.10 # raising the cap "stops paying" when doubling it brings less than 10% more tokens a second MAX_STEPS = 1_000_000 # a guard, not a limit: the longest honest run (cap 1) takes about 94,000 steps; a bug that stops chats finishing must fail, not hang def blocks_for(tokens: int) -> int: return -(-tokens // BLOCK_TOKENS) # --- the memory budget ---------------------------------------------------------------------------------------------------------------------- def memory_parts(share: int, card: int = GPU_BYTES) -> dict: """Bytes of the card in each part for a memory share given in whole percent. The pool is whole blocks; what is left over counts as unused.""" allowed = card * share // 100 blocks = max(0, allowed - WEIGHT_BYTES - RESERVE_BYTES) // BLOCK_BYTES pool = blocks * BLOCK_BYTES return {"weights": WEIGHT_BYTES, "reserve": RESERVE_BYTES, "pool": pool, "unused": card - WEIGHT_BYTES - RESERVE_BYTES - pool, "allowed": allowed, "blocks": blocks, "pool_tokens": blocks * BLOCK_TOKENS} def room(share: int, longest: int, reserved: bool = False) -> dict: """How many chats the pool holds. `starts` is False when it cannot hold even one chat of the longest length (vLLM then refuses to start). `typical` chats of TYPICAL_CHAT tokens fit when blocks are handed out as a chat grows; a server that reserves the full length for each chat fits `at_longest`.""" parts = memory_parts(share) per_longest = blocks_for(longest) per_typical = blocks_for(min(TYPICAL_CHAT, longest)) at_longest = parts["blocks"] // per_longest return {**parts, "starts": at_longest >= 1, "at_longest": at_longest, "typical": at_longest if reserved else parts["blocks"] // per_typical, "per_typical": per_typical, "per_longest": per_longest} # --- the invented chats --------------------------------------------------------------------------------------------------------------------- SEED = 2026 N_CHATS = 600 PROMPT_RANGE = (300, 6000) # tokens already in the chat when the next reply starts: equal steps on a log scale REPLY_RANGE = (40, 400) def draw_chats(seed: int = SEED, n: int = N_CHATS) -> dict: rng = random.Random(seed) prompts, replies = [], [] for _ in range(n): prompts.append(round(PROMPT_RANGE[0] * (PROMPT_RANGE[1] / PROMPT_RANGE[0]) ** rng.random())) replies.append(round(REPLY_RANGE[0] * (REPLY_RANGE[1] / REPLY_RANGE[0]) ** rng.random())) return {"prompts": prompts, "replies": replies} CHATS = draw_chats() def backlog(chats: dict = CHATS) -> list[Request]: """Every chat is waiting at time 0: a full queue, so the server is never short of work.""" return [Request(0.0, p, r) for p, r in zip(chats["prompts"], chats["replies"])] # --- the server ----------------------------------------------------------------------------------------------------------------------------- def serve(requests: list[Request], cap: int, pool_blocks: int, longest: int, reserved: bool = False, cost=L24.step_seconds) -> dict: """Lesson 24's continuous batching with a cache pool of `pool_blocks` blocks and at most `cap` chats running. A chat's id is its index. A step starts like this. Each running chat that needs a new block takes one; if none is free, the newest running chat is paused: its blocks are freed and it goes to the front of the queue (the oldest chats are served first). In a step where a chat was paused, nobody new joins. Otherwise, waiting chats join, oldest first, while fewer than `cap` run and the pool has the blocks the chat needs for the next token; the first one that does not fit holds up the rest. A chat that joins has its whole cache (prompt, plus the tokens it had written before it was paused) read in that step, as in lesson 24, and writes one token like every running chat. With `reserved`, a chat takes blocks for `longest` tokens when it joins, so it is never paused. Chats longer than `longest` are turned away. A server that cannot hold one chat of `longest` tokens does not start. The context used for the cost of a step is prompt + tokens written, as in lesson 24.""" assert pool_blocks >= blocks_for(longest), "the pool cannot hold one chat of the longest length" n = len(requests) refused = [i for i, r in enumerate(requests) if r.prompt + r.reply > longest] turned_away = set(refused) made, held = [0] * n, [0] * n first, last = [-1.0] * n, [-1.0] * n waiting, running = deque(), [] pending = [i for i in range(n) if i not in turned_away] nxt, free, now, done = 0, pool_blocks, 0.0, 0 steps = paused = rebuilt = 0 area, used_area, lowest_free = 0.0, 0.0, pool_blocks while done < len(pending): while nxt < len(pending) and requests[pending[nxt]].arrival <= now: waiting.append(pending[nxt]) nxt += 1 if not running and not waiting: now = requests[pending[nxt]].arrival # nothing to do: wait for the next chat continue paused_now = False k = 0 while k < len(running): i = running[k] if not reserved and (requests[i].prompt + made[i]) % BLOCK_TOKENS == 0: # the next token starts a new block if free == 0: j = running.pop() # the newest running chat gives up its blocks (when it is `i` itself, the scan ends) free += held[j] held[j] = 0 waiting.appendleft(j) paused += 1 paused_now = True continue free -= 1 held[i] += 1 k += 1 joined = [] while not paused_now and waiting and len(running) < cap: i = waiting[0] need = blocks_for(longest) if reserved else blocks_for(requests[i].prompt + made[i] + 1) if need > free: break waiting.popleft() free -= need held[i] = need running.append(i) joined.append(i) joining = set(joined) contexts = [requests[i].prompt + made[i] for i in running if i not in joining] reread = sum(requests[i].prompt + made[i] for i in joined) rebuilt += sum(requests[i].prompt + made[i] for i in joined if made[i] > 0) # a chat that was paused is read in again, in full dt = cost(contexts, reread) now += dt steps += 1 if steps > MAX_STEPS: raise RuntimeError(f"more than {MAX_STEPS:,} steps: some chat can never finish") area += len(running) * dt used_area += (pool_blocks - free) * dt lowest_free = min(lowest_free, free) still = [] for i in running: made[i] += 1 if first[i] < 0: first[i] = now if made[i] == requests[i].reply: last[i] = now free += held[i] held[i] = 0 done += 1 else: still.append(i) running = still served = [i for i in range(n) if i not in turned_away] tokens = per_word = 0.0 for i in served: tokens += requests[i].reply per_word += (last[i] - first[i]) / (requests[i].reply - 1) return {"seconds": now, "steps": steps, "served": len(served), "refused": len(refused), "paused": paused, "rebuilt_tokens": rebuilt, "tokens_per_second": tokens / now, "seconds_per_word": per_word / len(served), "running": area / now, "pool_used": used_area / (now * pool_blocks), "lowest_free": lowest_free, "first_token": first, "finished": last} # --- what the widget shows ----------------------------------------------------------------------------------------------------------------- def gb(nbytes: float) -> str: return f"{nbytes / 1e9:.1f}" def commas(n: int) -> str: return f"{n:,}" def percent(x: float) -> int: return math.floor(x * 100 + 0.5) def room_line(r: dict, longest: int, reserved: bool) -> str: if not r["starts"]: return f"The server will not start. One chat of {commas(longest)} tokens needs room for {commas(r['per_longest'] * BLOCK_TOKENS)} tokens; the pool holds {commas(r['pool_tokens'])}." if reserved: return f"Room for {commas(r['at_longest'])} chats at once, because each one holds {commas(longest)} tokens whether it uses them or not." return f"Room for {commas(r['typical'])} chats of {commas(min(TYPICAL_CHAT, longest))} tokens at once. {commas(r['at_longest'])} if every chat ran to {commas(longest)}." def run_line(run: dict, cap: int) -> str: return (f"Running on average: {commas(math.floor(run['running'] + 0.5))} chats (up to {commas(cap)}). Everyone together: {commas(math.floor(run['tokens_per_second'] + 0.5))} tokens a second. " f"One user waits {run['seconds_per_word'] * 1000:.1f} ms between words.") def extra_line(run: dict) -> str: bits = [f"The chats in progress hold {percent(run['pool_used'])}% of the cache pool on average."] if run["refused"]: bits.append(f"{commas(run['refused'])} of {commas(run['refused'] + run['served'])} chats turned away: too long.") if run["paused"]: bits.append(f"{commas(run['paused'])} chats paused to free room; {commas(run['rebuilt_tokens'])} tokens of cache rebuilt.") return " ".join(bits) def describe(share: int, longest: int, cap: int, reserved: bool = False, requests: list[Request] | None = None) -> dict: """Everything the widget shows for one setting of its controls: the four parts of the bar, the legend, how full the pool is, and three lines of words. `requests` defaults to the 600 invented chats; the widget test passes fewer to keep Python's side fast.""" r = room(share, longest, reserved) out = {"grow": [r["weights"], r["reserve"], r["pool"], r["unused"]], "fill": 0, "room": room_line(r, longest, reserved), "run": "", "extra": "", "legend": [f"Model weights {gb(r['weights'])} GB", f"Working reserve {gb(r['reserve'])} GB", f"Cache pool {gb(r['pool'])} GB", f"Left unused {gb(r['unused'])} GB"]} if not r["starts"]: out["extra"] = "Raise the memory share or lower the longest chat." return out run = serve(backlog() if requests is None else requests, cap, r["blocks"], longest, reserved) out.update(run=run_line(run, cap), extra=extra_line(run), fill=percent(run["pool_used"])) return out def sweep(share: int = OPENING["share"], longest: int = OPENING["longest"], caps: list[int] = CAPS) -> list[dict]: """One run per cap on the same chats and pool: the rows of the lesson's table.""" r = room(share, longest) rows = [] for cap in caps: run = serve(backlog(), cap, r["blocks"], longest) rows.append({"cap": cap, "tokens_per_second": run["tokens_per_second"], "seconds_per_word": run["seconds_per_word"], "running": run["running"], "paused": run["paused"], "rebuilt_tokens": run["rebuilt_tokens"]}) return rows def stops_paying(rows: list[dict]) -> int: """The first cap whose doubling brings under KNEE_GAIN times the tokens a second.""" by_cap = {row["cap"]: row["tokens_per_second"] for row in rows} for cap in sorted(by_cap): if 2 * cap in by_cap and by_cap[2 * cap] < KNEE_GAIN * by_cap[cap]: return cap return max(by_cap)