AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
01 Campus1 of 10

STARGATE CAMPUS, ABILENE · cutaway, roofs off, illustrative layout

WALKTHROUGH 1A data centre, top down
Gridhigh voltageTurbinessome campusesSubstation8 feedersDry coolerswarm waterBuilding 18 data halls
8 halls · 128 racks each
SIZED IN POWERfacility power · one square = 10 MW

Stargate, Abilenethe whole campus

1,200 MW

One Abilene buildingthe one we follow in

150 MW

xAI Colossus, 2024Memphis, a whole supercomputer

150 MW
Planned or designed capacity, not measured draw.
1.2 GWPLANNED CAMPUS
8BUILDINGS
esc
Full video on YouTube

01 Campus

Frontier AI campuses are sized in gigawatts of power, not floor space, split across a few buildings.

  1. Getting that connection approved and built is now often the slowest part of a project.

  2. Each feeder powers one building.

  3. Gas turbines start a site before the grid is ready; batteries smooth a training run's sharp swings in demand.

  4. One hall: 9,216 GPUs drawing 15.4 MW.

    Inside the data hall →
  5. Computing stores almost no energy, so all of it becomes heat. Dry coolers reject it with fans and no water; evaporative towers use less power but consume water.

  6. The campus plans eight of them.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • The Ultra-Scale Playbook: training on GPU clustersHugging Face · page, new tab
  • land
  • concrete
  • steel
  • copper (power)
  • water (warm)
  • the building we follow
  • power
  • buildings
  • cooling
  • heat out
  • the building we follow
Data hall →
  • land
  • concrete
  • steel
  • copper (power)
  • water (warm)
  • the building we follow
  • power
  • buildings
  • cooling
  • heat out
  • the building we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
02 Data hall2 of 10

ONE DATA HALL · 1 of 8 in the building · illustrative layout

Hall 1 of 8, Building 1↖ Campus
CDUCDU…… 16 rows, each with 8× NVL72 GPU racksRack 172 GPUsElectricity inbusway on every rowNetwork rowdetails in 03

POWER IN · grid to rack

Gridhigh voltage

↓

Substationsteps down

↓

Switchgearroutes, protects

↓

UPS + batteriesrides blips

↓

HEAT OUT · rack to air

CDUsloops never mix

↓

Facility water loopwarm out, cool back

↓

Chillerstandby · hottest days

↓

Dry coolers, rooffans, no water used

↓

Heat to the air

128RACKS
9,216GPUs
15.4 MWIT POWER
esc
Full video on YouTube

02 Data hall

Each building has data halls: rows of GPU racks which consume power to perform calculations. Water carries the heat away.

  1. Overhead copper bars with a tap-off per rack. The rack's own power shelves take it from there.

  2. Cool air enters at the front and leftover heat leaves into the hot aisle; the water arrives through the back manifolds.

    Inside one NVL72 rack →
  3. It pumps the racks' coolant and passes the heat to the building's water through a heat exchanger, so the two loops never mix.

  4. Fibre runs overhead to every rack.

    See also: Step 03 · Network

  5. Cold plates work with warm water, so on most days no chiller runs.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 3 HardwareBaseten · page, new tab
  • floor
  • racks (steel)
  • busway (copper)
  • coolant, cold
  • coolant, warm
  • CDU · Coolant Distribution Unit
  • the rack we follow
  • power
  • racks
  • network
  • cooling
  • heat out
  • the rack we follow
← CampusNetwork →
  • floor
  • racks (steel)
  • busway (copper)
  • coolant, cold
  • coolant, warm
  • CDU · Coolant Distribution Unit
  • the rack we follow
  • power
  • racks
  • network
  • cooling
  • heat out
  • the rack we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
03 Network3 of 10

ONE DATA HALL, CUT AWAY · 3 of 16 rows drawn, network row in front, illustrative

Network↖ Hall
…Corebuilding network roomto the other 7 hallscrosses racksGradient syncs · trainingPipeline activations · trainingKV caches · prefill → decode… 16 rowsRack 172 GPUsRack 872 GPUs…Front-end switchesno GPU traffic · storage, usersScale-up · NVLinkcopper, stays in the rackLeafrail-optimisedSpinetaller, top layerOverhead fibre

AS A GRAPH · 8 racks

ACROSSother sitesCORESPINELEAFRACKSSCALE-OUTSCALE-UP

Rail-optimised: each leaf takes one GPU slot from every rack. 8-rack pod: ≈32 leaf, 12 spine.

Scale-up · NVLinkcopper, inside the rack

900 GB/s (7,200 Gb/s)

Scale-out · networkfibre, between racks

50 GB/s (400 Gb/s)
Per GPU, each way · GB200 NVL72, ConnectX-7.
3SWITCH LAYERS
18×SLOWER THAN NVLINK
esc
Full video on YouTube

03 Scale-out network

