GDN planner: price the dense projection work per GDN-owned token; refit the chain overheads and gate

dev/trainer_rank_cost_fit.py

@@ -119,6 +119,48 @@ def _cell_key(row: dict[str, Any]) -> str:
     )
 
 
+def carried_certificate_cells(
+    certificate_path: Path,
+    evidence: list[Path],
+    *,
+    carry_shapes: set[str],
+    carry_cells: list[str] = (),
+) -> tuple[list[Candidate], list[dict[str, Any]], list[str]]:
+    """Certificate cells the evidence does not re-measure, as candidates plus
+    their recorded cell records, and the problems that forbid the carry.
+
+    Re-certifying after a runtime change re-measures the affected shapes; the
+    cells the change does not touch — named explicitly by shape in
+    ``carry_shapes`` (e.g. ``tp1cp1``, ``tp2cp1ep2``) or by cell-key substring
+    in ``carry_cells`` (e.g. ``|Qwen/Qwen3-4B|`` for an attention-only
+    geometry a GDN change cannot reach) — keep their certified aggregates,
+    which the certificate records exactly (medians, counts, spreads, features,
+    fingerprints). Every other certificate cell must be present in the
+    evidence: an affected cell that was not re-measured is a problem, never a
+    silent carry."""
+
+    measured = set(cell_fingerprints(evidence))
+    candidates, payload = load_certificate(certificate_path)
+    problems: list[str] = []
+    if not carry_shapes and not carry_cells:
+        problems.append(
+            "re-certification needs the configurations whose cells may be "
+            "carried (--carry-shapes and/or --carry-cells)"
+        )
+    records = []
+    for cell in payload["cells"]:
+        if cell["cell"] in measured:
+            continue
+        if _shape_label_of(cell["shape"]) in carry_shapes or any(
+            pattern in cell["cell"] for pattern in carry_cells
+        ):
+            records.append(cell)
+        else:
+            problems.append(f"affected cell not re-measured: {cell['cell']}")
+    carried = {cell["cell"] for cell in records}
+    return [c for c in candidates if c.cell in carried], records, problems
+
+
 def production_regret(
     candidates: list[Candidate], paths: list[Path]
 ) -> dict[str, dict[str, Any]]:
