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()