| """KL1 item 1: achieved bandwidth per kernel, per step, from both SP1 windows' traces (both ranks). |
| |
| PYTHONPATH=. python3 -m offload.kl1.bandwidth --window <SP1-W-dir> [--window ...] --output <json> |
| |
| Every byte count is ESTIMATED (offload/kl1/shapes.py states the formula and each assumption); every duration is |
| MEASURED by kineto in a profiled window (profile-affected). GB/s = ESTIMATED bytes / MEASURED time. The routed experts' |
| bytes depend on how many distinct experts a step's rows pick, which no trace records: they are given as a function of |
| that count, at the uniform-independent count (an upper bound) and at the counts 273 GB/s and the source's 240 GB/s |
| would allow. Output: numbers, kernel families and hashes only. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import statistics |
| import sys |
| from collections.abc import Sequence |
| from pathlib import Path |
| from typing import cast |
|
|
| from offload.gb10_budget.trace_budget import rules, sweep |
| from offload.kl1 import shapes as S |
| from offload.kl1 import trace_kernels as T |
| from offload.prefill_maxlen.evidence import Json, number, record |
|
|
| DECODE_CAPTURES = ('prof-short-b1:decode', 'prof-32k-b1:decode', 'prof-150k-b1:decode', 'prof-32k-b4:decode') |
| PREFILL_CAPTURES = ('prof-32k-b1:prefill-early', 'prof-150k-b1:prefill-late') |
| FULL_LAYERS = 20 |
| SUFFIX_ROWS = 128 |
|
|
|
|
| def _mean(xs: Sequence[float]) -> float: |
| if not xs: |
| raise ValueError('mean of nothing') |
| return statistics.fmean(xs) |
|
|
|
|
| def expert_mean_bytes() -> tuple[float, float]: |
| """(gate|up, down) bytes of an average routed expert on one rank: the K mix weighted by expert counts (assumes a |
| pick's K is distributed like the bank's; picks that favour high-K experts move more).""" |
| total = sum(S.K_MIX.values()) |
| gu = sum(n * S.expert_gu_bytes(k) for k, n in S.K_MIX.items()) / total |
| dn = sum(n * S.expert_down_bytes(k) for k, n in S.K_MIX.items()) / total |
| return gu, dn |
|
|
|
|
| def shared_bytes() -> float: |
| return S.expert_bytes(S.SHARED_K) |
|
|
|
|
| def smoe_bytes(unique: float, live_rows: int) -> float: |
| """One layer's smoe weight bytes: `unique` routed experts read once each (<= 8 rows each, exl3_smoe.cu RX) plus the |
| shared expert once per run of 8 live rows (no L2 reuse assumed between its runs).""" |
| gu, dn = expert_mean_bytes() |
| return unique * (gu + dn) + S.shared_runs(live_rows) * shared_bytes() |
|
|
|
|
| def unique_at(gbps: float, span_us: float, live_rows: int) -> float: |
| """Routed experts a layer could have read in span_us at gbps (shared runs first).""" |
| gu, dn = expert_mean_bytes() |
| return (gbps * 1e3 * span_us - S.shared_runs(live_rows) * shared_bytes()) / (gu + dn) |
|
|
|
|
| def sp1_split(step: T.Step) -> dict[str, float]: |
| """SP1's exclusive budget for this step (equal split among active buckets, trace_budget.sweep) twice: with SP1's |
| rules (l2pf_kernel is 'routed_moe') and with l2pf as its own bucket. ms.""" |
| out: dict[str, float] = {} |
| for tag, rename in (('sp1', False), ('l2pf_apart', True)): |
| spans = [(k.start, k.end, 'l2pf' if rename and k.family == 'l2pf' else k.bucket) for k in step.kernels] |
| excl, _, gaps = sweep(spans, step.lo, step.hi) |
| for b, v in excl.items(): |
| out[f'{tag}.{b}'] = v / 1e3 |
| out[f'{tag}.idle'] = sum(g1 - g0 for g0, g1 in gaps) / 1e3 |
| |
| excl, _, gaps = sweep([(k.start, k.end, k.bucket) for k in step.kernels if k.family != 'l2pf'], step.lo, step.hi) |
| for b, v in excl.items(): |
| out[f'no_l2pf.{b}'] = v / 1e3 |
| out['no_l2pf.idle'] = sum(g1 - g0 for g0, g1 in gaps) / 1e3 |
| return out |
|
|
|
|
| def moe_wall(step: T.Step) -> dict[str, float]: |
| """Wall time of the MoE phase: the union of router, routing top-k and smoe/TF expert kernels, ms per step.""" |
| fams = ('router', 'route_topk', 'smoe_rot', 'smoe_gu', 'smoe_down', 'smoe_fused', 'tf_gu', 'tf_down', 'tf_aux') |
| ks = [k for k in step.kernels if k.family in fams] |
| experts = [k for k in ks if k.family not in ('router', 'route_topk')] |
| return {'moe_union_ms': T.union(((k.start, k.end) for k in ks), step.lo, step.hi) / 1e3, |
| 'expert_union_ms': T.union(((k.start, k.end) for k in experts), step.lo, step.hi) / 1e3, |
| 'router_topk_union_ms': T.union(((k.start, k.end) for k in ks if k not in experts), step.lo, step.hi) / 1e3} |
|
|
|
|
| def dense_decode(step: T.Step) -> dict[str, Json]: |
| """exl3_mrow launches: four a layer in forward order (wq_a+wkv, wq_b, wo_a, wo_b), K from the template.""" |
| ks = [k for k in step.kernels if k.family == 'mrow'] |
| if len(ks) % len(S.DECODE_DENSE): |
| raise ValueError(f'step {step.index}: {len(ks)} mrow launches, not a multiple of 4') |
| rows: dict[str, Json] = {} |
| for j, lin in enumerate(S.DECODE_DENSE): |
| mine = ks[j::len(S.DECODE_DENSE)] |
| for k in mine: |
| if T.template_k(k.name) != lin.k: |
| raise ValueError(f'{lin.name}: template K {T.template_k(k.name)} != {lin.k}') |
| us = sum(k.dur for k in mine) |
| rows[lin.name] = {'launches': len(mine), 'us_per_step': us, 'mb_per_launch': lin.bytes / 1e6, |
| 'gbps': S.gbps(lin.bytes * len(mine), us), 'grid': cast(Json, list(mine[0].grid))} |
| us = sum(k.dur for k in ks) |
| total = sum(lin.bytes for lin in S.DECODE_DENSE) * len(ks) / len(S.DECODE_DENSE) |
| rows['all'] = {'launches': len(ks), 'us_per_step': us, 'mb_per_step': total / 1e6, 'gbps': S.gbps(total, us)} |
| head = [k for k in step.kernels if k.family == 'mrow_head' or (k.family == 'exl3_gemm' and T.template_k(k.name) == 6)] |
| if head: |
| hus = sum(k.dur for k in head) |
| rows['lm_head'] = {'launches': len(head), 'us_per_step': hus, 'mb_per_step': S.LM_HEAD.bytes / 1e6, |
| 'gbps': S.gbps(S.LM_HEAD.bytes, hus), 'kernel': head[0].family} |
| for fam, nbytes in (('router', S.ROUTER_BYTES), ('hc_prenorm', S.MHC_BYTES)): |
| ks2 = [k for k in step.kernels if k.family == fam] |
| if ks2: |
| us2 = sum(k.dur for k in ks2) |
| rows[fam] = {'launches': len(ks2), 'us_per_step': us2, 'mb_per_launch': nbytes / 1e6, |
| 'gbps': S.gbps(nbytes * len(ks2), us2)} |
| return rows |
|
|
|
|
| def smoe_decode(step: T.Step, requests: int) -> dict[str, Json]: |
| layers = T.moe_layers(step) |
| if len(layers) != S.LAYERS: |
| raise ValueError(f'step {step.index}: {len(layers)} smoe layers') |
| m = step.rows |
| slots = layers[0].gu.grid[0] // (S.I_LOCAL // 128) |
| span = sum(L.span for L in layers) |
| gu_us = sum(L.gu.dur for L in layers) |
| down_after = sum(L.down_after_gu for L in layers) |
| rot_us = sum(L.rot.dur for L in layers if L.rot is not None) |
| uni = S.unique_uniform(m) |
| per_layer_span = span / len(layers) |
| return {'rows': m, 'requests': requests, 'graph_slots': slots, 'graph_rows': slots // (S.TOP_K + 1), |
| 'span_ms': span / 1e3, 'gu_ms': gu_us / 1e3, 'down_after_gu_ms': down_after / 1e3, 'rot_ms': rot_us / 1e3, |
| 'shared_runs': S.shared_runs(m), |
| 'unique_uniform': uni, |
| 'gbps_if_uniform': smoe_bytes(uni, m) * len(layers) / span / 1e3, |
| 'gbps_if_6_per_request': smoe_bytes(min(uni, 6.0 * requests), m) * len(layers) / span / 1e3, |
| 'unique_max_at_273': unique_at(S.GB10_SPEC_GBPS, per_layer_span, m), |
| 'unique_max_at_240': unique_at(S.SOURCE_READ_GBPS[1], per_layer_span, m)} |
|
|
|
|
| def tf_prefill(step: T.Step) -> dict[str, Json]: |
| layers = T.moe_layers(step) |
| if len(layers) != S.LAYERS: |
| raise ValueError(f'prefill step {step.index}: {len(layers)} TF layers') |
| gu, dn = expert_mean_bytes() |
| out: dict[str, Json] = {} |
| for tag, sel, rows in (('full', layers[:FULL_LAYERS], step.rows), ('suffix', layers[FULL_LAYERS:], SUFFIX_ROWS)): |
| uni = S.unique_uniform(rows) |
| gus = sum(L.gu.dur for L in sel) |
| dns = sum(L.down.dur for L in sel) |
| out[tag] = {'layers': len(sel), 'rows': rows, 'unique_uniform': uni, 'gu_ms': gus / 1e3, 'down_ms': dns / 1e3, |
| 'gu_gbps_if_uniform': S.gbps(uni * gu * len(sel), gus), |
| 'down_gbps_if_uniform': S.gbps(uni * dn * len(sel), dns), |
| 'span_ms': sum(L.span for L in sel) / 1e3} |
| return out |
|
|
|
|
| def reconstruct_prefill(step: T.Step) -> dict[str, Json]: |
| ks = [k for k in step.kernels if k.family == 'reconstruct'] |
| by: dict[str, list[T.Kernel]] = {} |
| for k in ks: |
| by.setdefault('x'.join(map(str, S.reconstruct_shape(k.grid))), []).append(k) |
| out: dict[str, Json] = {} |
| tb = tus = 0.0 |
| for shape, group in sorted(by.items()): |
| kk = T.template_k(group[0].name) |
| if kk is None: |
| raise ValueError('reconstruct without template K') |
| nbytes = S.reconstruct_bytes(group[0].grid, kk) * len(group) |
| us = sum(k.dur for k in group) |
| tb += nbytes |
| tus += us |
| out[shape] = {'launches': len(group), 'K': kk, 'us': us, 'gbps': S.gbps(nbytes, us)} |
| if ks: |
| out['all'] = {'launches': len(ks), 'us': tus, 'gb': tb / 1e9, 'gbps': S.gbps(tb, tus)} |
| return out |
|
|
|
|
| def _fmean_rows(rows: Sequence[dict[str, Json]], keys: Sequence[str]) -> dict[str, float]: |
| return {k: _mean([number(r[k]) for r in rows]) for k in keys if all(k in r for r in rows)} |
|
|
|
|
| def _flatten(d: dict[str, Json], prefix: str = '') -> dict[str, float]: |
| out: dict[str, float] = {} |
| for k, v in d.items(): |
| if isinstance(v, dict): |
| out.update(_flatten(v, f'{prefix}{k}.')) |
| elif isinstance(v, (int, float)) and not isinstance(v, bool): |
| out[f'{prefix}{k}'] = float(v) |
| return out |
|
|
|
|
| def capture(paths: dict[int, Path], kind: str) -> dict[str, Json]: |
| ranks: dict[str, Json] = {} |
| for rank, path in sorted(paths.items()): |
| trace = T.load(path) |
| steps = [s for s in T.steps(trace) if s.phase == ('decode' if kind == 'decode' else 'prefill')] |
| per: list[dict[str, float]] = [] |
| for s in steps: |
| row: dict[str, Json] = {'period_ms': s.period / 1e3, 'rows': s.rows} |
| row.update(cast(dict[str, Json], sp1_split(s))) |
| row.update(cast(dict[str, Json], moe_wall(s))) |
| fam = T.family_totals(s) |
| row['family_union_ms'] = {f: v['union_us'] / 1e3 for f, v in fam.items()} |
| row['family_sum_ms'] = {f: v['sum_us'] / 1e3 for f, v in fam.items()} |
| if kind == 'decode': |
| row['dense'] = dense_decode(s) |
| row['smoe'] = smoe_decode(s, s.requests) |
| else: |
| if s.rows < 1024: |
| continue |
| row['tf'] = tf_prefill(s) |
| row['reconstruct'] = reconstruct_prefill(s) |
| per.append(_flatten(row)) |
| if not per: |
| raise ValueError(f'{path.name}: no complete {kind} steps') |
| keys = sorted(set().union(*per)) |
| ranks[str(rank)] = {'file_sha256': hashlib.sha256(path.read_bytes()).hexdigest(), 'steps': len(per), |
| 'rows_hist': cast(Json, _hist([p['rows'] for p in per])), |
| 'mean': cast(Json, {k: _mean([p.get(k, 0.0) for p in per]) for k in keys}), |
| 'per_step': cast(Json, per)} |
| return {'kind': kind, 'ranks': ranks} |
|
|
|
|
| def _hist(xs: Sequence[float]) -> dict[str, int]: |
| out: dict[str, int] = {} |
| for x in xs: |
| out[str(int(x))] = out.get(str(int(x)), 0) + 1 |
| return dict(sorted(out.items(), key=lambda kv: int(kv[0]))) |
|
|
|
|
| def window(root: Path) -> dict[str, Json]: |
| manifest = [record(x) for x in cast(list[Json], json.loads((root / 'traces-manifest.json').read_text()))] |
| groups: dict[str, dict[int, Path]] = {} |
| for row in manifest: |
| label = f"{row['cell']}:{row['capture']}" |
| path = (root / str(row['file'])).resolve() |
| if root.resolve() not in path.parents: |
| raise ValueError('manifest path leaves the window') |
| if hashlib.sha256(path.read_bytes()).hexdigest() != row['sha256']: |
| raise ValueError(f'{path.name}: sha256 differs from the manifest') |
| groups.setdefault(label, {})[int(number(row['rank']))] = path |
| out: dict[str, Json] = {} |
| for label in (*DECODE_CAPTURES, *PREFILL_CAPTURES): |
| if label not in groups: |
| continue |
| out[label] = capture(groups[label], 'decode' if label.endswith(':decode') else 'prefill') |
| return out |
|
|
|
|
| def constants() -> dict[str, Json]: |
| gu, dn = expert_mean_bytes() |
| return {'status': 'ESTIMATED bytes (offload/kl1/shapes.py); MEASURED kineto durations; UPSTREAM CLAIM bandwidths', |
| 'mean_k': S.mean_k(), 'expert_gu_mb': gu / 1e6, 'expert_down_mb': dn / 1e6, |
| 'shared_expert_mb': shared_bytes() / 1e6, |
| 'dense_mb_per_layer': {lin.name: lin.bytes / 1e6 for lin in S.DECODE_DENSE}, |
| 'lm_head_mb': S.LM_HEAD.bytes / 1e6, 'router_mb': S.ROUTER_BYTES / 1e6, 'mhc_mb': S.MHC_BYTES / 1e6, |
| 'gb10_spec_gbps': S.GB10_SPEC_GBPS, 'source_read_gbps': cast(Json, list(S.SOURCE_READ_GBPS)), |
| 'source_decode_gbps': cast(Json, list(S.SOURCE_DECODE_GBPS))} |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument('--window', type=Path, action='append', required=True) |
| ap.add_argument('--output', type=Path, required=True) |
| args = ap.parse_args(argv) |
| if args.output.exists(): |
| print(f'{args.output} exists; refusing to overwrite', file=sys.stderr) |
| return 1 |
| rules() |
| result: dict[str, Json] = {'type': 'kl1-bandwidth', 'constants': constants(), 'windows': {}} |
| for root in args.window: |
| cast(dict[str, Json], result['windows'])[root.name] = window(root) |
| args.output.write_text(json.dumps(result, indent=1, allow_nan=False) + '\n') |
| return 0 |
|
|
|
|
| if __name__ == '__main__': |
| raise SystemExit(main()) |
|
|
|
|
