prain: span profiler with flame graphs (PRAIN_PROFILE: engine step, every model module, MoE route/plan/copy/round, call counts; python -m prain.profile); prain web --host

README.md

@@ -35,8 +35,8 @@ The first start writes the bf16 pack (42.5 GiB, ~3 min); with `quant = "int8"` o
 then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files
 (`layerNNN.w13`, `.w2`, `.done`) are not read and can be deleted; a later bf16 or int8 start
 writes them again. Chat with `prain tui`, or `prain web --config FILE` for several users with
-drafts. Tool calls work through vLLM's `gemma4` parser, so coding agents and MCP clients can use
-the OpenAI-compatible API.
+drafts (`--host 0.0.0.0` serves it to other machines; it has no login). Tool calls work through
+vLLM's `gemma4` parser, so coding agents and MCP clients can use the OpenAI-compatible API.
 
 ## Expert formats
 
@@ -123,6 +123,14 @@ code at a negative log-likelihood of 5-9 per token on a correct engine.
 The simulator treats every disk-tier expert as a disk read, so it under-predicts bf16 (0.6 vs 1.04
 tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s).
 
+## Profile
+
+Serve with `PRAIN_PROFILE=/abs/dir`. Writing a number N into `dir/start` records the next N engine
+steps: spans around the engine step, every module of the model and prain's MoE phases (route,
+plan, copy: h2d/repack/d2d, round), each synced with the GPU. `python -m prain.profile dir`
+writes `decode.svg`/`prefill.svg` flame graphs and prints self time and calls per step. The
+syncs remove the CPU/GPU overlap: a decode step took 395 ms profiled, 307 ms without.
+
 ## vLLM patches
 
 - `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"

prain/cli.py

@@ -167,6 +167,8 @@ def main(argv=None):
                              "--scheduling-policy priority")
         if name == "web":
             ui.add_argument("--port", type=int, default=8080)
+            ui.add_argument("--host", default="127.0.0.1",
+                            help="address to bind; 0.0.0.0 for every interface (no login!)")
             ui.add_argument("--config", help="placement file: [serve] max_requests and the "
                                              "[draft] table (default: their defaults)")
     args = parser.parse_args(argv)
@@ -199,7 +201,8 @@ def main(argv=None):
         from prain.web import serve as web
 
         config = load_config(args.config) if args.config else Config(pack_dir="")  # defaults
-        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft)
+        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft,
+            host=args.host)
 
 
 if __name__ == "__main__":

prain/experts.py

@@ -32,7 +32,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
-from prain import drafts
+from prain import drafts, profile
 from prain.config import layer_index, load_config
 from prain.policy import SlotCache, write_stats
 from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4,
@@ -127,6 +127,8 @@ class TieredRoutedExperts(RoutedExperts):
             raise NotImplementedError("prain: one GPU only (no TP, DP or EP)")
         if vllm_config.use_v2_model_runner:
             _install_draft_hook()
+            if profile.DIR:
+                profile.install()
         elif not LAYERS:
             logger.warning("prain: drafts need vLLM's V2 model runner; with "
                            "VLLM_USE_V2_MODEL_RUNNER=0 every request runs with all its experts")
@@ -308,8 +310,12 @@ class TieredRoutedExperts(RoutedExperts):
         for dsts, srcs in zip(self._slots, self._source(expert)):
             for i, (dst, src) in enumerate(zip(dsts, srcs)):
                 if i == 0 and self.quant == "int4":  # the pack keeps the CPU kernel's order
-                    src = to_gpu_int4(src.to(dst.device, non_blocking=True))
-                dst[slot].copy_(src, non_blocking=True)
+                    with profile.span("h2d"):
+                        src = src.to(dst.device, non_blocking=True)
+                    with profile.span("repack"):
+                        src = to_gpu_int4(src)
+                with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
+                    dst[slot].copy_(src, non_blocking=True)
 
     def _load(self, expert, slot):
         """Demand copy on the compute stream."""

prain/profile.py

@@ -0,0 +1,277 @@
+"""Span profiler for the vLLM engine process: flame-graph data per engine step.
+
+Active only with PRAIN_PROFILE=DIR. Writing a number N into DIR/start records the next N engine
+steps: the step, the runner's execute_model and sample_tokens, every module forward of the model
+and prain's MoE phases become nested spans. Each span syncs the GPU at both ends, so GPU work is
+charged to the span that queued it (the overlap of CPU and GPU work is lost; a sync costs ~10 us).
+Then it writes DIR/decode.folded and DIR/prefill.folded ("frame;frame;frame microseconds" lines:
+time spent in the last frame itself, summed over the steps of that kind; a step with more than
+DECODE_TOKENS tokens is a prefill step), DIR/steps.json, and removes DIR/start.
+scripts/flamegraph.py draws them.
+"""
+
+import contextlib
+import json
+import os
+import threading
+import time
+
+import torch
+
+DIR = os.environ.get("PRAIN_PROFILE")
+SYNC = [torch.cuda.synchronize]  # tests replace it
+DECODE_TOKENS = 4
+
+
+class _Recorder:
+    def __init__(self):
+        self.active, self.left, self.thread = False, 0, None
+        self.stack = []  # [frame, start, time in children]
+        self.totals, self.steps, self.hooks = {}, {}, []
+        self.kind = "decode"
+
+
+_R = _Recorder()
+_NULL = contextlib.nullcontext()
+
+
+def _enter(frame):
+    SYNC[0]()
+    _R.stack.append([frame, time.perf_counter(), 0.0])
+
+
+def _exit():
+    SYNC[0]()
+    frame, start, children = _R.stack.pop()
+    took = time.perf_counter() - start
+    stack = (*(f[0] for f in _R.stack), frame)
+    _R.step[stack] = _R.step.get(stack, 0.0) + took - children
+    _R.step_calls[stack] = _R.step_calls.get(stack, 0) + 1
+    if _R.stack:
+        _R.stack[-1][2] += took
+
+
+def _recording():
+    return _R.active and threading.get_ident() == _R.thread
+
+
+def _pre_hook(frame):
+    def hook(module, args):  # must return None: anything else replaces the module's inputs
+        if _recording():
+            _enter(frame)
+    return hook
+
+
+def _post_hook(module, args, out):  # must return None, or it replaces the output
+    if _recording():
+        _exit()
+
+
+@contextlib.contextmanager
+def _span(name):
+    _enter(name)
+    try:
+        yield
+    finally:
+        _exit()
+
+
+def span(name):
+    """with span("copy"): ... -- a no-op unless a recording runs."""
+    return _span(name) if _recording() else _NULL
+
+
+def step_begin():
+    if not _R.active and DIR and os.path.exists(os.path.join(DIR, "start")):
+        with open(os.path.join(DIR, "start")) as f:
+            text = f.read().strip()
+        _R.left = int(text) if text.isdigit() else 1
+        _R.active, _R.thread, _R.totals, _R.steps = True, threading.get_ident(), {}, {}
+        _R.calls = {}
+    if _recording():
+        _R.step, _R.step_calls, _R.kind = {}, {}, "decode"
+        _enter("step")
+
+
+def note_step(model, tokens):
+    """In execute_model: the step's token count, and the module hooks on first use."""
+    if not _recording():
+        return
+    if tokens > DECODE_TOKENS:
+        _R.kind = "prefill"
+    if not _R.hooks:
+        for name, module in model.named_modules():
+            attr, cls = name.rsplit(".", 1)[-1], type(module).__name__
+            frame = cls if not name or attr.isdigit() else f"{attr}:{cls}"
+            _R.hooks.append(module.register_forward_pre_hook(_pre_hook(frame)))
+            _R.hooks.append(module.register_forward_hook(_post_hook))
+
+
+def step_end():
+    if not _recording():
+        return
+    while _R.stack:  # also closes spans an exception left open
+        _exit()
+    for stack, seconds in _R.step.items():
+        key = (_R.kind, *stack)
+        _R.totals[key] = _R.totals.get(key, 0.0) + seconds
+        _R.calls[key] = _R.calls.get(key, 0) + _R.step_calls[stack]
+    _R.steps[_R.kind] = _R.steps.get(_R.kind, 0) + 1
+    _R.left -= 1
+    if _R.left <= 0:
+        _finish()
+
+
+def _finish():
+    _R.active = False
+    for hook in _R.hooks:
+        hook.remove()
+    _R.hooks = []
+    for kind in _R.steps:
+        for suffix, values, fmt in (("folded", _R.totals, lambda s: round(s * 1e6)),
+                                    ("calls", _R.calls, int)):
+            with open(os.path.join(DIR, f"{kind}.{suffix}"), "w") as f:
+                for key, value in sorted(values.items()):
+                    if key[0] == kind:
+                        f.write(f"{';'.join(key[1:])} {fmt(value)}\n")
+    with open(os.path.join(DIR, "steps.json"), "w") as f:
+        json.dump({"steps": _R.steps, "decode_tokens": DECODE_TOKENS}, f)
+    with contextlib.suppress(OSError):
+        os.remove(os.path.join(DIR, "start"))
+
+
+def install():
+    """Spans around vLLM's engine step and the model runner (the vLLM process with the GPU)."""
+    from vllm.v1.engine.core import EngineCoreProc
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    def wrap(cls, name, around):
+        original = getattr(cls, name)
+        if not getattr(original, "prain_profile", False):
+            wrapped = around(original)
+            wrapped.prain_profile = True
+            setattr(cls, name, wrapped)
+
+    def engine_step(original):
+        def step(self, *args, **kwargs):
+            step_begin()
+            try:
+                return original(self, *args, **kwargs)
+            finally:
+                step_end()
+        return step
+
+    def execute_model(original):
+        def execute(self, scheduler_output, *args, **kwargs):
+            with span("execute_model"):
+                note_step(self.model, getattr(scheduler_output, "total_num_scheduled_tokens", 0))
+                return original(self, scheduler_output, *args, **kwargs)
+        return execute
+
+    def sample_tokens(original):
+        def sample(self, *args, **kwargs):
+            with span("sample_tokens"):
+                return original(self, *args, **kwargs)
+        return sample
+
+    wrap(EngineCoreProc, "_process_engine_step", engine_step)
+    wrap(GPUModelRunner, "execute_model", execute_model)
+    wrap(GPUModelRunner, "sample_tokens", sample_tokens)
+
+
+_COLORS = (("copy", "#e8603c"), ("round", "#f0a030"), ("expert", "#f0a030"), ("moe", "#f0a030"),
+           ("route", "#e8c440"), ("plan", "#e8c440"), ("degrade", "#e8c440"), ("cpu", "#e8c440"),
+           ("attn", "#5b9bd5"), ("rotary", "#5b9bd5"), ("mlp", "#6abf69"), ("gate_up", "#6abf69"),
+           ("down_proj", "#6abf69"), ("norm", "#a0a0a0"), ("logits", "#a070c0"),
+           ("lm_head", "#a070c0"), ("sample", "#a070c0"), ("embed", "#a070c0"))
+
+
+def _tree(lines):
+    root = {"name": "all", "self": 0.0, "kids": {}}
+    for line in lines:
+        stack, _, us = line.strip().rpartition(" ")
+        node = root
+        for frame in stack.split(";"):
+            node = node["kids"].setdefault(frame, {"name": frame, "self": 0.0, "kids": {}})
+        node["self"] += float(us)
+
+    def total(node):
+        node["total"] = node["self"] + sum(total(k) for k in node["kids"].values())
+        return node["total"]
+    total(root)
+    return root
+
+
+def top(lines, steps=1, n=25, calls=None):
+    """[(frame, ms per step, % of all, calls per step or None)] by time in the frame itself;
+    calls: the .calls lines of the same recording."""
+    own, count = {}, {}
+    for line in lines:
+        stack, _, us = line.strip().rpartition(" ")
+        frame = stack.rsplit(";", 1)[-1]
+        own[frame] = own.get(frame, 0.0) + float(us)
+    for line in calls or ():
+        stack, _, n_calls = line.strip().rpartition(" ")
+        frame = stack.rsplit(";", 1)[-1]
+        count[frame] = count.get(frame, 0) + int(n_calls)
+    whole = sum(own.values()) or 1.0
+    return [(frame, round(us / 1e3 / steps, 1), round(100 * us / whole, 1),
+             round(count[frame] / steps, 1) if frame in count else None)
+            for frame, us in sorted(own.items(), key=lambda kv: -kv[1])[:n]]
+
+
+def svg(lines, steps=1, title="", width=1200, row=18):
+    """A flame graph of folded stack lines (root on top, widest child first); each box's
+    tooltip gives its time per step, its share of all time and its own (self) time."""
+    from html import escape
+
+    root = _tree(lines)
+    whole = root["total"] or 1.0
+    boxes = []
+
+    def place(node, x, depth):
+        boxes.append((node, x, depth, node["total"] / whole * width))
+        for kid in sorted(node["kids"].values(), key=lambda k: -k["total"]):
+            place(kid, x, depth + 1)
+            x += kid["total"] / whole * width
+    place(root, 0.0, 0)
+    height = (max(d for _, _, d, _ in boxes) + 1) * row + 30
+    out = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
+           f'font-family="sans-serif" font-size="11">',
+           f'<text x="4" y="16" font-size="14">{escape(title)}: '
+           f'{whole / 1e3 / steps:.1f} ms per step over {steps} steps (hover for times)</text>']
+    for node, x, depth, w in boxes:
+        if w < 0.5:
+            continue
+        name, y = node["name"], 26 + depth * row
+        color = next((c for key, c in _COLORS if key in name.lower()), "#d8b890")
+        ms = node["total"] / 1e3 / steps
+        text = f"{name} {ms:.1f} ms" if ms >= 0.1 else name  # readable without hovering
+        label = text if len(text) * 6.5 < w - 4 else text[:max(int((w - 4) / 6.5) - 1, 0)] + "…"
+        out.append(
+            f'<g><title>{escape(name)}: {ms:.1f} ms per step '
+            f'({100 * node["total"] / whole:.1f}%), self {node["self"] / 1e3 / steps:.1f} ms</title>'
+            f'<rect x="{x:.1f}" y="{y}" width="{w:.1f}" height="{row - 1}" fill="{color}" rx="2"/>'
+            + (f'<text x="{x + 3:.1f}" y="{y + 12}">{escape(label)}</text>' if w > 14 else "")
+            + "</g>")
+    return "\n".join(out + ["</svg>"])
+
+
+if __name__ == "__main__":  # python -m prain.profile DIR: DIR/KIND.svg and a summary per kind
+    import sys
+
+    directory = sys.argv[1]
+    with open(os.path.join(directory, "steps.json")) as f:
+        counts = json.load(f)["steps"]
+    for kind, steps in counts.items():
+        with open(os.path.join(directory, f"{kind}.folded")) as f:
+            lines = f.read().splitlines()
+        calls_path = os.path.join(directory, f"{kind}.calls")
+        calls = open(calls_path).read().splitlines() if os.path.exists(calls_path) else None
+        with open(os.path.join(directory, f"{kind}.svg"), "w") as f:
+            f.write(svg(lines, steps, title=kind))
+        print(f"== {kind}: {steps} steps, {_tree(lines)['total'] / 1e3 / steps:.1f} ms per step"
+              f"  (self time per step, share, calls per step, ms per call)")
+        for frame, ms, share, n in top(lines, steps, calls=calls):
+            per_call = f"{n:8.1f} x {ms / n:7.3f} ms" if n else ""
+            print(f"{ms:9.1f} ms {share:5.1f}%  {per_call}  {frame}")

prain/tiered.py

@@ -3,6 +3,8 @@
 import torch
 import torch.nn.functional as F
 
+from prain import profile
+
 HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
 
 
@@ -119,21 +121,23 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
     raw = None
     skip = miss is not None and miss.min_weight < 1
     if draft is not None or skip:
-        if on_stats is not None:  # the trace records the routing, not what was skipped
-            raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
-        rows = draft if draft is not None else torch.zeros(topk_ids.shape[0], dtype=torch.bool,
-                                                           device=topk_ids.device)
-        need = min_weight
-        if skip:  # every row may skip; draft rows by the smaller share
-            need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
-            rows = torch.ones_like(rows)
-        available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
-                                      cache.num_experts)
-        _, kept = degrade(topk_weights, topk_ids, rows, available, need)
-        # experts some row copies anyway serve every row for free
-        available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
-        topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
-    stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # host sync
+        with profile.span("degrade"):
+            if on_stats is not None:  # the trace records the routing, not what was skipped
+                raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
+            rows = draft if draft is not None else torch.zeros(
+                topk_ids.shape[0], dtype=torch.bool, device=topk_ids.device)
+            need = min_weight
+            if skip:  # every row may skip; draft rows by the smaller share
+                need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
+                rows = torch.ones_like(rows)
+            available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
+                                          cache.num_experts)
+            _, kept = degrade(topk_weights, topk_ids, rows, available, need)
+            # experts some row copies anyway serve every row for free
+            available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
+            topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
+    with profile.span("route"):
+        stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # sync
     tokens = int(topk_ids.shape[0])
     if on_stats is not None:
         on_stats(raw if raw is not None else stats, tokens)
@@ -150,14 +154,20 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
         nonlocal total
         total = out if total is None else total.to(acc).add_(out)  # .to: a no-op once fp32
 
-    for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
+    with profile.span("plan"):
+        rounds = cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu)
+    for slot_of, loads in rounds:
         for expert, slot in loads:
-            load(expert, slot)
+            with profile.span("copy"):
+                load(expert, slot)
         if slot_of:  # each round computes only its own routes
-            table = slot_table(slot_of, cache.num_experts, x.device)
-            add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
+            with profile.span("round"):
+                table = slot_table(slot_of, cache.num_experts, x.device)
+                add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
     if after_enqueue is not None:
-        after_enqueue(cpu, stats)
+        with profile.span("prefetch"):
+            after_enqueue(cpu, stats)
     if cpu:  # runs while the GPU works through its rounds
-        add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
+        with profile.span("cpu experts"):
+            add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
     return torch.zeros_like(x) if total is None else total.to(x.dtype)

prain/web.py

@@ -3,8 +3,8 @@
 A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes
 {"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
 "messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
-"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
-machines.
+"refine" | "error"} events of its own. Binds to 127.0.0.1 unless --host says otherwise
+(0.0.0.0 serves every interface, without a login).
 
 Several browsers at once ([serve] and [draft] tables of the placement file):
 - At most draft.full_requests answers run with all their experts. An answer that arrives
@@ -267,8 +267,9 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
     return server
 
 
-def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()):
-    server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
-    print(f"prain web: http://127.0.0.1:{server.server_port}  (vLLM at {base_url}; "
+def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1"):
+    server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve,
+                         draft=draft)
+    print(f"prain web: http://{host}:{server.server_port}  (vLLM at {base_url}; "
           f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
     server.serve_forever()

tests/test_cli.py

@@ -61,6 +61,13 @@ class Cli(TmpCase):
             main(["serve", config, "/models/m", "--", "--max-num-seqs=4"])
         execvp.assert_not_called()
 
+    def test_web_binds_the_host_it_is_given(self):
+        with mock.patch("prain.web.serve") as web:
+            main(["web", "--port", "8081"])
+            self.assertEqual(web.call_args.kwargs["host"], "127.0.0.1")
+            main(["web", "--host", "0.0.0.0"])
+            self.assertEqual(web.call_args.kwargs["host"], "0.0.0.0")
+
     def test_synth_sim_and_tune(self):
         hf = self.write("config.json", json.dumps(GEMMA4))
         trace = os.path.join(self.dir, "trace.jsonl.gz")

tests/test_profile.py

@@ -0,0 +1,100 @@
+import json
+import os
+import threading
+import time
+
+import torch
+
+from prain import profile
+from tests.helpers import TmpCase
+
+
+class Inner(torch.nn.Module):
+    def forward(self, x):
+        time.sleep(0.002)
+        return x
+
+
+class Outer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([Inner(), Inner()])
+        self.head = Inner()
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        with profile.span("copy"):
+            time.sleep(0.001)
+        return self.head(x)
+
+
+class TestProfile(TmpCase):
+    def setUp(self):
+        super().setUp()
+        self.syncs = []
+        saved = profile.DIR, profile.SYNC[0]
+        profile.DIR, profile.SYNC[0] = self.dir, lambda: self.syncs.append(1)
+        self.addCleanup(lambda: (setattr(profile, "DIR", saved[0]),
+                                 profile.SYNC.__setitem__(0, saved[1])))
+
+    def steps(self, model, tokens):
+        for n in tokens:
+            profile.step_begin()
+            with profile.span("execute_model"):
+                profile.note_step(model, n)
+                model(torch.zeros(1))
+            profile.step_end()
+
+    def folded(self, kind):
+        with open(os.path.join(self.dir, f"{kind}.folded")) as f:
+            return {stack: int(us) for stack, us in (l.rsplit(" ", 1) for l in f.read().splitlines())}
+
+    def test_records_the_requested_steps_as_folded_stacks(self):
+        model = Outer()
+        self.steps(model, [1])  # no start file: nothing recorded, nothing synced
+        self.assertEqual(self.syncs, [])
+        self.write("start", "3")
+        self.steps(model, [1, 1, 300, 1])  # records 2 decode steps and 1 prefill, not the 4th
+        self.assertFalse(os.path.exists(os.path.join(self.dir, "start")))
+        with open(os.path.join(self.dir, "steps.json")) as f:
+            self.assertEqual(json.load(f)["steps"], {"decode": 2, "prefill": 1})
+        decode = self.folded("decode")
+        self.assertEqual(set(decode), {"step", "step;execute_model", "step;execute_model;Outer",
+                                       "step;execute_model;Outer;Inner",  # both layers, merged
+                                       "step;execute_model;Outer;copy",
+                                       "step;execute_model;Outer;head:Inner"})
+        for stack, ms in (("Outer;Inner", 8), ("Outer;head:Inner", 4), ("Outer;copy", 2)):
+            self.assertGreaterEqual(decode[f"step;execute_model;{stack}"], ms * 1000)
+            self.assertLess(decode[f"step;execute_model;{stack}"], ms * 1000 * 3)
+        self.assertGreater(self.folded("prefill")["step;execute_model;Outer;Inner"], 4000)
+        with open(os.path.join(self.dir, "decode.calls")) as f:
+            calls = dict(line.rsplit(" ", 1) for line in f.read().splitlines())
+        self.assertEqual(calls["step;execute_model;Outer;Inner"], "4")  # 2 layers x 2 steps
+        self.assertEqual(calls["step"], "2")
+        recorded = len(self.syncs)
+        self.steps(model, [1])  # finished: the module hooks are gone
+        self.assertEqual(len(self.syncs), recorded)
+        self.assertFalse(model.head._forward_hooks or model.head._forward_pre_hooks)
+
+    def test_hooks_leave_other_threads_inputs_and_outputs_alone(self):
+        model = Outer()
+        self.write("start", "1")
+        profile.step_begin()
+        profile.note_step(model, 1)  # installs the hooks
+        out = []
+        thread = threading.Thread(target=lambda: out.append(model(torch.ones(1))))
+        thread.start()
+        thread.join()
+        self.assertTrue(torch.equal(out[0], torch.ones(1)))
+        profile.step_end()
+
+    def test_svg_shows_time_per_step(self):
+        lines = ["step 500", "step;attn 3000", "step;attn;copy 1000"]
+        svg = profile.svg(lines, steps=2, title="decode")
+        self.assertTrue(svg.startswith("<svg"))
+        self.assertIn("attn: 2.0 ms per step (88.9%)", svg)  # 4000 us of 4500, over 2 steps
+        self.assertIn("copy: 0.5 ms per step", svg)
+        self.assertEqual(profile.top(lines, steps=2)[0], ("attn", 1.5, 66.7, None))  # self time
+        calls = ["step;attn 10", "step;attn;copy 4", "step 2"]
+        self.assertEqual(profile.top(lines, steps=2, calls=calls)[0], ("attn", 1.5, 66.7, 5.0))