feat: implement edge reconnection, unified handle management, and persistent routing layout flow

Author: apstenku123

cppmega_mlx/runtime/path_c_physical_abi.py

@@ -463,8 +463,18 @@ def validate_physical_abi_runtime_bindings(
     mapping: Mapping[str, Any],
     bank_shapes: Mapping[str, Any],
     buffers: Mapping[str, Any] | None,
+    *,
+    allow_superset: bool = False,
 ) -> dict[str, Any]:
-    """Validate caller-supplied physical bank buffers without allocating."""
+    """Validate caller-supplied physical bank buffers without allocating.
+
+    ``allow_superset`` (opt-in; default exact-equality) accepts a caller/model-owned
+    bank whose flat element count is >= the required extent and whose dtype matches.
+    This is for bindings against a single schedule-template region when the model owns
+    banks sized from the MERGED launcher+generated-stage ABI (a strict superset of any
+    one region's need); the kernel-arg builder slices a contiguous prefix view to the
+    exact per-region shape at call time. All other callers keep strict equality.
+    """
 
     bridge_plan = plan_physical_abi_runtime_bridge(mapping, bank_shapes)
     required_banks = list(bridge_plan["required_bank_buffers"])
@@ -498,10 +508,19 @@ def validate_physical_abi_runtime_bindings(
         expected_shape = tuple(int(dim) for dim in tuple(bank_shapes[bank]))
         actual_shape = _shape_of(provided[bank])
         if actual_shape != expected_shape:
-            shape_mismatches.append(bank)
-            errors.append(
-                f"{bank}: shape {actual_shape} does not match expected {expected_shape}"
+            expected_numel = int(prod(expected_shape)) if expected_shape else 1
+            actual_numel = (
+                int(prod(actual_shape)) if actual_shape else 0
             )
+            if allow_superset and actual_shape is not None and actual_numel >= expected_numel:
+                # Larger-but-sufficient caller-owned bank (merged-ABI superset); the
+                # kernel-arg builder slices a contiguous prefix view to the exact shape.
+                pass
+            else:
+                shape_mismatches.append(bank)
+                errors.append(
+                    f"{bank}: shape {actual_shape} does not match expected {expected_shape}"
+                )
         expected_dtype = bank_dtypes.get(bank)
         actual_dtype = _dtype_of(provided[bank])
         if actual_dtype != expected_dtype:

cppmega_mlx/training/compiled.py

@@ -130,6 +130,11 @@ class PathCFusedTrainBlockCallableArtifact:
 
     hidden_packing_performed = False
     no_hidden_allocation_policy = True
+    # A base (non-staged) artifact wraps ONE compiled TileLang/Metal kernel that
+    # lowers the whole forward and the full recurrent backward internally and only
+    # gates on ``path_c_run_backward``; it is dispatched grid-chunked (launched once
+    # per forward / backward). Multi-kernel stage artifacts override this to True.
+    generated_stage_artifact = False
 
     def __init__(
         self,

outputs/webwright_canvas/probe_drag.py

@@ -0,0 +1,37 @@
+import time
+from playwright.sync_api import sync_playwright
+URL="http://127.0.0.1:8765"; PRESET="deepseek_v4_flash"
+with sync_playwright() as p:
+    b=p.firefox.launch(headless=True); pg=b.new_page(viewport={"width":2200,"height":1300})
+    pg.goto(URL,wait_until="networkidle"); time.sleep(2)
+    sels=pg.locator("select")
+    for i in range(sels.count()):
+        if any(PRESET in o for o in sels.nth(i).locator("option").all_inner_texts()):
+            sels.nth(i).select_option(label=PRESET); break
+    time.sleep(1.5)
+    g=pg.get_by_role("button",name="Generate Architecture")
+    if g.count()>0: g.first.click()
+    pg.wait_for_selector(".react-flow__node",timeout=15000); time.sleep(3)
+    pg.locator(".react-flow__controls-fitview").first.click(); time.sleep(1)
+    before = pg.locator('.react-flow__edge').count()
+    # Select first edge, grab one of its updater anchors, drag onto its target node top handle.
+    edge = pg.locator('.react-flow__edge').first
+    box = edge.bounding_box()
+    pg.mouse.click(box['x']+box['width']/2, box['y']+box['height']/2); time.sleep(0.4)
+    anchors = pg.locator('.react-flow__edgeupdater')
+    n = anchors.count()
+    moved = False
+    if n >= 1:
+        a = anchors.first.bounding_box()
+        # drag the anchor a bit (just to exercise the reconnect drag path); drop near a node center
+        node = pg.locator('.react-flow__node').nth(1).bounding_box()
+        pg.mouse.move(a['x']+a['width']/2, a['y']+a['height']/2)
+        pg.mouse.down()
+        pg.mouse.move(node['x']+node['width']/2, node['y']+5, steps=15)  # near top edge of a node
+        pg.mouse.up(); time.sleep(0.6)
+        moved = True
+    after = pg.locator('.react-flow__edge').count()
+    print(f"edges before={before} after={after}  anchors={n}  dragPerformed={moved}")
+    print("RESULT edges_survived_drag:", after >= before-1)
+    pg.screenshot(path="/tmp/route_geo_folded.png")
+    b.close()

outputs/webwright_canvas/probe_folded_tb.py

@@ -0,0 +1,20 @@
+import time
+from playwright.sync_api import sync_playwright
+URL="http://127.0.0.1:8765"; PRESET="deepseek_v4_flash"
+with sync_playwright() as p:
+    b=p.firefox.launch(headless=True); pg=b.new_page(viewport={"width":2200,"height":1300})
+    pg.goto(URL,wait_until="networkidle"); time.sleep(2)
+    sels=pg.locator("select")
+    for i in range(sels.count()):
+        if any(PRESET in o for o in sels.nth(i).locator("option").all_inner_texts()):
+            sels.nth(i).select_option(label=PRESET); break
+    time.sleep(1.5)
+    g=pg.get_by_role("button",name="Generate Architecture")
+    if g.count()>0: g.first.click()
+    pg.wait_for_selector(".react-flow__node",timeout=15000); time.sleep(3)
+    pg.locator(".react-flow__controls-fitview").first.click(); time.sleep(1)
+    zin=pg.locator(".react-flow__controls-zoomin").first
+    zin.click();time.sleep(0.3);zin.click();time.sleep(0.6)
+    pg.screenshot(path="/tmp/route_folded_tb.png")
+    print("saved")
+    b.close()

outputs/webwright_canvas/probe_folded_v2.py

@@ -0,0 +1,37 @@
+import time, sys
+from playwright.sync_api import sync_playwright
+URL="http://127.0.0.1:8765"
+PRESETS = ["gemma4","llama3_8b","deepseek_v3","gpt_oss_20b","minimax_m2"]
+VIOLATION_JS = open("probe_routing.py").read().split('VIOLATION_JS = r"""')[1].split('"""')[0]
+def measure(pg):
+    return pg.evaluate("()=>{"+VIOLATION_JS.split("() => {",1)[1])
+with sync_playwright() as p:
+    b=p.firefox.launch(headless=True); pg=b.new_page(viewport={"width":2400,"height":1400})
+    grand=0
+    for preset in PRESETS:
+        for scen in ["folded","unpackall"]:
+            pg.goto(URL,wait_until="networkidle"); time.sleep(1.5)
+            sels=pg.locator("select"); ok=False
+            for i in range(sels.count()):
+                if any(preset==o.strip() for o in sels.nth(i).locator("option").all_inner_texts()):
+                    sels.nth(i).select_option(label=preset); ok=True; break
+            if not ok: print(f"  {preset}: NOT FOUND"); continue
+            time.sleep(1.2)
+            g=pg.get_by_role("button",name="Generate Architecture")
+            if g.count()>0: g.first.click()
+            try: pg.wait_for_selector(".react-flow__node",timeout=15000)
+            except: print(f"  {preset}/{scen}: no nodes"); continue
+            time.sleep(1.5)
+            if scen=="unpackall":
+                ub=pg.locator('[data-testid^="unpack-btn-"]')
+                if ub.count()>0:
+                    ub.first.click(); time.sleep(0.4)
+                    ua=pg.locator('[data-testid^="confirm-unpack-all-"]')
+                    if ua.count()>0: ua.first.click(); time.sleep(2.2)
+            pg.get_by_role("button",name="Auto Align Graph").first.click(); time.sleep(4)
+            m=measure(pg)
+            grand+=len(m['violations'])
+            print(f"  {preset:16s}/{scen:9s}: edges={m['edges']:3d} nodes={m['nodes']:3d}  OVER-BOXES={len(m['violations'])}")
+            for v in m['violations'][:4]: print(f"        ! {v}")
+    print(f"\nGRAND TOTAL edges-over-boxes = {grand}")
+    b.close()

outputs/webwright_canvas/probe_multi.py

@@ -0,0 +1,38 @@
+import time
+from playwright.sync_api import sync_playwright
+URL="http://127.0.0.1:8765"; PRESET="deepseek_v4_flash"
+with sync_playwright() as p:
+    b=p.firefox.launch(headless=True); pg=b.new_page(viewport={"width":2200,"height":1300})
+    pg.goto(URL,wait_until="networkidle"); time.sleep(2)
+    sels=pg.locator("select")
+    for i in range(sels.count()):
+        if any(PRESET in o for o in sels.nth(i).locator("option").all_inner_texts()):
+            sels.nth(i).select_option(label=PRESET); break
+    time.sleep(1.5)
+    g=pg.get_by_role("button",name="Generate Architecture")
+    if g.count()>0: g.first.click()
+    pg.wait_for_selector(".react-flow__node",timeout=15000); time.sleep(3)
+    pg.locator(".react-flow__controls-fitview").first.click(); time.sleep(1)
+    # Count edges + check reconnect anchors exist (react-flow__edgeupdater) when hovering an edge.
+    info = pg.evaluate(r"""
+() => {
+  const edges = document.querySelectorAll('.react-flow__edge').length;
+  // selecting an edge should render edgeupdater anchors; check the DOM supports them
+  const hasUpdaterClass = !!document.querySelector('.react-flow__edge');
+  return { edges };
+}
+""")
+    print("edges:", info['edges'])
+    # Click an edge to select it, then check for edgeupdater anchors
+    edge = pg.locator('.react-flow__edge').first
+    box = edge.bounding_box()
+    if box:
+        pg.mouse.click(box['x']+box['width']/2, box['y']+box['height']/2)
+        time.sleep(0.5)
+    anchors = pg.locator('.react-flow__edgeupdater').count()
+    print("edgeupdater anchors after selecting an edge:", anchors)
+    # Programmatic reconnect test: pick the De-Tokenizer incoming edge, drag its
+    # target anchor to a different handle. Simpler: verify the anchors are draggable
+    # by checking they exist (>0 means reconnection UI is active).
+    print("RECONNECT-UI-ACTIVE:", anchors > 0)
+    b.close()