diff --git a/README.md b/README.md
index 7c284bf..f768d35 100644
--- a/README.md
+++ b/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"
diff --git a/prain/cli.py b/prain/cli.py
index 7aed86c..e44b023 100644
--- a/prain/cli.py
+++ b/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__":
diff --git a/prain/experts.py b/prain/experts.py
index 102c242..a91ff45 100644
--- a/prain/experts.py
+++ b/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."""
diff --git a/prain/profile.py b/prain/profile.py
new file mode 100644
index 0000000..22c7544
--- /dev/null
+++ b/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'"])
+
+
+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}")
diff --git a/prain/tiered.py b/prain/tiered.py
index 3386d40..c316267 100644
--- a/prain/tiered.py
+++ b/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)
diff --git a/prain/web.py b/prain/web.py
index 4e3040e..0ea9a9b 100644
--- a/prain/web.py
+++ b/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()
diff --git a/tests/test_cli.py b/tests/test_cli.py
index aa2bec7..a286c4c 100644
--- a/tests/test_cli.py
+++ b/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")
diff --git a/tests/test_profile.py b/tests/test_profile.py
new file mode 100644
index 0000000..a274a63
--- /dev/null
+++ b/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("