Communication between two GPU racks is slow, over a switched network: any GPU can reach any other.

  1. It is separate from the copper NVLink inside the rack.

  2. Each leaf takes the same GPU slot from every rack, so same-slot GPUs are one leaf hop apart (the graph's gold path).

  3. GPUs in different slots go up through a spine (the cut-away's gold path); many equal paths spread the load.

  4. Smaller ones stop at leaf and spine.

  5. Network speeds are sold in bits: 400 Gb/s is 50 GB/s.

  6. Training keeps tensor-parallel and expert traffic inside a rack; gradient syncs and pipeline activations cross here.

    See also: Step 05 · Switch trays

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • NVL72 AI factory: network logical architectureNVIDIA · page, new tab
  • The Ultra-Scale Playbook: training on GPU clustersHugging Face · page, new tab
  • front-end network
  • network (leaf, spine, core)
  • NVLink
  • the rack we follow
← Data hallNVL72 rack →
  • front-end network
  • network (leaf, spine, core)
  • NVLink
  • the rack we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
04 NVL72 rack4 of 10

GB200 NVL72 · side view and front view, top to bottom

Rack 1 of 128↖ Hall

SIDE VIEW

FRONT VIEW

FRONT

BACK · coolant, power

Power shelves

10 compute trays4 GPUs each

9 switch trays

8 compute trays4 GPUs each

Power shelves

←NVLink spinecopper, rack back

The tray we followcompute tray 1, opened at step 06

SCALE-UP · NVLINK · copper, every GPU to every switch chip· every GPU to every switch chip

18 NVSwitch chips

GPUs · one GPU's 18 links drawn

130 TB/s across the whole spine

POWER AND COOLING · at the back

Power shelvesAC → DC
→
Busbar
→
Every tray
CDUrow end
→
Manifoldscold inhot out
→
Quick-connects
72GPUs
36GRACE CPUs
≈120 kWPER RACK
1.36 tWEIGHT
esc
Full video on YouTube

04 GB200 NVL72 rack

Communication inside one GPU rack is fast: 72 GPUs joined so tightly they work like one.

  1. Each holds four GPUs.

    Inside the compute tray →
  2. They hold the 18 NVSwitch chips; the copper NVLink spine at the back joins every GPU to them. This is the scale-up network: all 72 GPUs act as one.

    Inside the switch trays →
  3. A copper busbar carries it down the back; each tray clips on when it slides in.

  4. Every tray has a quick-connect pair, so it slides out without draining the rack.

    See also: Step 02 · Data hall

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GB200 NVL72NVIDIA · page, new tab
  • NVL72 AI factory: network logical architectureNVIDIA · page, new tab
  • Inference Engineering · 3 HardwareBaseten · page, new tab
  • steel
  • circuit boards
  • copper
  • coolant in
  • coolant out
  • fibre (network)
  • the tray we follow
  • compute
  • NVLink
  • network
  • power
  • coolant in
  • coolant out
  • the tray we follow
← NetworkSwitch trays →
  • steel
  • circuit boards
  • copper
  • coolant in
  • coolant out
  • fibre (network)
  • the tray we follow
  • compute
  • NVLink
  • network
  • power
  • coolant in
  • coolant out
  • the tray we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
05 Switch trays5 of 10

NVLINK SWITCH TRAY · top view, back on the left

Switch tray 1 of 9↖ Rack
SHARPNVSwitch chip 1SHARPNVSwitch chip 2PowerinMgmtBMCcoolant →

Link fanswitch chips above, GPUs below

GPU 1GPU 2GPU 3GPU 44 of the 72 GPUs

Spine connectorsback edge, to the copper spine

Two NVSwitch chipsa port for each of the 72 GPUs

IN-SWITCH SUMone all-reduce, 4 of the 72 GPUs

g₁g₂g₃g₄
each GPU's gradients, sent once
→
SHARPg₁ + g₂ + g₃ + g₄
→
ΣgΣgΣgΣg
the sum, back to every GPU
9SWITCH TRAYS
130 TB/sACROSS THE RACK
esc
Full video on YouTube

05 NVLink switch trays

Single-rack speed comes from switch chips in the middle of the rack: one hop between any two GPUs on that rack.

  1. Over 5,000 copper cables carry all 1,296 GPU links (72 GPUs × 18). Copper needs no optical transceivers, saving power and cost, but only reaches across one rack.

  2. Each switch chip has a port for every one of the 72 GPUs.

  3. So every GPU in the rack connects to all 18 switch chips, and traffic between any two GPUs crosses exactly one switch chip.

  4. In training, every GPU adds its gradients to everyone else's (an all-reduce); the switch does the sums on the way through, so less data crosses the links.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GB200 NVL72NVIDIA · page, new tab
  • silicon
  • circuit board
  • copper
  • coolant in
  • coolant out
  • one GPU's links
  • NVLink
  • sums in the switch
  • power
  • coolant in
  • coolant out
  • one GPU's links
← NVL72 rackCompute tray →
  • silicon
  • circuit board
  • copper
  • coolant in
  • coolant out
  • one GPU's links
  • NVLink
  • sums in the switch
  • power
  • coolant in
  • coolant out
  • one GPU's links
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
06 Compute tray6 of 10

COMPUTE TRAY · top view · back on the left, front on the right· top view

Tray 1 of 18↖ Rack
BACK
FRONT
Coolantin, out
NVLink72 links
out
GB200 SUPERCHIP
Grace CPU72 Arm cores
LPDDR5X480 GB
B200
B200
GB200 SUPERCHIP
Grace CPU72 Arm cores
LPDDR5X480 GB
B200
B200
ConnectX-7400 Gb/s
ConnectX-7400 Gb/s
ConnectX-7400 Gb/s
ConnectX-7400 Gb/s
BlueField-32 DPUs
Fans + PSU
NVLink-C2CNVLink-C2C

cold in

warm out

NVLink to
the spine

to other
racks

storage,
front end

WITHIN ONE RACK

B200 this GPU
→
NVLink connector tray's back edge
→
Copper spine rack back, behind trays
→
Switch tray any of the 9
→
B200 another tray, via spine

BETWEEN TWO RACKS

B200
→
ConnectX-7 front
→
Scale-out network fibre, step 03
→
GPU in another rack
NVLink spineback edge
GB200 SUPERCHIP
Grace CPU
B200186 GB
B200186 GB
GB200 SUPERCHIP
Grace CPU
B200186 GB
B200186 GB
NIC
NIC
NIC
NIC
DPU
PSU

back at the top · front at the bottom

744 GBHBM3E ON THE TRAY
72NVLINKS OUT THE BACK
esc
Full video on YouTube

06 Compute tray

Around the 9 switch trays sit 18 compute trays, each holding four GPUs and two CPUs.

  1. They share one board, joined by NVLink-C2C.

    Inside the superchip →
  2. Each has its own cold plate fed from the blue line.

    Inside the B200 GPU →
  3. 18 per GPU, to connect that GPU to each of the 18 switch chips in the rack's switch trays.

    See also: Step 05 · Switch trays

  4. 400 Gb/s each: the inter-rack, scale-out network path.

    See also: Step 03 · Network

  5. They keep that traffic off the GPUs' network cards.

  6. Coolant enters cold from the rack manifold and leaves warm; air alone can't remove ≈6 kW from a 1U tray.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 3 HardwareBaseten · page, new tab
  • GPU Glossary · GPU RAMModal · page, new tab
  • NVL72 AI factory: network logical architectureNVIDIA · page, new tab
  • silicon
  • circuit board
  • copper
  • steel
  • coolant in
  • coolant out
  • the GPU we follow
  • compute
  • memory
  • NVLink
  • network
  • power
  • coolant in
  • coolant out
  • the GPU we follow
← Switch traysSuperchip →
  • silicon
  • circuit board
  • copper
  • steel
  • coolant in
  • coolant out
  • the GPU we follow
  • compute
  • memory
  • NVLink
  • network
  • power
  • coolant in
  • coolant out
  • the GPU we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
07 Superchip7 of 10

GB200 SUPERCHIP · one board, top view

Superchip 1 of 2↖ Tray
GB200 SUPERCHIP
Grace CPU
LPDDR5X
NVLink-C2C
B200
B200

HOW MUCHmemory one GPU can reach

HBM3eown package

186 GB

Grace LPDDR5Xover C2C

480 GBshared by 2

0250500GB · LINEAR SCALE

HOW FASTread speed from one GPU

HBM3e

8,000 GB/s

Grace LPDDR5X

450 GB/s

x86 serverCPU RAM over PCIe 5

64 GB/s

101001,00010,000GB/S · LOG SCALE

One GB200 GPU reading; per direction where it matters.

≈18×SLOWER THAN HBM
esc
Full video on YouTube

07 GB200 Superchip

Each CPU pairs with two GPUs in one unified memory space, so they can use its larger, slower LPDDR5X.

  1. It runs the operating system and the serving program: it tokenises requests, batches them, launches GPU kernels and moves data. Arm cores use a simpler, more power-efficient instruction set than a laptop's x86 cores, so 72 fit in the power of a few; they do no model maths themselves.

  2. Each keeps its own HBM (memory) inside its package: the weights, KV cache and activations every pass reads. Only memory millimetres away, on thousands of wires, reaches 8 TB/s; LPDDR5X over C2C manages 450 GB/s.

    See also: Step 08 · B200 GPU

  3. This is unified memory: either side reads the other’s memory as its own, with no copies managed by software: a B200 reads Grace’s LPDDR5X, and Grace reads the B200’s HBM3e.

  4. Overflow KV cache, paused chats and optimizer state go here.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GB200 NVL72NVIDIA · page, new tab
  • NVIDIA Grace CPU Superchip Architecture In DepthNVIDIA · page, new tab
  • GPU Glossary · GPU RAMModal · page, new tab
  • silicon
  • circuit board
  • HBM
  • copper
  • compute
  • memory, fast
  • memory, overflow
  • C2C link
← Compute trayB200 GPU →
  • silicon
  • circuit board
  • HBM
  • copper
  • compute
  • memory, fast
  • memory, overflow
  • C2C link
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
08 B200 GPU8 of 10

B200 GPU PACKAGE · top view, drawn with its real counts

GPU 1 of 4↖ Tray
L2 cache
L2 cache
10 TB/s
26 × 33 mm

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

HBM3e24 GB

GB100 die

GB100 die

die-to-die link · 10 TB/s

SUBSTRATE · INTERPOSER

NVLINK OUTto the other 71 GPUs

  1. back connector
  2. ↓
  3. copper spine
  4. ↓
  5. switch tray, 1 of 9
  6. ↓
  7. any GPU in the rack

From this GPU: 18 links900 GB/s each way

OTHER RACKS PCIe → ConnectX-7 → leaf switch

186 GBHBM3E
8 TB/sREAD SPEED
208 BTRANSISTORS
126 MBL2
esc
Full video on YouTube

08 B200 GPU

Within each GPU are two silicon dies working as one, fed by stacks of High Bandwidth Memory (HBM) right beside them.

  1. The weights and the KV cache live here, and every decode pass (one per new token) reads them again.

    Inside the HBM3e stack →
  2. L2 cache is shared by every SM, so data one SM fetched needn’t come from HBM again (a cache hit). It is on the die, so it is faster than HBM.

  3. All identical, all working at once.

    Inside one SM →

    See also: Step 10 · Inside one tensor core

  4. When a model is split across GPUs, partial results go to the switch chips and on to the other GPU.

    See also: Step 05 · Switch trays

  5. They skip the switch trays: PCIe to the ConnectX-7 at the tray front, then fibre to a leaf switch.

    See also: Step 03 · Network

  6. A scanner prints at most 26 × 33 mm at a time, so Blackwell joins two dies with a 10 TB/s link and software sees one GPU.

  7. Same two dies and eight HBM stack positions; more HBM, more FP4 maths.

    Inside B300 vs B200 →

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GPU Glossary · CUDA device architectureModal · page, new tab
  • GPU Glossary · Graphics processing clusterModal · page, new tab
  • Microbenchmarking NVIDIA’s Blackwell ArchitecturearXiv · page, new tab
  • Inference Engineering · 3 HardwareBaseten · page, new tab
  • silicon die
  • HBM stack
  • interposer
  • substrate
  • copper links
  • reticle limit
  • the SM we follow
  • compute
  • memory
  • interposer
  • substrate
  • NVLink
  • reticle limit
  • the SM we follow
← SuperchipSM →
  • silicon die
  • HBM stack
  • interposer
  • substrate
  • copper links
  • reticle limit
  • the SM we follow
  • compute
  • memory
  • interposer
  • substrate
  • NVLink
  • reticle limit
  • the SM we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
09 SM9 of 10

ONE SM · four partitions around shared memory

SM 1 of 148↖ GPU
L2 cacheoutside the SM · shared by all 148 SMs · ~126 MB
weights in from HBM ←
SM
TMAcopies tiles in
Shared memory / L1
Warp scheduler
Registers64 KB
32 CUDA cores
Tensor core5th gen
Tensor memory64 KB
Warp scheduler
Registers64 KB
32 CUDA cores
Tensor core5th gen
Tensor memory64 KB
Warp scheduler
Registers64 KB
32 CUDA cores
Tensor core5th gen
Tensor memory64 KB
Warp scheduler
Registers64 KB
32 CUDA cores
Tensor core5th gen
Tensor memory64 KB
2,048THREADS IN FLIGHT
256 KBREGISTERS
256 KBTENSOR MEMORY
256 KBL1 + SHARED
esc
Full video on YouTube

09 Streaming multiprocessor (SM)

Inside a die, 148 SMs work in parallel. Moving data (model weights) from HBM to these SMs is usually the speed bottleneck (not the actual maths itself).

  1. The kernel (e.g. a cuBLAS matmul) cuts each matrix into tiles and hands them to SMs. A tile comes from HBM through L2, which all SMs share.

  2. One instruction moves a whole tile. Shared memory (L1) sits inside the SM, even faster than L2, so kernels reuse the tile many times before going back.

  3. Threads run in groups of 32, called warps. Each cycle the scheduler issues one instruction for a warp that is ready; a warp waiting on memory is skipped, so the cores never sit idle.

  4. Almost all of a model’s maths (the matmuls) happens in these.

    Inside the tensor core →
  5. New in Blackwell: a store beside the tensor cores, so running totals don’t tie up registers.

  6. Activations, normalisation, softmax.

  7. Each weight read serves one token per request, so HBM bandwidth, not the tensor cores, sets the pace.

    See also: Walkthrough 2 · Passes

  8. In prefill each weight serves every prompt token, so the tensor cores, not HBM, set the pace.

    See also: Walkthrough 2 · Passes

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GPU Glossary · Streaming multiprocessorModal · page, new tab
  • GPU Glossary · Tensor coreModal · page, new tab
  • Nvidia’s B200: Keeping the CUDA Juggernaut RollingChips and Cheese · page, new tab
  • memory
  • compute
  • control
  • the tensor core we follow
← B200 GPUTensor core →
  • memory
  • compute
  • control
  • the tensor core we follow
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core
10 Tensor core10 of 10

ONE TILE MULTIPLY · C = A × B, 4 × 4, illustrative values

16 cells · 0.25 s per cell
Tensor core 1 of 4↖ SM
A · activations
1201013220111102×B · weights
2101102103101212=C · output
16 cells × 4 multiply-adds, in one instruction
Where tile B comes from

One 4,096 × 4,096 weight matrix in Llama 3.1 8B, cut into 256 tiles of 256 × 256 · tile size illustrative

the tile drawn above, shrunk to 4 × 4

BF16 2.5 · FP8 5 · FP4 10 PFLOPS per GPU, dense↳ Number formats
64MULTIPLY-ADDS IN THIS TILE
esc
Full video on YouTube

10 Tensor core

At the lowest level, tensor cores multiply tiles of matrices. Fewer bits per number means more multiplications per second.

  1. Rows are tokens, columns are features.

  2. A weight matrix is far too big for one multiply, so it is cut into tiles. The tensor cores take them one after another.

  3. C₂₃ is row 2 of A times column 3 of B: four products, added up.

  4. Real Blackwell tiles are far bigger; the player only slows it down.

  5. Each halving roughly doubles throughput and halves the bytes per weight.

    Inside: Number formats →

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GPU Glossary · Tensor coreModal · page, new tab
  • GPU Glossary · Streaming multiprocessorModal · page, new tab
  • A · activations
  • B · weights
  • C · output
  • the cell being computed
← SM2 · Inside a transformer →
  • A · activations
  • B · weights
  • C · output
  • the cell being computed
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU↳ HBM3e stack09 SM10 Tensor core
08 B200 GPU ↳ HBM3e stack8 of 10

HBM3E STACK · side section, not to scale vertically

Stack 1 of 8↖ GPU
HBM3E STACK
B200 diecompute

DRAM dies
3 GB each

TSVs
copper, vertical

base die
talks to the GPU

interposer

the bus, sideways to the die

HBM3eone stack

1,024 wires

DDR5one DIMM

64 wires
05121,024DATA WIRES PER DEVICE

Interface width only, not speed.

MADE BYSK hynixSamsungMicron
24 GBPER STACK
≈1 TB/sPER STACK
8STACKS PER B200
esc
Full video on YouTube

HBM3e stack

Those HBM stacks are fast because of width: a very wide bus, millimetres from the dies that read it.

  1. Each holds 3 GB. 8 stacks × 24 GB = 192 GB built; 186 GB usable on GB200.

  2. Through-silicon vias replace a memory module’s long traces; each path is well under a millimetre.

  3. It holds no memory and drives the interface out through the interposer.

  4. A DDR5 module has 64. Width, not clock speed, is where the bandwidth comes from.

  5. SK hynix, Samsung and Micron. Its supply limits how many GPUs can be built.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GPU Glossary · GPU RAMModal · page, new tab
  • DRAM die
  • base die
  • interposer
  • TSVs (copper)
  • microbumps
  • memory
  • memory wiring
  • compute
  • packaging
← Back to B200 GPU
  • DRAM die
  • base die
  • interposer
  • TSVs (copper)
  • microbumps
  • memory
  • memory wiring
  • compute
  • packaging
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU↳ B300 vs B20009 SM10 Tensor core
08 B200 GPU ↳ B300 vs B2008 of 10

B200 AND B300 · same package, per-GPU figures

B300↖ GPU
die 1die 2

ONE STACK, SIDE VIEW

B300: 12 high

B200: 8 high

base die

GPU

HBMmemory per GPU

B200186 GB
B300278 GB
0300GB

FP44-bit, AI inference

B20010 PFLOPS
B30015 PFLOPS
015PFLOPS

FP6464-bit, scientific simulation

B20040 TFLOPS
B3001.4 TFLOPS
040TFLOPS

NetworkNIC per GPU

B20050 GB/s each way
B300100 GB/s each way
0100GB/S EACH WAY

NVLinkto the switch

B200900 GB/s each way
B300900 GB/s each way
0900GB/S EACH WAY

Per GPU, dense, from NVIDIA’s rack specs ÷ 72.

LINEAGE· HBM per GPU, by year of first shipment

H1002022

80 GB

HBM3
→

H2002024

141 GB

HBM3e
→

B2002024

186 GB

HBM3e · 8-high
→

B3002025

278 GB

HBM3e · 12-high
→

Rubin2026

288 GB

HBM4 · announced
esc
Full video on YouTube

B300 vs B200

The next GPU keeps those two dies and HBM stack positions, with taller HBM stacks and more FP4 maths, but little FP64.

  1. The package outline doesn’t change.

  2. More layers add capacity; the interface stays 1,024 bits wide.

  3. Dense FP4 goes from 10 to 15 PFLOPS, good for inference, which tolerates 4-bit numbers. FP64 (64-bit, used by weather and physics simulations) drops from 40 to 1.4 TFLOPS: the B300 gives up science work for AI.

  4. ConnectX-8 (800 Gb/s) replaces ConnectX-7 (400 Gb/s), so traffic between racks, through the leaf switches, gets twice the bandwidth.

  5. H100 in 2022, H200 and B200 in 2024, B300 in 2025; Rubin is announced for 2026.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GB300 NVL72NVIDIA · page, new tab
  • GB200 NVL72NVIDIA · page, new tab
  • silicon die
  • HBM stack
  • interposer
  • substrate
  • compute
  • memory
  • packaging
← Back to B200 GPU
  • silicon die
  • HBM stack
  • interposer
  • substrate
  • compute
  • memory
  • packaging
AI STACK·
A data centre, top down ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Campus02 Data hall03 Network04 NVL72 rack05 Switch trays06 Compute tray07 Superchip08 B200 GPU09 SM10 Tensor core↳ Number formats
10 Tensor core ↳ Number formats10 of 10

ONE MODEL, SIX FORMATS · Llama 3.1 8B, 8.03 B weights · bits drawn to scale

Number formats↖ Tensor core
FORMATBITS PER NUMBER · to scalePER WEIGHT8B MODEL

FP3232 bits · E8M23

4 bytes
32.1 GB

BF1616 bits · E8M7

2 bytes
16.1 GB

FP88 bits · E4M3

1 byte
8.0 GB

FP44 bits · E2M1

+ one 8-bit scale per 16 weights
0.56 bytes
4.5 GB
GGUF FILES· whole numbers in blocks, each block with a scale

Q8_08.5 bits avg

1.06 bytes
8.54 GB

Q4_K_M4.9 bits avg

the file W3 runs on the 5090
0.61 bytes
4.92 GB

EVERY POSITIVE FP4 VALUE · before the scale

00.511.52346
7.1×SMALLER · FP32 → FP4
4×MATHS PER SECOND · BF16 → FP4
15VALUES IN FP4
esc
Full video on YouTube

Number formats

Fewer bits per weight shrinks a model and speeds it up; each halving gives up range or precision.

  1. Four bits hold eight values from 0 to 6, either sign. One 8-bit scale per 16 weights stretches that grid to fit the block.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • OCP Microscaling Formats (MX) Specification v1.0Open Compute Project · page, new tab
  • llama.cpp quantize READMEggml-org · page, new tab
  • sign
  • exponent · range
  • mantissa · precision
  • whole number
  • scale
  • the file we follow
← Back to Tensor core
  • sign
  • exponent · range
  • mantissa · precision
  • whole number
  • scale
  • the file we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens1 of 14

TOKENIZER · text in on the left, ids out on the right

TEXT · 21 CHARACTERS
Tom opened a new bank
TOKENIZER
cut intoknown pieces
TOKENS
IDS
Tom
Tom
25763
·opened
·opened
9107
·a
·a
264
·new
·new
502
·bank
·bank
6201

HOW OTHER TEXT GETS CUT· the same tokenizer, each piece with its id

Common words: one piece
·the279
·account2759
1 token each
Rare words: several pieces
·Ou69327
ag351
ad329
oug49404
ou283
·Ouagadougou · 5 tokens
Numbers: up to 3 digits
·220
17811256
924
·1789 · 3 tokens
Same word, different ids
·bank6201
Bank26913
·Bank8715
space and case change the id
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

5TOKENS
128,256IN THE VOCABULARY

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

01 Tokenization

A language model never sees text: the prompt is cut into tokens, and each token gets a unique number, its token ID.

  1. At the end, the model answers ·account. Llama also puts a start token first; it is left out of every diagram.

  2. A fixed list of text pieces, learned once before training; frequent chunks get their own entry.

    Inside: Byte pairs →
  3. ·opened after a word and opened at the start of a line are different tokens.

  4. They are all the model receives. ·bank (6201) is the one whose prediction becomes ·account.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 2 ModelsBaseten · page, new tab
  • text
  • tokens and ids
  • the token we follow
Embedding →
  • text
  • tokens and ids
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
02 Embedding2 of 14

EMBEDDING LOOKUP · real shapes · matrix not to scale, broken below row 10,000

EMBEDDING MATRIX⋮row 0⋮↓ row 128,255[128,256 × 4,096]1 row per token in the model’s vocabularyrow 25,763 · Tomrow 9,107 · ·openedrow 264 · ·arow 502 · ·newrow 6,201 · ·bankTOKEN ID25,7639,1072645026,201TOKENDIM 0DIM 1DIM 2DIM 3…DIM 4,095Tom−0.01860.0105−0.0093−0.0236…0.0009·opened0.0007−0.00830.01350.0041…0.0016·a−0.0015−0.0006−0.0009−0.0008…0.0001·new0.0017−0.00300.0033−0.0084…−0.0025·bank−0.01930.0057−0.0096−0.0002…0.0038each row [1 × 4,096] · values from the published weightsSTREAMSx₀x₀x₀x₀x₀ONE ROW, TWO MEANINGS· 2 of 4,096 directions · illustrativeMONEYRIVERaccountloandepositcashrivershorewaterfishbankone row sits between both;context moves it (dashed: after attention)EMBEDDING MATRIX · A LOOKUP IS A ONE-HOT MULTIPLY10 0 00 0 0one-hot(6201)[1 × 128,256]·E[128,256 × 4,096]=row 6,201 · ·bank[1 × 4,096]Row 6,201 equals a one-hot vector times E: 128,255 zeros and one 1, atposition 6,201.So nothing is multiplied: the GPU copies 4,096 numbers.Training nudges every row; ·bank gets one row for all its meanings(below).128,256 × 4,096 = 525 M weights · 6.5% of the model.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

4,096NUMBERS PER TOKEN
525 MNUMBERS IN THE MATRIX
6.5%OF THE MODEL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

02 Word embedding

Each token ID selects the corresponding row vector from the embedding matrix, which contains one row per token in the model's vocabulary.

  1. From the tokenizer, in prompt order.

  2. Each row holds 4,096 numbers: the embedding dimension. It is learned in training like every other weight.

  3. Id 6,201 means “take row 6,201”.

  4. No single number means anything: ·bank starts 0.0193 below zero in its first dimension, and nobody chose that.

  5. But note, ·bank has one row for money and river alike.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • tokens (ids and vectors)
  • weights (the matrix)
  • the token we follow
← TokensStream →
  • tokens (ids and vectors)
  • weights (the matrix)
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream3 of 14

RESIDUAL STREAM · isometric · blocks run left to right, tokens go into the page

block 1block 32Attention 1all 5 communicateMLP 1each on its ownAttention 32all 5 communicateMLP 32each on its ownTom· · ··opened· · ··a· · ··new· · ··bank· · ·x₀x0 = E[id]embedding row·thex₁x1 = x0 + ΔΔ = attention·ax₂x2 = x1 + ΔΔ = MLP·thex₆₃x63 = x62 + ΔΔ = attention·accountx₆₄x64 = x63 + ΔΔ = MLP·account· · ·blocks 2–31· · ·IF THE MODEL STOPPED HERE, ·bank WOULD PREDICTlogit lens: each x sent straight to the unembed · illustrativeunembed→·account4,096 → 14,336 → 4,0964,096 → 14,336 → 4,096
HOW TO READ THIS
  • Dark arcs: ·bank reads Tom, opened, a and new.
  • Faint arcs: each earlier vector reads the ones before it.
  • Both planes are big matrix multiplies: the GPU runs all 5 vectors in parallel.
  • Each step adds a Δ to the stream.
  • Attention: every vector reads the earlier ones (arcs)
  • MLP: every vector, each on its own
  • one column = one token’s vector, [4,096 × 1]
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

5VECTORS · ONE PER TOKEN
4,096NUMBERS PER VECTOR d_model = EMBEDDING DIM
64ADDITIONS · 32 BLOCKS × 2

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

03 Residual stream

That stream of vectors runs the length of the model: every layer reads it and adds a small change, nudging each vector's meaning.

  1. The ·bank column ends as the vector that predicts ·account.

    Inside: Attention residual →
  2. It reads every column and writes back onto each.

  3. It reads and writes one column at a time.

  4. Because the original is kept, gradients flow straight through in training, which is what lets many blocks be stacked.

    Inside: Hyper-connections →Inside: Engram →

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • stream
  • attention
  • MLP
  • the token we follow
← EmbeddingBlock →
  • stream
  • attention
  • MLP
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
04 Block4 of 14

ONE BLOCK · block 1 of the stream in 03, zoomed in · attention first, then the MLP

block 1RMSNormrescaleAttention32 heads, all 5 communicateRMSNormrescaleMLP · SwiGLUeach vector on its ownTom++·opened++·a++·new++·bank+14,336WupWdown4,0964,096+x₀x₁x₂block 2same steps,own weightsn = RMSNorm(x0)rescaled copy of x0x1 = x0 + Attn(n)Δ from attentionn = RMSNorm(x1)rescaled copy of x1x2 = x1 + MLP(n)Δ from the MLPSwiGLU
HOW TO READ THIS
WHERE THE BLOCK’S WEIGHTS ARE · 218 M per block
attn 19%MLP 81% · 176 M
attention 41.9 M · the 2 norms hold 8,192, too few to draw
  • Both run all 5 vectors in parallel on the GPU. Each step adds a Δ.
  • RMSNorm: rescale each vector first
  • Attention: vectors read the earlier ones (arcs) · 32 heads, 4 drawn
  • MLP: each vector on its own · Wup, Wgate widen, Wdown narrows
  • one column = one token’s vector, [4,096 × 1]
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

41.9 MATTENTION PARAMETERS
176 MMLP PARAMETERS
2ADDITIONS TO THE STREAM

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

04 One transformer block

The repeating unit of an LLM: one attention layer, then one MLP (or feed-forward) layer. The block repeats 32 times, so each vector is nudged twice per block.

  1. RMSNorm divides by the vector’s average size, then applies a learned scale.

    Inside: Normalisation →
  2. The only step where tokens share information.

    See also: Step 05 · QKV

  3. The stream itself is only ever added to.

  4. Widen, gate, narrow.

    See also: Step 08 · MLP

  5. The next block repeats all of it with its own weights.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • stream
  • norm
  • attention
  • MLP
  • the token we follow
← StreamQKV →
  • stream
  • norm
  • attention
  • MLP
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
05 QKV5 of 14

QKV · one token, three projections, each cut into heads · hover: head 1

E·bank[4,096 × 1]WQ[4,096 × 4,096]32 heads, 128 matrix rows per headqWK[1,024 × 4,096]kWV[1,024 × 4,096]v32 QUERY HEADS · q cut into pieces of 128head 1328kv8 KV HEADS · a k and a v each, shared by 4 query headsheads 1–4 score against KV head 1’s k and blend its vk and v are kept for every token: the KV cache, 12E · ·bank, RESCALED BY RMSNormx÷ rms(x)× gE[4,096 × 1]the same E goesinto WQ, WK, WVrms(x): the square root of the mean of x², one number per token.g: 4,096 learned gains, one per number.Every token enters attention at a similar size.RMSNorm has 4,096 weights; GPT-3 uses LayerNorm.WQ · HEAD 1’S QUERY · WHAT THIS TOKEN LOOKS FORWQ: Queryproj. matrix[4,096 × 4,096]32 heads × 128 rowsHead 1 of WQ:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 4,096]each row: 4,096 weights·E: tokenvector⋮[4,096 × 1]=q1: head 1’squery vectorR1·ER2·E⋮R128·E[128 × 1]Each row R is dotted with E: 4,096 products, one number.128 rows give q1: what head 1 looks for in other tokens.Head 2 uses rows 129–256, and so on: 32 heads.One multiply makes all 32 heads; q is then cut into pieces.WQ: 16.8 M weights, 4,096 × 4,096.WK · HEAD 1’S KEY · WHAT THIS TOKEN OFFERSWK: Keyproj. matrix[1,024 × 4,096]8 heads × 128 rowsHead 1 of WK:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 4,096]each row: 4,096 weights·E: tokenvector⋮[4,096 × 1]=k1: head 1’skey vectorR1·ER2·E⋮R128·E[128 × 1]Same slice, own weights: k1 is what this token offers, matched against every q1 (06).k1 and v1 serve query heads 1–4: only 8 KV heads (GQA).k is kept for every token: the KV cache (12).WK: 4.2 M weights, 1,024 × 4,096.WV · HEAD 1’S VALUE · WHAT THIS TOKEN HANDS OVERWV: Valueproj. matrix[1,024 × 4,096]8 heads × 128 rowsHead 1 of WV:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 4,096]each row: 4,096 weights·E: tokenvector⋮[4,096 × 1]=v1: head 1’svalue vectorR1·ER2·E⋮R128·E[128 × 1]v1 is what this token hands over when head 1 picks it (07).k1 and v1 serve query heads 1–4: only 8 KV heads (GQA).v is kept next to k in the KV cache (12).WV: 4.2 M weights, 1,024 × 4,096.32 QUERY HEADS IN 8 GROUPS1–45–89–1213–1617–2021–2425–2829–32Four q heads (plum) share one k and one v head (blue)Head h scores each earlier token j: qh·kj.kj is from h’s group’s KV head: heads 1–4 use KV head 1, …Softmax turns the scores into weights on the vj (06, 07).Each token’s k, v are cached (12): 8 of each, not 32.WQ: 16.8 M weights; the 32 heads are slices of it.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

128NUMBERS PER HEAD
32QUERY HEADS
8KV HEADS

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

05 QKV

Attention is the only place tokens share information (they talk to each other). It starts by giving each token a query, key and value.

  1. 4,096 numbers, one column.

  2. In this head, ·bank’s q asks for a word that says which kind of bank is meant.

  3. ·opened’s k matches ·bank’s q best: their dot product q·k is the largest. So ·bank learns it is a bank that can be opened, not a river bank.

  4. k sets how much ·bank attends to ·opened; v is what it then takes from it. Two vectors, so what gets a token found and what it passes on can differ.

  5. Each piece is a head. 06b shows why many beat one.

    Inside: Fewer keys →

    See also: Step 12 · KV cache

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • token vector
  • weights
  • query
  • key, value
  • head 1
← BlockScores →
  • token vector
  • weights
  • query
  • key, value
  • head 1
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
06 Scores6 of 14

SCORES · how much each token listens to each earlier one · one head

WEIGHTS · HEAD 1Tomopenedanewbankrows: queries q · columns: keys kTomopenedanewbank1.00.35.65.10.60.30.05.30.15.50.03.52.05.32.08each row sums to 1grey: later tokens, masked out (−∞)ROW ·bank, WORKED THROUGHs = qbank·k / √128scoresoftmaxw = eˢ / Σ eˢweightTomopenedanewbank0.493.351.002.861.47.03.52.05.32.08Σ = 1.00 · same as the gold rowillustrative valuesq·k for every earlier token÷ √128mask the futuresoftmaxweights → 07 blends vThis is head 1 of 32 · 06b shows the other 31.HEAD 1 · WHERE THE 5 × 5 GRID COMES FROMTomopenedanewbankq · 128 numbers →Q1 · queries[5 × 128]·TomTomopenedopenedaanewnewbankbankk·128↓K1ᵀ · keys, on their side[128 × 5]=.80.401.02.201.991.30.101.891.202.41.493.351.002.861.47S1 = Q1·K1ᵀ ÷ √128[5 × 5] · scoresmask, softmax →06’s weightsOne cell: ·bank’s q · ·opened’s k = 128 products added = 37.9 ÷ √128 = 3.35 → softmax .525 × 5: rows and columns are tokens. 5 tokens in, 5 × 5 out; a 1,000-token prompt makes a 1,000 × 1,000 grid.128: the length of each q and k. It hides inside every cell (the dot product), not on the grid’s sides.32: head 1 uses numbers 1–128 of each token’s q and k, head 2 the next 128, and so on: 32 grids, all 5 × 5 (06b).values illustrative
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

25SCORES IN THIS HEAD
÷√128SCALING
−∞ → 0MASKED

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

06 Attention scores

Each token's query vector is scored against every earlier token's key vector; softmax turns the scores into weights that sum to 1. A higher score means that earlier token is more relevant to it.

  1. ·opened (0.52) and ·new (0.32) together take 84% of it.

    See also: Step 07 · Blend

  2. A dot product q·k, large when the two vectors point the same way; dividing by √128 keeps softmax from going all-or-nothing.

    See also: Step 05 · QKV

  3. Each score is exponentiated, then divided by the total (54.66); a score 0.5 higher gets about 1.6× the weight.

  4. Their scores are set to −∞ before softmax, and exp(−∞) is 0. Training uses the same mask.

    See also: Step 12 · KV cache

  5. Each cell is one q·k of 128 numbers. The card shows the whole grid as one matrix product, Q·Kᵀ.

    See also: Heads

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • Inference Engineering, chapter 2: ModelsBaseten · page, new tab
  • scores and weights, low → high
  • masked
  • the token we follow
← QKVHeads →
  • scores and weights, low → high
  • masked
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
06b Heads7 of 14

HEADS · 06 showed head 1 · there are 32, each with its own slice of q and k · ·bank’s row outlined

06’S GRID · HEAD 1 OF 32rows: queries q · columns: keys kTom·opened·a·new·bankTom·opened·a·new·bank1.00.35.65.10.60.30.05.30.15.50.03.52.05.32.08·bank’s row: one mix ofthe 5 tokens, summing to 1123432… each k is cut the same wayALL 32 HEADS · 128 NUMBERS EACH·bank’s q (4,096 numbers), cut into 32 slices of 128head 1·bank →·openedhead 2previous tokenhead 3first tokenhead 4itselfKV head 1: one k, one v per tokenall four heads score its k and blend its v+ 7 more groups like this: heads 5–32, KV heads 2–8Why 32? One row gives ·bank one mix. Each head’s row sums to 1 on its own, so ·bankcan lean on ·opened in head 1 and on Tom in head 3 at once. 07 blends all 32 back together.Same work as one 4,096-wide head: 32 × 128 = 4,096 multiplies per cell · heads 2–4 illustrative
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

32GRIDS PER LAYER
128NUMBERS PER HEAD
8KV HEADS

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

06b Attention heads

Attention runs as 32 heads side by side, so ·bank gets 32 rows of weights, not one.

  1. Its weights sum to 1, so ·opened and ·new share one budget.

    See also: Step 06 · Scores

  2. 32 × 128 = 4,096, so the work is the same as one wide head. Heads 2–4 here are made up; real heads learn their own patterns.

  3. Each head keeps its own q, so its own grid. Only k and v are shared, and they are what the KV cache (12) stores: 131 KB per token, not 524 KB.

    See also: Step 12 · KV cache

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • weights, low → high
  • k and v (a KV head)
  • the token we follow
← ScoresBlend →
  • weights, low → high
  • k and v (a KV head)
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
07 Blend8 of 14

BLENDING VALUES · the row for ·bank, 3 of 128 numbers shown, illustrative

HEAD 1 · EACH v × ITS WEIGHT FROM 06, THEN ADDED UPv: first 3 of 128 numbers per token · illustrativew: weight (06)v1 · from KV head 1w × v1Tom.03×0.2 −0.1 0.0 …=0.01 0.00 0.00 …opened.52×−0.2 1.4 0.6 …=−0.10 0.73 0.31 …a.05×0.1 0.2 0.1 …=0.00 0.01 0.00 …new.32×1.1 0.3 −0.4 …=0.35 0.10 −0.13 …bank.08×0.3 −0.2 0.2 …=0.02 −0.02 0.02 …o1 · add each column=0.28 0.82 0.21 …head 1’s output[1 × 128]Biggest weights win: ·opened (.52) and ·new (.32) make up most of o1,so head 1 pulls them into ·bank. The other tokens add almost nothing.JOIN THE 32 HEADS, MIX, ADDWO[4,096 × 4,096]·o₁o₃₂[4,096 × 1]32 × 128=Δ[4,096 × 1]·bank+Δ=·bankthe stream (03)MLP (08)Every head does this with its own w (v shared by 4):32 outputs of 128 numbers each.Joined end to end: 4,096 again.WO mixes the heads into one change, Δ.Δ is added onto ·bank; nothing is replaced.KV cache (12): k and v of past tokensare kept, so a new token only adds its own.WO · JOIN THE 32 HEADS, MIX, ADDWO: Output projection matrix[4,096 × 4,096]lit: head 1 and its block of WO·=+o: heads’ output vec[4,096 × 1]Δ: change vec[4,096 × 1]x: token vec[4,096 × 1]x′: new token vec[4,096 × 1]Join: the 32 head outputs (128 each) stack into o.Mix: WO lets every head reach every number.Block i of WO only meets head i: Δ = Σ WO,i · oi.Add: Δ goes back onto the stream, x′ = x + Δ.WO: 16.8 M weights. Attention in all: 41.9 M per block.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

84%FROM ·opened AND ·new

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

07 Blend

Those weights mix the values, so ·bank now carries the context it needs: mostly ·opened and ·new.

  1. From the Wᵥ projection.

    See also: Step 05 · QKV

  2. Tom and ·a get 8% between them, so they add almost nothing.

  3. 128 numbers; the first three are drawn.

  4. Real heads show patterns like the previous token, the first token, or the subject.

  5. Wₒ turns the 32 heads’ outputs back into 4,096 numbers.

    Inside: Exclusive attention →

    See also: Step 08 · MLP

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • values, head outputs, Δ
  • attention weights (06)
  • model weights
  • the token we follow, in the stream
← HeadsMLP →
  • values, head outputs, Δ
  • attention weights (06)
  • model weights
  • the token we follow, in the stream
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP9 of 14

MLP · SwiGLU · one token, left to right · 4,096 = 52 px, to scale · hover a W for its maths

E·bank[4,096 × 1]Wup[14,336 × 4,096] · 58.7 M weightsWgate[14,336 × 4,096] · 58.7 M weightsSiLU(gate)[14,336 × 1]up[14,336 × 1]⊗elementwiseh · [14,336 × 1]a few units fire, the rest near 0Wdown[4,096 × 14,336]58.7 M weightsΔE[4,096 × 1]added tothe streamWup · EACH ROW ASKS ONE QUESTIONWup: Up proj. matrix— R1 —— R2 —— R14,336 —⋮[14,336 × 4,096]·E: tokenvector⋮[4,096 × 1]=up: widenedvectorR1·ER2·ER14,336·E⋮[14,336 × 1]no biasEach row R is a direction in 4,096-number space.R·E: how far ·bank points along it. One number per row.14,336 rows, so 14,336 numbers: 3.5× wider than E.Wgate asks a second question with its own rows; SiLU(gate) sets how much of eachup number passes (⊗).58.7 M weights here, and the same again in the gate matrix.Wdown · EACH COLUMN WRITES ONE ANSWERWdown: Down proj. matrixc1c2⋯c14,336[4,096 × 14,336]·h: hiddenvectorh1h2h14,336⋮[14,336 × 1]=ΔE: changevector⋮[4,096 × 1]ΔE = h1c1 + h2c2 + … + h14,336c14,336Each column c is a direction in 4,096-number space: one hidden unit’s answer.Unit i adds hi × ci; units near 0 add almost nothing.The sum ΔE is added onto ·bank in the stream.58.7 M weights. The three MLP matrices: 176 M per block.SiLU(gate) · NEGATIVES SQUISHED TO ~0x · σ(x)dashed: y = xgate, beforeafter SiLUEach of the 14,336 gate numbers goes through x · σ(x).Positive: passes almost unchanged. Negative: near 0.The result scales the up numbers (⊗): near 0 switches a unitoff.12 of 14,336 shown, illustrative. GPT-3: GELU.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

176 MMLP WEIGHTS PER BLOCK
5.6 BIN ALL 32 MLPS
70%OF THE MODEL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

08 Multilayer perceptron (MLP)

Then the MLP works on each token alone. It holds 70% of the model’s parameters.

  1. 4,096 numbers, now carrying context from the others.

  2. Each hidden unit is a row of weights that tests for one pattern.

  3. SiLU is x · sigmoid(x): near 0 for negative inputs, near x for positive; multiplying by it lets a unit switch the other off.

    Inside: Gated units →
  4. Each unit that fired writes its own direction back.

    See also: Step 09 · Stacking

  5. A router picks a few per token.

    Inside: Experts →Inside: 896 experts →

    See also: Step 13 · Tricks

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • token vector
  • weights
  • hidden units, dark = fired · illustrative
← BlendStacking →
  • token vector
  • weights
  • hidden units, dark = fired · illustrative
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
09 Stacking10 of 14

THE STACK · the residual stream runs left to right

EARLY · 1–8word pieces, grammarMIDDLE · 9–24facts looked upLATE · 25–32answer moved into ·bankTomopenedanewbank12891524252732⋯⋯⋯⋯⋯blockMLPs on ·opened recall·opened a ·bank → ·accountMeng et al. 2022 · ROMEattention copies itinto ·banklast ·bank→ Output (10)EVERY BLOCKx ← x + Attention(x) · then · x ← x + MLP(x)The lanes are the residual stream: one vector per token. Attention and MLP sit on it and add to it.Roles of early, middle and late blocks: rough, from interpretability studies · illustrative.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

32BLOCKS
8.03 BPARAMETERS IN ALL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

09 Stacking

Stacked up, the blocks divide the work: early ones handle local patterns, middle ones facts, late ones the prediction.

  1. Bigger models stack 80 or more.

  2. The previous token, short phrases.

  3. ·opened ·a ·new ·bank stop being separate tokens and becomes one idea, a money bank, carried onto ·bank.

  4. ROME showed this in GPT-2 XL and GPT-J: editing those MLPs changes the fact the model states.

  5. By block 32 the ·bank vector points at ·account.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Locating and Editing Factual Associations in GPTMeng et al. · paper, new tab
  • ·bank, the token we predict from
  • attention: spans all tokens
  • MLP: each token on its own
← MLPOutput →
  • ·bank, the token we predict from
  • attention: spans all tokens
  • MLP: each token on its own
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
10 Output11 of 14

OUTPUT HEAD · the last vector scored against every token, illustrative logits· illustrative logits

UNEMBEDDING MATRIX⋮row 2,759 · accountrow 0↓ row 128,255[128,256 × 4,096]1 row per tokenin the model’s vocabulary··bank[4,096× 1]=⋮21.0−3.14.2⋮0.7z[128,256 × 1]logits: one scoreper tokenTOP 5 OF 128,256logit zp = eᶻ / Σ eᶻsoftmaxprobability paccountbranchloancardmanager128,251others21.018.217.617.116.56.0 eachillustrative values5.2%2.9%1.7%1.0%3.4% togetherΣ = 100%85.9%SAMPLEpick one token at random, weighted by p → ·account about 86 times in 100Temperature T divides z first: e^(z/T). T below 1 sharpens toward account, above 1 flattens. Here T = 1.WU · ONE DOT PRODUCT PER TOKENW_U[128,256 × 4,096]·h=z[128,256 × 1]one score per tokenh: [4,096 × 1]zaccount = 21.0zaccount = row 2,759 · h = Σ wihi over 4,096 = 21.0illustrative valuesEach row of WU is one token’s output direction.128,256 rows, so 128,256 dot products, all at once.h points most along the account row, so zaccount is largest.Softmax turns z into next-token odds.525 M weights. Llama keeps WU separate; GPT-3 reuses E.
·bank
→
128,256 rows × 4,096
one score per token
PROBABILITIES · top five, then the rest
·account85.9%
·branch5.2%
·loan2.9%
·card1.7%
·manager1.0%
others3.4% together
050100
PROBABILITY, %
Temperature T divides z first: e^(z/T). T below 1 sharpens toward account, above 1 flattens. Here T = 1.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

86%·account AT T = 1

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

10 Output

At the top, the last token's vector becomes probabilities for every token, and the sampler picks one.

  1. After 32 blocks it holds everything worked out about the prompt.

  2. Each logit is a dot product with that token’s row; ·account scores 21.0.

  3. A lead of 2.8 in score becomes a 16× lead in probability.

    Inside: Softmax collapse →
  4. Below 1 the favourite gains; at 0 it always wins.

  5. Then it is appended for the next pass.

    See also: Step 11 · Passes

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • last vector, chosen token
  • weights
  • logits
  • probabilities
← StackingPasses →
  • last vector, chosen token
  • weights
  • logits
  • probabilities
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
11 Passes12 of 14

FOUR PASSES · each row is one run of the whole model

TOKENS IN · dashed: already run, keys and values reused
READS EVERY WEIGHT
OUT
PASS 1t = 0.0 msTom·opened·a·new·bankall weights→·account
PASS 2t = 2.0 msTom·opened·a·new·bank·accountall weights→.
PASS 3t = 4.0 msTom·opened·a·new·bank·account.all weights→·He
PASS 4t = 6.0 msTom·opened·a·new·bank·account.·Heall weights→·deposited
pass 1 · 2.0 mspass 2 · 2.0 mspass 3 · 2.0 mspass 4 · 2.0 ms012345678TIME, MSOne B200, batch 1, BF16 weights, ceiling not measured.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

16.06 GBREAD PER PASS
2.0 msPER PASS AT BEST

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

11 Passes

That pick is appended and the model runs again: one pass per token, so speed is tokens per second.

  1. They are all known, so they go through side by side: prefill.

  2. The sampler’s pick becomes the last input token.

  3. Earlier tokens’ keys k and values v are reused.

    See also: Step 12 · KV cache

  4. A pass can’t start until the previous one has picked its token.

  5. On one B200 at batch 1, that caps speed at ≈500 tokens/s.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 0 InferenceBaseten · page, new tab
  • prompt tokens
  • newest token
  • already run
  • the model
  • token picked
  • time
← OutputKV cache →
  • prompt tokens
  • newest token
  • already run
  • the model
  • token picked
  • time
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
12 KV cache13 of 14

KV CACHE · one column per token, 32 layers each

KEYS AND VALUES SAVED · one column per token
layer 32layer 1Tom·opened·a·new·bank·accountONE CELLone token's key and value
in one layer
ONE COLUMNone token’s keys and values, all 32 layersNEWEST COLUMN·account computes only its own k and v,
then the cache grows by one column
ONE CHAT'S CACHE BY CONTEXT LENGTH · the weights for scale
KV cache · 4,096 tokens0.54 GB
KV cache · 8,192 tokens1.07 GB
KV cache · 32,768 tokens4.29 GB
KV cache · 131,072 tokens17.2 GB
Weights · shared by all chats16.06 GB
05101520MEMORY, GBLlama 3.1 8B, BF16, one chat.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

4×SMALLER WITH 8 KV HEADS
0.524 MBPER TOKEN IF ALL 32 HEADS KEPT THEIR OWN

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

12 KV cache

To avoid recomputing, earlier tokens' keys and values are saved; each pass adds only its own, so memory keeps growing.

  1. Computed once, kept in GPU memory, read by every later token.

  2. 32 layers × 8 KV heads × 128 numbers × k and v × 2 bytes.

    See also: Step 05 · QKV

  3. At 131,072 tokens one chat’s cache outweighs the model.

    Inside: Linear attention →Inside: Compressed cache →
  4. That is why cache memory, not weights, usually limits how many users fit on a GPU.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • cached keys and values
  • newest token’s K and V
  • weights
  • memory, GB
← PassesTricks →
  • cached keys and values
  • newest token’s K and V
  • weights
  • memory, GB
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
13 Tricks14 of 14

FOUR GROUPS · by what they fix

BF16 + cache16-bit weights16.6 GB
FP8 + cache8-bit weights8.6 GB
4-bit (Q4_K_M) + cache4-bit weights5.5 GB
Same, 32 KV headsno sharing: 4× the cache7.1 GB
01020GB READ PER GENERATED TOKENLlama 3.1 8B, batch 1, 4,096-token context. Teal is weights, purple is the KV cache.
MEMORYweights and KV cache
NAMED TECHNIQUES
FP8 and 4-bit weightsGQAMLA
MORE PER PASSexperts, guesses
NAMED TECHNIQUES
Mixture of expertsSpeculative decoding
ATTENTIONfewer pairs scored
NAMED TECHNIQUES
FlashAttentionSliding window
OUTPUTcontrol, not speed
NAMED TECHNIQUES
Constrained decodingSampling settings
BF164-bitsame weights, ¼ of the bits32 KV heads8 KV headsheads share keys and valuesDeepSeek-V3: 37 B of 671 Bparameters active per token✓✓✓✗a small model guesses,one big-model pass checksall earlierlast 3 onlyscore fewer pairs,or do the same work fasteronly tokens that fitthe format can be pickedsame speed, valid output
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

13 Tricks

Generation is limited by memory, so almost every trick reads fewer bytes, at home or in a data centre.

  1. 4-bit weights take a third of BF16’s bytes, and usually cost a little accuracy.

  2. Grouped-query attention stores 8 KV heads for 32 query heads; DeepSeek’s latent attention compresses them further.

    See also: Step 12 · KV cache

  3. DeepSeek-V3 reads 37 B of its 671 B parameters per token.

    See also: Step 08 · MLP

  4. A cheap drafter guesses; the big model verifies all of them for one read of the weights.

    Inside: Draft and verify →
  5. FlashAttention never writes the score grid to HBM; sliding windows stop the cache growing.

  6. Constrained decoding zeroes tokens that would break JSON or a grammar.

    See also: Step 10 · Output

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 5 TechniquesBaseten · page, new tab
  • GPU Glossary · Memory-boundModal · page, new tab
  • weights
  • KV cache
  • attention work
  • output control
  • skipped or rejected
← KV cache1 · A data centre, top down →
  • weights
  • KV cache
  • attention work
  • output control
  • skipped or rejected
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens↳ Byte pairs02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens ↳ Byte pairs1 of 14

BYTE-PAIR MERGES · "·opened", illustrative merge order

Tokenizer↖ Tokens
ROWSYMBOLS AFTER THE MERGE
start · 7·opened
merge 1 · e+n·opened
merge 2 · e+d·opened
merge 3 · ·+o·opened
merge 4 · ·o+p·opened
merge 5 · en+ed·opened
merge 6·openeda token in Llama 3.1’s list

Merge order and pair counts illustrative. · is the leading space.

TOP PAIRS · before merge 1

e n9.4 M
e d9.0 M
· o8.6 M
o p7.9 M
n e6.1 M
p e5.2 M

05 M9.4 M

TIMES SEEN, MILLIONS

The winner becomes merge 1. Then every count is redone.

128,256ENTRIES
4.2CHARACTERS PER TOKEN, THIS PROMPT
esc
Full video on YouTube

Byte pairs

The tokenizer's pieces come from a list built by merging the most frequent pair of symbols, again and again.

  1. "·opened" begins as seven symbols, the space included.

  2. The most frequent pair becomes one new symbol.

  3. Each merge adds one entry to the list; training stops at the target size.

  4. Rare words stay in pieces.

    See also Step 01 · Tokens

  5. The prompt’s 21 characters become 5 tokens.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Neural Machine Translation of Rare Words with Subword UnitsSennrich et al. · paper, new tab
  • text
  • tokens
  • the pair merged in each row
← Back to TokensBytes, not tokens →
  • text
  • tokens
  • the pair merged in each row
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens↳ Byte pairs02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens ↳ Byte pairs1 of 14

BYTE PATCHES · the same prompt, patch edges illustrative

Tokenizer↖ Tokens
NEXT-BYTE ENTROPY, BITS024cut above 2.5 bitsTom·opened·a·new·bank4 bytes25244one per patch: 6 here, against 21 bytes3 bytes7245one per token: 5 here, against 21 bytesBIG-MODEL STEPS

Entropy from a small byte model; values illustrative.

Token edges from Llama 3.1’s tokenizer; · is the leading space.

CUT BY
21BYTES
8BLARGEST BLT IN THE STUDY
esc
Full video on YouTube

Bytes, not tokens

Byte-level models skip that list and group raw bytes into patches by difficulty; either way, the model gets numbers.

  1. No list to learn, and no word is ever unknown.

  2. Predictable bytes score low; the start of a new word scores high.

  3. Easy stretches become long patches; hard ones get more, shorter patches.

  4. Compute follows difficulty instead of a fixed vocabulary.

  5. It was trained on up to 4 trillion bytes.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Byte Latent Transformer: Patches Scale Better Than TokensPagnoni et al. · paper, new tab
  • bytes
  • next-byte surprise (entropy)
  • patches
  • tokens
  • the patch holding ·bank
  • the token ·bank
← Byte pairs← Back to Tokens
  • bytes
  • next-byte surprise (entropy)
  • patches
  • tokens
  • the patch holding ·bank
  • the token ·bank
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Engram04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Engram3 of 14

ENGRAM · hashed n-gram lookup beside the stream, illustrative rows

Engram↖ Stream
TOKENS SO FARTomopenedanewbankkey: “opened a new”hash × 3key → row numberTABLE 1TABLE 2TABLE 3memorygatecontextstream+for “bank”opened a newBENCHMARK GAIN OVER AN EQUAL MOE, POINTSMMLU+3.0BBH+5.0HumanEval+3.00246

27B models, same parameters and FLOPs; gains from the paper's Table 1, NIAH from Table 2. Rows and hash count illustrative.

84.2MULTI-QUERY NIAH, EQUAL MoE
97.0MULTI-QUERY NIAH, WITH ENGRAM
esc
Full video on YouTube

Engram

Layers spend effort rebuilding common facts in the stream; Engram looks them up in a table keyed by word sequences.

  1. At ·bank, the key could be ·opened ·a ·new.

  2. Several hash functions point into several tables, so one collision does little harm.

  3. The table can be huge and still cheap to use.

  4. If it doesn't fit, the gate closes and little is added.

  5. At 27B, Engram beat an equal-size MoE on knowledge and reasoning.

    See also Step 08 · MLP

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Conditional Memory via Scalable LookupCheng et al. · paper, new tab
  • tokens
  • lookup tables
  • gate
  • the token we follow
← Back to Stream
  • tokens
  • lookup tables
  • gate
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Hyper-connections04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Hyper-connections3 of 14

HYPER-CONNECTIONS · one token, four streams, Llama width for scale

Hyper-connections↖ Stream
from Embedding4 streams, 4,096 numbers eachMIXING WEIGHTS, 4 × 4learned for each layerdarker = stronger linkBlock 1attention + MLPread a blendwrite to all fourBlock 2attention + MLPread a blendwrite to all fourMixing weights illustrative.from EmbeddingBlock 1attention + MLPread+add, weight 1Block 2attention + MLPread+add, weight 11carried on with weight 1
Residual
4STREAMS
16,384NUMBERS PER TOKEN
esc
Full video on YouTube

Hyper-connections

The stream is one vector; hyper-connections widen it into several copies, with learned mixing between them.

  1. Each block adds its output to it, weight one.

  2. For a Llama-width model that is 16,384 numbers per token.

  3. The network learns how strongly each depth feeds the next.

  4. Layers can effectively be reordered or run side by side.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Hyper-ConnectionsZhu et al. · paper, new tab
  • the stream for ·bank
  • block
  • mixing weights
← Back to StreamKeeping the mix tame →
  • the stream for ·bank
  • block
  • mixing weights
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Hyper-connections04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Hyper-connections3 of 14

HC · each mixing matrix free to take any values

Keeping the mix tame↖ Stream
ONE MIXING MATRIX, 4 × 40.900.600.200.300.501.100.400.200.100.301.200.500.400.200.600.802.002.202.102.001.902.202.401.80row sumscolumn sumsMatrix values illustrative.LARGEST SIGNAL GAIN ACROSS DEPTH, × (LOG SCALE)1101001,00010,000Maximum composite gain reported in the mHC paper.HC≈3,000×mHC≈1.6×
Mixing
27BLARGEST MODEL TRAINED
6.7%EXTRA TIME
esc
Full video on YouTube

Keeping the mix tame

That mixing can amplify signals 3,000-fold, so mHC makes each mix keep the total, and the stream stays tame.

  1. Small gains compound across depth; in the paper's runs they reached about 3,000.

  2. A mix can move signal between streams but never create it.

  3. Divide rows by their sums, then columns, and repeat.

  4. Training stayed stable up to 27B, for 6.7% more time.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • mHC: Manifold-Constrained Hyper-ConnectionsXie et al. · paper, new tab
  • unstable gain (HC)
  • mixing weights; stable gain (mHC)
← Hyper-connections← Back to Stream
  • unstable gain (HC)
  • mixing weights; stable gain (mHC)
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Attention residual04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Attention residual3 of 14

ATTENTION OVER DEPTH · one token at one layer, weights illustrative

Attention residual↖ Stream
STREAM OF THE TOKEN WE FOLLOWnext layerevery output added with weight 11 block = 8 outputsWEIGHT ON EACH EARLIER OUTPUT, SHARE OF 100.10.2116324864Earlier output, 1 to 64, oldest first (attention, MLP, attention, …)HOW ONE LAYER MAKES ITS WEIGHTS1 · SCOREsᵢ = w · RMSNorm(vᵢ)One learned query w per layerscores every earlier output vᵢ.2 · SOFTMAX OVER DEPTHαᵢ = exp(sᵢ) ÷ Σⱼ exp(sⱼ)Turns the 64 scores into weightsthat are positive and sum to 1.3 · MIXh = Σᵢ αᵢ · vᵢThe layer's input is theweighted sum of earlier outputs.HOW A PLAIN STREAM MIXES1 · SCOREnoneNo query: nothing looks at whatthe earlier outputs contain.2 · WEIGHTSevery weight = 1Every output is added once,an equal 1/64 share of the sum.3 · MIXh = v₁ + v₂ + … + v₆₄The layer's input is theplain sum of earlier outputs.→→
equal, 1/64
Depth weights
64EARLIER OUTPUTS IN LLAMA 8B
1.4TTRAINING TOKENS, KIMI'S 48B TEST
esc
Full video on YouTube

Attention residual

The stream adds every block's output with weight one; attention residual lets each layer choose which earlier outputs to use.

  1. Deep in the model, each new output is a small share of a growing sum.

  2. One learned query per layer scores every earlier output; a softmax over depth turns the scores into weights.

  3. Attend over block summaries, not every layer.

  4. Output sizes and gradients stayed more even across depth.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Attention ResidualsKimi Team · paper, new tab
  • stream of the token we follow
  • weight on an earlier output
  • block group
← Back to Stream
  • stream of the token we follow
  • weight on an earlier output
  • block group
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block↳ Normalisation05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
04 Block ↳ Normalisation4 of 14

RMSNORM · the ·bank vector before and after, first 4 of 4,096 numbers

RMSNorm↖ Block
−3−1.50+1.5+3Value, no unitLlama 3.1 8B, layer 1, before the gain.As it enterstiny numbersRMS of all 4,096÷ RMSAfter dividingRMS is now 1× gain: one learned weight per number
−0.0193
dim 0
0.0057
dim 1
−0.0096
dim 2
−0.0002
dim 3
−2.13
dim 0
0.63
dim 1
−1.06
dim 2
−0.02
dim 3
0.0090RMS IN
1.0RMS OUT
2PER BLOCK
esc
Full video on YouTube

Normalisation

Each block first rescales its copy of the stream: RMSNorm brings every token's vector to a fixed size.

  1. The root of the mean square: 0.0090 for ·bank as it enters.

  2. The vector keeps its direction; its RMS becomes 1.

  3. 4,096 weights let the model restore what it needs.

  4. Training becomes unstable.

    See also Step 03 · Stream

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Root Mean Square Layer NormalizationZhang & Sennrich · paper, new tab
  • ·bank, the token we follow
  • learned gain
← Back to BlockWithout normalising →
  • ·bank, the token we follow
  • learned gain
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block↳ Normalisation05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
04 Block ↳ Normalisation4 of 14

INPUT → OUTPUT · one number, curves illustrative

RMSNorm↖ Block
−4−2024−101Input value, no unitOUTPUT VALUE, NO UNITScatter illustrative; formulas from the papers.α, s: learned
Dynamic Tanh (DyT)tanh(α·x)
Derferf(α·x + s)
Curve
esc
Full video on YouTube

Without normalising

A tanh curve can replace that normalisation, squashing large values instead of measuring the vector.

  1. In trained transformers, the points form an S: straight in the middle, flat at the ends.

  2. Dynamic Tanh: no mean, no RMS, just a per-number curve.

  3. Usually with no new tuning.

  4. A search over curve shapes found it beat both RMSNorm and DyT.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Transformers without NormalizationZhu et al. · paper, new tab
  • Stronger Normalization-Free TransformersChen et al. · paper, new tab
  • norm layer input → output
  • the curve
  • one number of ·bank
← Normalisation← Back to Block
  • norm layer input → output
  • the curve
  • one number of ·bank
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV↳ Fewer keys06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
05 QKV ↳ Fewer keys5 of 14

KV CACHE PER TOKEN · Llama 3.1 8B shape, BF16

KV heads↖ QKV
32 QUERY HEADS32 KEY–VALUE HEADS00.10.20.30.40.5KV cache per token, MBLlama 3.1 8B shape, 32 layers, BF16Multi-head · 320.524 MBGrouped · 8 · Llama0.131 MBMulti-query · 10.016 MBLatent · MLAone latent per token · DeepSeek-V2, no bar: a different model
KV heads
4×SMALLER THAN MULTI-HEAD (LLAMA)
−93.3%DEEPSEEK-V2 CACHE
esc
Full video on YouTube

Fewer keys

Those keys and values are kept for every past token, so models share them across heads, or compress them.

  1. One for each query head: 0.524 MB per token.

  2. Sixteen kilobytes, with some quality lost.

  3. Close to multi-head quality, close to multi-query speed.

  4. Keys and values are rebuilt from it; DeepSeek-V2 cut its cache 93.3%.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GQA: Training Generalized Multi-Query Transformer ModelsAinslie et al. · paper, new tab
  • DeepSeek-V2DeepSeek-AI · paper, new tab
  • query head
  • key–value head
  • Llama 3.1 8B's choice
← Back to QKVRead fewer keys →
  • query head
  • key–value head
  • Llama 3.1 8B's choice
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV↳ Fewer keys06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
05 QKV ↳ Fewer keys5 of 14

SPARSE ATTENTION · one query over a long context, illustrative scores

KV heads↖ QKV
" bank"0326496128kept above the cutINDEXER SCORE, NO UNITPosition in context, thousands of tokens
2,048KEYS READ
1.6%OF 128K
esc
Full video on YouTube

Read fewer keys

