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:
@@ -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]:
|
||||||
|
|||||||
Reference in New Issue
Block a user