Skip to content
C.W.K.
Stream
Lesson 06 of 06 · published

The Same Token on CUDA

~14 min · journey, cuda, pcie, vram, mlx-cuda-backend, physics

Level 0Spec-Sheet Skimmer
0 XP0/91 lessons0/19 achievements
0/100 XP to next level100 XP to go0% complete
"On a card the journey has one more step, taken once. Everything else is the same division with a bigger number on top — until the weights stop fitting, and then the extra step is taken on every token."

The Journey, Step by Step, on a Discrete GPU

This quest measured no CUDA machine; every number below is a vendor figure or a division of one, and labeled so. The path is still worth walking, because it is the path this track's other five lessons are silently compared against. Step 1, file to host pages: identical — the checkpoint is memory-mapped into the CPU's memory, faulted in from the page cache or the SSD. Step 2, the extra step: the weights cross PCI Express into the card's own memory, once, at load — a cudaMemcpy from a page-locked staging buffer, at about 63 GB/s on sixteen lanes of Gen 5, so the 27B's 16 GB take a quarter of a second. Step 3, decode: exactly the physics track's loop, against the card's memory instead of the pool: the same 14.42 GB per token, divided by 1,792 GB/s on an RTX 5090 or 3,350 on an H100, which is why the interface lesson conceded that a model that fits the card decodes faster there. Step 4, the token comes home: each sampled id is a few bytes copied back over the bus, microseconds; runtimes hide even that by keeping the loop on the device. Step 5, the cache: the KV cache lives in the card's memory too, so context spends VRAM, not host memory — 64 KB a token for the 27B, 6.4 GB at a hundred thousand tokens, out of 32.

StepMac, MLXMac, PyTorch MPSCard, CUDAEvidence
file → memorymmap, lazy; pages wired at first forwardmmap, lazymmap into host memorymeasured (Mac) / physics
into the compute unit's memorynone — one poola CPU memcpy inside the pool: 110–450 ms for 2.47 GB (22 GB/s warm)PCIe copy once: ~63 GB/s, 0.25 s for 16 GBmeasured (Mac) / vendor-claim (bus)
decode, per token14.42 GB ÷ 638 measured → 44 tok/s ceilingsame bytes, same bus, bigger fixed term14.42 GB ÷ 1,792 → 124 (RTX 5090); ÷ 3,350 → 232 (H100)measured / derived from vendor
token id back to the hosta sync, same poola sync, same poola bus crossing of a few bytes; hidden by CUDA graphsphysics
KV cachein the pool, wiredin the poolin VRAM: 6.4 GB per 100k tokens on the 27Bmeasured (size) / physics

The Step That Repeats When the Weights Do Not Fit

The card's memory is the wall the interface lesson priced. A 32 GB card holds the 27B at 4 bits with room for a long context; it does not hold a 70B at 4 bits, which is about 40 GB, and the runtime's answer is to keep part of the model on the host and stream it across the bus every token. Then step 2 stops being taken once: the division becomes bytes-per-token over 63 GB/s, and the 27B would decode at 4.4 tokens per second on a bus that can move it, against 44 from the pool. That number is the whole reason a 512 GB Mac Studio exists in a household that could have bought cards, and the whole reason the rivals track spends a lesson on the multi-GPU tax: two cards make the wall taller and add a second bus between them. The GB10 in the same table is NVIDIA's own answer to the wall, a pool of 128 GB with a card's software and a laptop's bandwidth — 273 GB/s, a ceiling of 19 tokens per second on the 27B — a third of the Studio's 57 at spec, and 43% of the 44 a kernel actually streams there. The rivals track measures that trade in both directions.

The Inversion

The last fact of this track is the odd one. MLX — the framework Apple built for the one-pool machine, 263 days after the community got there first — now ships a CUDA backend: pip install mlx[cuda12] on Linux, for cards of architecture SM 7.5 and up. Code written against unified memory runs on the modular machine, with the copy step hidden inside the framework where the Mac never needed it. For twenty years the porting ran the other way, from CUDA toward everything else. The mlx quest's production track covers the backend as engineering; here it is the mouse's last move: the framework born of the accident reaching the machine the accident was measured against. The Mac card's entry for this lesson is one line — the five steps, with step 2 marked once or every token for the machine you actually own.

Code

journey_any_device.py — the same loop on whatever accelerator you have; step 2 is the line that differs·python
#!/usr/bin/env python3
"""The same token loop on whatever accelerator this machine has. The one line that
differs between a Mac and a CUDA box is the .to(device): on a Mac it is a copy inside
one pool; on a discrete GPU it is the PCIe crossing, once, at load. Everything after it
is bytes per token against that device's own memory bandwidth."""
import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "unsloth/Llama-3.2-1B-Instruct"
dev = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
sync = {"cuda": lambda: torch.cuda.synchronize(), "mps": lambda: torch.mps.synchronize()}.get(dev, lambda: None)

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).eval()      # step 1: file -> host pages (mmap, lazy)
nbytes = sum(p.numel() * p.element_size() for p in model.parameters())

t0 = time.perf_counter(); model.to(dev); sync(); t_move = time.perf_counter() - t0     # step 2: THE extra step
ids = tok.apply_chat_template([{"role": "user", "content": "Name three uses of unified memory."}],
                              add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to(dev)
with torch.no_grad():
    model.generate(ids, max_new_tokens=4, do_sample=False)                             # warm-up
    sync(); t0 = time.perf_counter(); model.generate(ids, max_new_tokens=1, do_sample=False); sync(); ttft = time.perf_counter() - t0
    sync(); t0 = time.perf_counter(); out = model.generate(ids, max_new_tokens=128, min_new_tokens=128, do_sample=False); sync(); dt = time.perf_counter() - t0
n = out.shape[1] - ids.shape[1]
print(f"device {dev}: {nbytes/1e9:.2f} GB bf16")
print(f"  step 2, .to({dev}): {t_move*1e3:.0f} ms = {nbytes/t_move/1e9:.1f} GB/s  <- one pool on a Mac; PCIe on a card (x16 Gen5 ~63 GB/s)")
print(f"  step 3, decode:     {n/(dt-ttft):.1f} tok/s = {nbytes*n/(dt-ttft)/1e9:.0f} GB/s effective; TTFT {ttft*1e3:.0f} ms on {ids.shape[1]} prompt tokens")

# office, M3 Ultra, torch 2.11.0, 2026-09-15 (a Mac: the step exists in the framework, not in the hardware):
#   step 2, .to(mps): 1133 ms in a cold process; 451 ms in a warm one; 110-115 ms on a repeat (22 GB/s) -- a CPU memcpy inside one pool
#   step 3, decode:   44.2 tok/s = 109 GB/s effective; TTFT 26 ms on 42 prompt tokens
#   for comparison, mlx-lm on the same bytes copies nothing at step 2 and decodes at 190 tok/s (previous lessons)
cuda_derive.py — the 27B on five memory systems, vendor numbers in, divisions out·python
#!/usr/bin/env python3
"""The same 27B 4-bit checkpoint (14.42 GB per token, 64 KB of KV per token) on
five memory systems: vendor bandwidth figures in, derived ceilings and copy times out.
Nothing here is measured -- every input is a vendor number, every output a division."""
BYTES_PER_TOKEN = 14.42e9          # Qwen3.5-27B-4bit, read per token (measurements file)
KV_PER_TOKEN = 64e3                # 2 x 16 full-attention layers x 4 KV heads x 256 x 2 bytes
PCIE5_X16 = 63e9                   # ~63 GB/s per direction (T2 interface lesson)

systems = [  # name, memory GB, bandwidth GB/s, weights cross a bus at load?
    ("M3 Ultra (Mac Studio)",        512, 819,  False),
    ("M2 Ultra (Mac Studio)",        192, 800,  False),
    ("RTX 5090 (32 GB GDDR7)",        32, 1792, True),
    ("RTX PRO 6000 Blackwell (96 GB)", 96, 1792, True),
    ("H100 SXM (80 GB HBM3)",         80, 3350, True),
    ("DGX Spark GB10 (128 GB LPDDR5X)", 128, 273, False),
]
print(f"{'system':34} {'fits 27B+100k ctx?':20} {'ceiling tok/s':>14} {'load copy over PCIe5':>22}")
for name, mem, bw, bus in systems:
    need = (BYTES_PER_TOKEN + 1.6e9 + KV_PER_TOKEN * 100_000) / 1e9      # weights + vision/embed + 100k-token cache
    fits = need <= mem * 0.9
    ceiling = bw * 1e9 / BYTES_PER_TOKEN
    copy = f"{16.05e9 / PCIE5_X16:.2f} s" if bus else "none (one pool)"
    print(f"{name:34} {('yes' if fits else 'NO') + f' ({need:.1f} GB)':20} {ceiling:14.0f} {copy:>22}")

print("\nIf the weights do not fit the card, every token re-crosses the bus:")
print(f"  27B at 4-bit over PCIe 5 x16: {PCIE5_X16 / BYTES_PER_TOKEN:.1f} tok/s ceiling -- the bus becomes the memory")

# system                             fits 27B+100k ctx?    ceiling tok/s   load copy over PCIe5
# M3 Ultra (Mac Studio)              yes (22.4 GB)                    57        none (one pool)
# M2 Ultra (Mac Studio)              yes (22.4 GB)                    55        none (one pool)
# RTX 5090 (32 GB GDDR7)             yes (22.4 GB)                   124                 0.25 s
# RTX PRO 6000 Blackwell (96 GB)     yes (22.4 GB)                   124                 0.25 s
# H100 SXM (80 GB HBM3)              yes (22.4 GB)                   232                 0.25 s
# DGX Spark GB10 (128 GB LPDDR5X)    yes (22.4 GB)                    19        none (one pool)
# If the weights do not fit the card: 4.4 tok/s over PCIe 5 x16

External links

Exercise

Run journey_any_device.py on your Mac and, if you have access to any CUDA machine, there too. Put the step-2 time and the decode rate for each on your card, then run cuda_derive.py with your own card's memory and bandwidth added as a row. Answer on the card: for the largest model you actually run, is step 2 taken once or every token on each machine?
Hint
Once if the 4-bit weights plus the cache you use fit in the card's memory with margin; every token if not. On the Mac the answer is always once, for anything inside the working set — and never at all on the MLX path. If your card's decode is below its ceiling by more than half with a model that fits, the fixed term is the framework's, as in the MPS lesson.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.