From f9d0077472853938ccd803e515d3ed1bc7834ff5 Mon Sep 17 00:00:00 2001
From: Mukesh Garg <mukeshgarg81@gmail.com>
Date: Mon, 27 Jul 2026 23:39:39 -0700
Subject: [PATCH] router: memoize Dijkstra path lookups
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per-instance cache keyed by (id(adj), start, goal). All 4 adj dicts
(_adj, _adj_all, _adj_local, _adj_mcpu_dma) are built in __init__ and
never mutated (topology is static per ADR-0006 / SPEC §0.1), so id(adj)
is stable for the router's lifetime. Cache is populated on the
successful-return paths; RoutingError paths intentionally re-run each
call (rare, keeps error semantics unchanged).

Motivation: cProfile of Case 4 decode at S_kv=8K showed 1,142
_run_dijkstra_with_dist calls consuming ~1.95s tottime. Paths depend
only on (adj, src, dst) so ~99% of those calls are recomputing the same
result.

Verified:
- tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass
- 128K decode wall: 237.28s -> 233.71s (-1.5%)
- Modeled kernel latency: 461.13us (byte-identical before/after)
- op_log_len: 3057 (unchanged)

Small win but no risk: memoization returns byte-identical results and
paths cannot change during a sim run. Larger event-count reductions
require touching the per-hop hot path (zero-latency chain collapse
etc.).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 src/kernbench/policy/routing/router.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/kernbench/policy/routing/router.py b/src/kernbench/policy/routing/router.py
index 0869079..d4ec39b 100644
--- a/src/kernbench/policy/routing/router.py
+++ b/src/kernbench/policy/routing/router.py
@@ -111,6 +111,13 @@ class PathRouter:
         self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list)
         self._adj_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list)
         self._adj_local: dict[str, list[tuple[str, float]]] = defaultdict(list)
+        # Memoize path lookups: adj dicts are built once here and never
+        # mutated (topology is static per ADR-0006 / SPEC §0.1), so
+        # (id(adj), start, goal) is a stable cache key for the router's
+        # lifetime. Callers use the returned path list read-only.
+        self._path_cache: dict[
+            tuple[int, str, str], tuple[list[str], float]
+        ] = {}
         for e in graph.edges:
             w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm
             self._adj_all[e.src].append((e.dst, w))
@@ -185,8 +192,14 @@ class PathRouter:
         start: str,
         goal: str,
     ) -> tuple[list[str], float]:
+        cache_key = (id(adj), start, goal)
+        cached = self._path_cache.get(cache_key)
+        if cached is not None:
+            return cached
         if start == goal:
-            return [start], 0.0
+            result = ([start], 0.0)
+            self._path_cache[cache_key] = result
+            return result
         best: dict[str, float] = {start: 0.0}
         prev: dict[str, str] = {}
         heap: list[tuple[float, str]] = [(0.0, start)]
@@ -200,7 +213,9 @@ class PathRouter:
                     cur = prev[cur]
                 path.append(start)
                 path.reverse()
-                return path, d
+                result = (path, d)
+                self._path_cache[cache_key] = result
+                return result
             if d > best.get(node, float("inf")):
                 continue
             for neighbor, edge_dist in adj[node]: