Why split reading the prompt from writing the reply? A toy you can run

On this page

In one minute
Explore the numbers
How it works
Your turn
What this means when you build
Go deeperoptional
Skills this lesson builds
Check yourself
Further reading
Codeoptional

An AI chat server does two jobs for every request: it reads your prompt, then writes its reply a word at a time. When one GPU, the chip that does the arithmetic, does both, a long prompt from someone else stops it writing. In this lesson’s toy, a prompt of 2000 tokens, which are short pieces of text, takes 103.0 ms to read, and every reply in progress waits: a word due 7.9 ms after the last comes 110.8 ms after it. At 4 requests a second only 13.5% of requests then get a first word within 0.5 s and no later gap over 0.1 s. Give the reading to one GPU and the writing to another and 98.7% do. The price is that the reading GPU’s notes on the prompt (the cache, 262 MB for 2000 tokens) must cross the link between the two GPUs.

In one minute

Picture a shop with two clerks who answer letters by hand. A letter is read in full before its answer starts, and the answer is written one word at a time, so a clerk is usually in the middle of several answers. When a long letter arrives, a clerk who both reads and writes puts every answer down to read it, and her customers see a long pause mid-sentence. The shop can put one clerk on reading and one on writing. The writer never stops, but the notes the reader takes must be carried across the room, and a slow corridor can cost more than the pause it removed.

The reader does the prefill, one big pass over the prompt, and the writer does the decode, one small step per word. The notes are the KV cache of lesson 22 and the corridor is the link between the GPUs. The first word comes later (a longer time to first token), because the prompt waits its turn at one reader and then crosses the link. So the split pays only for long prompts on a fast link, and goodput, the share of requests that meet every target, tells you when.

Explore the numbers

Open the widget. It opens on the toy’s default: 4 requests a second, prompts of 2000 tokens, replies of 200 words and a 25 GB/s link. The sentence on top says that shared GPUs get 13.5% of requests right and the split 98.7%. Drag Requests a second to 10: shared falls to 1.7% and the split to 32.2%. Drag Link speed to 0.25 GB/s: the split falls to 0.0%, because no prompt’s cache crosses in time. Drag Prompt tokens to 200: the two tie at 100.0%.

How it works

The picture, step by step

A toy with two GPUs serves 600 requests, laid out two ways. In the shared layout both GPUs read prompts and write replies, and requests take turns. In the other, one GPU reads and the other writes, which is called a prefill-decode split, or disaggregated serving. A request is good if its first word comes within 0.5 s and no gap between its words is over 0.1 s (invented targets). Figure 1 shows one stall.

Four rows of time. On a shared GPU, reply A waits 110.8 ms for a word while a 2000-token prompt B is read. On the split pair the writer GPU gives A a word every 7.9 ms throughout, and B’s cache lands 10.5 ms after its read.

Figure 1. One long prompt stops a shared GPU’s replies and does not stop a writer GPU.

Requests arrive. 600 requests come at 4 a second, a random time apart, with prompts of 1000 to 3000 tokens and replies of 100 to 300 words.
A shared GPU reads whole. A 2000-token prompt takes 103.0 ms to read, and no reply on that GPU moves meanwhile. A word due 7.9 ms after the last comes 110.8 ms after it.
Measure the shared pair. Half of all gaps between words are 7.9 ms (the p50), but the slowest 1 in 100 (the p99, the 99th percentile) is 110.8 ms. Only 13.5% of requests are good.
Give the reading to one GPU. The reader reads prompts in order, one at a time, with 0.40 s of reading to do each second. The writer only writes, 7.9 ms a step.
Carry the cache across. A 2000-token prompt’s notes are 2000 x 131,072 bytes = 262.1 MB. At 25 GB/s that takes 10.5 ms, against 103.0 ms to read.
Count good requests. On the split no gap passes 7.9 ms and 98.7% of requests are good, against 13.5%. The first word comes later: p50 143.0 ms against 106.6 ms.

Each step in code

One real line of the program does each step. The files are in code/, and the line numbers are those of this lesson’s files.

Step	The code line	What it does
1. Space the arrivals	`pd_sim.py:41` `t += rng.expovariate(rps)`	Each request comes a random time after the last, 1 / rps seconds on average.
2. Read the whole prompt	`pd_sim.py:57` `t += read`	The GPU’s clock jumps by the whole read before any reply gets another word.
3. Record the gap	`pd_sim.py:72` `gaps[a[0]].append(t - last[a[0]])`	The gap before a word is the time now minus the time of that reply’s last word, so a read in between shows up here.
4. Read in order	`pd_sim.py:98` `free = max(free, r.arrival) + r.prompt * READ_S_PER_TOKEN`	The reader starts a prompt when it is free and the request has arrived. `free` is when it finishes.
5. Price the hand-over	`pd_sim.py:80` `return prompt_tokens * CACHE_BYTES_PER_TOKEN / (link_gbps * 1e9)`	The cache’s bytes divided by the link’s bytes per second.
6. Count good requests	`pd_sim.py:112` `good = sum(1 for i, r in enumerate(requests) if first[i] - r.arrival <= first_target and max(gaps[i], default=0.0) <= gap_target)`	A request counts when its first word came within the target and its largest gap is within the target.

A real run

trace.py runs the six steps with pd_sim.py’s own functions and prints one line for each, then one line for the limit. Run it from the repository root:

cd lessons/why-split-reading-the-prompt-from-writing-the-reply/code && python3 trace.py

Phone readers: slide the box sideways to see every column.

step 1  600 requests in 148.9 s, 4 a second; prompts of 1000, 2000, 3000 tokens, replies of 100, 200, 300 words
step 2  shared: a 2000-token prompt takes 103.0 ms to read and no reply on that GPU moves meanwhile; a word due 7.9 ms after the last comes 110.8 ms after it
step 3  shared: first word p50 106.6 ms, p99 267.0 ms; gap between words p50 7.9 ms, p99 110.8 ms
step 4  split: the reader GPU has 0.40 s of prompts to read each second; the writer GPU only writes, 7.9 ms a step
step 5  hand-over: 2000 tokens x 131072 bytes = 262.1 MB; at 25 GB/s that takes 10.5 ms, against 103.0 ms to read (equal at 2.55 GB/s)
step 6  split: first word p50 143.0 ms, p99 616.3 ms; gap p50 7.9 ms, p99 7.9 ms; goodput 98.7% against 13.5% shared
limit   goodput of the split by link speed, GB/s: 100 98.7%, 25 98.7%, 12.5 98.7%, 1.25 86.3%, 0.125 0.0%

Read it against the steps.

Steps 2 and 3. The word waits 110.8 ms, 14 times the usual 7.9. A 200-word reply lasts about 1.6 s, long enough for a prompt to land on its GPU, so 86.5% of replies have a gap over 0.1 s and 13.5% are good. Figure 2 shows the tail grow with load: the shared p99 gap passes 100 ms at 4 requests a second and is 213.8 ms at 10, while the split stays at 7.9 ms.

The p99 gap in ms against 1 to 10 requests a second. Shared GPUs rise from 7.9 to 213.8, passing the 100 ms limit at 4; the split stays at 7.9.

Figure 2. The shared p99 gap passes the limit at 4 requests a second. The split’s does not move.

Step 4. One reader can read about 9.7 prompts a second (1 / 0.103 s). Past that a queue builds: the split’s first-word p99 climbs from 0.27 s at 1 request a second to 3.72 s at 10.
Step 5. The hand-over equals the read at 2.55 GB/s for any prompt length, because the cache and the read both grow with the prompt.
Step 6. The first word is slower, p99 616.3 ms against 267.0 ms, because a prompt waits for one reader and crosses the link, while the shared pair reads on two GPUs. Goodput counts both targets, so it shows which layout keeps requests inside both.
The limit. Slow the link and the split fails: 86.3% at 1.25 GB/s and 0.0% at 0.125. It beats the shared layout on links faster than about 0.3 GB/s (Figure 3). Below that even the shortest prompt, 1000 tokens, has a 131 MB cache that cannot be read and cross within 0.5 s, so no request is good.

Goodput against link speed. The split falls from 98.7% at 12.5 GB/s to 0.0% at 0.125; shared GPUs stay at 13.5%. The lines cross at about 0.3 GB/s.

Figure 3. The split beats shared GPUs on goodput on links faster than about 0.3 GB/s, and loses below it.

Hands on

Your turn

The link was 25 GB/s, fast enough to carry the cache in a tenth of the read. See what a link 100 times slower does.

In code/trace.py, change line 8 to LINK_GBPS = 0.25.
Run python3 trace.py again in the code folder.

Look at steps 5 and 6. How long does the hand-over of a 2000-token prompt take, and what share of requests is good?

What you should see

Steps 5 and 6 change and the other lines stay as they were:

step 1  600 requests in 148.9 s, 4 a second; prompts of 1000, 2000, 3000 tokens, replies of 100, 200, 300 words
step 2  shared: a 2000-token prompt takes 103.0 ms to read and no reply on that GPU moves meanwhile; a word due 7.9 ms after the last comes 110.8 ms after it
step 3  shared: first word p50 106.6 ms, p99 267.0 ms; gap between words p50 7.9 ms, p99 110.8 ms
step 4  split: the reader GPU has 0.40 s of prompts to read each second; the writer GPU only writes, 7.9 ms a step
step 5  hand-over: 2000 tokens x 131072 bytes = 262.1 MB; at 0.25 GB/s that takes 1048.6 ms, against 103.0 ms to read (equal at 2.55 GB/s)
step 6  split: first word p50 1151.5 ms, p99 1976.0 ms; gap p50 7.9 ms, p99 7.9 ms; goodput 0.0% against 13.5% shared
limit   goodput of the split by link speed, GB/s: 100 98.7%, 25 98.7%, 12.5 98.7%, 1.25 86.3%, 0.125 0.0%

The hand-over takes 1048.6 ms, about ten times the read, and no request is good. Even the shortest prompt, 1000 tokens, needs 524 ms to cross, which is over the 0.5 s allowed for the first word. The gaps stay at 7.9 ms, because the writer is never held up. The split fails on the corridor, not on the writing. The limit line does not move, because it tries its own link speeds.

What this means when you build

Watch the p99 gap, not the middle. In the toy the middle gap is 7.9 ms on both layouts. The slowest 1 in 100 is 110.8 ms shared and 7.9 ms split (lesson 27).
Know your traffic before you split. In the toy, prompts of 600 tokens or fewer nearly tie, and at 2000 tokens the split leads 98.7% to 13.5%. A pause between words matters to a person who watches the words appear, not to a program that waits for the whole reply.
Price the hand-over first. The cache is the prompt’s tokens times the bytes per token, and the link must carry it inside your first-word target. In the toy the split stops paying below about 0.3 GB/s. Real figures depend on the model and the hardware.
Time each leg separately. In pull request 3201 of vllm-project/aibrix (merged 2026-10-20; read 2026-10-20), a request set to pd:1,least-request:1 was routed as pd, but two places compared the raw setting with pd. Routing time then ran past the prefill leg, and the response side recorded none of its prefill, KV-transfer or decode counters. The merge resolves the setting before both checks, at gateway_req_body.go:214 and gateway_rsp_body.go:698 (line numbers at the merge commit 3469dad). A metric that mixes two legs, or skips one, hides where the delay is.
On one GPU, try chunked prefill first. If you cannot add a second pool or a fast link, mixing pieces of a prompt into each writing step shortens the stall with no hand-over (Go deeper).

Optionalabout 17 min

Go deeper

Optional. This part is for readers who write Python. The code is in the code/ folder next to this lesson. From inside that folder, python3 trace.py prints the run above, python3 make_figs.py redraws the figures and python3 -B -m pytest -q -p no:cacheprovider runs the tests.

