router: memoize Dijkstra path lookups

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>
This commit is contained in:
2026-07-27 23:39:39 -07:00
parent 65f358fdfa
commit f9d0077472
+17 -2
View File
@@ -111,6 +111,13 @@ class PathRouter:
self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list) 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_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list)
self._adj_local: 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: for e in graph.edges:
w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm 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)) self._adj_all[e.src].append((e.dst, w))
@@ -185,8 +192,14 @@ class PathRouter:
start: str, start: str,
goal: str, goal: str,
) -> tuple[list[str], float]: ) -> 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: 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} best: dict[str, float] = {start: 0.0}
prev: dict[str, str] = {} prev: dict[str, str] = {}
heap: list[tuple[float, str]] = [(0.0, start)] heap: list[tuple[float, str]] = [(0.0, start)]
@@ -200,7 +213,9 @@ class PathRouter:
cur = prev[cur] cur = prev[cur]
path.append(start) path.append(start)
path.reverse() path.reverse()
return path, d result = (path, d)
self._path_cache[cache_key] = result
return result
if d > best.get(node, float("inf")): if d > best.get(node, float("inf")):
continue continue
for neighbor, edge_dist in adj[node]: for neighbor, edge_dist in adj[node]: