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