Models can also read fewer keys: a cheap scorer picks each query's top few thousand, and attention runs on those.

  1. Few heads, low precision: far less work than real attention.

  2. Every other past key is skipped for this token.

  3. Cost per token stops growing with context length.

  4. The indexer was trained first, then everything together.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • DeepSeek-V3.2DeepSeek-AI · paper, new tab
  • key read
  • key skipped
  • the query, "·bank"
  • indexer score
← Fewer keys← Back to QKV
  • key read
  • key skipped
  • the query, "·bank"
  • indexer score
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend↳ Exclusive attention08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
07 Blend ↳ Exclusive attention8 of 14

EXCLUSIVE SELF-ATTENTION · "·bank", first 3 of 128 numbers, illustrative values

Head output↖ Blend
y · attention outputv · own valuealong v · removedz · the rest = new outputy0.55−0.050.40v0.30−0.200.20z0.100.250.10GAIN OVER STANDARD ATTENTION0.00.51.01.50.7B+0.261.3B+1.032.7B+1.36Average accuracy gain, points8 tasks, context 2,048, same training
Attention
2LINES OF CODE
esc
Full video on YouTube

Exclusive attention

That blend often echoes the token's own value; subtracting the echo keeps what others added, and improved small models.

  1. Part of what attention returns is the token's own information, echoed back.

  2. Project y onto v, then subtract that part from y.

  3. The MLP still sees the token itself through the stream.

    See also Step 03 · Stream

  4. Up to +1.36 points average accuracy at 2.7B.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Exclusive Self AttentionZhai · paper, new tab
  • y · attention output
  • v · own value
  • part along v
  • z · the rest
← Back to Blend
  • y · attention output
  • v · own value
  • part along v
  • z · the rest
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ Gated units09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ Gated units9 of 14

ACTIVATIONS · output against input, one hidden unit

ACTIVATION
Gate↖ MLP
4
2
0
−4
−2
2
4
INPUT, NO UNIT
OUTPUT
ReLU
SiLU
INSIDE THE MLP · one token, 4,096 numbers in
x
W gate4,096 × 14,336
W up4,096 × 14,336
SiLU
ReLU
×
W down
4,096 out
gate · decides how much passes
SwiGLU3 matrices · Llama14,336 units
2-matrix MLPsame parameters21,504 units
0
5,000
10,000
15,000
20,000
25,000
HIDDEN WIDTH, UNITS
Llama 3.1 8B, one block.
3MATRICES
176 MPARAMETERS PER BLOCK
esc
Full video on YouTube

Gated units

Inside that MLP, Llama multiplies two projections, one passed through SiLU, instead of applying a single activation.

  1. Below zero the output is flat at zero; above, it is the input.

  2. x · sigmoid(x): a small negative dip, then nearly x.

  3. One decides how much passes; the other carries the content.

  4. The hidden width shrinks to keep parameters level: 14,336, not 21,504.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GLU Variants Improve TransformerShazeer · paper, new tab
  • token vector
  • weights
  • activation
  • not used in this setting
  • not used in this setting
← Back to MLP
  • token vector
  • weights
  • activation
  • not used in this setting
  • not used in this setting
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ Experts09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ Experts9 of 14

MIXTURE OF EXPERTS · DeepSeek-V3 layer, one token

LAYER
MLP↖ Block
Router7,168 × 256
ROUTER SCORE · one per expert, no unit
expert 1256
256 ROUTED EXPERTS · each one an MLP · 8 gold run
8 outputs × score
+
One MLP4,096 → 14,336 → 4,096every weight works for every token
x
7,168 out
Shared expertruns for every token
PARAMETERS, B · per token
37 B work671 B in the model
WORK PER TOKEN, B · same scale
Llama 3.1 8B8
DeepSeek-V337
8 of 256EXPERTS PER TOKEN
5.5%OF PARAMETERS ACTIVE
esc
Full video on YouTube

Experts

Bigger models split the MLP into many experts, and a router runs only a few of them per token.

  1. A small matrix: one score per expert.

  2. Their outputs are added, weighted by the router's scores.

  3. It holds what all tokens need.

  4. A large model at a small model's cost per token.

    See also Step 13 · Tricks

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Outrageously Large Neural NetworksShazeer et al. · paper, new tab
  • DeepSeek-V3 Technical ReportDeepSeek-AI · paper, new tab
  • token vector
  • weights
  • runs for this token
← Back to MLPKeeping experts busy →
  • token vector
  • weights
  • runs for this token
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ Experts09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ Experts9 of 14

EXPERT LOAD · share of tokens per expert, illustrative

BALANCING
MLP↖ Block
0
10
20
30
40
SHARE OF TOKENS, %
- - - even share, 12.5%
34
expert 1
29
expert 2
9
expert 3
8
expert 4
7
expert 5
5
expert 6
4
expert 7
4
expert 8
34 with no balancing
ROUTER CORRECTION
Aux losspenalty on uneven load · off
Biasadded to each score · off
↻ more tokens → more training → higher scoresThe loop feeds experts 1 and 2; six others starve.
GPU busy, %
GPU 1 · 100
GPU 2 · 85
GPU 3 · 26
GPU 4 · 24
GPU 5 · 21
GPU 6 · 15
GPU 7 · 12
GPU 8 · 12
8 experts on 8 GPUs, one batch; loads illustrative. Busy relative to the busiest GPU.
esc
Full video on YouTube

Keeping experts busy

Left alone, that router keeps picking the same few experts, so balancing keeps every expert learning and every GPU busy.

  1. The loop collapses onto a few.

  2. It works, but pulls against the main training goal.

  3. Overloaded experts' scores drop a little each step; idle ones rise.

  4. Experts live on different GPUs.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Switch TransformersFedus et al. · paper, new tab
  • overloaded
  • expert load
  • barely used
  • with no balancing
  • with no balancing
← Experts← Back to MLP
  • overloaded
  • expert load
  • barely used
  • with no balancing
  • with no balancing
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ 896 experts09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ 896 experts9 of 14

LATENT MoE · Kimi K3, one token through one layer

MLP↖ Block
x · 7,168
3,584
down · 7,168 → 3,584
up · 3,584 → 7,168
896 ROUTED EXPERTS · 28 × 32 · 16 gold run
16 outputs × score
+
7,168 out
Shared expert 1full width, 7,168
Shared expert 2full width, 7,168
always run
PARAMETERS · per token
104 B active2.78 T in the model
16 of 896EXPERTS PER TOKEN
3.7%OF PARAMETERS, PER TOKEN
esc
Full video on YouTube

896 experts

Kimi K3 pushes experts further: hundreds of them, each working at half width, only a few running per token.

  1. 7,168 numbers become 3,584.

  2. Each expert is cheaper, so there can be more of them.

  3. 2.78 T parameters, 104 B active per token.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Kimi K3: Open Frontier IntelligenceKimi Team · paper, new tab
  • LatentMoEElango et al. · paper, new tab
  • token vector
  • weights
  • runs for this token
  • latent, half width
← Back to MLPStable at 896 →
  • token vector
  • weights
  • runs for this token
  • latent, half width
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ 896 experts09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ 896 experts9 of 14

STABLE LATENT MoE · three guards, Kimi K3

MLP↖ Block
FOUR MATRIX MULTIPLIES IN A ROW · down, expert in, expert out, up
x
down7,168 → 3,584
expert in×16 experts
expert out×16 experts
+
RMSNorm
up3,584 → 7,168
7,168 out
soft cap≤ 100
Router896 scores
quantile bias
resets the size
EXPERT OUTPUT · against input, one unit, illustrative
+100
0
−100
−4
0
4
INPUT, NO UNIT
- - - cap at ±100
without a cap
with the cap
QUANTILE BALANCING · one expert, router scores over a batch, illustrative
bias set at this score
routed
here
low score
high score
TOKENS PER SCORE BAND · each expert should get 1 in 56 tokens
2.78 TPARAMETERS IN KIMI K3
esc
Full video on YouTube

Stable at 896

At that scale training can blow up, so a norm, capped outputs and quantile-based balancing keep 896 experts stable.

  1. Down, expert in, expert out, up.

  2. See also Normalisation

  3. tanh limits both branches; no output exceeds 100.

  4. The bias is set at the score that sends this expert exactly its share, in one step.

    See also Experts, screen 2 · Keeping experts busy

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Kimi K3: Open Frontier IntelligenceKimi Team · paper, new tab
  • token vector
  • weights
  • expert output, capped
  • without the cap
← 896 experts← Back to MLP
  • token vector
  • weights
  • expert output, capped
  • without the cap
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output↳ Softmax collapse11 Passes12 KV cache13 Tricks
10 Output ↳ Softmax collapse11 of 14

LOGIT MARGIN · target logit minus log-sum-exp of the others

PRECISION
Softmax↖ Output
|GRADIENT| · target token, log scale
1
10−6
10−12
10−18
exactly 0
0
10
20
30
40
LOGIT MARGIN m, no unit
Hard-target cross-entropy; curve shapes from the paper's Figure 1.
BF16 · ≈5.5
FP32 · 16.6
FP64 · 36.7
standard softmax: rounded to 0
fixed: keeps shrinking
“·account” · 1.8
m = 1.8
0
TARGET GRADIENT
at the marker
WORKED EXAMPLElogits 21 and 0 · BF16
a e−21 ≈ 7.6 × 10−10the other token's term (the leader's is 1)
→
b 1 + 7.6 × 10−10 = 1BF16's step near 1 is 2−7 ≈ 0.008
→
c p = 1, so 1 − p = 0the target's gradient is exactly 0
esc
Full video on YouTube

Softmax collapse

Those probabilities come from softmax, which can round a training gradient to exactly zero when one token leads by enough.

  1. For ·account here it is 1.8.

  2. In FP32 this happens past the cliff line. The worked example below shows one case in BF16.

  3. Gradients stop summing to zero, which causes loss spikes with Adam.

  4. Compute the target's gradient as minus the sum of the others'.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • A Fix for Softmax and Cross-EntropyBarron et al. · paper, new tab
  • target gradient, fixed softmax
  • standard softmax, lost past the cliff
  • precision cliff, selected
  • the token we follow, ·account
← Back to Output
  • target gradient, fixed softmax
  • standard softmax, lost past the cliff
  • precision cliff, selected
  • the token we follow, ·account
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache↳ Linear attention13 Tricks
12 KV cache ↳ Linear attention13 of 14

MEMORY PER CHAT · a Llama-shaped model, if its attention were linear

CONTEXT
Cache↖ KV cache
SOFTMAX ATTENTION · one column per token
keys above, values below
token 1 → … grows by one per token
newest
drop softmax,
regroup the sums
LINEAR STATE · one per head
128 × 128 numbers
same size at token 10 or 1,000,000
NEW TOKEN
kvq
adds k·vᵀ
q reads the state
MEMORY PER CHAT · GB, log scale
KV cache
one chat, 128K tokens
Linear state
any length
17.2 GB
0.034 GB · the cache passes this at 256 tokens
0.01
0.1
1
10
100
GB, BF16
Llama 3.1 8B shape, BF16; the linear version is hypothetical. The cache bar follows the context control.
O(N)WORK FOR N TOKENS, WAS O(N²)
esc
Full video on YouTube

Linear attention

Instead of a cache that keeps growing, linear attention keeps a fixed-size summary of every past key and value.

  1. The cache grows with every token.

  2. Keys times values can be summed once into one matrix per head.

  3. Same cost at token 10 or token 1,000,000.

  4. The real cache passes that at 256 tokens.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Transformers are RNNsKatharopoulos et al. · paper, new tab
  • cached keys and values
  • linear state
  • newest token
  • memory, GB
← Back to KV cacheForgetting on purpose →
  • cached keys and values
  • linear state
  • newest token
  • memory, GB
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache↳ Linear attention13 Tricks
12 KV cache ↳ Linear attention13 of 14

STATE UPDATE · one head, illustrative values

UPDATE
Cache↖ KV cache
BEFORE
state S, 8 × 8 of 128 × 128 shown
darker = larger number, illustrative
AFTER THE GATE
every cell × α, α = 0.8
AFTER THE DELTA WRITE
key k’s row replaced
gold row = what key k returns
× α
write v for key k
S ← α·S·(I − β·k·kᵀ) + β·v·kᵀ
α fades · β·k·kᵀ removes the old value · β·v·kᵀ writes the new
MIXED STACK · Kimi Linear, layers in order
Three linear layers, then one full-attention layer. Only the full layers keep a growing cache.
linear (KDA)
and full (MLA)
layer 1 →
full
… repeats
−75%KV CACHE
6×DECODING AT 1M TOKENS
esc
Full video on YouTube

Forgetting on purpose

A fixed summary fills up, so newer designs fade it and overwrite old values, often alongside a few full-attention layers.

  1. Old associations pile up and blur.

  2. Recent tokens weigh more than distant ones.

  3. Read what the key returns, then correct it toward the new value.

  4. Kimi Linear cut the cache by up to 75%.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Gated Delta NetworksYang et al. · paper, new tab
  • Kimi LinearKimi Team · paper, new tab
  • linear state
  • the key being written
  • full-attention layer, keeps a cache
  • linear layer
← Linear attention← Back to KV cache
  • linear state
  • the key being written
  • full-attention layer, keeps a cache
  • linear layer
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache↳ Compressed cache13 Tricks
12 KV cache ↳ Compressed cache13 of 14

CSA AND HCA · one million tokens of context, DeepSeek-V4-Pro

LAYER TYPE
Cache↖ KV cache
CONTEXT · 1,000,000 tokens, not to scale
token 1 →
999,872 older tokens: compressed below
last 128 tokens, kept whole
read in full by every layer
now
CSA · compressed sparse
250,000 entries; an indexer picks the top 1,024 (0.4%)
group → 1 entry
dark = read · light = skipped · 70 of 250,000 drawn
HCA · heavily compressed
7,813 entries, all read
group → 1 entry
14 of 7,813 drawn
LAYERS · in order
CSA and HCA alternate
CSA
HCA
CSA
HCA
CSA
HCA
CSA
HCA
CSA
HCA
CSA
HCA
… repeats
DeepSeek-V4-Pro at one million tokens; entry counts from the DeepSeek-V4 report
4TOKENS PER CSA ENTRY
128PER HCA ENTRY
esc
Full video on YouTube

Compressed cache

The cache can also be compressed: DeepSeek-V4 squeezes groups of tokens into single entries and reads only the most relevant.

  1. Nearby detail matters most.

  2. Then reads only the top 1,024 of 250,000.

  3. Then reads all 7,813 of them.

  4. Fine detail sparsely, coarse overview densely.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • DeepSeek-V4DeepSeek-AI · paper, new tab
  • cache entry, read
  • entry skipped
  • older tokens, before compression
  • current token
← Back to KV cacheHow much smaller →
  • cache entry, read
  • entry skipped
  • older tokens, before compression
  • current token
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache↳ Compressed cache13 Tricks
12 KV cache ↳ Compressed cache13 of 14

KV CACHE AT ONE MILLION TOKENS · relative to a grouped-query baseline

Cache↖ KV cache
CACHE SIZE · % of a BF16 grouped-query cache
Grouped-query
BF16 GQA8, head dim 128
100%
DeepSeek-V3.2
latent attention
≈20%
derived figure
DeepSeek-V4-Pro
CSA + HCA, FP8 entries
≈2%
1/50 of the baseline
0%
25%
50%
75%
100%
% of BF16 grouped-query cache, one million tokens
COMPUTE PER TOKEN AT 1M · % of DeepSeek-V3.2
DeepSeek-V3.2
100%
DeepSeek-V4-Pro
27%
0%
25%
50%
75%
100%
Per the DeepSeek-V4 report, one million tokens of context. V3.2 cache bar derived from its 10% figure.
esc
Full video on YouTube

How much smaller

At a million tokens, that cache is about 2% of a standard grouped-query cache: far less memory per chat.

  1. The report's baseline: BF16 GQA8, head dimension 128.

  2. Compression, sparse reads and FP8 storage combined.

  3. V3.2 already used latent attention.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • DeepSeek-V4DeepSeek-AI · paper, new tab
  • grouped-query baseline
  • DeepSeek-V3.2 cache
  • DeepSeek-V4-Pro cache
  • compute per token
← Compressed cache← Back to KV cache
  • grouped-query baseline
  • DeepSeek-V3.2 cache
  • DeepSeek-V4-Pro cache
  • compute per token
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks↳ Draft and verify
13 Tricks ↳ Draft and verify14 of 14

SPECULATIVE DECODING · one big-model pass, illustrative draft

More per pass↖ Tricks
DRAFT · a small model guesses ahead
Draft modelsmall and cheap
·account
.
·He
·deposited
all four enter together
Big model · one pass, all four at oncereads 16.06 GB of weights once
·account
.
·He
·deposited
CHECKED
·account
✓
.
✓
·He
✓
·deposited
✗
·has
big model's pick
4 tokens from one pass · draft words illustrative
TIME · ms, to produce four tokens
Draft and verify
One token per pass
2.0 ms · 4 tokens
8.0 ms
0
2
4
6
8
Time, ms
One B200, batch 1, BF16; draft cost not drawn
2–3×FASTER ON T5-XXL
esc
Full video on YouTube

Draft and verify

One trick gets several tokens per weight read: a small model guesses, the big model checks them in one pass.

  1. Cheap: it reads a fraction of the bytes.

  2. The weights are read once, the same as for one token.

  3. The big model's own pick replaces the first wrong one.

  4. On T5-XXL this ran 2–3× faster.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Fast Inference from Transformers via Speculative DecodingLeviathan et al. · paper, new tab
  • weights, one pass
  • guessed token
  • rejected
  • the token we follow
← Back to TricksSame answers →
  • weights, one pass
  • guessed token
  • rejected
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks↳ Draft and verify
13 Tricks ↳ Draft and verify14 of 14