@@ -167,7 +209,9 @@ def production_regret(
     return report
 
 
-def validate_completeness(paths: list[Path], *, repeat: int) -> list[str]:
+def validate_completeness(
+    paths: list[Path], *, repeat: int, carried: list[dict[str, Any]] = ()
+) -> list[str]:
     """Every mandatory candidate of every cell must have ``repeat`` usable rows.
 
     The runners record failures, and the fitter silently skips thin candidates,
@@ -204,6 +248,12 @@ def validate_completeness(paths: list[Path], *, repeat: int) -> list[str]:
                 ):
                     usable[(key, label)] += 1
     gaps: list[str] = []
+    for record in carried:
+        for c in record["candidates"]:
+            if c["label"] != "automatic" and int(c["n"]) < repeat:
+                gaps.append(
+                    f"{record['cell']}: {c['label']} carried with {c['n']} rows (< {repeat})"
+                )
     for cell, labels in expected.items():
         for label in labels:
             if (cell, label) in failed:
@@ -250,6 +300,7 @@ def validate_manifest(
     *,
     excluded: list[str],
     exact: bool = False,
+    carried: list[dict[str, Any]] = (),
 ) -> tuple[list[str], list[str]]:
     """Exact cell identities: every expected cell present unless excluded, no
     unexpected cells, exclusions listed in the manifest, and one execution
@@ -261,6 +312,8 @@ def validate_manifest(
     expected = {cell["key"] for cell in manifest["cells"]}
     listed_exclusions = {cell["key"] for cell in manifest["excluded"]}
     present = cell_fingerprints(paths)
+    for record in carried:
+        present[record["cell"]].add(_fingerprint(record))
     problems: list[str] = []
     excluded_keys = {
         key
@@ -297,6 +350,7 @@ def export_certificate(
     table_id: str = "",
     reranked: list[Candidate] | None = None,
     reranker: dict[str, Any] | None = None,
+    carried: list[dict[str, Any]] = (),
 ) -> None:
     """Write the compact, reproducible record binding the table to its data.
 
@@ -328,6 +382,23 @@ def export_certificate(
                     "geometry": row.get("geometry"),
                     "shape": list(_shape(row)),
                 }
+    for record in carried:
+        fingerprints[record["cell"]] = {
+            key: record.get(key)
+            for key in (
+                "requests_sha256",
+                "source",
+                "workload",
+                "model",
+                "device",
+                "device_capability",
+                "device_memory_class",
+                "param_dtype",
+                "hidden_size",
+                "geometry",
+                "shape",
+            )
+        }
     by_cell: dict[str, list[Candidate]] = defaultdict(list)
     for candidate in [*candidates, *(reranked or [])]:
         by_cell[candidate.cell].append(candidate)
@@ -567,7 +638,11 @@ def evaluate_groups(
 
 
 def _shape_label(candidate: Candidate) -> str:
-    tp, cp, ep, etp = candidate.shape
+    return _shape_label_of(candidate.shape)
+
+
+def _shape_label_of(shape: tuple[int, ...] | list[int]) -> str:
+    tp, cp, ep, etp = (list(shape) + [1, 1])[:4]
     return (
         f"tp{tp}cp{cp}"
         + (f"ep{ep}" if ep > 1 else "")
@@ -1264,6 +1339,25 @@ def main() -> None:
         default="",
         help="fit from a checked-in certificate's aggregates instead of raw evidence",
     )
+    parser.add_argument(
+        "--carry-shapes",
+        default="",
+        help=(
+            "re-certification (--from-certificate with evidence): comma-separated "
+            "parallel shapes (tp1cp1, tp2cp1ep2, ...) the change does not touch; "
+            "their certificate cells are carried as recorded aggregates, every "
+            "other certificate cell must be re-measured by the evidence"
+        ),
+    )
+    parser.add_argument(
+        "--carry-cells",
+        default="",
+        help=(
+            "re-certification: comma-separated cell-key substrings the change "
+            "does not touch (e.g. an attention-only geometry for a GDN planner "
+            "change); their certificate cells are carried like --carry-shapes"
+        ),
+    )
     parser.add_argument(
         "--export-certificate",
         default="",
@@ -1329,6 +1423,24 @@ def main() -> None:
     )
     arguments = parser.parse_args()
     terms = tuple(t for t in arguments.terms.split(",") if t)
+    # Re-certification: cells the evidence does not re-measure are carried
+    # from the previous certificate as their recorded aggregates.
+    carried: list[Candidate] = []
+    carried_records: list[dict[str, Any]] = []
+    carry_shapes = sorted({s for s in arguments.carry_shapes.split(",") if s})
+    carry_cells = sorted({s for s in arguments.carry_cells.split(",") if s})
+    if arguments.from_certificate and arguments.evidence:
+        carried, carried_records, carry_problems = carried_certificate_cells(
+            Path(arguments.from_certificate),
+            arguments.evidence,
+            carry_shapes=set(carry_shapes),
+            carry_cells=carry_cells,
+        )
+        if carry_problems:
+            print("re-certification refused:", file=sys.stderr)
+            for problem in carry_problems:
+                print("  -", problem, file=sys.stderr)
+            raise SystemExit(1)
     if arguments.exclude_cells == "@manifest":
         # Exactly the manifest's listed exclusions (each carries its reason).
         if not arguments.manifest:
@@ -1350,6 +1462,7 @@ def excluded_cell(cell: str) -> bool:
             Path(arguments.manifest),
             excluded=excluded,
             exact=exact,
+            carried=carried_records,
         )
         if problems:
             print("manifest validation failed:", file=sys.stderr)
@@ -1370,7 +1483,9 @@ def excluded_cell(cell: str) -> bool:
         gaps = [
             gap
             for gap in validate_completeness(
-                arguments.evidence, repeat=arguments.require_complete
+                arguments.evidence,
+                repeat=arguments.require_complete,
+                carried=carried_records,
             )
             if not excluded_cell(gap.split(":", 1)[0])
         ]
@@ -1382,7 +1497,16 @@ def excluded_cell(cell: str) -> bool:
         print(
             f"evidence complete: every mandatory candidate has >= {arguments.require_complete} rows"
         )
-    if arguments.from_certificate:
+    if arguments.from_certificate and arguments.evidence:
+        candidates = sorted(
+            load_candidates(arguments.evidence) + carried,
+            key=lambda c: (c.cell, c.label),
+        )
+        print(
+            f"carried {len(carried_records)} cells from {Path(arguments.from_certificate).name}"
+            f" ({len({c.cell for c in candidates}) - len(carried_records)} re-measured)"
+        )
+    elif arguments.from_certificate:
         candidates, _certificate = load_certificate(Path(arguments.from_certificate))
     else:
         candidates = load_candidates(arguments.evidence)
@@ -1591,6 +1715,8 @@ def held_out(cell: str) -> bool:
                 "objective": arguments.objective,
                 "holdout": arguments.holdout,
                 "exclude_cells": excluded,
+                "carry_shapes": carry_shapes,
+                "carry_cells": carry_cells,
                 "require_complete": arguments.require_complete,
                 "rerank_shapes": sorted(rerank_shapes),
                 "shortlist_size": arguments.shortlist_size,
@@ -1607,6 +1733,7 @@ def held_out(cell: str) -> bool:
             table_id=arguments.table_id,
             reranked=reranked,
             reranker=report.get("reranker"),
+            carried=carried_records,
         )
         print("certificate written:", arguments.export_certificate)