Add WASD/arrow key navigation and fix unused variable

Perfetto-style keyboard controls: A/D or left/right to pan,
W/S or up/down to zoom in/out. Also removes unused active_by_addr
variable flagged by ruff.

Author: drisspg

transformer_nuggets/utils/benchmark.py

@@ -265,6 +265,7 @@ def save_memory_snapshot(file_path: Path | str, viz: str = "torch"):
                 html = torch.cuda._memory_viz.trace_plot(s)  # type: ignore
             case "d3":
                 from transformer_nuggets.utils.memory_viz import generate_memory_html
+
                 html = generate_memory_html(s)
             case _:
                 raise ValueError(f"Unknown viz backend: {viz!r}, expected 'torch' or 'd3'")
@@ -334,6 +335,7 @@ def oom_observer(device, alloc, device_alloc, device_free):
                     html = torch.cuda._memory_viz.trace_plot(snapshot)  # type: ignore
                 case "d3":
                     from transformer_nuggets.utils.memory_viz import generate_memory_html
+
                     html = generate_memory_html(snapshot)
                 case _:
                     html = torch.cuda._memory_viz.trace_plot(snapshot)  # type: ignore

transformer_nuggets/utils/memory_viz.py

@@ -7,7 +7,10 @@ def _extract_frames(frames: list[dict]) -> list[str]:
         fn = f.get("filename", "")
         name = f.get("name", "")
         line = f.get("line", 0)
-        if not name or name in ("torch::unwind::unwind()", "torch::CapturedTraceback::gather(bool, bool, bool)"):
+        if not name or name in (
+            "torch::unwind::unwind()",
+            "torch::CapturedTraceback::gather(bool, bool, bool)",
+        ):
             continue
         if fn and fn != "??" and fn != "":
             result.append(f"{fn}:{line} {name}")
@@ -48,7 +51,6 @@ def get_stack_idx(frames: list[dict]) -> int:
     reserved = 0
     hwm = 0
     timeline: list[dict] = []
-    active_by_addr: dict[int, dict] = {}
     last_time = 0
 
     current_stack: list[int] = []
@@ -70,12 +72,14 @@ def get_stack_idx(frames: list[dict]) -> int:
                 allocated += size
                 offset = allocated - size
                 alloc_id = len(alloc_polys)
-                alloc_polys.append({
-                    "si": si,
-                    "s": size,
-                    "ts": [timestep],
-                    "offsets": [offset],
-                })
+                alloc_polys.append(
+                    {
+                        "si": si,
+                        "s": size,
+                        "ts": [timestep],
+                        "offsets": [offset],
+                    }
+                )
                 current_stack.append(alloc_id)
                 alloc_id_by_addr[addr] = alloc_id
                 timestep += 1
@@ -110,10 +114,17 @@ def get_stack_idx(frames: list[dict]) -> int:
                 pass
 
         hwm = max(hwm, allocated)
-        timeline.append({
-            "t": time_us, "a": allocated, "r": reserved,
-            "h": hwm, "act": action, "s": size, "si": si,
-        })
+        timeline.append(
+            {
+                "t": time_us,
+                "a": allocated,
+                "r": reserved,
+                "h": hwm,
+                "act": action,
+                "s": size,
+                "si": si,
+            }
+        )
 
     for alloc_id in current_stack:
         poly = alloc_polys[alloc_id]
@@ -149,8 +160,7 @@ def generate_memory_html(
     }
 
     return (
-        _MEMORY_VIZ_TEMPLATE
-        .replace("__TITLE__", title)
+        _MEMORY_VIZ_TEMPLATE.replace("__TITLE__", title)
         .replace("__TIMELINE__", json.dumps(timeline))
         .replace("__ALLOCS__", json.dumps(alloc_polys))
         .replace("__STACKS__", json.dumps(stacks))
@@ -566,18 +576,60 @@ def generate_memory_html(
   .extent([[0, 0], [width, height]])
   .on('zoom', (event) => updateChart(event.transform));
 
-chartArea.append('rect')
+const zoomRect = chartArea.append('rect')
   .attr('width', width).attr('height', height)
   .attr('fill', 'none').attr('pointer-events', 'all')
   .call(zoom);
 
 // Raise polys above the zoom rect
 polysG.raise();
 
+// WASD / arrow key navigation (Perfetto-style)
+const PAN_STEP = 0.15;
+const ZOOM_STEP = 1.3;
+
+document.addEventListener('keydown', function(event) {
+  if (event.target.tagName === 'INPUT') return;
+  const k = event.key.toLowerCase();
+  let t = currentTransform;
+
+  switch (k) {
+    case 'a':
+    case 'arrowleft':
+      t = t.translate(width * PAN_STEP, 0);
+      break;
+    case 'd':
+    case 'arrowright':
+      t = t.translate(-width * PAN_STEP, 0);
+      break;
+    case 'w':
+    case 'arrowup':
+      t = d3.zoomIdentity
+        .translate(width / 2, 0)
+        .scale(t.k * ZOOM_STEP)
+        .translate(-width / 2, 0)
+        .translate(t.x / (t.k * ZOOM_STEP), 0);
+      break;
+    case 's':
+    case 'arrowdown':
+      t = d3.zoomIdentity
+        .translate(width / 2, 0)
+        .scale(t.k / ZOOM_STEP)
+        .translate(-width / 2, 0)
+        .translate(t.x / (t.k / ZOOM_STEP), 0);
+      break;
+    default:
+      return;
+  }
+
+  event.preventDefault();
+  zoomRect.call(zoom.transform, t);
+});
+
 document.getElementById('hwm-toggle').onchange = function() {
   hwmG.style('display', this.checked ? null : 'none');
 };
 </script>
 </body>
 </html>
-""";
+"""