ACCEPT RULE · big model p against draft q, one position

More per pass↖ Tricks
PROBABILITY · %, big model p against draft q
0%
25%
50%
75%
100%
85.9%
60%
·account
5.2%
15%
·branch
2.9%
10%
·loan
1.7%
5%
·card
1%
4%
·manager
p from the W2 output head at T = 1; q illustrative; other tokens not drawn
ACCEPT TEST · keep with chance min(1, p ÷ q)
·account
p 0.859 ≥ q 0.60
chance 1.00 · always kept
·branch
p 0.052 ÷ q 0.15 = 0.35
chance 0.35 · roll 0.71 lands outside
so ·branch is rejected
ON REJECTION · resample from the excess
·branch
·account
only ·account has p above q here
RESULT
Each token comes out at its p:
exactly as if the big model ran alone.
roll value illustrative
esc
Full video on YouTube

Same answers

A probabilistic accept rule keeps the output distribution exactly the big model's: faster, with no change in quality.

  1. Big model p, draft q.

  2. "·account": 0.859 against 0.60, always kept.

  3. Overconfident guesses are thinned out.

  4. The result matches sampling from the big model alone.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Accelerating Large Language Model Decoding with Speculative SamplingChen et al. · paper, new tab
  • big model p
  • draft q
  • excess p − q, the token we follow
  • chance of keeping
  • rejected
← Draft and verify← Back to Tricks
  • big model p
  • draft q
  • excess p − q, the token we follow
  • chance of keeping
  • rejected
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens1 of 14

TOKENIZER · text in on the left, ids out on the right

TEXT · 21 CHARACTERS
Tom opened a new bank
TOKENIZER
cut intoknown pieces
TOKENS
IDS
Tom
Tom
13787
·opened
·opened
4721
·a
·a
257
·new
·new
649
·bank
·bank
3331

HOW OTHER TEXT GETS CUT· the same tokenizer, each piece with its id

Common words: one piece
·the262
·account1848
1 token each
Rare words: several pieces
·Ou42201
ag363
ad324
oug20805
ou280
·Ouagadougou · 5 tokens
Numbers: cut unevenly
·171596
894531
·1789 · 2 tokens
Same word, different ids
·bank3331
Bank28650
·Bank5018
space and case change the id
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

5TOKENS
50,257IN THE VOCABULARY

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

01 Tokenization

A language model never sees text: the prompt is cut into tokens, and each token gets a unique number, its token ID.

  1. At the end, the model answers ·account.

  2. A fixed list of text pieces, learned once before training; frequent chunks get their own entry.

    Inside: Byte pairs →
  3. ·opened after a word and opened at the start of a line are different tokens.

  4. They are all the model receives. ·bank (3331) is the one whose prediction becomes ·account.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 2 ModelsBaseten · page, new tab
  • text
  • tokens and ids
  • the token we follow
Embedding →
  • text
  • tokens and ids
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
02 Embedding2 of 14

EMBEDDING LOOKUP · real shapes · matrix not to scale, broken below row 10,000

EMBEDDING MATRIX⋮row 0⋮↓ row 50,256[50,257 × 12,288]1 row per token in the model’s vocabularyrow 13,787 · Tomrow 4,721 · ·openedrow 257 · ·arow 649 · ·newrow 3,331 · ·bankTOKEN ID13,7874,7212576493,331TOKENDIM 0DIM 1DIM 2DIM 3…DIM 12,287Tom−0.02310.0164−0.00580.0297…−0.0112·opened0.0142−0.02070.00890.0035…0.0186·a0.00460.0013−0.0071−0.0024…0.0052·new−0.01190.02510.0167−0.0083…−0.0204·bank0.0268−0.0095−0.01830.0121…0.0147each row [1 × 12,288] · values illustrativeSTREAMSx₀x₀x₀x₀x₀ONE ROW, TWO MEANINGS· 2 of 12,288 directions · illustrativeMONEYRIVERaccountloandepositcashrivershorewaterfishbankone row sits between both;context moves it (dashed: after attention)EMBEDDING MATRIX · A LOOKUP IS A ONE-HOT MULTIPLY10 0 00 0 0one-hot(3331)[1 × 50,257]·E[50,257 × 12,288]=row 3,331 · ·bank[1 × 12,288]Row 3,331 equals a one-hot vector times E: 50,256 zeros and one 1, atposition 3,331.So nothing is multiplied: the GPU copies 12,288 numbers.Training nudges every row; ·bank gets one row for all its meanings(below).50,257 × 12,288 = 618 M weights · 0.35% of the model.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

12,288NUMBERS PER TOKEN
618 MNUMBERS IN THE MATRIX
0.35%OF THE MODEL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

02 Word embedding

Each token ID selects the corresponding row vector from the embedding matrix, which contains one row per token in the model's vocabulary.

  1. From the tokenizer, in prompt order.

  2. Each row holds 12,288 numbers: the embedding dimension. It is learned in training like every other weight.

  3. Id 3,331 means “take row 3,331”.

  4. No single number means anything; these are illustrative, as GPT-3’s weights were never published.

  5. But note, ·bank has one row for money and river alike.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • tokens (ids and vectors)
  • weights (the matrix)
  • the token we follow
← TokensStream →
  • tokens (ids and vectors)
  • weights (the matrix)
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream3 of 14

RESIDUAL STREAM · isometric · blocks run left to right, tokens go into the page

block 1block 96Attention 1all 5 communicateMLP 1each on its ownAttention 96all 5 communicateMLP 96each on its ownTom· · ··opened· · ··a· · ··new· · ··bank· · ·x₀x0 = E[id] + P[pos]token + position·thex₁x1 = x0 + ΔΔ = attention·ax₂x2 = x1 + ΔΔ = MLP·thex₁₉₁x191 = x190 + ΔΔ = attention·accountx₁₉₂x192 = x191 + ΔΔ = MLP·account· · ·blocks 2–95· · ·IF THE MODEL STOPPED HERE, ·bank WOULD PREDICTlogit lens: each x sent straight to the unembed · illustrativeunembed→·account12,288 → 49,152 → 12,28812,288 → 49,152 → 12,288
HOW TO READ THIS
  • Dark arcs: ·bank reads Tom, opened, a and new.
  • Faint arcs: each earlier vector reads the ones before it.
  • Both planes are big matrix multiplies: the GPU runs all 5 vectors in parallel.
  • Each step adds a Δ to the stream.
  • Attention: every vector reads the earlier ones (arcs)
  • MLP: every vector, each on its own
  • one column = one token’s vector, [12,288 × 1]
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

5VECTORS · ONE PER TOKEN
12,288NUMBERS PER VECTOR d_model = EMBEDDING DIM
192ADDITIONS · 96 BLOCKS × 2

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

03 Residual stream

That stream of vectors runs the length of the model: every layer reads it and adds a small change, nudging each vector's meaning.

  1. The ·bank column ends as the vector that predicts ·account.

    Inside: Attention residual →
  2. It reads every column and writes back onto each.

  3. It reads and writes one column at a time.

  4. Because the original is kept, gradients flow straight through in training, which is what lets many blocks be stacked.

    Inside: Hyper-connections →Inside: Engram →

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • stream
  • attention
  • MLP
  • the token we follow
← EmbeddingBlock →
  • stream
  • attention
  • MLP
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
04 Block4 of 14

ONE BLOCK · block 1 of the stream in 03, zoomed in · attention first, then the MLP

block 1LayerNormcentre, rescaleAttention96 heads, all 5 communicateLayerNormcentre, rescaleMLP · GELUeach vector on its ownTom++·opened++·a++·new++·bank+49,152WupWdown12,28812,288+x₀x₁x₂block 2same steps,own weightsn = LayerNorm(x0)normalised copy of x0x1 = x0 + Attn(n)Δ from attentionn = LayerNorm(x1)normalised copy of x1x2 = x1 + MLP(n)Δ from the MLPGELU
HOW TO READ THIS
WHERE THE BLOCK’S WEIGHTS ARE · 1.81 B per block
attn 33%MLP 67% · 1.21 B
attention 604 M · the 2 norms hold 49,152, too few to draw
  • Both run all 5 vectors in parallel on the GPU. Each step adds a Δ.
  • LayerNorm: centre and rescale each vector
  • Attention: vectors read the earlier ones (arcs) · 96 heads, 4 drawn
  • MLP: each vector on its own · Wup widens, Wdown narrows
  • one column = one token’s vector, [12,288 × 1]
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

604 MATTENTION PARAMETERS
1.21 BMLP PARAMETERS
2ADDITIONS TO THE STREAM

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

04 One transformer block

The repeating unit of an LLM: one attention layer, then one MLP (or feed-forward) layer. The block repeats 96 times, so each vector is nudged twice per block.

  1. LayerNorm subtracts the mean and divides by the spread, then applies a learned scale and bias.

    Inside: Normalisation →
  2. The only step where tokens share information.

    See also: Step 05 · QKV

  3. The stream itself is only ever added to.

  4. Widen, apply GELU, narrow.

    See also: Step 08 · MLP

  5. The next block repeats all of it with its own weights.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • stream
  • norm
  • attention
  • MLP
  • the token we follow
← StreamQKV →
  • stream
  • norm
  • attention
  • MLP
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
05 QKV5 of 14

QKV · one token, three projections, each cut into heads · hover: head 1

E·bank[12,288 × 1]⋮WQ[12,288 × 12,288]96 heads, 128 matrix rows per head⋮q⋮WK[12,288 × 12,288]⋮k⋮WV[12,288 × 12,288]⋮v96 QUERY HEADS · q cut into pieces of 128⋯⋯head 19696kv96 KV HEADS · a k and a v each, one per query headhead 1 scores against KV head 1’s k and blends its vk and v are kept for every token: the KV cache, 12E · ·bank, RESCALED BY LayerNormx÷ std× g + bE[12,288 × 1]the same E goesinto WQ, WK, WVmean and std: over the numbers of x, one pair per token.g, b: 12,288 learned gains and 12,288 biases.Every token enters attention at a similar size.LayerNorm has 24,576 weights; Llama uses RMSNorm.WQ · HEAD 1’S QUERY · WHAT THIS TOKEN LOOKS FORWQ: Queryproj. matrix⋮[12,288 × 12,288]96 heads × 128 rowsHead 1 of WQ:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 12,288]each row: 12,288 weights·E: tokenvector⋮[12,288 × 1]=q1: head 1’squery vectorR1·ER2·E⋮R128·E[128 × 1]Each row R is dotted with E: 12,288 products, one number.128 rows give q1: what head 1 looks for in other tokens.Head 2 uses rows 129–256, and so on: 96 heads.One multiply makes all 96 heads; q is then cut into pieces.WQ: 151 M weights, 12,288 × 12,288, plus 12,288 biases.WK · HEAD 1’S KEY · WHAT THIS TOKEN OFFERSWK: Keyproj. matrix⋮[12,288 × 12,288]96 heads × 128 rowsHead 1 of WK:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 12,288]each row: 12,288 weights·E: tokenvector⋮[12,288 × 1]=k1: head 1’skey vectorR1·ER2·E⋮R128·E[128 × 1]Same slice, own weights: k1 is what this token offers, matched against every q1 (06).k1 and v1 are head 1’s alone: one KV head per query head.k is kept for every token: the KV cache (12).WK: 151 M weights, 12,288 × 12,288, plus 12,288 biases.WV · HEAD 1’S VALUE · WHAT THIS TOKEN HANDS OVERWV: Valueproj. matrix⋮[12,288 × 12,288]96 heads × 128 rowsHead 1 of WV:rows 1–128— R1 —— R2 —⋮— R128 —[128 × 12,288]each row: 12,288 weights·E: tokenvector⋮[12,288 × 1]=v1: head 1’svalue vectorR1·ER2·E⋮R128·E[128 × 1]v1 is what this token hands over when head 1 picks it (07).k1 and v1 are head 1’s alone: one KV head per query head.v is kept next to k in the KV cache (12).WV: 151 M weights, 12,288 × 12,288, plus 12,288 biases.96 QUERY HEADS, EACH WITH ITS OWN k AND v12345678910111296⋯each q head (plum) has its own k and v head (blue)Head h scores each earlier token j: qh·kj.kj is from h’s own KV head: head 1 uses KV head 1, …Softmax turns the scores into weights on the vj (06, 07).Each token’s k, v are cached (12): 96 of each per block.WQ: 151 M weights; the 96 heads are slices of it.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

128NUMBERS PER HEAD
96QUERY HEADS
96KV HEADS

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

05 QKV

Attention is the only place tokens share information (they talk to each other). It starts by giving each token a query, key and value.

  1. 12,288 numbers, one column.

  2. In this head, ·bank’s q asks for a word that says which kind of bank is meant.

  3. ·opened’s k matches ·bank’s q best: their dot product q·k is the largest. So ·bank learns it is a bank that can be opened, not a river bank.

  4. k sets how much ·bank attends to ·opened; v is what it then takes from it. Two vectors, so what gets a token found and what it passes on can differ.

  5. Each piece is a head. 06b shows why many beat one.

    Inside: Fewer keys →

    See also: Step 12 · KV cache

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • token vector
  • weights
  • query
  • key, value
  • head 1
← BlockScores →
  • token vector
  • weights
  • query
  • key, value
  • head 1
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
06 Scores6 of 14

SCORES · how much each token listens to each earlier one · one head

WEIGHTS · HEAD 1Tomopenedanewbankrows: queries q · columns: keys kTomopenedanewbank1.00.35.65.10.60.30.05.30.15.50.03.52.05.32.08each row sums to 1grey: later tokens, masked out (−∞)ROW ·bank, WORKED THROUGHs = qbank·k / √128scoresoftmaxw = eˢ / Σ eˢweightTomopenedanewbank0.493.351.002.861.47.03.52.05.32.08Σ = 1.00 · same as the gold rowillustrative valuesq·k for every earlier token÷ √128mask the futuresoftmaxweights → 07 blends vThis is head 1 of 96 · 06b shows the other 95.HEAD 1 · WHERE THE 5 × 5 GRID COMES FROMTomopenedanewbankq · 128 numbers →Q1 · queries[5 × 128]·TomTomopenedopenedaanewnewbankbankk·128↓K1ᵀ · keys, on their side[128 × 5]=.80.401.02.201.991.30.101.891.202.41.493.351.002.861.47S1 = Q1·K1ᵀ ÷ √128[5 × 5] · scoresmask, softmax →06’s weightsOne cell: ·bank’s q · ·opened’s k = 128 products added = 37.9 ÷ √128 = 3.35 → softmax .525 × 5: rows and columns are tokens. 5 tokens in, 5 × 5 out; a 1,000-token prompt makes a 1,000 × 1,000 grid.128: the length of each q and k. It hides inside every cell (the dot product), not on the grid’s sides.96: head 1 uses numbers 1–128 of each token’s q and k, head 2 the next 128, and so on: 96 grids, all 5 × 5 (06b).values illustrative
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

25SCORES IN THIS HEAD
÷√128SCALING
−∞ → 0MASKED

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

06 Attention scores

Each token's query vector is scored against every earlier token's key vector; softmax turns the scores into weights that sum to 1. A higher score means that earlier token is more relevant to it.

  1. ·opened (0.52) and ·new (0.32) together take 84% of it.

    See also: Step 07 · Blend

  2. A dot product q·k, large when the two vectors point the same way; dividing by √128 keeps softmax from going all-or-nothing.

    See also: Step 05 · QKV

  3. Each score is exponentiated, then divided by the total (54.66); a score 0.5 higher gets about 1.6× the weight.

  4. Their scores are set to −∞ before softmax, and exp(−∞) is 0. Every other GPT-3 layer also masks tokens outside a local band.

    See also: Step 12 · KV cache

  5. Each cell is one q·k of 128 numbers. The card shows the whole grid as one matrix product, Q·Kᵀ.

    See also: Heads

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • Inference Engineering, chapter 2: ModelsBaseten · page, new tab
  • scores and weights, low → high
  • masked
  • the token we follow
← QKVHeads →
  • scores and weights, low → high
  • masked
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
06b Heads7 of 14

HEADS · 06 showed head 1 · there are 96, each with its own slice of q and k · ·bank’s row outlined

06’S GRID · HEAD 1 OF 96rows: queries q · columns: keys kTom·opened·a·new·bankTom·opened·a·new·bank1.00.35.65.10.60.30.05.30.15.50.03.52.05.32.08·bank’s row: one mix ofthe 5 tokens, summing to 11–496… each k is cut the same wayALL 96 HEADS · 128 NUMBERS EACH·bank’s q (12,288 numbers), cut into 96 slices of 128head 1·bank →·openedhead 2previous tokenhead 3first tokenhead 4itselfKV heads 1–4: one k, one v eacheach head scores its own k and blends its own v+ 92 more heads like this, KV heads 5–96Why 96? One row gives ·bank one mix. Each head’s row sums to 1 on its own, so ·bankcan lean on ·opened in head 1 and on Tom in head 3 at once. 07 blends all 96 back together.Same work as one 12,288-wide head: 96 × 128 = 12,288 multiplies per cell · heads 2–4 illustrative
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

96GRIDS PER LAYER
128NUMBERS PER HEAD
96KV HEADS

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

06b Attention heads

Attention runs as 96 heads side by side, so ·bank gets 96 rows of weights, not one.

  1. Its weights sum to 1, so ·opened and ·new share one budget.

    See also: Step 06 · Scores

  2. 96 × 128 = 12,288, so the work is the same as one wide head. Heads 2–4 here are made up; real heads learn their own patterns.

  3. Every head’s k and v go into the KV cache (12): 4719 KB per token. Later models share them across heads to shrink that.

    See also: Step 12 · KV cache

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • weights, low → high
  • k and v (a KV head)
  • the token we follow
← ScoresBlend →
  • weights, low → high
  • k and v (a KV head)
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
07 Blend8 of 14

