File size: 7,757 Bytes
3de4238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | import json
import os
import statistics
from collections import defaultdict
from typing import Dict, List, Optional
import torch
class _FirstTokenTimer:
"""Streamer compatible with HuggingFace generate(streamer=...).
Records the time of the first put() call as the TTFT end point; subsequent calls only accumulate the token count.
On end(), after a CUDA sync, the elapsed time is computed and recorded into LATENCY_TRACKER.
"""
def __init__(self, tracker: "LatencyTracker", subtask: str):
self.tracker = tracker
self.subtask = subtask
self.t_start = torch.cuda.Event(enable_timing=True)
self.t_first = torch.cuda.Event(enable_timing=True)
self.t_end = torch.cuda.Event(enable_timing=True)
self.first_token_seen = False
self.prompt_echo_skipped = False
self.n_tokens = 0
self.t_start.record()
def put(self, value):
# HF generate() calls streamer.put(input_ids) once before generation
# (transformers/generation/utils.py line 180). Skip that prompt echo
# so TTFT is measured to the first *generated* token.
if not self.prompt_echo_skipped:
self.prompt_echo_skipped = True
return
n = value.numel() if isinstance(value, torch.Tensor) else 1
if not self.first_token_seen:
self.t_first.record()
self.first_token_seen = True
self.n_tokens += n
def end(self):
self.t_end.record()
torch.cuda.synchronize()
if not self.first_token_seen:
return
ttft_ms = self.t_start.elapsed_time(self.t_first)
e2e_ms = self.t_start.elapsed_time(self.t_end)
decode_ms = max(e2e_ms - ttft_ms, 1e-6)
self.tracker.record(
subtask=self.subtask,
ttft_ms=ttft_ms,
decode_ms=decode_ms,
e2e_ms=e2e_ms,
n_tokens=self.n_tokens,
)
class LatencyTracker:
def __init__(self):
self.enabled: bool = False
self.warmup_batches: int = 0
self.down_t: int = 4
self._batch_idx: int = 0
self._current_subtask: str = "unknown"
self._records: Dict[str, List[Dict[str, float]]] = defaultdict(list)
def enable(self, warmup_batches: int = 5, down_t: int = 4):
self.enabled = True
self.warmup_batches = int(warmup_batches)
self.down_t = int(down_t)
self._batch_idx = 0
self._current_subtask = "unknown"
self._records.clear()
def disable(self):
self.enabled = False
def mark_batch_start(self):
if not self.enabled:
return
self._batch_idx += 1
def set_current_subtask(self, name: str):
self._current_subtask = str(name)
def make_streamer(self) -> Optional[_FirstTokenTimer]:
if not self.enabled:
return None
return _FirstTokenTimer(self, self._current_subtask)
def record(self, subtask: str, ttft_ms: float, decode_ms: float,
e2e_ms: float, n_tokens: int):
if not self.enabled:
return
if self._batch_idx <= self.warmup_batches:
return
if n_tokens <= 0:
return
tps = (n_tokens - 1) / (decode_ms / 1000.0) if n_tokens > 1 else 0.0
self._records[subtask].append({
"ttft_ms": ttft_ms,
"decode_ms": decode_ms,
"e2e_ms": e2e_ms,
"n_tokens": n_tokens,
"tps": tps,
})
@staticmethod
def _stats(values: List[float], with_p99: bool = True) -> Dict[str, float]:
if not values:
return {}
s = sorted(values)
n = len(s)
def pct(p: float) -> float:
if n == 1:
return s[0]
k = (n - 1) * p
f = int(k)
c = min(f + 1, n - 1)
return s[f] + (s[c] - s[f]) * (k - f)
out = {
"mean": statistics.fmean(values),
"p50": pct(0.50),
"p95": pct(0.95),
"min": s[0],
"max": s[-1],
}
if with_p99:
out["p99"] = pct(0.99)
return out
def _build_entry(self, recs: List[Dict[str, float]]) -> dict:
ttft = [r["ttft_ms"] for r in recs]
tps = [r["tps"] for r in recs if r["n_tokens"] > 1]
e2e = [r["e2e_ms"] for r in recs]
motion_fps = [r["tps"] * self.down_t for r in recs if r["n_tokens"] > 1]
return {
"count": len(recs),
"ttft_ms": self._stats(ttft, with_p99=True),
"tps": self._stats(tps, with_p99=False),
"motion_fps": self._stats(motion_fps, with_p99=False),
"e2e_ms": self._stats(e2e, with_p99=True),
"total_tokens": sum(r["n_tokens"] for r in recs),
}
def summary(self) -> dict:
per_subtask = {}
all_records: List[Dict[str, float]] = []
for name, recs in self._records.items():
all_records.extend(recs)
per_subtask[name] = self._build_entry(recs)
overall = self._build_entry(all_records)
device = "cuda:" + str(torch.cuda.current_device()) if torch.cuda.is_available() else "cpu"
return {
"warmup_batches": self.warmup_batches,
"down_t": self.down_t,
"motion_fps_formula": "tps * down_t (motion frames per second per single generated sequence)",
"device": device,
"per_subtask": per_subtask,
"overall": overall,
}
def print_summary(self):
s = self.summary()
line_w = 156
print("=" * line_w)
print("LM Generate: TTFT / TPS / Motion FPS")
print(f"Warm-up batches skipped: {s['warmup_batches']} "
f"down_t: {s['down_t']} Device: {s['device']}")
print(f"motion_fps = tps * down_t (motion frames/sec per single generated sequence)")
print("-" * line_w)
header1 = (f"{'Subtask':<18}{'n':>6} "
f"{'TTFT(ms)':<36}{'TPS(tok/s)':<30}{'motion_fps':<24}{'E2E(ms)':<30}{'tokens':>10}")
header2 = (f"{'':<18}{'':>6} "
f"{'mean':>6} {'p50':>6} {'p95':>6} {'p99':>6} {'min':>6} {'max':>6} "
f"{'mean':>6} {'p50':>6} {'p95':>6} {'min':>6} "
f"{'mean':>6} {'p50':>6} {'min':>6} "
f"{'mean':>6} {'p50':>6} {'p95':>6} {'p99':>6} {'max':>6} "
f"{'':>10}")
print(header1)
print(header2)
print("-" * line_w)
def _fmt_row(name: str, d: dict):
ttft = d["ttft_ms"] or {}
tps = d["tps"] or {}
mfps = d["motion_fps"] or {}
e2e = d["e2e_ms"] or {}
def g(x, k):
return f"{x.get(k, 0.0):>6.1f}" if x else f"{'-':>6}"
return (f"{name:<18}{d['count']:>6} "
f"{g(ttft,'mean')} {g(ttft,'p50')} {g(ttft,'p95')} {g(ttft,'p99')} {g(ttft,'min')} {g(ttft,'max')} "
f"{g(tps,'mean')} {g(tps,'p50')} {g(tps,'p95')} {g(tps,'min')} "
f"{g(mfps,'mean')} {g(mfps,'p50')} {g(mfps,'min')} "
f"{g(e2e,'mean')} {g(e2e,'p50')} {g(e2e,'p95')} {g(e2e,'p99')} {g(e2e,'max')} "
f"{d['total_tokens']:>10}")
for name, d in s["per_subtask"].items():
print(_fmt_row(name, d))
print("-" * line_w)
print(_fmt_row("overall", s["overall"]))
print("=" * line_w)
def dump_json(self, path: str):
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
with open(path, "w") as f:
json.dump(self.summary(), f, indent=2)
LATENCY_TRACKER = LatencyTracker()
|