pd_sim.py is 127 lines, with its comments and blank lines. It uses only the standard library, reads no file, opens no connection and loads no model. Two of its costs echo lesson 23’s best case for one chip, the chip of lesson 19: 51.5 microseconds to read a prompt token and 7.877 ms for a writing step. The 131,072 cache bytes per token echo lesson 22. The link speeds, the 32 slots a GPU writes at once and the two targets are invented.

The simulation

    waiting, active, first, gaps, last, t = sorted(jobs), [], {}, {}, {}, 0.0
    while waiting or active:
        if not active and waiting[0][0] > t:
            t = waiting[0][0]                                         # an idle GPU waits for the next job
        ready = list(takewhile(lambda j: j[0] <= t, waiting))         # `waiting` is sorted, so the jobs that are ready are its first few
        if ready and len(active) < SLOTS and ready[0][1] > 0:         # step 2: a prompt waits and a slot is free, so read the whole prompt now
            ready_t, read, left, i = ready[0]
            waiting.remove(ready[0])
            t += read
            first[i], gaps[i], last[i] = t, [], t
            if left > 0:
                active.append([i, left])
            continue

serve is one GPU. A job is its ready time, the seconds of reading it needs, the words left to write and an id. When a prompt is waiting and a slot is free, the whole prompt is read at once: t jumps, the first word is stamped at the end, and continue returns to the top without running a writing step. There is no chunking. That is the stall.

        t += STEP_S                                                   # one writing step: every reply in progress gets one word
        for a in active:
            gaps[a[0]].append(t - last[a[0]])
            last[a[0]], a[1] = t, a[1] - 1
        active = [a for a in active if a[1] > 0]
    return first, gaps

A writing step costs STEP_S for any number of replies in progress, up to 32, and gives each one a word. The gap recorded is the time since that reply’s last word, so a read that ran between two steps is inside it. shared sends request 0 to the first GPU, request 1 to the second, and so on, each with its own serve.

    free, jobs = 0.0, []
    for i, r in enumerate(requests):
        free = max(free, r.arrival) + r.prompt * READ_S_PER_TOKEN
        jobs.append((free + handover_seconds(r.prompt, link_gbps), 0.0, r.reply - 1, i))
    return serve(jobs)

split makes the reader a single queue. free is when the reader finishes each prompt, and the job reaches the writer at free plus the hand-over, with 0.0 seconds of reading to do. serve treats such a job as one that arrives with its first word made, so the writer GPU never reads and its steps are never interrupted. The link has no limit on how many caches cross at once.

Why the two jobs differ

READ_S_PER_TOKEN = 2 * PARAMS / A100_PEAK_FLOPS         # reading the prompt: 2 FLOPs per weight per token, limited by arithmetic (lesson 23)
STEP_S = 2 * PARAMS / A100_BANDWIDTH                    # one writing step: every weight is read once, limited by memory, the same for any number of replies (lesson 23)

Reading a prompt does about 2 operations per weight for every token, so a long prompt keeps the arithmetic busy and a token costs READ_S_PER_TOKEN. A writing step has to fetch every weight from memory once, whatever the number of replies, so it costs STEP_S however full the slots are (lessons 19 and 23). Splitwise states the same pair, prompt computation compute-intensive and token generation memory-intensive. One GPU doing both serves each badly: the prompt holds the writers back, and the writers hold the arithmetic back.

The cost of the hand-over

The cache is prompt tokens x 131,072 bytes, so it grows with the prompt just as the read does. For a 2000-token prompt, 262.1 MB crosses, and the time depends only on the link:

Link	Hand-over of 2000 tokens	Against the 103.0 ms read
100 GB/s (800 Gbit/s)	2.6 ms	2.5%
25 GB/s (200 Gbit/s)	10.5 ms	10.2%
12.5 GB/s (100 Gbit/s)	21.0 ms	20.4%
1.25 GB/s (10 Gbit/s)	209.7 ms	2.0 times
0.125 GB/s (1 Gbit/s)	2097.2 ms	20 times

The hand-over is slower than the read below 2.55 GB/s, but the split still wins goodput long after that: at 2 GB/s the split is good for 94.8% of requests, because the first-word target is 0.5 s and the hand-over only has to fit inside what is left after the read. The split stops paying at about 0.3 GB/s (0.2922 in the code: at 0.29 no request is good and at 0.293, 22.3% are), where the 1000-token prompts run out of time: 51.5 ms to read plus 448.5 ms to cross fills the 0.5 s (Figure 3). The sweep is in make_figs.py, which finds the crossover by bisection on the same compare function.

Where the split does not pay

Figure 4 holds the link at 25 GB/s and the load at 4 requests a second, and varies the prompt.

Goodput against typical prompt length from 100 to 6000 tokens on a log scale. Shared GPUs fall from 100.0% at 200 to 94.8% at 1000, 13.5% at 2000 and 3.2% at 4000. The split stays at 100.0% to 1000, then 98.7% at 2000, 44.0% at 4000 and 1.2% at 6000.

Figure 4. The two layouts nearly tie up to 600 tokens, and the split pulls ahead after that. At 6000 tokens the single reader is overloaded and both collapse.

With short prompts, a read is short, so the stall is short: at 200 tokens the read takes 10.3 ms, a stall is about 18 ms and the 0.1 s gap target is never missed. The split then pays for a link it does not need. At the other end, 6000-token prompts need 0.31 s of reading each, so one reader falls behind (the split’s first-word p99 is 32.5 s) and the shared pair, which reads on two GPUs, falls behind too. The split has a window between these: at this load the two nearly tie up to 600 tokens, and the split leads from 800 to 4000.

What the papers report

DistServe (Zhong et al., OSDI 2024, arXiv 2401.09670; read 2026-10-20) finds that putting prefill and decoding on the same GPUs causes “strong prefill-decoding interferences”, and assigns them to different GPUs. It reports that it can serve 7.4 times more requests or meet a 12.6 times tighter latency target than the systems it compares with, “while staying within latency constraints for > 90% of requests”, and on one chat workload a 2.0 to 4.6 times higher request rate than vLLM. Its goodput is a rate, the most requests per second per GPU that meet the latency targets for over 90% of requests. On the transfer, it puts the cache of one 512-token request to a 66-billion-parameter model at about 1.13 GB, which at 10 requests a second is 11.3 GB/s (90 Gbit/s), and reports the transfer under 0.1% of total latency for a 175-billion-parameter model and under 30 ms for over 95% of requests. Its discussion says that a throughput-only offline job may favour chunked prefill, that one or a few GPUs may suit a non-split system, and that disaggregation “remains promising in long-context serving”: the cache, and with it the transfer, grows “linearly with the prompt length” but “the prefill computation grows quadratically, so the relative duration of transmission and prefill job decreases”.

Splitwise (Patel et al., ISCA 2024, arXiv 2311.18677, version 2; read 2026-10-20) reports up to 1.4 times higher throughput at 20% lower cost than current designs, or 2.35 times more throughput under the same power and cost budgets. It sends the cache layer by layer while the prompt is still being read, which leaves a constant transfer time that is not hidden of about 8 ms on A100 GPUs with a 200 Gbit/s link and about 5 ms on H100 GPUs with 400 Gbit/s. For small prompts it sends the cache in one piece, below 512 tokens on the H100 setup.

Sarathi-Serve (Agrawal et al., OSDI 2024, arXiv 2403.02310; read 2026-10-20) keeps reading and writing on the same GPUs and “splits a prefill request into near equal sized chunks”, mixing a chunk into each step so that writing is not stalled. It reports 2.6 times higher serving capacity for Mistral-7B on one A100 and up to 3.7 times for Yi-34B on two A100s, both against vLLM, and up to 5.6 times for Falcon-180B with pipeline parallelism. It notes that a split needs the cache moved, which is hard without a fast link.

What the toy leaves out

One reader and one writer. Real systems run several of each and choose the ratio. The toy’s single reader saturates near 9.7 requests a second.
No chunking and no batching of prompts. The reader takes one prompt at a time, whole, and the shared GPU does the same.
An uncontended link. Any number of caches cross at once. A real link is shared, and the cache can cross in layers as the prompt is read (Splitwise).
A fixed split of traffic. The shared layout alternates requests between two GPUs. A real scheduler can send a request to the less busy one.
A share, not a rate. The toy’s goodput is the share of requests that meet both targets. DistServe’s goodput is a rate per GPU.
Toy costs. The read and write times are best-case figures for one chip, not measurements, and every request is served to its last word with no failures.

Proof it works

test_pd_sim.py checks that the copied costs equal the ones in lesson 23’s code, that a read stalls a shared GPU’s replies for the read plus one step, that a hand-over is bytes over link speed, and that the split removes the stall and reproduces the crossover. test_lesson.py checks that every number this lesson, the trace and the figures quote comes from the code, that each code block and each file:line in the table matches the real file, that the figures are the ones make_figs.py draws, and that trace.py prints the run above and the Your turn change prints the lines shown. test_widget.py checks that the numbers inside the widget are the ones show.py computes from pd_sim.py, that the page’s JavaScript gives the same words as show.py for every setting of its six controls, that each state named in “Explore the numbers” is reachable and matches a fresh run of the simulation, and that the chart’s key stays above the plot and its text stays large enough to read on a phone.

Where this is used in practice

The papers above describe systems that serve with the reading and the writing apart. The open-source AIBrix gateway has a routing strategy named pd for it, and the pull request in the section above is a fix to how that strategy is timed (read 2026-10-20). This lesson shows the idea on a toy and measures no real server.

Skills this lesson builds

Say what a prefill-decode split changes and what it costs.
Read goodput and the p99 gap instead of an average.
Work out a cache’s size and its hand-over time from bytes per token and link speed.
Find the link speed or prompt length at which a split stops paying.
Time the legs of a split request separately.

Portfolio project: extend the simulation with a second reader GPU and with a shared layout that reads prompts in chunks. Chart goodput against requests a second for all three layouts, and add a short note on which one your own traffic would need and why. Publish the code, the chart and the note.

Check yourself

On shared GPUs the middle gap between words is 7.9 ms, the same as on the split. Why do they differ so much in goodput?

Answer

The middle gap hides the stalls. A reply lasts about 1.6 s, long enough for a prompt to land on its GPU, and one gap over 0.1 s makes a request fail: 86.5% of the shared requests have one. The split’s gaps never pass 7.9 ms.
At a 2 GB/s link the hand-over of a 2000-token prompt takes 131 ms, longer than the 103 ms read. Does the split lose?

Answer

No. The split is good for 94.8% of requests at 2 GB/s, against 13.5% shared. Goodput is set by the targets, not by comparing the hand-over with the read: the first word has 0.5 s, and a stall of 110.8 ms on shared GPUs breaks the gap target while a hand-over delays only the first word.
With prompts of 200 tokens the two layouts both reach 100.0%. Why does the split give nothing here?

Answer

A 200-token prompt takes 10.3 ms to read, so a shared GPU stalls its replies for about 18 ms, well inside the 0.1 s gap target. There is no stall for the split to remove, and it still adds a hand-over to the first word.
At 10 requests a second the split’s first-word p99 reaches 3.72 s. What limit has it met?

Answer

The single reader. An average 2000-token prompt takes 0.103 s to read, so one reader handles about 9.7 prompts a second. Above that, prompts queue and each waits for those ahead of it. The writer GPU is not the limit, and a second reader would be the fix.
Why would counting the prefill leg inside a request’s “routing time” mislead an operator?

Answer

Routing time should end where the prefill leg starts. For the pd strategy the router sends the prefill leg from inside routing, so a routing time that runs to the end of routing holds the whole prefill: a slow prefill looks like slow routing, and the operator looks in the wrong place. Pull request 3201 of vllm-project/aibrix made a weighted setting that resolves to pd end routing time at the prefill start too.

Code