BLENDING VALUES · the row for ·bank, 3 of 128 numbers shown, illustrative

HEAD 1 · EACH v × ITS WEIGHT FROM 06, THEN ADDED UPv: first 3 of 128 numbers per token · illustrativew: weight (06)v1 · head 1’s ownw × v1Tom.03×0.2 −0.1 0.0 …=0.01 0.00 0.00 …opened.52×−0.2 1.4 0.6 …=−0.10 0.73 0.31 …a.05×0.1 0.2 0.1 …=0.00 0.01 0.00 …new.32×1.1 0.3 −0.4 …=0.35 0.10 −0.13 …bank.08×0.3 −0.2 0.2 …=0.02 −0.02 0.02 …o1 · add each column=0.28 0.82 0.21 …head 1’s output[1 × 128]Biggest weights win: ·opened (.52) and ·new (.32) make up most of o1,so head 1 pulls them into ·bank. The other tokens add almost nothing.JOIN THE 96 HEADS, MIX, ADDWO[12,288 × 12,288]·o₁o₉₆[12,288 × 1]96 × 128=Δ[12,288 × 1]·bank+Δ=·bankthe stream (03)MLP (08)Every head does this with its own w and v:96 outputs of 128 numbers each.Joined end to end: 12,288 again.WO mixes the heads into one change, Δ.Δ is added onto ·bank; nothing is replaced.KV cache (12): k and v of past tokensare kept, so a new token only adds its own.WO · JOIN THE 96 HEADS, MIX, ADDWO: Output projection matrix[12,288 × 12,288]lit: head 1 and its block of WO·=+o: heads’ output vec[12,288 × 1]Δ: change vec[12,288 × 1]x: token vec[12,288 × 1]x′: new token vec[12,288 × 1]Join: the 96 head outputs (128 each) stack into o.Mix: WO lets every head reach every number.Block i of WO only meets head i: Δ = Σ WO,i · oi.Add: Δ goes back onto the stream, x′ = x + Δ.WO: 151 M weights. Attention in all: 604 M per block.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

84%FROM ·opened AND ·new

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

07 Blend

Those weights mix the values, so ·bank now carries the context it needs: mostly ·opened and ·new.

  1. From the Wᵥ projection.

    See also: Step 05 · QKV

  2. Tom and ·a get 8% between them, so they add almost nothing.

  3. 128 numbers; the first three are drawn.

  4. Real heads show patterns like the previous token, the first token, or the subject.

  5. Wₒ turns the 96 heads’ outputs back into 12,288 numbers.

    Inside: Exclusive attention →

    See also: Step 08 · MLP

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • LLM Visualization: step through a GPT in 3DBrendan Bycroft · page, new tab
  • values, head outputs, Δ
  • attention weights (06)
  • model weights
  • the token we follow, in the stream
← HeadsMLP →
  • values, head outputs, Δ
  • attention weights (06)
  • model weights
  • the token we follow, in the stream
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP9 of 14

MLP · GELU · one token, left to right · 12,288 = 52 px, to scale · hover a W for its maths

E·bank[12,288 × 1]Wup[49,152 × 12,288]604 M weightsup + b[49,152 × 1]GELUh · [49,152 × 1]a few units fire, the rest near 0Wdown[12,288 × 49,152]604 M weightsΔE[12,288 × 1]added tothe streamWup · EACH ROW ASKS ONE QUESTIONWup: Up proj. matrix— R1 —— R2 —— R49,152 —⋮[49,152 × 12,288]·E: tokenvector⋮[12,288 × 1]+b: biasvector⋮[49,152 × 1]=up + b:widened vectorR1·E+b1R2·E+b2R49,152·E+b49,152⋮[49,152 × 1]Each row R is a direction in 12,288-number space.R·E: how far ·bank points along it. One number per row.49,152 rows, so 49,152 numbers: 4× wider than E.Plus a bias b per row; GELU then passes the units that answer yes.604 M weights here, plus 49,152 biases.Wdown · EACH COLUMN WRITES ONE ANSWERWdown: Down proj. matrixc1c2⋯c49,152[12,288 × 49,152]·h: hiddenvectorh1h2h49,152⋮[49,152 × 1]=ΔE: changevector⋮[12,288 × 1]ΔE = h1c1 + h2c2 + … + h49,152c49,152 + bEach column c is a direction in 12,288-number space: one hidden unit’s answer.Unit i adds hi × ci; units near 0 add almost nothing.The sum ΔE is added onto ·bank in the stream.604 M weights. The two MLP matrices: 1.2 B per block.GELU · NEGATIVES SQUISHED TO ~0x · Φ(x)dashed: y = xup + b, beforeafter GELUEach of the 49,152 hidden numbers goes through x · Φ(x).Positive: passes almost unchanged. Negative: near 0.The result is h: near 0 switches a unit off.12 of 49,152 shown, illustrative. Llama: SiLU.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

1.2 BMLP WEIGHTS PER BLOCK
116 BIN ALL 96 MLPS
66%OF THE MODEL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

08 Multilayer perceptron (MLP)

Then the MLP works on each token alone. It holds 66% of the model’s parameters.

  1. 12,288 numbers, now carrying context from the others.

  2. Each hidden unit is a row of weights that tests for one pattern.

  3. GELU is near x for positive inputs and near 0 for negative ones, so most units go quiet.

    Inside: Gated units →
  4. Each unit that fired writes its own direction back.

    See also: Step 09 · Stacking

  5. A router picks a few per token.

    Inside: Experts →Inside: 896 experts →

    See also: Step 13 · Tricks

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • token vector
  • weights
  • hidden units, dark = fired · illustrative
← BlendStacking →
  • token vector
  • weights
  • hidden units, dark = fired · illustrative
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
09 Stacking10 of 14

THE STACK · the residual stream runs left to right

EARLY · 1–24word pieces, grammarMIDDLE · 25–72facts looked upLATE · 73–96answer moved into ·bankTomopenedanewbank1224254572738196⋯⋯⋯⋯⋯blockMLPs on ·opened recall·opened a ·bank → ·accountMeng et al. 2022 · ROMEattention copies itinto ·banklast ·bank→ Output (10)EVERY BLOCKx ← x + Attention(x) · then · x ← x + MLP(x)The lanes are the residual stream: one vector per token. Attention and MLP sit on it and add to it.Roles and band edges (scaled from 32 blocks): rough, from interpretability studies · illustrative.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

96BLOCKS
174.6 BPARAMETERS IN ALL

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

09 Stacking

Stacked up, the blocks divide the work: early ones handle local patterns, middle ones facts, late ones the prediction.

  1. Smaller ones stack fewer: Llama 3.1 8B has 32.

  2. The previous token, short phrases.

  3. ·opened ·a ·new ·bank stop being separate tokens and becomes one idea, a money bank, carried onto ·bank.

  4. ROME showed this in GPT-2 XL and GPT-J: editing those MLPs changes the fact the model states.

  5. By block 96 the ·bank vector points at ·account.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Locating and Editing Factual Associations in GPTMeng et al. · paper, new tab
  • ·bank, the token we predict from
  • attention: spans all tokens
  • MLP: each token on its own
← MLPOutput →
  • ·bank, the token we predict from
  • attention: spans all tokens
  • MLP: each token on its own
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
10 Output11 of 14

OUTPUT HEAD · the last vector scored against every token, illustrative logits· illustrative logits

EMBEDDING MATRIX, REUSED⋮row 1,848 · accountrow 0↓ row 50,256[50,257 × 12,288]1 row per tokenin the model’s vocabulary··bank[12,288× 1]=⋮21.0−3.14.2⋮0.7z[50,257 × 1]logits: one scoreper tokenTOP 5 OF 50,257logit zp = eᶻ / Σ eᶻsoftmaxprobability paccountbranchloancardmanager50,252others21.018.217.617.116.56.0 eachillustrative values5.3%2.9%1.8%1.0%1.3% togetherΣ = 100%87.6%SAMPLEpick one token at random, weighted by p → ·account about 88 times in 100Temperature T divides z first: e^(z/T). T below 1 sharpens toward account, above 1 flattens. Here T = 1.WU = E · ONE DOT PRODUCT PER TOKENE[50,257 × 12,288]·h=z[50,257 × 1]one score per tokenh: [12,288 × 1]zaccount = 21.0zaccount = row 1,848 · h = Σ wihi over 12,288 = 21.0illustrative valuesEach row of E is one token’s output direction.50,257 rows, so 50,257 dot products, all at once.h points most along the account row, so zaccount is largest.Softmax turns z into next-token odds.No extra weights: GPT-3 reuses E (618 M). Llama keeps WU (525 M).
·bank
→
50,257 rows × 12,288
one score per token
PROBABILITIES · top five, then the rest
·account87.6%
·branch5.3%
·loan2.9%
·card1.8%
·manager1.0%
others1.3% together
050100
PROBABILITY, %
Temperature T divides z first: e^(z/T). T below 1 sharpens toward account, above 1 flattens. Here T = 1.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

88%·account AT T = 1

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

10 Output

At the top, the last token's vector becomes probabilities for every token, and the sampler picks one.

  1. After 96 blocks it holds everything worked out about the prompt.

  2. Each logit is a dot product with that token’s row; ·account scores 21.0.

  3. A lead of 2.8 in score becomes a 16× lead in probability.

    Inside: Softmax collapse →
  4. Below 1 the favourite gains; at 0 it always wins.

  5. Then it is appended for the next pass.

    See also: Step 11 · Passes

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • last vector, chosen token
  • weights
  • logits
  • probabilities
← StackingPasses →
  • last vector, chosen token
  • weights
  • logits
  • probabilities
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
11 Passes12 of 14

FOUR PASSES · each row is one run of the whole model

TOKENS IN · dashed: already run, keys and values reused
READS EVERY WEIGHT
OUT
PASS 1t = 0.0 msTom·opened·a·new·bankall weights→·account
PASS 2t = 43.6 msTom·opened·a·new·bank·accountall weights→.
PASS 3t = 87.3 msTom·opened·a·new·bank·account.all weights→·He
PASS 4t = 130.9 msTom·opened·a·new·bank·account.·Heall weights→·deposited
pass 1 · 43.6 mspass 2 · 43.6 mspass 3 · 43.6 mspass 4 · 43.6 ms0255075100125150175200TIME, MSBatch 1, BF16 read at one B200’s 8 TB/s; the weights need 2 B200s. Not measured.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

349.2 GBREAD PER PASS
43.6 msPER PASS AT 8 TB/S

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

11 Passes

That pick is appended and the model runs again: one pass per token, so speed is tokens per second.

  1. They are all known, so they go through side by side: prefill.

  2. The sampler’s pick becomes the last input token.

  3. Earlier tokens’ keys k and values v are reused.

    See also: Step 12 · KV cache

  4. A pass can’t start until the previous one has picked its token.

  5. At 8 TB/s, reading 349.2 GB takes ≈44 ms, ≈23 tokens/s. The weights need 2 B200s; this is arithmetic, not a measurement.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 0 InferenceBaseten · page, new tab
  • prompt tokens
  • newest token
  • already run
  • the model
  • token picked
  • time
← OutputKV cache →
  • prompt tokens
  • newest token
  • already run
  • the model
  • token picked
  • time
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
12 KV cache13 of 14

KV CACHE · one column per token, 96 layers each, 3 per stripe

KEYS AND VALUES SAVED · one column per token
layer 96layer 1Tom·opened·a·new·bank·accountONE CELLone token's keys and values
in 3 layers
ONE COLUMNone token’s keys and values, all 96 layersNEWEST COLUMN·account computes only its own k and v,
then the cache grows by one column
ONE CHAT'S CACHE BY CONTEXT LENGTH · the weights for scale
KV cache · 256 tokens1.21 GB
KV cache · 512 tokens2.42 GB
KV cache · 1,024 tokens4.83 GB
KV cache · 2,048 tokens9.66 GB
Weights · shared by all chats349.2 GB
02.557.510MEMORY, GBGPT-3 175B, BF16, one chat. The weights bar is cut: 349.2 GB is off the scale.
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

96KV HEADS, ONE PER QUERY HEAD
9.66 GBONE CHAT AT 2,048 TOKENS

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

12 KV cache

To avoid recomputing, earlier tokens' keys and values are saved; each pass adds only its own, so memory keeps growing.

  1. Computed once, kept in GPU memory, read by every later token.

  2. 96 layers × 96 KV heads × 128 numbers × k and v × 2 bytes.

    See also: Step 05 · QKV

  3. At its 2,048-token limit one chat’s cache is 9.66 GB, 3% of the weights.

    Inside: Linear attention →Inside: Compressed cache →
  4. That is why cache memory, not weights, usually limits how many users fit on a GPU.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • cached keys and values
  • newest token’s K and V
  • weights
  • memory, GB
← PassesTricks →
  • cached keys and values
  • newest token’s K and V
  • weights
  • memory, GB
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
13 Tricks14 of 14

FOUR GROUPS · by what they fix

BF16 + cache16-bit weights358.9 GB
FP8 + cache8-bit weights184.3 GB
4-bit + cacheestimate, never released116.7 GB
Same, 8 KV headsshared: 1/12 the cache107.8 GB
0200400GB READ PER GENERATED TOKENGPT-3 175B, batch 1, 2,048-token context. Teal is weights, purple is the KV cache. 4-bit is an estimate.
MEMORYweights and KV cache
NAMED TECHNIQUES
FP8 and 4-bit weightsGQAMLA
MORE PER PASSexperts, guesses
NAMED TECHNIQUES
Mixture of expertsSpeculative decoding
ATTENTIONfewer pairs scored
NAMED TECHNIQUES
FlashAttentionSliding window
OUTPUTcontrol, not speed
NAMED TECHNIQUES
Constrained decodingSampling settings
BF164-bitsame weights, ¼ of the bits96 KV heads8 KV headsheads share keys and valuesDeepSeek-V3: 37 B of 671 Bparameters active per token✓✓✓✗a small model guesses,one big-model pass checksall earlierlast 3 onlyscore fewer pairs,or do the same work fasteronly tokens that fitthe format can be pickedsame speed, valid output
MODEL HYPERPARAMETERS · SAME ON EVERY SCREEN
Llama 3.1 8BGPT-3 175B
layers3296
embed dim d4,09612,288
attention heads3296
KV heads896
head dim128128
MLP width (W_up rows)14,336 · 3.5 d49,152 · 4 d
vocabulary128,25650,257
context (tokens)131,0722,048
normRMSNormLayerNorm
activationSwiGLUGELU
positionsRoPElearned
parameters8.0 B175 B
  • Llama’s MLP is 3.5 d, not 4 d: SwiGLU uses 3 matrices, so each is made narrower.
  • The model toggle (top right) swaps the whole column.

Every number on every screen is computed from this table.

INSIDE A TRANSFORMER · Vaswani et al. 2017, Fig. 1 and 2

esc
Add & NormFeed forwardAdd & NormMulti-head attentionEmbeddingSource textN×EncoderTranslation only.Chat models drop it.Add & NormMulti-head attentionreads the encoderOutput probabilitiesSoftmaxLinear10Add & NormFeed forward08Add & NormMasked multi-headattention05·06·07PositionalencodingEmbedding02Inputs01040312DecoderN×0911Scaled dot-product attentionMatMul07SoftMaxMask (opt.)ScaleMatMul06QKVMulti-head attentionLinearConcatScaled dot-productattentionh06bLinearLinearLinear05VKQ
esc
Full video on YouTube

13 Tricks

Generation is limited by memory, so almost every trick reads fewer bytes, at home or in a data centre.

  1. 4-bit weights would take about a third of BF16’s bytes. GPT-3’s were never released, so its 107 GB is an estimate.

  2. GPT-3 keeps a KV head for each of its 96 heads; grouped-query attention shares a few. DeepSeek’s latent attention compresses them further.

    See also: Step 12 · KV cache

  3. DeepSeek-V3 reads 37 B of its 671 B parameters per token.

    See also: Step 08 · MLP

  4. A cheap drafter guesses; the big model verifies all of them for one read of the weights.

    Inside: Draft and verify →
  5. FlashAttention never writes the score grid to HBM; sliding windows stop the cache growing.

  6. Constrained decoding zeroes tokens that would break JSON or a grammar.

    See also: Step 10 · Output

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Inference Engineering · 5 TechniquesBaseten · page, new tab
  • GPU Glossary · Memory-boundModal · page, new tab
  • weights
  • KV cache
  • attention work
  • output control
  • skipped or rejected
← KV cache1 · A data centre, top down →
  • weights
  • KV cache
  • attention work
  • output control
  • skipped or rejected
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens↳ Byte pairs02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens ↳ Byte pairs1 of 14

BYTE-PAIR MERGES · "·opened", illustrative merge order

Tokenizer↖ Tokens
ROWSYMBOLS AFTER THE MERGE
start · 7·opened
merge 1 · e+n·opened
merge 2 · e+d·opened
merge 3 · ·+o·opened
merge 4 · ·o+p·opened
merge 5 · en+ed·opened
merge 6·openeda token in GPT-3’s list

Merge order and pair counts illustrative. · is the leading space.

TOP PAIRS · before merge 1

e n9.4 M
e d9.0 M
· o8.6 M
o p7.9 M
n e6.1 M
p e5.2 M

05 M9.4 M

TIMES SEEN, MILLIONS

The winner becomes merge 1. Then every count is redone.

50,257ENTRIES
4.2CHARACTERS PER TOKEN, THIS PROMPT
esc
Full video on YouTube

Byte pairs

The tokenizer's pieces come from a list built by merging the most frequent pair of symbols, again and again.

  1. "·opened" begins as seven symbols, the space included.

  2. The most frequent pair becomes one new symbol.

  3. Each merge adds one entry to the list; training stops at the target size.

  4. Rare words stay in pieces.

    See also Step 01 · Tokens

  5. The prompt’s 21 characters become 5 tokens.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Neural Machine Translation of Rare Words with Subword UnitsSennrich et al. · paper, new tab
  • text
  • tokens
  • the pair merged in each row
