278 lines
10 KiB
Python
278 lines
10 KiB
Python
"""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}")
|