Each file runs on its own. Download it, or read it here.

embed_widget.py29 linesDownload

Show code

"""Keep the table of numbers inside ../widget.html equal to the one show.py computes from pd_sim.py.

    python3 embed_widget.py          rewrite the data block in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

import show

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
CONFIG_BLOCK = re.compile(r'(<script type="application/json" id="pd-config">)(.*?)(</script>)', re.S)


def embed(html):
    """html with the data block rewritten from show.py."""
    settings = json.dumps(show.widget_config(), ensure_ascii=False, separators=(", ", ": "))
    return CONFIG_BLOCK.sub(lambda m: m.group(1) + "\n" + settings + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data block now matches show.py")

make_figs.py268 linesDownload

Show code

"""Draw the four figures of lesson 33 from pd_sim.py, so that no number in a figure is typed by hand.

    python3 make_figs.py          rewrite ../figs/*.svg
    python3 make_figs.py --check  exit 1 if a figure on disk differs (test_lesson.py runs the same check)

All four figures are 540 units wide. Text is at least 22 units (20 for axis ticks) in ink or grey on white or a light fill; orange and green are used on marks only,
and every mark that colour tells apart also differs in shape or dash and is named in words.
"""
import math
import sys
from functools import lru_cache
from pathlib import Path
from xml.sax.saxutils import escape

import pd_sim as sim

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"
PAD = 14                                            # room between the axis and the first mark
RPS, PROMPT, REPLY, LINK = 4, 2000, 200, 25         # the lesson's default run
RPS_GRID = tuple(range(1, 11))
PROMPT_GRID = (100, 200, 400, 600, 800, 1000, 1500, 2000, 3000, 4000, 6000)
LINK_GRID = (0.1, 0.125, 0.15, 0.2, 0.25, 0.29, 0.3, 0.4, 0.5, 0.75, 1, 1.25, 1.5, 2, 2.5, 3, 5, 12.5, 25, 100)   # includes the five speeds of sim.LINKS_GBPS


def num(x: float) -> str:
    return f"{x:.2f}".rstrip("0").rstrip(".")


# The SVG namespace is a name that every standalone .svg file must carry, not an address this program visits. It is written in two pieces so that a plain
# search of code/ for the start of a web address finds nothing (test_lesson.py runs that search).
NAMESPACE = "ht" + "tp://www.w3.org/2000/svg"


def frame(width: int, height: int, title: str, desc: str, body: list[str]) -> str:
    head = (f'<svg xmlns="{NAMESPACE}" viewBox="0 0 {width} {height}" width="{width}" height="{height}" role="img" '
            f'font-family="system-ui, sans-serif" font-size="22" fill="{GREY}">')
    return "\n".join([head, f"  <title>{escape(title)}</title>", f"  <desc>{escape(desc)}</desc>", '  <rect width="100%" height="100%" fill="#ffffff"/>',
                      *("  " + line for line in body), "</svg>", ""])


def text(x: float, y: float, content: str, anchor: str = "start", color: str = INK, bold: bool = False, size: int = 22) -> str:
    extra = (f' text-anchor="{anchor}"' if anchor != "start" else "") + (' font-weight="bold"' if bold else "") + (f' font-size="{size}"' if size != 22 else "")
    return f'<text x="{num(x)}" y="{num(y)}"{extra} fill="{color}">{escape(content)}</text>'


def tick(x: float, y: float, label: str, anchor: str = "middle") -> str:
    return text(x, y, label, anchor, GREY, size=20)


def circle(x: float, y: float, r: float = 6, color: str = GREEN) -> str:
    return f'<circle cx="{num(x)}" cy="{num(y)}" r="{num(r)}" fill="{color}" stroke="#ffffff" stroke-width="1"/>'


def square(x: float, y: float, half: float = 6) -> str:
    return f'<rect x="{num(x - half)}" y="{num(y - half)}" width="{num(2 * half)}" height="{num(2 * half)}" fill="#ffffff" stroke="{GREEN}" stroke-width="3"/>'


def triangle(x: float, y: float) -> str:
    return f'<path d="M{num(x)} {num(y - 9)}L{num(x + 9)} {num(y + 7)}L{num(x - 9)} {num(y + 7)}z" fill="{ORANGE}" stroke="#ffffff" stroke-width="1"/>'


def block(x0: float, x1: float, y: float, height: float, stroke: str, dash: str = "") -> str:
    return (f'<rect x="{num(x0)}" y="{num(y)}" width="{num(x1 - x0)}" height="{num(height)}" fill="{FILL}" stroke="{stroke}" stroke-width="3"'
            + (f' stroke-dasharray="{dash}"' if dash else "") + "/>")


def ms(seconds: float) -> str:
    return f"{seconds * 1000:.1f}"


def path(points: list[tuple[float, float]], color: str, dash: str = "") -> str:
    return f'<path d="M{" L".join(f"{num(x)} {num(y)}" for x, y in points)}" fill="none" stroke="{color}" stroke-width="4"' + (f' stroke-dasharray="{dash}"' if dash else "") + "/>"


def key(x: float, y: float, label: str, kind: str) -> list[str]:
    """One line of a key: a short sample of the mark, then its name. kind is shared (orange line, filled circle), split (green dashed line, hollow square) or rule (grey dash)."""
    if kind == "rule":
        return [f'<path d="M{num(x - 18)} {num(y - 7)}H{num(x + 18)}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>', text(x + 30, y, label)]
    line = path([(x - 18, y - 7), (x + 18, y - 7)], ORANGE if kind == "shared" else GREEN, "" if kind == "shared" else DASH)
    return [line, circle(x, y - 7, 7, ORANGE) if kind == "shared" else square(x, y - 7), text(x + 30, y, label)]


# --- the numbers: each is computed once from the simulation ------------------------------------------------------------------------

@lru_cache(maxsize=None)
def runs(rps: float, prompt: int, reply: int, link: float) -> dict:
    return sim.compare(rps, prompt, reply, link)


@lru_cache(maxsize=None)
def crossover_gbps(rps: float = RPS, prompt: int = PROMPT, reply: int = REPLY) -> float:
    """The slowest link at which the split's goodput is still above the shared layout's, found by halving the gap between a link that loses and one that wins."""
    lose, win = 0.1, 100.0
    assert runs(rps, prompt, reply, lose)["split"]["goodput"] <= runs(rps, prompt, reply, win)["shared"]["goodput"] < runs(rps, prompt, reply, win)["split"]["goodput"]
    for _ in range(40):
        middle = math.sqrt(lose * win)
        if runs(rps, prompt, reply, middle)["split"]["goodput"] > runs(rps, prompt, reply, middle)["shared"]["goodput"]:
            win = middle
        else:
            lose = middle
    return win


@lru_cache(maxsize=None)
def stall_scene() -> dict:
    """Figure 1: reply A is in progress when B, a long prompt, arrives; the same two requests on one shared GPU and on a split pair."""
    a_read, b_read, arrival = 100 * sim.READ_S_PER_TOKEN, PROMPT * sim.READ_S_PER_TOKEN, 0.045
    first, gaps = sim.serve([(0.0, a_read, 40, 0), (arrival, b_read, 5, 1)])
    requests = [sim.Request(0.0, 100, 41), sim.Request(arrival, PROMPT, 6)]
    split_first, split_gaps = sim.split(requests, LINK)

    def words(start: float, steps: list[float]) -> list[float]:
        times = [start]
        for gap in steps:
            times.append(times[-1] + gap)
        return times

    return {"shared_words": words(first[0], gaps[0]), "read_end": first[1], "read_start": first[1] - b_read, "split_words": words(split_first[0], split_gaps[0]),
            "arrival": arrival, "reader_start": arrival, "reader_end": arrival + b_read, "landed": split_first[1], "hand_over": sim.handover_seconds(PROMPT, LINK)}


# --- figure 1: one long prompt, on a shared GPU and on a split pair -----------------------------------------------------------------

def fig_stall() -> str:
    s = stall_scene()
    x0, window, eps = 20, 0.24, 1e-9
    px = lambda t: x0 + t / window * 500
    words = [t for t in s["shared_words"] if t <= window]
    before = max(t for t in words if t <= s["read_start"] + eps)
    after = min(t for t in s["shared_words"] if t >= s["read_end"] + sim.STEP_S - eps)
    body = [text(x0, 28, "One shared GPU", bold=True)]
    body.append(block(px(s["read_start"]), px(s["read_end"]), 40, 36, ORANGE))
    body.append(text((px(s["read_start"]) + px(s["read_end"])) / 2, 64, f"reads B: {ms(s['read_end'] - s['read_start'])} ms", "middle"))
    body += [circle(px(t), 58) for t in words if t <= s["read_start"] + eps or t >= after - eps]
    body.append(f'<path d="M{num(px(before))} 102H{num(px(after))}M{num(px(before))} 95V109M{num(px(after))} 95V109" stroke="{INK}" stroke-width="3"/>')
    body.append(text((px(before) + px(after)) / 2, 132, f"{ms(after - before)} ms with no word from A", "middle"))
    body.append(triangle(px(s["read_end"]), 88))
    body.append(text(x0, 168, "Split: reader GPU", bold=True))
    body.append(block(px(s["reader_start"]), px(s["reader_end"]), 180, 36, ORANGE))
    body.append(text((px(s["reader_start"]) + px(s["reader_end"])) / 2, 204, f"reads B: {ms(s['reader_end'] - s['reader_start'])} ms", "middle"))
    body.append(text(x0, 248, "Split: link", bold=True))
    body.append(block(px(s["reader_end"]), px(s["landed"]), 260, 28, GREY, DASH))
    body.append(text(px(s["reader_end"]) - 12, 282, f"B's cache crosses: {ms(s['hand_over'])} ms", "end"))
    body.append(text(x0, 330, "Split: writer GPU", bold=True))
    body.append(text(x0 + 500, 330, f"A: a word every {ms(sim.STEP_S)} ms", "end"))
    body += [circle(px(t), 358) for t in s["split_words"] if t <= window]
    body.append(triangle(px(s["landed"]), 388))
    body.append(f'<path d="M{x0} 412H{x0 + 500}" stroke="{GREY}" stroke-width="2"/>')
    for t in (0, 0.05, 0.10, 0.15, 0.20):
        body.append(tick(px(t), 440, f"{round(t * 1000)}"))
    body.append(text(x0 + 500, 470, "time, in ms", "end", GREY))
    body.append(circle(x0 + 8, 463, 7))
    body.append(text(x0 + 24, 470, "A's word"))
    body.append(triangle(x0 + 150, 463))
    body.append(text(x0 + 168, 470, "B's first word"))
    desc = (f"Time in milliseconds on four rows. One shared GPU: reply A gets a word every {ms(sim.STEP_S)} ms until B's {PROMPT}-token prompt arrives at {ms(s['arrival'])} ms. "
            f"The GPU then reads B from {ms(s['read_start'])} to {ms(s['read_end'])} ms ({ms(s['read_end'] - s['read_start'])} ms), and A gets no word for {ms(after - before)} ms. "
            f"Split: the reader GPU reads B from {ms(s['reader_start'])} to {ms(s['reader_end'])} ms, B's cache crosses the link in {ms(s['hand_over'])} ms and lands at {ms(s['landed'])} ms, "
            f"and the writer GPU gives A a word every {ms(sim.STEP_S)} ms throughout.")
    return frame(540, 486, f"A {PROMPT}-token prompt stops a shared GPU's replies for {ms(after - before)} ms, and does not stop a writer GPU", desc, body)


# --- figures 2 to 4: lines against a setting -----------------------------------------------------------------------------------------

def axes(x0: float, x1: float, top: float, base: float, yticks: list[tuple[float, str]], header: str, xlabel: str) -> list[str]:
    body = [text(20, 36, header, color=GREY)]
    for y, label in yticks:
        body.append(f'<path d="M{x0} {num(y)}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(tick(x0 - 8, y + 7, label, "end"))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    body.append(text((x0 + x1) / 2, base + 64, xlabel, "middle", GREY))
    return body


def fig_gap_by_load() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    px = lambda r: x0 + PAD + (r - RPS_GRID[0]) / (RPS_GRID[-1] - RPS_GRID[0]) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v / 0.25 * (base - top)
    got = {r: runs(r, PROMPT, REPLY, LINK) for r in RPS_GRID}
    gap = lambda layout: [got[r][layout]["gap_p99"] for r in RPS_GRID]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.1), "100"), (py(0.2), "200")], "p99 gap between words, in ms", "requests per second")
    body.append(f'<path d="M{x0} {num(py(sim.WORD_GAP_TARGET_S))}H{x1}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    for r in RPS_GRID:
        body.append(tick(px(r), base + 30, str(r)))
    body.append(path([(px(r), py(g)) for r, g in zip(RPS_GRID, gap("shared"))], ORANGE))
    body.append(path([(px(r), py(g)) for r, g in zip(RPS_GRID, gap("split"))], GREEN, DASH))
    body += [circle(px(r), py(g), 7, ORANGE) for r, g in zip(RPS_GRID, gap("shared"))]
    body += [square(px(r), py(g)) for r, g in zip(RPS_GRID, gap("split"))]
    body += key(x0 + 40, top + 12, "shared GPUs", "shared") + key(x0 + 40, top + 40, "split", "split") + key(x0 + 40, top + 68, "100 ms limit", "rule")
    at = lambda layout, r: ms(got[r][layout]["gap_p99"])
    over = min(r for r in RPS_GRID if got[r]["shared"]["gap_p99"] > sim.WORD_GAP_TARGET_S)
    desc = (f"p99 gap between words in ms against requests per second, from 1 to 10, with {PROMPT}-token prompts and {REPLY}-word replies (before the mix of lengths). "
            f"Shared GPUs (orange line, filled circles): {at('shared', 1)} at 1, {at('shared', 3)} at 3, {at('shared', 4)} at 4, {at('shared', 6)} at 6 and {at('shared', 10)} at 10. "
            f"Split (green dashed line, hollow squares): {at('split', 1)} at 1 and {at('split', 10)} at 10. A grey dashed line marks the 100 ms limit.")
    return frame(540, base + 84, f"On shared GPUs the p99 gap passes {sim.WORD_GAP_TARGET_S * 1000:.0f} ms at {over} requests a second; on a split pair it stays at {at('split', over)} ms", desc, body)


def fig_goodput_by_link() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    low, high = math.log10(LINK_GRID[0]), math.log10(LINK_GRID[-1])
    px = lambda g: x0 + PAD + (math.log10(g) - low) / (high - low) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v * (base - top)
    got = {g: runs(RPS, PROMPT, REPLY, g) for g in LINK_GRID}
    cut = crossover_gbps()
    shared_goodput = got[LINK]["shared"]["goodput"]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.5), "50%"), (py(1), "100%")], "goodput: requests that meet both targets", "link speed, in GB/s (log scale)")
    for g in (0.1, 1, 10, 100):
        body.append(tick(px(g), base + 30, str(g if g < 1 else int(g))))
    body.append(f'<path d="M{num(px(cut))} {top}V{base}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    body.append(path([(x0, py(shared_goodput)), (x1, py(shared_goodput))], ORANGE))
    body.append(path([(px(g), py(got[g]["split"]["goodput"])) for g in LINK_GRID], GREEN, DASH))
    body += [square(px(g), py(got[g]["split"]["goodput"])) for g in sim.LINKS_GBPS]
    body += key(x0 + 190, top + 80, "split", "split") + key(x0 + 190, top + 114, f"shared GPUs, {shared_goodput:.1%}", "shared") + key(x0 + 190, top + 148, "crossover,", "rule")
    body.append(text(x0 + 220, top + 174, f"about {cut:.1f} GB/s"))                                         # the rule's label runs to a second line so it fits the figure's width
    at = lambda g: f"{got[g]['split']['goodput']:.1%}"
    desc = (f"Goodput against link speed from 0.1 to 100 GB/s, on a log scale, at {RPS} requests a second with {PROMPT}-token prompts and {REPLY}-word replies. "
            f"Split (green dashed line, hollow squares at the five speeds on offer): {at(100)} at 100, {at(25)} at 25, {at(12.5)} at 12.5, {at(1.25)} at 1.25 and {at(0.125)} at 0.125 GB/s. "
            f"Shared GPUs (orange line) stay at {shared_goodput:.1%}. A grey dashed line marks the crossover at about {cut:.1f} GB/s: below it no request meets the 0.5 s target for its first word.")
    return frame(540, base + 84, f"A split pair beats shared GPUs on goodput on links faster than about {cut:.1f} GB/s, and loses below it", desc, body)


def fig_goodput_by_prompt() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    low, high = math.log10(PROMPT_GRID[0]), math.log10(PROMPT_GRID[-1])
    px = lambda p: x0 + PAD + (math.log10(p) - low) / (high - low) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v * (base - top)
    got = {p: runs(RPS, p, REPLY, LINK) for p in PROMPT_GRID}
    series = lambda layout: [got[p][layout]["goodput"] for p in PROMPT_GRID]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.5), "50%"), (py(1), "100%")], "goodput: requests that meet both targets", "tokens in a typical prompt (log scale)")
    for p in (100, 200, 500, 1000, 2000, 5000):
        body.append(tick(px(p), base + 30, f"{p:,}"))
    body.append(path([(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("shared"))], ORANGE))
    body.append(path([(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("split"))], GREEN, DASH))
    body += [circle(px(p), py(v), 7, ORANGE) for p, v in zip(PROMPT_GRID, series("shared"))]
    body += [square(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("split"))]
    body += key(x0 + 30, top + 160, "shared GPUs", "shared") + key(x0 + 30, top + 194, "split", "split")
    at = lambda layout, p: f"{got[p][layout]['goodput']:.1%}"
    desc = (f"Goodput against the typical prompt length, from {PROMPT_GRID[0]} to {PROMPT_GRID[-1]} tokens on a log scale, at {RPS} requests a second with {REPLY}-word replies and a {LINK} GB/s link. "
            f"Shared GPUs (orange line, filled circles): {at('shared', 200)} at 200, {at('shared', 1000)} at 1,000, {at('shared', 2000)} at 2,000 and {at('shared', 4000)} at 4,000. "
            f"Split (green dashed line, hollow squares): {at('split', 200)} at 200, {at('split', 1000)} at 1,000, {at('split', 2000)} at 2,000, {at('split', 4000)} at 4,000 and {at('split', 6000)} at 6,000.")
    tie = max(p for p in PROMPT_GRID if all(abs(got[q]["split"]["goodput"] - got[q]["shared"]["goodput"]) < 0.01 for q in PROMPT_GRID if q <= p))     # the layouts tie up to here
    return frame(540, base + 84, f"The two layouts nearly tie on prompts up to {tie} tokens, and the split pulls ahead after that", desc, body)


FIGURES = {"stall.svg": fig_stall, "gap-by-load.svg": fig_gap_by_load, "goodput-by-link.svg": fig_goodput_by_link, "goodput-by-prompt.svg": fig_goodput_by_prompt}


def main() -> int:
    stale = []
    for name, draw in FIGURES.items():
        target, svg = FIGS / name, draw()
        if "--check" in sys.argv:
            if not target.exists() or target.read_text(encoding="utf-8") != svg:
                stale.append(name)
        else:
            FIGS.mkdir(exist_ok=True)
            target.write_text(svg, encoding="utf-8")
    if "--check" in sys.argv:
        print("figures differ from pd_sim.py: " + ", ".join(stale) if stale else "figures match pd_sim.py")
        return 1 if stale else 0
    print("wrote " + ", ".join(FIGURES))
    return 0


if __name__ == "__main__":
    sys.exit(main())

pd_sim.py127 linesDownload

Show code

"""A toy simulation of two GPUs that serve chats, laid out two ways: both GPUs do everything ("shared"), or one reads prompts and the other writes replies ("split").

Every cost is a toy value. Two of them echo lesson 23 (why-ai-waits-before-the-first-word): the time to read one prompt token and the time of one writing step are its
best-case numbers for Llama 3.1 8B with 16-bit weights on an NVIDIA A100 80GB SXM (the chip of lesson 19). The cache bytes per token echo lesson 22 (paged-kv-cache). The link
speeds, the slot count and the two targets are invented. Nothing here is a measurement, and test_pd_sim.py checks the copied values against lesson 23's code.

Time moves in whole steps. A GPU that has a waiting prompt and a free slot reads that whole prompt first (nothing else runs; there is no chunking) and the prompt's first word
comes out at the end. Otherwise it runs one writing step: every reply in progress gets one more word. A new reply joins at the start of a step. Python standard library only.
The widget (../widget.html) shows a table that show.py builds from this file, and test_widget.py checks the page against it.
"""
from __future__ import annotations

import random
from itertools import takewhile
from typing import NamedTuple

# --- toy values: the chip, copied from lesson 23 (prefill_sim.py) and lesson 19 (tiled-matmul-roofline)
PARAMS = 8_030_261_248                                  # Llama 3.1 8B, lesson 23
A100_PEAK_FLOPS, A100_BANDWIDTH = 312e12, 2039e9        # A100 80GB SXM data sheet, lesson 19: 312 TFLOPS, 2,039 GB/s (GB is 10**9 bytes)
READ_S_PER_TOKEN = 2 * PARAMS / A100_PEAK_FLOPS         # reading the prompt: 2 FLOPs per weight per token, limited by arithmetic (lesson 23)
STEP_S = 2 * PARAMS / A100_BANDWIDTH                    # one writing step: every weight is read once, limited by memory, the same for any number of replies (lesson 23)
CACHE_BYTES_PER_TOKEN = 131_072                         # 2 x 32 layers x 8 key-value heads x 128 x 2 bytes, lesson 22
# --- toy values: invented
SLOTS = 32                                              # replies one GPU writes at once
LINKS_GBPS = (100, 25, 12.5, 1.25, 0.125)               # link speeds offered, in 10**9 bytes per second (800, 200, 100, 10 and 1 gigabits per second)
FIRST_WORD_TARGET_S, WORD_GAP_TARGET_S = 0.5, 0.1       # a request is good if its first word comes within 0.5 s and no gap between words is over 0.1 s
MIX = (0.5, 1.0, 1.5)                                   # each prompt and each reply is its stated length times one of these, chosen at random
N_REQUESTS, SEED = 600, 7


class Request(NamedTuple):
    arrival: float      # seconds since the first request
    prompt: int         # tokens in the prompt
    reply: int          # words in the reply, at least 1


def make_requests(rps: float, prompt: int, reply: int, n: int = N_REQUESTS, seed: int = SEED) -> list[Request]:
    """Step 1: seeded Poisson arrivals (gaps between requests are exponential, mean 1 / rps)."""
    rng, t, out = random.Random(seed), 0.0, []
    for _ in range(n):
        t += rng.expovariate(rps)
        out.append(Request(t, round(prompt * rng.choice(MIX)), max(1, round(reply * rng.choice(MIX)))))
    return out


def serve(jobs: list[tuple]) -> tuple[dict, dict]:
    """One GPU. A job is (ready_time, read_seconds, words_left_to_write, id). Returns {id: time of first word} and {id: [gaps between its words]}.
    A job with read_seconds 0 arrives with its cache and first word already made (the split's writing GPU); the others read their prompt here."""
    waiting, active, first, gaps, last, t = sorted(jobs), [], {}, {}, {}, 0.0
    while waiting or active:
        if not active and waiting[0][0] > t:
            t = waiting[0][0]                                         # an idle GPU waits for the next job
        ready = list(takewhile(lambda j: j[0] <= t, waiting))         # `waiting` is sorted, so the jobs that are ready are its first few
        if ready and len(active) < SLOTS and ready[0][1] > 0:         # step 2: a prompt waits and a slot is free, so read the whole prompt now
            ready_t, read, left, i = ready[0]
            waiting.remove(ready[0])
            t += read
            first[i], gaps[i], last[i] = t, [], t
            if left > 0:
                active.append([i, left])
            continue
        for j in ready[: SLOTS - len(active)]:                        # a job that comes with its cache joins at the start of a step
            if j[1] == 0:
                waiting.remove(j)
                first[j[3]], gaps[j[3]], last[j[3]] = j[0], [], j[0]
                if j[2] > 0:
                    active.append([j[3], j[2]])
        if not active:
            continue
        t += STEP_S                                                   # one writing step: every reply in progress gets one word
        for a in active:
            gaps[a[0]].append(t - last[a[0]])
            last[a[0]], a[1] = t, a[1] - 1
        active = [a for a in active if a[1] > 0]
    return first, gaps


def handover_seconds(prompt_tokens: int, link_gbps: float) -> float:
    """Step 5: the prompt's cache is prompt tokens x bytes per token, and it crosses the link at the link's speed."""
    return prompt_tokens * CACHE_BYTES_PER_TOKEN / (link_gbps * 1e9)


def shared(requests: list[Request]) -> tuple[dict, dict]:
    """Step 2: two GPUs that each read prompts and write replies. Requests take turns: 0 goes to the first GPU, 1 to the second, 2 to the first."""
    first, gaps = {}, {}
    for g in (0, 1):
        f, w = serve([(r.arrival, r.prompt * READ_S_PER_TOKEN, r.reply - 1, i) for i, r in enumerate(requests) if i % 2 == g])
        first.update(f)
        gaps.update(w)
    return first, gaps


def split(requests: list[Request], link_gbps: float) -> tuple[dict, dict]:
    """Step 4: one GPU only reads prompts, one at a time in order; each prompt's cache crosses the link (step 5); the other GPU only writes. The first word reaches the
    reader of the reply when the cache lands on the writing GPU. The link carries any number of caches at once (no contention)."""
    free, jobs = 0.0, []
    for i, r in enumerate(requests):
        free = max(free, r.arrival) + r.prompt * READ_S_PER_TOKEN
        jobs.append((free + handover_seconds(r.prompt, link_gbps), 0.0, r.reply - 1, i))
    return serve(jobs)


def percentile(values: list[float], p: int) -> float:
    """Nearest rank, as in lesson 27 (why-average-speed-is-misleading)."""
    ordered = sorted(values)
    return ordered[-(-p * len(ordered) // 100) - 1]


def goodput(requests: list[Request], run: tuple[dict, dict], first_target: float = FIRST_WORD_TARGET_S, gap_target: float = WORD_GAP_TARGET_S) -> float:
    """Step 6: the share of requests whose first word came within `first_target` seconds and whose every gap between words was at most `gap_target`."""
    first, gaps = run
    good = sum(1 for i, r in enumerate(requests) if first[i] - r.arrival <= first_target and max(gaps[i], default=0.0) <= gap_target)
    return good / len(requests)


def measure(requests: list[Request], run: tuple[dict, dict]) -> dict:
    """Steps 3 and 6: time to first word per request, the gaps between words pooled over all requests, and goodput at the default targets."""
    first, gaps = run
    waits = [first[i] - r.arrival for i, r in enumerate(requests)]
    pooled = [g for i in range(len(requests)) for g in gaps[i]]
    return {"first_p50": percentile(waits, 50), "first_p99": percentile(waits, 99), "gap_p50": percentile(pooled, 50), "gap_p99": percentile(pooled, 99), "goodput": goodput(requests, run)}


def compare(rps: float, prompt: int, reply: int, link_gbps: float) -> dict:
    """The same requests through both layouts."""
    requests = make_requests(rps, prompt, reply)
    return {"shared": measure(requests, shared(requests)), "split": measure(requests, split(requests, link_gbps))}

show.py67 linesDownload

Show code

"""What the widget (../widget.html) shows, worked out from pd_sim.py. Run it from this folder: python3 show.py

The widget cannot run Python, so this file runs the toy once for every setting its controls offer, keeps how many of the 600 requests were good and the slowest 1 in 100 times,
and embed_widget.py writes those numbers into widget.html. The words the page shows are built here and again in the page; test_widget.py checks the two agree for every setting.
A control is an index into one of the tuples below (the first stop is 0).
"""
from __future__ import annotations

from functools import cache

import pd_sim as sim

RPS = (2, 4, 10)                            # requests a second
PROMPTS = (200, 1000, 2000)                 # tokens in a typical prompt
REPLIES = (100, 200, 400)                   # words in a typical reply
LINKS = (0.25, 0.5, 1.25, 2.5, 12.5, 25)    # GB/s, 10**9 bytes per second (2, 4, 10, 20, 100 and 200 gigabits per second)
FIRST_TARGETS = (0.25, 0.5, 1.0)            # seconds to the first word
GAP_TARGETS = (0.05, 0.1, 0.2)              # seconds between two words
OPENING = {"rps": 1, "prompt": 2, "link": 5, "reply": 1, "first": 1, "gap": 1}    # 4 a second, 2000 tokens, 25 GB/s, 200 words, 0.5 s, 0.1 s: the toy's default
N = sim.N_REQUESTS


def summarize(requests: list[sim.Request], run: tuple[dict, dict]) -> dict:
    """For one layout: the requests that were good under each pair of targets (first word target varies slowest), then the slowest 1 in 100 first word and gap, in ms."""
    first, gaps = run
    waits = [first[i] - r.arrival for i, r in enumerate(requests)]
    longest = [max(gaps[i], default=0.0) for i in range(len(requests))]
    pooled = [g for i in range(len(requests)) for g in gaps[i]]
    return {"good": [sum(1 for w, g in zip(waits, longest) if w <= f and g <= t) for f in FIRST_TARGETS for t in GAP_TARGETS],
            "first": round(sim.percentile(waits, 99) * 1000, 1), "gap": round(sim.percentile(pooled, 99) * 1000, 1)}


@cache
def table() -> list:
    """cells[rps][prompt][reply] = {"shared": summary, "split": [summary for each link speed]}."""
    cells = []
    for rps in RPS:
        row = []
        for prompt in PROMPTS:
            column = []
            for reply in REPLIES:
                requests = sim.make_requests(rps, prompt, reply)
                column.append({"shared": summarize(requests, sim.shared(requests)), "split": [summarize(requests, sim.split(requests, link)) for link in LINKS]})
            row.append(column)
        cells.append(row)
    return cells


def widget_config() -> dict:
    """The block of numbers embedded in widget.html."""
    return {"n": N, "rps": RPS, "prompt": PROMPTS, "reply": REPLIES, "link": LINKS, "first": FIRST_TARGETS, "gap": GAP_TARGETS, "opening": OPENING, "cells": table()}


def describe(rps: int, prompt: int, link: int, reply: int, first: int, gap: int) -> dict:
    """The words for one setting and the counts the chart draws: the shared layout's good requests, and the split's at every link speed."""
    cell, k = table()[rps][prompt][reply], first * len(GAP_TARGETS) + gap
    a, b, pct = cell["shared"], cell["split"][link], lambda count: f"{count / N * 100:.1f}%"
    return {"head": f"Both layouts get {pct(a['good'][k])} of requests right." if a["good"][k] == b["good"][k]
            else f"Shared GPUs get {pct(a['good'][k])} of requests right and the split gets {pct(b['good'][k])}.",
            "rule": f"Right means the first word within {FIRST_TARGETS[first]:g} s and no gap between words over {GAP_TARGETS[gap]:g} s.",
            "tails": f"Slowest 1 in 100: first word {a['first']:.1f} ms shared, {b['first']:.1f} ms split; gap between words {a['gap']:.1f} ms shared, {b['gap']:.1f} ms split.",
            "shared": a["good"][k], "split": [s["good"][k] for s in cell["split"]]}


if __name__ == "__main__":
    shown = describe(*(OPENING[name] for name in ("rps", "prompt", "link", "reply", "first", "gap")))
    print(shown["head"], shown["rule"], shown["tails"], sep="\n")

trace.py34 linesDownload

Show code

"""Follow one run of the toy through the lesson's six steps. Run it from this folder: python3 trace.py

It uses pd_sim.py's own functions and adds no rules of its own. It prints one line for each step, then one line for the limit: how the split does at each link speed on offer.
"""
from pd_sim import CACHE_BYTES_PER_TOKEN, LINKS_GBPS, READ_S_PER_TOKEN, STEP_S, compare, handover_seconds, make_requests, measure, shared, split

RPS, PROMPT, REPLY = 4, 2000, 200   # requests per second, tokens in a typical prompt, words in a typical reply
LINK_GBPS = 25      # the link between the two GPUs, in 10**9 bytes per second

ms = lambda seconds: f"{seconds * 1000:.1f}"

requests = make_requests(RPS, PROMPT, REPLY)
print(f"step 1  {len(requests)} requests in {requests[-1].arrival:.1f} s, {RPS} a second; prompts of {', '.join(str(p) for p in sorted({r.prompt for r in requests}))} tokens, "
      f"replies of {', '.join(str(w) for w in sorted({r.reply for r in requests}))} words")

read = PROMPT * READ_S_PER_TOKEN
print(f"step 2  shared: a {PROMPT}-token prompt takes {ms(read)} ms to read and no reply on that GPU moves meanwhile; a word due {ms(STEP_S)} ms after the last comes {ms(read + STEP_S)} ms after it")

both = measure(requests, shared(requests))
print(f"step 3  shared: first word p50 {ms(both['first_p50'])} ms, p99 {ms(both['first_p99'])} ms; gap between words p50 {ms(both['gap_p50'])} ms, p99 {ms(both['gap_p99'])} ms")

load = sum(r.prompt for r in requests) * READ_S_PER_TOKEN / requests[-1].arrival
print(f"step 4  split: the reader GPU has {load:.2f} s of prompts to read each second; the writer GPU only writes, {ms(STEP_S)} ms a step")

cache = PROMPT * CACHE_BYTES_PER_TOKEN
print(f"step 5  hand-over: {PROMPT} tokens x {CACHE_BYTES_PER_TOKEN} bytes = {cache / 1e6:.1f} MB; at {LINK_GBPS} GB/s that takes {ms(handover_seconds(PROMPT, LINK_GBPS))} ms, "
      f"against {ms(read)} ms to read (equal at {CACHE_BYTES_PER_TOKEN / READ_S_PER_TOKEN / 1e9:.2f} GB/s)")

apart = measure(requests, split(requests, LINK_GBPS))
print(f"step 6  split: first word p50 {ms(apart['first_p50'])} ms, p99 {ms(apart['first_p99'])} ms; gap p50 {ms(apart['gap_p50'])} ms, p99 {ms(apart['gap_p99'])} ms; "
      f"goodput {apart['goodput']:.1%} against {both['goodput']:.1%} shared")

by_link = ", ".join(f"{link} {compare(RPS, PROMPT, REPLY, link)['split']['goodput']:.1%}" for link in LINKS_GBPS)
print(f"limit   goodput of the split by link speed, GB/s: {by_link}")

test_lesson.py163 linesDownload

Show code

"""Tests that tie lesson.md, trace.py and the four figures to pd_sim.py. Run from this folder: python3 -B -m pytest -q -p no:cacheprovider"""
import io
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
from contextlib import redirect_stdout
from functools import lru_cache
from pathlib import Path

import make_figs
import pd_sim as sim

HERE = Path(__file__).resolve().parent
LESSON = HERE.parent
SOURCE = (HERE / "pd_sim.py").read_text(encoding="utf-8")
RPS, PROMPT, REPLY, LINK = make_figs.RPS, make_figs.PROMPT, make_figs.REPLY, make_figs.LINK
ms = lambda seconds: f"{seconds * 1000:.1f}"
pct = lambda share: f"{share:.1%}"
READ = PROMPT * sim.READ_S_PER_TOKEN


@lru_cache(maxsize=None)
def reading():
    prose = (LESSON / "lesson.md").read_text(encoding="utf-8")
    return prose, {name: (LESSON / "figs" / name).read_text(encoding="utf-8") for name in make_figs.FIGURES}


def test_every_number_the_lesson_and_its_figures_quote_comes_from_the_code():
    """Guards prose, trace and figure numbers drifting from the code: each fragment is built from a value computed here and must appear in the lesson or the figure.
    The numbers quoted from the papers and the pull request are not computed here."""
    prose, figures = reading()
    run = make_figs.runs(RPS, PROMPT, REPLY, LINK)
    both, apart = run["shared"], run["split"]
    at = lambda rps=RPS, prompt=PROMPT, link=LINK: make_figs.runs(rps, prompt, REPLY, link)
    requests = sim.make_requests(RPS, PROMPT, REPLY)
    _, gaps = sim.shared(requests)
    stalled = sum(1 for i in range(len(requests)) if max(gaps[i], default=0.0) > sim.WORD_GAP_TARGET_S) / len(requests)
    load = sum(r.prompt for r in requests) * sim.READ_S_PER_TOKEN / requests[-1].arrival
    cache = PROMPT * sim.CACHE_BYTES_PER_TOKEN
    break_even = sim.CACHE_BYTES_PER_TOKEN / sim.READ_S_PER_TOKEN / 1e9
    stall = READ + sim.STEP_S
    cross = make_figs.crossover_gbps()
    shortest = round(PROMPT * min(sim.MIX))
    over = next(rps for rps in make_figs.RPS_GRID if at(rps)["shared"]["gap_p99"] > 0.1)
    slow = at(link=0.25)["split"]
    scene = make_figs.stall_scene()
    wanted = {"lesson.md": [
        f"a prompt of {PROMPT} tokens", f"takes {ms(READ)} ms to read", f"a word due {ms(sim.STEP_S)} ms after the last comes {ms(stall)} ms after it",
        f"At {RPS} requests a second only {pct(both['goodput'])} of requests", f"within {sim.FIRST_WORD_TARGET_S} s and no later gap over {sim.WORD_GAP_TARGET_S} s",
        f"and {pct(apart['goodput'])[:-1]}% do", f"The price is that the reading GPU's notes on the prompt (the cache, {cache / 1e6:.0f} MB for {PROMPT} tokens) must cross the link between the two GPUs",
        f"{RPS} requests a second, prompts of {PROMPT} tokens, replies of {REPLY} words and a {LINK} GB/s link", f"shared GPUs get {pct(both['goodput'])} of requests right and the split {pct(apart['goodput'])}",
        f"Drag Requests a second to 10: shared falls to {pct(at(10)['shared']['goodput'])} and the split to {pct(at(10)['split']['goodput'])}",
        f"Drag Link speed to 0.25 GB/s: the split falls to {pct(slow['goodput'])}", f"Drag Prompt tokens to 200: the two tie at {pct(at(prompt=200)['split']['goodput'])}",
        f"serves {sim.N_REQUESTS} requests", f"600 requests come at {RPS} a second",
        f"prompts of {round(PROMPT * min(sim.MIX))} to {round(PROMPT * max(sim.MIX))} tokens and replies of {round(REPLY * min(sim.MIX))} to {round(REPLY * max(sim.MIX))} words",
        f"Half of all gaps between words are {ms(both['gap_p50'])} ms (the p50), but the slowest 1 in 100 (the p99, the 99th percentile) is {ms(both['gap_p99'])} ms",
        f"Only {pct(both['goodput'])} of requests are good", f"with {load:.2f} s of reading to do each second",
        f"{PROMPT} x {sim.CACHE_BYTES_PER_TOKEN:,} bytes = {cache / 1e6:.1f} MB", f"At {LINK} GB/s that takes {ms(sim.handover_seconds(PROMPT, LINK))} ms, against {ms(READ)} ms to read",
        f"no gap passes {ms(apart['gap_p99'])} ms and {pct(apart['goodput'])} of requests are good, against {pct(both['goodput'])}",
        f"p50 {ms(apart['first_p50'])} ms against {ms(both['first_p50'])} ms", f"The word waits {ms(stall)} ms, {round(stall / sim.STEP_S)} times the usual {ms(sim.STEP_S)}",
        f"A 200-word reply lasts about {REPLY * sim.STEP_S:.1f} s", f"so {pct(stalled)} of replies have a gap over 0.1 s and {pct(both['goodput'])} are good",
        f"passes 100 ms at {over} requests a second and is {ms(at(10)['shared']['gap_p99'])} ms at 10", f"One reader can read about {1 / READ:.1f} prompts a second (1 / {READ:.3f} s)",
        f"from {at(1)['split']['first_p99']:.2f} s at 1 request a second to {at(10)['split']['first_p99']:.2f} s at 10", f"equals the read at {break_even:.2f} GB/s for any prompt length",
        f"p99 {ms(apart['first_p99'])} ms against {ms(both['first_p99'])} ms", f"{pct(at(link=1.25)['split']['goodput'])} at 1.25 GB/s and {pct(at(link=0.125)['split']['goodput'])} at 0.125",
        f"on links faster than about {cross:.1f} GB/s", f"the shortest prompt, {shortest} tokens, has a {shortest * sim.CACHE_BYTES_PER_TOKEN / 1e6:.0f} MB cache that cannot be read and cross within {sim.FIRST_WORD_TARGET_S} s",
        "change line 8 to `LINK_GBPS = 0.25`", f"The hand-over takes {ms(sim.handover_seconds(PROMPT, 0.25))} ms, about ten times the read",
        f"Even the shortest prompt, {shortest} tokens, needs {sim.handover_seconds(shortest, 0.25) * 1000:.0f} ms to cross",
        f"The slowest 1 in 100 is {ms(both['gap_p99'])} ms shared and {ms(apart['gap_p99'])} ms split", f"at {PROMPT} tokens the split leads {pct(apart['goodput'])} to {pct(both['goodput'])}",
        f"the split stops paying below about {cross:.1f} GB/s", f"(0.2922 in the code: at 0.29 no request is good and at 0.293, {pct(at(link=0.293)['split']['goodput'])} are)",
        f"{ms(shortest * sim.READ_S_PER_TOKEN)} ms to read plus {ms(sim.handover_seconds(shortest, cross))} ms to cross fills the {sim.FIRST_WORD_TARGET_S} s",
        f"`pd_sim.py` is {len(SOURCE.splitlines())} lines", f"{sim.READ_S_PER_TOKEN * 1e6:.1f} microseconds to read a prompt token and {sim.STEP_S * 1000:.3f} ms for a writing step",
        f"The {sim.CACHE_BYTES_PER_TOKEN:,} cache bytes per token", f"the {sim.SLOTS} slots", f"For a {PROMPT}-token prompt, {cache / 1e6:.1f} MB crosses",
        f"at 2 GB/s the split is good for {pct(at(link=2)['split']['goodput'])} of requests", f"at 200 tokens the read takes {ms(200 * sim.READ_S_PER_TOKEN)} ms, a stall is about {round((200 * sim.READ_S_PER_TOKEN + sim.STEP_S) * 1000)} ms",
        f"6000-token prompts need {6000 * sim.READ_S_PER_TOKEN:.2f} s of reading each", f"(the split's first-word p99 is {at(prompt=6000)['split']['first_p99']:.1f} s)",
        f"takes {sim.handover_seconds(PROMPT, 2) * 1000:.0f} ms, longer than the {READ * 1000:.0f} ms read", f"{pct(at(link=2)['split']['goodput'])} of requests at 2 GB/s, against {pct(both['goodput'])} shared",
        f"{pct(stalled)} of the shared requests have one", f"A 200-token prompt takes {ms(200 * sim.READ_S_PER_TOKEN)} ms to read", f"about {round((200 * sim.READ_S_PER_TOKEN + sim.STEP_S) * 1000)} ms",
        f"first-word p99 reaches {at(10)['split']['first_p99']:.2f} s", f"An average {PROMPT}-token prompt takes {READ:.3f} s to read, so one reader handles about {1 / READ:.1f} prompts a second",
    ], "stall.svg": [
        f"{ms(stall)} ms with no word from A", f"reads B: {ms(READ)} ms",
        f"B's cache crosses: {ms(sim.handover_seconds(PROMPT, LINK))} ms", f"A: a word every {ms(sim.STEP_S)} ms",
    ], "gap-by-load.svg": [f"{ms(at(rps)['shared']['gap_p99'])} at {rps}" for rps in (1, 3, 4, 6, 10)] + [f"{ms(at(10)['split']['gap_p99'])} at 10"],
        "goodput-by-link.svg": [f"{pct(at(link=link)['split']['goodput'])} at {link}" for link in sim.LINKS_GBPS] + [pct(both["goodput"]), f"crossover at about {cross:.1f} GB/s", f"about {cross:.1f} GB/s</text>"],
        "goodput-by-prompt.svg": [f"{pct(at(prompt=p)['shared']['goodput'])} at {p:,}" for p in (200, 1000, 2000, 4000)] + [f"{pct(at(prompt=p)['split']['goodput'])} at {p:,}" for p in (200, 1000, 2000, 4000, 6000)],
    }
    texts = {"lesson.md": prose, **figures}
    missing = [f"{name}: {fragment!r}" for name, fragments in wanted.items() for fragment in fragments if fragment not in texts[name]]
    assert not missing, "not found: " + "; ".join(missing)
    assert len(SOURCE.splitlines()) < 150 and f"{cross:.4f}" == "0.2922" and at(link=0.29)["split"]["goodput"] == 0.0 and stalled == 1 - both["goodput"]


def test_the_lessons_claims_about_where_the_split_pays_hold():
    """Guards the shape claims: the layouts tie up to 600 tokens and the split leads from 800 to 4000 and both collapse at 6000; a link under the crossover leaves no request
    able to meet the first-word target even for the shortest prompt, and one just over it wins; the split loses the first word at every link speed on offer."""
    for prompt in (100, 200, 400, 600):
        c = make_figs.runs(RPS, prompt, REPLY, LINK)
        assert abs(c["split"]["goodput"] - c["shared"]["goodput"]) < 0.01, prompt
    for prompt in (800, 1000, 1500, 2000, 3000, 4000):
        c = make_figs.runs(RPS, prompt, REPLY, LINK)
        assert c["split"]["goodput"] - c["shared"]["goodput"] > 0.02, prompt
    both = make_figs.runs(RPS, 6000, REPLY, LINK)
    assert max(both["shared"]["goodput"], both["split"]["goodput"]) < 0.05
    cross, shortest = make_figs.crossover_gbps(), round(PROMPT * min(sim.MIX))
    assert abs(shortest * sim.READ_S_PER_TOKEN + sim.handover_seconds(shortest, cross) - sim.FIRST_WORD_TARGET_S) < 0.002
    assert make_figs.runs(RPS, PROMPT, REPLY, cross * 0.97)["split"]["goodput"] == 0.0
    assert make_figs.runs(RPS, PROMPT, REPLY, cross * 1.03)["split"]["goodput"] > make_figs.runs(RPS, PROMPT, REPLY, LINK)["shared"]["goodput"]
    assert all(make_figs.runs(RPS, PROMPT, REPLY, link)["split"]["first_p50"] > make_figs.runs(RPS, PROMPT, REPLY, link)["shared"]["first_p50"] for link in sim.LINKS_GBPS)


def test_the_figures_are_the_ones_the_code_draws_and_follow_the_house_style():
    """Guards a figure going stale or off-palette: the files on disk equal what make_figs draws, the lesson shows each one, and each parses, has a title and a description,
    a white ground, no script, no web address, only the house colours and no text under 20 units."""
    prose, figures = reading()
    assert set(figures) == set(re.findall(r"\]\(figs/([\w-]+\.svg)\)", prose)) == set(make_figs.FIGURES)
    allowed = {make_figs.INK, make_figs.GREY, make_figs.ORANGE, make_figs.GREEN, make_figs.FILL, make_figs.PANEL, "#ffffff"}
    web_address = re.compile("ht" + r"tps?://")
    for name, svg in figures.items():
        assert svg == make_figs.FIGURES[name](), name
        root = ET.fromstring(svg)
        assert root.find("{%s}title" % make_figs.NAMESPACE) is not None and root.find("{%s}desc" % make_figs.NAMESPACE) is not None and 'fill="#ffffff"' in svg, name
        assert "<script" not in svg and not web_address.search(svg.replace(make_figs.NAMESPACE, "")) and root.get("width") == "540", name
        assert {c.lower() for c in re.findall(r"#[0-9a-fA-F]{6}\b", svg)} <= {c.lower() for c in allowed}, name
        assert all(int(size) >= 20 for size in re.findall(r'font-size="(\d+)"', svg)), name
    for path in HERE.glob("*.py"):
        assert not web_address.search(path.read_text(encoding="utf-8").replace("ht" + "tp://www.w3.org", "")), path.name


def table_rows(prose):
    return re.findall(r"^\| \d\. [^|]+\| `(\w+\.py):(\d+)`<br>`(.+?)` \|", prose, re.M)


def test_code_blocks_are_slices_of_the_real_file_and_the_table_points_at_its_lines():
    """Guards a `file:line` in the code table pointing at the wrong line, and a code block in Go deeper that no longer matches the file."""
    prose, _ = reading()
    blocks = re.findall(r"```python\n(.*?)\n```", prose, re.S)
    assert len(blocks) == 4 and all(block in SOURCE for block in blocks)
    rows = table_rows(prose)
    assert [(name, int(n)) for name, n, _ in rows] == [("pd_sim.py", 41), ("pd_sim.py", 57), ("pd_sim.py", 72), ("pd_sim.py", 98), ("pd_sim.py", 80), ("pd_sim.py", 112)]
    for name, number, code in rows:
        line = (HERE / name).read_text(encoding="utf-8").splitlines()[int(number) - 1]
        assert re.sub(r"\s+# .*$", "", line).strip() == code, (name, number)


def run_trace(source):
    out = io.StringIO()
    with redirect_stdout(out):
        exec(compile(source, "trace.py", "exec"), {"__name__": "__main__"})
    return out.getvalue()


def test_trace_is_the_block_in_the_lesson_and_the_your_turn_answer_is_right():
    """Guards the pasted run drifting from `python3 trace.py` (the real command, from the code folder) and the Your turn answer: changing line 8 to LINK_GBPS = 0.25 must
    give the lines in the details block."""
    prose, _ = reading()
    real = subprocess.run([sys.executable, "-B", "trace.py"], cwd=HERE, capture_output=True, text=True, check=True).stdout
    pasted = re.search(r"python3 trace\.py\n```\n\nPhone readers: slide the box sideways to see every column\.\n\n```text\n(.*?)\n```", prose, re.S).group(1)
    assert real == pasted + "\n"
    trace = (HERE / "trace.py").read_text(encoding="utf-8")
    assert trace.splitlines()[7].startswith("LINK_GBPS = 25 ") and "change line 8 to `LINK_GBPS = 0.25`" in prose
    changed = run_trace(trace.replace("LINK_GBPS = 25 ", "LINK_GBPS = 0.25 ", 1))
    answer = re.search(r"<summary>What you should see</summary>.*?```text\n(.*?)\n```", prose, re.S).group(1)
    assert answer + "\n" == changed and "goodput 0.0% against 13.5% shared" in answer
    assert [a == b for a, b in zip(real.splitlines(), changed.splitlines())] == [True, True, True, True, False, False, True]

test_pd_sim.py55 linesDownload

Show code

"""Guards for pd_sim.py: the stall shows on the shared layout and the split removes it, the hand-over costs bytes over link speed, the lesson's crossover reproduces,
and the copied toy values still equal lessons 22 and 23."""
import importlib.util
import sys
from pathlib import Path

import pytest

import pd_sim as sim

LESSON_23 = Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py"


def lesson_23():
    spec = importlib.util.spec_from_file_location("prefill_sim_23", LESSON_23)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module            # a dataclass in lesson 23's file looks its module up here
    spec.loader.exec_module(module)
    return module


def test_toy_values_still_equal_lessons_22_and_23():
    ps = lesson_23()
    assert (sim.PARAMS, sim.A100_PEAK_FLOPS, sim.A100_BANDWIDTH) == (ps.PARAMS, ps.A100_PEAK_FLOPS, ps.A100_BANDWIDTH)
    assert sim.CACHE_BYTES_PER_TOKEN == ps.KV_BYTES_PER_TOKEN == 131_072
    assert sim.READ_S_PER_TOKEN == pytest.approx(ps.prefill(1000).math_s / 1000) and sim.STEP_S == pytest.approx(ps.decode_step(1).memory_s)


def test_a_long_prompt_stalls_every_reply_on_a_shared_gpu():
    # Reply A (100-token prompt, 10 words) is mid-reply when B (2,000-token prompt) arrives: A's next word waits for B's whole read plus one step.
    jobs = [(0.0, 100 * sim.READ_S_PER_TOKEN, 9, 0), (0.03, 2000 * sim.READ_S_PER_TOKEN, 3, 1)]
    first, gaps = sim.serve(jobs)
    assert first[0] == pytest.approx(100 * sim.READ_S_PER_TOKEN) and max(gaps[0]) == pytest.approx(2000 * sim.READ_S_PER_TOKEN + sim.STEP_S)
    assert min(gaps[0]) == pytest.approx(sim.STEP_S)


def test_the_hand_over_is_cache_bytes_over_link_speed():
    assert sim.handover_seconds(2000, 25) == 2000 * 131_072 / 25e9
    one = [sim.Request(1.0, 2000, 4)]
    first, gaps = sim.split(one, 1.25)
    assert first[0] - 1.0 == pytest.approx(2000 * sim.READ_S_PER_TOKEN + sim.handover_seconds(2000, 1.25))      # first word reaches the reader with the cache
    assert gaps[0] == pytest.approx([sim.STEP_S] * 3)


def test_split_removes_the_stall_and_the_lessons_crossover_reproduces():
    requests = sim.make_requests(4, 2000, 200)
    shared = sim.measure(requests, sim.shared(requests))
    fast = sim.measure(requests, sim.split(requests, 25))
    assert shared["gap_p99"] > 0.1 > 10 * fast["gap_p99"] and fast["gap_p99"] <= 2 * sim.STEP_S and shared["gap_p50"] < 0.01      # p99 stalls, p50 does not
    assert fast["goodput"] > 0.9 > 0.2 > shared["goodput"] and fast["first_p99"] > shared["first_p99"]                           # the split wins goodput, loses first word
    slow = sim.measure(requests, sim.split(requests, 0.125))
    assert slow["goodput"] < shared["goodput"]                                                                                   # a slow enough link ends the win
    short = sim.make_requests(4, 200, 200)
    s, f = sim.measure(short, sim.shared(short)), sim.measure(short, sim.split(short, 25))
    assert s["goodput"] == f["goodput"] == 1.0 and f["first_p99"] >= s["first_p99"] * 0.7                                        # short prompts: a tie

test_widget.py185 linesDownload

Show code

"""Guards for ../widget.html: its table of numbers, that its words and chart equal show.py's for every setting, that the states the lesson names are reachable and right
against pd_sim.py itself, and what a reader sees and can press. Run from this folder: python3 -B -m pytest -q -p no:cacheprovider"""
import json
import re
import subprocess
from itertools import product

import pd_sim as sim
import show
from embed_widget import CONFIG_BLOCK, WIDGET, embed

HTML = WIDGET.read_text(encoding="utf-8")
NAMES = ("rps", "prompt", "link", "reply", "first", "gap")
STOPS = {"rps": show.RPS, "prompt": show.PROMPTS, "link": show.LINKS, "reply": show.REPLIES, "first": show.FIRST_TARGETS, "gap": show.GAP_TARGETS}
# Every setting the six sliders offer, as the stop each is on.
GRID = list(product(*(range(len(STOPS[name])) for name in NAMES)))

# Reads the model and the picture out of the widget between its marker comments, runs every setting with Node and compares the words and counts with show.py, exactly. For
# each setting it also reads the chart back: the key sits above everything in the plot, nothing leaves the viewBox, and the marks for the link in use sit on the lines.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const cfg = JSON.parse(html.match(/<script type="application\/json" id="pd-config">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the model ('), html.indexOf('/* ---------- display'));
const api = new Function('cfg', "'use strict';" + js + '; return {describe, chart};')(cfg);
const bad = [];
const note = (args, what, got, want) => { if (bad.length < 3) bad.push([args.join('/'), what, JSON.stringify(got).slice(0, 300), JSON.stringify(want).slice(0, 300)]); };
const pairs = (d) => [...d.matchAll(/[ML]([\d.]+) ([\d.]+)/g)].map((m) => [+m[1], +m[2]]);
// The box a group of marks covers, in SVG units. Text is a guess: 15 units tall, 8.25 wide a character, placed by its anchor.
function box(group) {
  const b = { top: 1e9, bottom: -1e9, left: 1e9, right: -1e9 };
  const grow = (x0, y0, x1, y1) => { b.left = Math.min(b.left, x0); b.right = Math.max(b.right, x1); b.top = Math.min(b.top, y0); b.bottom = Math.max(b.bottom, y1); };
  for (const m of group.matchAll(/<text class="[^"]*" x="([\d.]+)" y="([\d.]+)"(?: text-anchor="(\w+)")?>([^<]*)<\/text>/g)) {
    const w = 8.25 * m[4].length, x = +m[1];
    const [x0, x1] = m[3] === 'middle' ? [x - w / 2, x + w / 2] : m[3] === 'end' ? [x - w, x] : [x, x + w];
    grow(x0, +m[2] - 11, x1, +m[2] + 4);
  }
  for (const m of group.matchAll(/<circle cx="([\d.]+)" cy="([\d.]+)" r="([\d.]+)"/g)) grow(+m[1] - +m[3], +m[2] - +m[3], +m[1] + +m[3], +m[2] + +m[3]);
  for (const m of group.matchAll(/<rect x="([\d.]+)" y="([\d.]+)" width="([\d.]+)" height="([\d.]+)"/g)) grow(+m[1], +m[2], +m[1] + +m[3], +m[2] + +m[4]);
  for (const m of group.matchAll(/<path d="([^"]+)"/g)) pairs(m[1]).forEach(([x, y]) => grow(x, y, x, y));
  return b;
}
const near = (p, q) => Math.abs(p[0] - q[0]) < 0.15 && Math.abs(p[1] - q[1]) < 0.15;
fx.cases.forEach((c) => {
  const got = api.describe(...c.args);
  if (JSON.stringify(got) !== JSON.stringify(c.want)) return note(c.args, 'words', got, c.want);
  const svg = api.chart(got, c.args[2]);
  const [w, h] = svg.match(/viewBox="0 0 (\d+) (\d+)"/).slice(1).map(Number);
  const key = box(svg.match(/<g class="pd-key">(.*?)<\/g>/)[1]), plot = svg.match(/<g class="pd-plot">(.*?)<\/g><\/svg>/)[1], p = box(plot);
  if (!(key.bottom + 4 < p.top)) note(c.args, 'the key reaches the plot', key, p);
  if (Math.min(key.left, p.left, key.top, p.top) < 0 || Math.max(key.right, p.right) > w || Math.max(key.bottom, p.bottom) > h) note(c.args, 'outside the viewBox', [w, h], [key, p]);
  const line = (color) => plot.match(new RegExp('<path d="M([^"]+)" fill="none" stroke="var\\(--' + color + '\\)"'))[1].split('L').map((q) => q.split(' ').map(Number));
  const green = line('green'), orange = line('orange'), n = cfg.link.length;
  const square = plot.match(/<rect x="([\d.]+)" y="([\d.]+)" width="11.0"/), dot = plot.match(/<circle cx="([\d.]+)" cy="([\d.]+)" r="6"/);
  const k = c.args[2], ok = green.length === n && orange.length === n && orange.every((q) => q[1] === orange[0][1]) && square && dot
    && near([+square[1] + 5.5, +square[2] + 5.5], green[k]) && near([+dot[1], +dot[2]], orange[k])
    && (plot.match(/width="6.0"/g) || []).length === n && (plot.match(/r="3.5"/g) || []).length === n;
  if (!ok) note(c.args, 'marks are off the lines', green, orange);
});
console.log(JSON.stringify({ checks: fx.cases.length, bad }));
process.exit(bad.length ? 1 : 0);
"""

# Runs the page script against a small stand-in for the DOM, moves the sliders as a reader would and reads back what the page shows.
PAGE = r"""
import fs from 'node:fs';
const [widget, fixture] = process.argv.slice(2);
const html = fs.readFileSync(widget, 'utf8');
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const config = html.match(/<script type="application\/json" id="pd-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];
function el(tag) {
  const e = { tag, listeners: {}, value: '', attrs: {}, textContent: '', innerHTML: '' };
  e.setAttribute = (k, v) => { e.attrs[k] = v; };
  e.addEventListener = (type, f) => (e.listeners[type] = e.listeners[type] || []).push(f);
  return e;
}
const page = { root: el('html'), messages: [], parent: { other: true } };
[...html.matchAll(/id="(pd-[^"]+)"/g)].forEach((m) => { page[m[1]] = el('x'); });
page['pd-config'].textContent = config;
const doc = { getElementById: (id) => page[id], documentElement: page.root };
const win = { addEventListener: (type, f) => page.messages.push(f) };
new Function('document', 'window', 'parent', code)(doc, win, page.parent);
const names = ['rps', 'prompt', 'link', 'reply', 'first', 'gap'];
const fire = (e, type) => (e.listeners[type] || []).forEach((f) => f({ target: e }));
const bad = [];
const want = (what, got, expected) => { if (JSON.stringify(got) !== JSON.stringify(expected)) bad.push([what, JSON.stringify(got).slice(0, 300), JSON.stringify(expected).slice(0, 300)]); };
const shown = () => [page['pd-line'].textContent, page['pd-rule'].textContent, page['pd-tails'].textContent];
const set = (stops) => names.forEach((n, i) => { page['pd-' + n].value = String(stops[i]); });
want('opening controls', names.map((n) => page['pd-' + n].value), fx.opening.controls);
want('opening state', shown(), fx.opening.shown);
want('nothing spoken yet', page['pd-live'].textContent, '');
want('opening chart', page['pd-chart'].innerHTML.match(/<path d="M306\.0 82L/) !== null && /viewBox="0 0 320 244"/.test(page['pd-chart'].innerHTML), true);
fx.states.forEach((s) => {   // the states the lesson tells the reader to reach, worked out from pd_sim.py itself
  set(s.controls);
  fire(page['pd-' + s.moved], 'input');
  want('state ' + s.moved, page['pd-line'].textContent, s.head);
});
set(fx.dragged.controls);
fire(page['pd-rps'], 'input');
want('dragging shows the new state', shown(), fx.dragged.shown);
want('dragging stays silent', page['pd-live'].textContent, '');
fire(page['pd-rps'], 'change');
want('release speaks', page['pd-live'].textContent, fx.dragged.shown[0]);
want('words on the controls', names.map((n) => page['pd-' + n + '-out'].textContent).concat(names.map((n) => page['pd-' + n].attrs['aria-valuetext'])), fx.dragged.words);
const theme = (source, name) => page.messages.forEach((f) => f({ data: { type: 'spicyshark:theme', theme: name }, source }));
theme({ other: 'frame' }, 'dark');
want('theme from another sender is ignored', page.root.attrs['data-theme'], undefined);
theme(page.parent, 'dark');
want('theme from the host page', page.root.attrs['data-theme'], 'dark');
console.log(JSON.stringify({ bad: bad.slice(0, 3) }));
process.exit(bad.length ? 1 : 0);
"""


def run_node(tmp_path, script, *args):
    path = tmp_path / "check.mjs"
    path.write_text(script, encoding="utf-8")
    done = subprocess.run(["node", str(path), *args], capture_output=True, text=True, timeout=120)
    assert done.returncode == 0, done.stdout + done.stderr
    return json.loads(done.stdout)


def sentence(rps, prompt, link, reply=200):
    """The top line the page must show, worked out from pd_sim.py without show.py's table."""
    got = sim.compare(rps, prompt, reply, link)
    a, b = f"{got['shared']['goodput']:.1%}", f"{got['split']['goodput']:.1%}"
    return f"Both layouts get {a} of requests right." if a == b else f"Shared GPUs get {a} of requests right and the split gets {b}."


def test_widget_data_block_matches_the_python_module():
    assert embed(HTML) == HTML
    assert json.loads(CONFIG_BLOCK.search(HTML).group(2)) == json.loads(json.dumps(show.widget_config()))
    for name in NAMES:                                                                                  # each slider walks the stops the table covers and starts on the opening one
        assert re.search(r'id="pd-%s" min="0" max="%d" step="1" value="%d"' % (name, len(STOPS[name]) - 1, show.OPENING[name]), HTML), name
    assert [STOPS[name][show.OPENING[name]] for name in NAMES] == [4, 2000, 25, 200, sim.FIRST_WORD_TARGET_S, sim.WORD_GAP_TARGET_S]      # the opening is the toy's default
    requests = sim.make_requests(4, 2000, 200)                                                          # show.py counts good requests itself: it must agree with pd_sim.goodput
    for run, summary in ((sim.shared(requests), show.table()[1][2][1]["shared"]), (sim.split(requests, 0.25), show.table()[1][2][1]["split"][0])):
        assert summary["good"] == [round(sim.goodput(requests, run, f, g) * sim.N_REQUESTS) for f in show.FIRST_TARGETS for g in show.GAP_TARGETS]
        assert (summary["first"], summary["gap"]) == tuple(round(sim.measure(requests, run)[key] * 1000, 1) for key in ("first_p99", "gap_p99"))


def test_widget_gives_the_same_words_and_a_clean_chart_as_python_for_every_setting(tmp_path):
    cases = [{"args": list(g), "want": show.describe(*g)} for g in GRID]
    heads = " ".join(c["want"]["head"] for c in cases)
    assert "Both layouts get" in heads and "Shared GPUs get" in heads and "the split gets 0.0%" in heads    # the grid reaches every kind of sentence
    (tmp_path / "fx.json").write_text(json.dumps({"cases": cases}), encoding="utf-8")
    assert run_node(tmp_path, PARITY, str(tmp_path / "fx.json"), str(WIDGET))["checks"] == len(GRID)


def test_widget_page_opens_on_the_point_and_reaches_the_states_the_lesson_names(tmp_path):
    opening = [show.OPENING[name] for name in NAMES]
    head = show.describe(*opening)["head"]
    assert head == "Shared GPUs get 13.5% of requests right and the split gets 98.7%." == sentence(4, 2000, 25)      # the first screen makes the point
    states = [{"moved": "rps", "controls": [2, 2, 5, 1, 1, 1], "head": sentence(10, 2000, 25)},              # Requests a second 10
              {"moved": "link", "controls": [1, 2, 0, 1, 1, 1], "head": sentence(4, 2000, 0.25)},            # Link speed 0.25 GB/s
              {"moved": "prompt", "controls": [1, 0, 5, 1, 1, 1], "head": sentence(4, 200, 25)}]             # Prompt tokens 200
    assert "1.7% of requests right and the split gets 32.2%" in states[0]["head"] and states[1]["head"].endswith("the split gets 0.0%.") and states[2]["head"] == "Both layouts get 100.0% of requests right."
    dragged = (2, 1, 3, 2, 0, 2)
    words = ["10", "1000", "2.5 GB/s", "400", "0.25 s", "0.2 s", "10 requests a second", "1000 tokens in a prompt", "link speed 2.5 gigabytes a second", "replies of 400 words",
             "first word within 0.25 seconds", "no gap between words over 0.2 seconds"]
    fixture = {"opening": {"controls": [str(v) for v in opening], "shown": [show.describe(*opening)[key] for key in ("head", "rule", "tails")]}, "states": states,
               "dragged": {"controls": list(dragged), "shown": [show.describe(*dragged)[key] for key in ("head", "rule", "tails")], "words": words}}
    (tmp_path / "fx.json").write_text(json.dumps(fixture), encoding="utf-8")
    run_node(tmp_path, PAGE, str(WIDGET), str(tmp_path / "fx.json"))


def test_widget_is_small_themed_and_uses_only_the_site_palette():
    body = HTML[HTML.index('<div class="pd-root">'):HTML.index("<script")]
    visible = re.sub(r"<details.*?</details>", "", body, flags=re.S)
    assert len(re.findall(r"<input|<select|<button", visible)) == 3 and "More settings" in body and "Toy values and 600 invented requests" in visible   # three controls on show, caveat always on show
    longest = {"rps": "10", "prompt": "2000", "link": "12.5 GB/s", "reply": "400", "first": "0.25 s", "gap": "0.05 s"}
    labels = re.findall(r'<label class="pd-lab" for="pd-(\w+)">(.*?)</label>', body)
    assert [name for name, _ in labels] == list(NAMES) and [re.sub(r"<[^>]+>", "", text).split(":")[0] for name, text in labels[:3]] == ["Requests a second", "Prompt tokens", "Link speed"]
    assert all(f'for="pd-{name}"' in visible for name in NAMES[:3]) and all(f'for="pd-{name}"' in body[body.index("<details"):] for name in NAMES[3:])     # the first three are on show, the rest are in More settings
    for name, label in labels:                                                                          # a label and its longest value fit one line at 340 px
        text = re.sub(r"<[^>]+>", "", label) + longest[name]
        assert len(text) <= 22, text
    assert len(HTML.splitlines()) < 300 and "light-dark(#ffffff, #0a0f17)" in HTML and "spicyshark:theme" in HTML and "prefers-reduced-motion" in HTML
    assert "forced-colors: active" in HTML and 'role="status" aria-live="polite"' in HTML
    assert not re.search(r"0072B2|56B4E9|D55E00|ht" r"tps?://|<img|@import", HTML, re.I)
    sizes = [int(px) for px in re.findall(r"\.pd-[kt] \{ font: (?:600 )?(\d+)px", HTML)]
    assert len(sizes) == 2 and min(sizes) >= 15 and 'viewBox="0 0 320 ' in HTML                         # chart text 15 units or more of a 320-unit chart: about 12.5 px in the 290 px frame at 340

In one minute

Explore the numbers

How it works

The picture, step by step

Each step in code

A real run

Your turn

What this means when you build

Go deeper

The simulation

Why the two jobs differ

The cost of the hand-over

Where the split does not pay

What the papers report

What the toy leaves out

Proof it works

Where this is used in practice

Skills this lesson builds

Check yourself

Further reading

Code