← Back to TokensBytes, not tokens →
  • text
  • tokens
  • the pair merged in each row
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens↳ Byte pairs02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
01 Tokens ↳ Byte pairs1 of 14

BYTE PATCHES · the same prompt, patch edges illustrative

Tokenizer↖ Tokens
NEXT-BYTE ENTROPY, BITS024cut above 2.5 bitsTom·opened·a·new·bank4 bytes25244one per patch: 6 here, against 21 bytes3 bytes7245one per token: 5 here, against 21 bytesBIG-MODEL STEPS

Entropy from a small byte model; values illustrative.

Token edges from GPT-3’s tokenizer; · is the leading space.

CUT BY
21BYTES
8BLARGEST BLT IN THE STUDY
esc
Full video on YouTube

Bytes, not tokens

Byte-level models skip that list and group raw bytes into patches by difficulty; either way, the model gets numbers.

  1. No list to learn, and no word is ever unknown.

  2. Predictable bytes score low; the start of a new word scores high.

  3. Easy stretches become long patches; hard ones get more, shorter patches.

  4. Compute follows difficulty instead of a fixed vocabulary.

  5. It was trained on up to 4 trillion bytes.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Byte Latent Transformer: Patches Scale Better Than TokensPagnoni et al. · paper, new tab
  • bytes
  • next-byte surprise (entropy)
  • patches
  • tokens
  • the patch holding ·bank
  • the token ·bank
← Byte pairs← Back to Tokens
  • bytes
  • next-byte surprise (entropy)
  • patches
  • tokens
  • the patch holding ·bank
  • the token ·bank
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Hyper-connections04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Hyper-connections3 of 14

HYPER-CONNECTIONS · one token, four streams, GPT-3 width for scale

Hyper-connections↖ Stream
from Embedding4 streams, 12,288 numbers eachMIXING WEIGHTS, 4 × 4learned for each layerdarker = stronger linkBlock 1attention + MLPread a blendwrite to all fourBlock 2attention + MLPread a blendwrite to all fourMixing weights illustrative.from EmbeddingBlock 1attention + MLPread+add, weight 1Block 2attention + MLPread+add, weight 11carried on with weight 1
Residual
4STREAMS
49,152NUMBERS PER TOKEN
esc
Full video on YouTube

Hyper-connections

The stream is one vector; hyper-connections widen it into several copies, with learned mixing between them.

  1. Each block adds its output to it, weight one.

  2. For a GPT-3-width model that is 49,152 numbers per token.

  3. The network learns how strongly each depth feeds the next.

  4. Layers can effectively be reordered or run side by side.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Hyper-ConnectionsZhu et al. · paper, new tab
  • the stream for ·bank
  • block
  • mixing weights
← Back to StreamKeeping the mix tame →
  • the stream for ·bank
  • block
  • mixing weights
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream↳ Attention residual04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
03 Stream ↳ Attention residual3 of 14

ATTENTION OVER DEPTH · one token at one layer, weights illustrative

Attention residual↖ Stream
STREAM OF THE TOKEN WE FOLLOWnext layerevery output added with weight 11 block = 24 outputsWEIGHT ON EACH EARLIER OUTPUT, SHARE OF 100.10.214896144192Earlier output, 1 to 192, oldest first (attention, MLP, attention, …)HOW ONE LAYER MAKES ITS WEIGHTS1 · SCOREsᵢ = w · RMSNorm(vᵢ)One learned query w per layerscores every earlier output vᵢ.2 · SOFTMAX OVER DEPTHαᵢ = exp(sᵢ) ÷ Σⱼ exp(sⱼ)Turns the 192 scores into weightsthat are positive and sum to 1.3 · MIXh = Σᵢ αᵢ · vᵢThe layer's input is theweighted sum of earlier outputs.HOW A PLAIN STREAM MIXES1 · SCOREnoneNo query: nothing looks at whatthe earlier outputs contain.2 · WEIGHTSevery weight = 1Every output is added once,an equal 1/192 share of the sum.3 · MIXh = v₁ + v₂ + … + v₁₉₂The layer's input is theplain sum of earlier outputs.→→
equal, 1/192
Depth weights
192EARLIER OUTPUTS IN GPT-3 175B
1.4TTRAINING TOKENS, KIMI'S 48B TEST
esc
Full video on YouTube

Attention residual

The stream adds every block's output with weight one; attention residual lets each layer choose which earlier outputs to use.

  1. Deep in the model, each new output is a small share of a growing sum.

  2. One learned query per layer scores every earlier output; a softmax over depth turns the scores into weights.

  3. Attend over block summaries, not every layer.

  4. Output sizes and gradients stayed more even across depth.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Attention ResidualsKimi Team · paper, new tab
  • stream of the token we follow
  • weight on an earlier output
  • block group
← Back to Stream
  • stream of the token we follow
  • weight on an earlier output
  • block group
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block↳ Normalisation05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
04 Block ↳ Normalisation4 of 14

LAYERNORM · the ·bank vector before and after, first 4 of 12,288 numbers, illustrative

LayerNorm↖ Block
−8−40+4+8Value, no unitGPT-3 175B, layer 1, before the gain. Values illustrative: weights unpublished.As it enterstiny numbersMean and SD of all 12,288− mean, ÷ SDAfter normalisingmean 0, SD 1× gain + bias: two learned weights per numberNEWER: RMSNORMno mean, no bias:÷ RMS, then × gain
0.0268
dim 0
−0.0095
dim 1
−0.0183
dim 2
0.0121
dim 3
4.85
dim 0
−2.00
dim 1
−3.66
dim 2
2.08
dim 3
0.0053SD IN (ILLUSTRATIVE)
1.0SD OUT
2PER BLOCK
esc
Full video on YouTube

Normalisation

Each block first rescales its copy of the stream: GPT-3’s LayerNorm centres each token’s vector and fixes its spread.

  1. Mean 0.0011, standard deviation 0.0053 for ·bank as it enters (illustrative).

  2. The numbers now average 0, with a standard deviation of 1.

  3. 24,576 weights let the model restore what it needs.

  4. Training becomes unstable.

    See also Step 03 · Stream

  5. It is cheaper and trains about as well; Llama 3.1 uses it.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Layer NormalizationBa, Kiros & Hinton · paper, new tab
  • Root Mean Square Layer NormalizationZhang & Sennrich · paper, new tab
  • ·bank, the token we follow
  • learned gain and bias
← Back to BlockWithout normalising →
  • ·bank, the token we follow
  • learned gain and bias
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV↳ Fewer keys06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks
05 QKV ↳ Fewer keys5 of 14

KV CACHE PER TOKEN · GPT-3 175B shape, BF16

KV heads↖ QKV
96 QUERY HEADS96 KEY–VALUE HEADS012345KV cache per token, MBGPT-3 175B shape, 96 layers, BF16Multi-head · 96 · GPT-34.719 MBGrouped · 80.393 MBMulti-query · 10.049 MBLatent · MLAone latent per token · DeepSeek-V2, no bar: a different model
KV heads
4.72 MBKV CACHE PER TOKEN (GPT-3)
−93.3%DEEPSEEK-V2 CACHE
esc
Full video on YouTube

Fewer keys

Keys and values are kept for every past token; GPT-3 keeps one per head, newer models share or compress them.

  1. One for each query head: 4.719 MB per token.

  2. About 49 kilobytes, with some quality lost.

  3. In GPT-3's shape that is 0.393 MB, close to multi-head quality.

  4. Keys and values are rebuilt from it; DeepSeek-V2 cut its cache 93.3%.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • GQA: Training Generalized Multi-Query Transformer ModelsAinslie et al. · paper, new tab
  • DeepSeek-V2DeepSeek-AI · paper, new tab
  • query head
  • key–value head
  • GPT-3 175B's choice
← Back to QKVRead fewer keys →
  • query head
  • key–value head
  • GPT-3 175B's choice
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ Gated units09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ Gated units9 of 14

ACTIVATIONS · output against input, one hidden unit

ACTIVATION
Activation↖ MLP
4
2
0
−4
−2
2
4
INPUT, NO UNIT
OUTPUT
ReLU
GELU
INSIDE THE MLP · one token, 12,288 numbers in
x
W gate12,288 × 49,152
W up12,288 × 49,152
SiLU
GELU
×
W down
12,288 out
gate branch · SwiGLU only
GELU MLP2 matrices · GPT-349,152 units
SwiGLUsame parameters32,768 units
0
10,000
20,000
30,000
40,000
50,000
HIDDEN WIDTH, UNITS
GPT-3 175B, one block.
2MATRICES
1208 MPARAMETERS PER BLOCK
esc
Full video on YouTube

Gated units

GPT-3's MLP applies one activation, GELU, between two matrices; newer models multiply two projections instead, a gated unit.

  1. Below zero the output is flat at zero; above, it is the input.

  2. x · Φ(x), Φ the normal curve's area below x: a small dip, then nearly x.

  3. One, through SiLU, decides how much passes; the other carries the content.

  4. The hidden width shrinks from 49,152 to 32,768.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Gaussian Error Linear Units (GELUs)Hendrycks & Gimpel · paper, new tab
  • GLU Variants Improve TransformerShazeer · paper, new tab
  • token vector
  • weights
  • activation
  • not used in this setting
  • not used in this setting
← Back to MLP
  • token vector
  • weights
  • activation
  • not used in this setting
  • not used in this setting
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP↳ Experts09 Stacking10 Output11 Passes12 KV cache13 Tricks
08 MLP ↳ Experts9 of 14

MIXTURE OF EXPERTS · DeepSeek-V3 layer, one token

LAYER
MLP↖ Block
Router7,168 × 256
ROUTER SCORE · one per expert, no unit
expert 1256
256 ROUTED EXPERTS · each one an MLP · 8 gold run
8 outputs × score
+
One MLP12,288 → 49,152 → 12,288every weight works for every token
x
7,168 out
Shared expertruns for every token
PARAMETERS, B · per token
37 B work671 B in the model
WORK PER TOKEN, B · same scale
GPT-3 175B175
DeepSeek-V337
8 of 256EXPERTS PER TOKEN
5.5%OF PARAMETERS ACTIVE
esc
Full video on YouTube

Experts

Bigger models split the MLP into many experts, and a router runs only a few of them per token.

  1. A small matrix: one score per expert.

  2. Their outputs are added, weighted by the router's scores.

  3. It holds what all tokens need.

  4. A large model at a small model's cost per token.

    See also Step 13 · Tricks

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Outrageously Large Neural NetworksShazeer et al. · paper, new tab
  • DeepSeek-V3 Technical ReportDeepSeek-AI · paper, new tab
  • token vector
  • weights
  • runs for this token
← Back to MLPKeeping experts busy →
  • token vector
  • weights
  • runs for this token
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output↳ Softmax collapse11 Passes12 KV cache13 Tricks
10 Output ↳ Softmax collapse11 of 14

LOGIT MARGIN · target logit minus log-sum-exp of the others

PRECISION
Softmax↖ Output
|GRADIENT| · target token, log scale
1
10−6
10−12
10−18
exactly 0
0
10
20
30
40
LOGIT MARGIN m, no unit
Hard-target cross-entropy; curve shapes from the paper's Figure 1.
BF16 · ≈5.5
FP32 · 16.6
FP64 · 36.7
standard softmax: rounded to 0
fixed: keeps shrinking
“·account” · 2.0
m = 2
0
TARGET GRADIENT
at the marker
WORKED EXAMPLElogits 21 and 0 · BF16
a e−21 ≈ 7.6 × 10−10the other token's term (the leader's is 1)
→
b 1 + 7.6 × 10−10 = 1BF16's step near 1 is 2−7 ≈ 0.008
→
c p = 1, so 1 − p = 0the target's gradient is exactly 0
esc
Full video on YouTube

Softmax collapse

Those probabilities come from softmax, which can round a training gradient to exactly zero when one token leads by enough.

  1. For ·account here it is 2.0.

  2. In FP32 this happens past the cliff line. The worked example below shows one case in BF16.

  3. Gradients stop summing to zero, which causes loss spikes with Adam.

  4. Compute the target's gradient as minus the sum of the others'.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • A Fix for Softmax and Cross-EntropyBarron et al. · paper, new tab
  • target gradient, fixed softmax
  • standard softmax, lost past the cliff
  • precision cliff, selected
  • the token we follow, ·account
← Back to Output
  • target gradient, fixed softmax
  • standard softmax, lost past the cliff
  • precision cliff, selected
  • the token we follow, ·account
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache↳ Linear attention13 Tricks
12 KV cache ↳ Linear attention13 of 14

MEMORY PER CHAT · a GPT-3-shaped model, if its attention were linear

CONTEXT
Cache↖ KV cache
SOFTMAX ATTENTION · one column per token
keys above, values below
token 1 → … grows by one per token
newest
drop softmax,
regroup the sums
LINEAR STATE · one per head
128 × 128 numbers
same size at token 10 or 1,000,000
NEW TOKEN
kvq
adds k·vᵀ
q reads the state
MEMORY PER CHAT · GB, log scale
KV cache
one chat, 2K tokens
Linear state
any length
9.66 GB
0.3 GB · the cache passes this at 64 tokens
0.1
1
10
100
1,000
10,000
GB, BF16
GPT-3 175B shape, BF16; the linear version is hypothetical. 8K and up are beyond GPT-3's 2,048.
O(N)WORK FOR N TOKENS, WAS O(N²)
esc
Full video on YouTube

Linear attention

Instead of a cache that keeps growing, linear attention keeps a fixed-size summary of every past key and value.

  1. The cache grows with every token.

  2. Keys times values can be summed once into one matrix per head.

  3. Same cost at token 10 or token 1,000,000.

  4. The real cache passes that at 64 tokens.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Transformers are RNNsKatharopoulos et al. · paper, new tab
  • cached keys and values
  • linear state
  • newest token
  • memory, GB
← Back to KV cacheForgetting on purpose →
  • cached keys and values
  • linear state
  • newest token
  • memory, GB
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks↳ Draft and verify
13 Tricks ↳ Draft and verify14 of 14

SPECULATIVE DECODING · one big-model pass, illustrative draft

More per pass↖ Tricks
DRAFT · a small model guesses ahead
Draft modelsmall and cheap
·account
.
·He
·deposited
all four enter together
Big model · one pass, all four at oncereads 349.2 GB of weights once
·account
.
·He
·deposited
CHECKED
·account
✓
.
✓
·He
✓
·deposited
✗
·has
big model's pick
4 tokens from one pass · draft words illustrative
TIME · ms, to produce four tokens
Draft and verify
One token per pass
43.6 ms · 4 tokens
174.4 ms
0
50
100
150
200
Time, ms
At one B200's 8 TB/s, though 349 GB needs several GPUs; batch 1, BF16; draft cost not drawn
2–3×FASTER ON T5-XXL
esc
Full video on YouTube

Draft and verify

One trick gets several tokens per weight read: a small model guesses, the big model checks them in one pass.

  1. Cheap: it reads a fraction of the bytes.

  2. The weights are read once, the same as for one token.

  3. The big model's own pick replaces the first wrong one.

  4. On T5-XXL this ran 2–3× faster.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Fast Inference from Transformers via Speculative DecodingLeviathan et al. · paper, new tab
  • weights, one pass
  • guessed token
  • rejected
  • the token we follow
← Back to TricksSame answers →
  • weights, one pass
  • guessed token
  • rejected
  • the token we follow
AI STACK·
Inside a transformer ▾
1A data centre, top down2Inside a transformer
All walkthroughsAll
01 Tokens02 Embedding03 Stream04 Block05 QKV06 Scores06b Heads07 Blend08 MLP09 Stacking10 Output11 Passes12 KV cache13 Tricks↳ Draft and verify
13 Tricks ↳ Draft and verify14 of 14

ACCEPT RULE · big model p against draft q, one position

More per pass↖ Tricks
PROBABILITY · %, big model p against draft q
0%
25%
50%
75%
100%
87.6%
60%
·account
5.3%
15%
·branch
2.9%
10%
·loan
1.8%
5%
·card
1%
4%
·manager
p from the W2 output head at T = 1; q illustrative; other tokens not drawn
ACCEPT TEST · keep with chance min(1, p ÷ q)
·account
p 0.876 ≥ q 0.60
chance 1.00 · always kept
·branch
p 0.053 ÷ q 0.15 = 0.36
chance 0.36 · roll 0.71 lands outside
so ·branch is rejected
ON REJECTION · resample from the excess
·branch
·account
only ·account has p above q here
RESULT
Each token comes out at its p:
exactly as if the big model ran alone.
roll value illustrative
esc
Full video on YouTube

Same answers

A probabilistic accept rule keeps the output distribution exactly the big model's: faster, with no change in quality.

  1. Big model p, draft q.

  2. "·account": 0.876 against 0.60, always kept.

  3. Overconfident guesses are thinned out.

  4. The result matches sampling from the big model alone.

Click a part to read about it. ← → steps, ↑ ↓ parts.

  • Accelerating Large Language Model Decoding with Speculative SamplingChen et al. · paper, new tab
  • big model p
  • draft q
  • excess p − q, the token we follow
  • chance of keeping
  • rejected
← Draft and verify← Back to Tricks
  • big model p
  • draft q
  • excess p − q, the token we follow
  • chance of keeping
  • rejected
KEYBOARDesc or ? to close
←→
Previous or next step
↑↓
Previous or next part; it expands in place
enter
Go inside the expanded part, if it has its own screen
esc
Close what's open, collapse the part, or leave a detour
?
Show or hide this list

MOUSE

Hover a part to highlight it. Click it to read about it.