From 58653b5eaeebeb3fbadfe556de3be8e2170b21e4 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sat, 23 Aug 2025 16:19:44 +0300
Subject: [PATCH] viz: store memory scale (#11794)

---
 tinygrad/viz/js/index.js | 4 ++--
 tinygrad/viz/serve.py    | 8 +++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index deda1d322b..21a94ce3d2 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -151,7 +151,7 @@ async function renderProfiler() {
   // layout once!
   if (data != null) return;
   const profiler = d3.select(".profiler").html("");
-  const { layout, dur } = await (await fetch("/get_profile")).json();
+  const { layout, dur, peak } = await (await fetch("/get_profile")).json();
   // place devices on the y axis and set vertical positions
   const [tickSize, padding] = [10, 8];
   const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
@@ -163,7 +163,7 @@ async function renderProfiler() {
   // color by key (name/category/device)
   const colorMap = new Map();
   data = {tracks:new Map(), axes:{}};
-  const heightScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxheight=100]);
+  const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
   for (const [k, v] of Object.entries(layout)) {
     if (v.shapes.length === 0) continue;
     const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px");
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index e391b2ecba..e179b07638 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -146,7 +146,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int)
     shapes.append({"name":name, "ref":ref, "st":st-start_ts, "dur":dur, "depth":depth, "cat":cat, "info":info})
   return {"shapes":shapes, "maxDepth":len(levels)}
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int]) -> dict:
   step, peak, mem = 0, 0, 0
   shps:dict[int, dict] = {}
   temp:dict[int, dict] = {}
@@ -173,6 +173,7 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_
     v["x"].append(step)
     v["y"].append(v["y"][-1])
   timestamps.append(end_ts-start_ts)
+  peaks.append(peak)
   return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
@@ -190,11 +191,12 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   if start_ts is None: return None
   # return layout of per device events
   layout:dict[str, dict] = {}
+  peaks:list[int] = []
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
     layout[k] = timeline_layout(v, start_ts)
-    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts))
-  return json.dumps({"layout":layout, "dur":unwrap(end_ts)-start_ts}).encode("utf-8")
+    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks)
+  return json.dumps({"layout":layout, "dur":unwrap(end_ts)-start_ts, "peak":max(peaks, default=0)}).encode("utf-8")
 
 def get_runtime_stats(key) -> list[dict]:
   ret:list[dict] = []