feat(v7-f56b,f53): DimEnvEditor for live H/nh/head_dim mutation

apstenku123 · apstenku123 · commit 999136f78f58 · 2026-05-23T02:11:00.000+02:00
vbgui/src/components/DimEnvEditor.tsx:
- Five-input row (H, nh, head_dim, B, S) + Apply button, rendered
  above the canvas. Each input has data-testid='dim-env-{key}'.
- Inline mismatch indicator (data-testid='dim-env-inline-mismatch')
  appears live as the user types — without waiting for Apply — when
  nh*head_dim ≠ H. Renders "nh*head_dim = N ≠ H = M".
- Numeric inputs; non-finite drafts are ignored on Apply.

vbgui/src/App.tsx:
- MINI_DIM_ENV constant was previously inlined into 3 RPC call sites
  (verify, suggest_sharding, buildVerifyParams default). Promoted
  to dimEnv useState, fed back into wireSpecRef and into the verify
  debouncer dep list (dimEnvKey) so a user mismatch immediately
  re-runs verify and surfaces the F56b warning badge.
- Canvas wrapper switched to flexDirection='column' so the editor
  stacks above FlowCanvas without occluding it.
- DimEnvEditor and the existing F56b badge co-render: the editor's
  inline indicator surfaces *during* typing; the badge surfaces
  *after* verify roundtrip completes.

vbgui/tests/DimEnvEditor.test.tsx — 5 unit tests:
- Inputs populate from value prop.
- Apply forwards parsed numeric draft.
- Mismatch indicator visible when nh*head_dim ≠ H.
- Mismatch indicator hidden when consistent.
- Indicator re-evaluates live (before Apply) on edit.

App.integration tests still pass 7/7; vitest total 12/12.

Unlocks F56b visual e2e (type mismatch into editor → assert badge)
and F53 scaling sweep (programmatic H ∈ {64,128,256,512} via Apply).

Ref: cppmega-mlx-j3qa.2, cppmega-mlx-j3qa.4
diff --git a/vbgui/src/App.tsx b/vbgui/src/App.tsx
@@ -4,6 +4,7 @@ import {
 } from "@xyflow/react";
 
 import { FlowCanvas } from "@/components/FlowCanvas";
+import { DimEnvEditor } from "@/components/DimEnvEditor";
 import { Palette } from "@/components/Palette";
 import { Sidebar } from "@/components/Sidebar";
 import { TopBar, type RunMode } from "@/components/TopBar";
@@ -147,11 +148,19 @@ export function App(): JSX.Element {
   // list that was missing 5 architectures.
   const PRESETS = usePresets(rpc);
 
+  // V7-F56b / V7-F53: editable dim_env (was a constant MINI_DIM_ENV).
+  // Threaded into the verify trigger so changes to H/nh/head_dim
+  // re-run verify and surface the F56b mismatch warning.
+  const [dimEnv, setDimEnv] = useState<Record<string, number>>(
+    { ...MINI_DIM_ENV });
+
   // Keep one stable spec snapshot for the verify debouncer to read.
-  const wireSpecRef = useRef({ nodes, edges, spec, availableSideChannels });
+  const wireSpecRef = useRef({ nodes, edges, spec, availableSideChannels,
+                               dimEnv });
   useEffect(() => {
-    wireSpecRef.current = { nodes, edges, spec, availableSideChannels };
-  }, [nodes, edges, spec, availableSideChannels]);
+    wireSpecRef.current = { nodes, edges, spec, availableSideChannels,
+                            dimEnv };
+  }, [nodes, edges, spec, availableSideChannels, dimEnv]);
 
   useEffect(() => {
     const available = new Set(availableSideChannels);
@@ -163,6 +172,7 @@ export function App(): JSX.Element {
     if (snap.nodes.length === 0) return;
     const params = buildVerifyParams(
       snap.nodes, snap.edges, snap.spec, snap.availableSideChannels,
+      snap.dimEnv,
     );
     try {
       const r = await rpc.call<{
@@ -227,10 +237,14 @@ export function App(): JSX.Element {
   const rewriterKey = spec.rewriters.map((r) => r.name).join(",");
   const sideChannelKey = JSON.stringify(spec.side_channels);
   const availableSideChannelKey = availableSideChannels.join(",");
+  // V7-F56b/F53: re-verify when the dim_env editor mutates H, nh, or
+  // head_dim — otherwise the F56b badge wouldn't appear after a user
+  // types a mismatch.
+  const dimEnvKey = JSON.stringify(dimEnv);
   useEffect(() => { scheduleVerify(); },
            [nodesKey, edgesKey, lossKey, optimKey, shardingKey,
             rewriterKey, sideChannelKey, availableSideChannelKey,
-            scheduleVerify]);
+            dimEnvKey, scheduleVerify]);
 
   // V7-H03: snapshot the (nodes, edges, spec) triple every time a
   // structural key changes — that's the same set of user-meaningful
@@ -343,7 +357,7 @@ export function App(): JSX.Element {
                      }[] } }[];
       }>("suggest_sharding", {
         graph: nodesToGraph(snap.nodes, snap.edges),
-        dim_env: MINI_DIM_ENV,
+        dim_env: snap.dimEnv,
         loss: { kind: snap.spec.loss.kind,
                 head_outputs: snap.spec.loss.head_outputs },
         optim: { kind: snap.spec.optim.kind,
@@ -626,8 +640,10 @@ export function App(): JSX.Element {
           {activeTab === "canvas" && (
             <>
               <Palette />
-              <div style={{ flex: 1, position: "relative", display: "flex",
+              <div style={{ flex: 1, position: "relative",
+                            display: "flex", flexDirection: "column",
                             minHeight: 0 }}>
+                <DimEnvEditor value={dimEnv} onApply={setDimEnv} />
                 {spec.gotchas.some(
                   (g) => g.id === "v7_f56b_dim_env_mismatch") && (
                   <div
@@ -821,10 +837,11 @@ function buildVerifyParams(
   edges: Edge[],
   spec: SpecState,
   availableSideChannels: string[] = ["doc_ids", "token_ids"],
+  dimEnv: Record<string, number> = MINI_DIM_ENV,
 ) {
   return {
     graph: nodesToGraph(nodes, edges),
-    dim_env: MINI_DIM_ENV,
+    dim_env: dimEnv,
     loss: { kind: spec.loss.kind, head_outputs: spec.loss.head_outputs,
             params: spec.loss.params },
     optim: { kind: spec.optim.kind,
diff --git a/vbgui/src/components/DimEnvEditor.tsx b/vbgui/src/components/DimEnvEditor.tsx
@@ -0,0 +1,74 @@
+// Minimal dim_env editor surfaced above the canvas. Unlocks the F56b
+// visual-warning flow (mismatch surfaces a banner in the canvas) and
+// the F53 dimension-scaling sweep (sweeps H ∈ {64,128,256,512} by
+// calling onApply per H).
+
+import { useState } from "react";
+
+export interface DimEnvEditorProps {
+  value: Record<string, number>;
+  onApply: (next: Record<string, number>) => void;
+}
+
+const EDITABLE_KEYS = ["H", "nh", "head_dim", "B", "S"] as const;
+type EditableKey = typeof EDITABLE_KEYS[number];
+
+export function DimEnvEditor({ value, onApply }: DimEnvEditorProps): JSX.Element {
+  const [draft, setDraft] = useState<Record<string, string>>(() => {
+    const out: Record<string, string> = {};
+    for (const k of EDITABLE_KEYS) out[k] = String(value[k] ?? "");
+    return out;
+  });
+
+  const mismatch = (() => {
+    const H = Number(draft.H);
+    const nh = Number(draft.nh);
+    const hd = Number(draft.head_dim);
+    if (!Number.isFinite(H) || !Number.isFinite(nh) ||
+        !Number.isFinite(hd)) return null;
+    if (nh * hd === H) return null;
+    return `nh*head_dim = ${nh * hd} ≠ H = ${H}`;
+  })();
+
+  return (
+    <div data-testid="dim-env-editor"
+         style={{ display: "flex", alignItems: "center", gap: 8,
+                  padding: "4px 8px", background: "#f9fafb",
+                  borderBottom: "1px solid #e5e7eb",
+                  fontFamily: "system-ui, sans-serif", fontSize: 12 }}>
+      <strong data-testid="dim-env-editor-label">dim_env:</strong>
+      {EDITABLE_KEYS.map((k: EditableKey) => (
+        <label key={k} style={{ display: "inline-flex", gap: 4 }}>
+          {k}
+          <input
+            data-testid={`dim-env-${k}`}
+            type="number"
+            value={draft[k]}
+            onChange={(e) => setDraft({ ...draft, [k]: e.target.value })}
+            style={{ width: 64 }}
+          />
+        </label>
+      ))}
+      <button
+        data-testid="dim-env-apply"
+        onClick={() => {
+          const next: Record<string, number> = { ...value };
+          for (const k of EDITABLE_KEYS) {
+            const n = Number(draft[k]);
+            if (Number.isFinite(n)) next[k] = n;
+          }
+          onApply(next);
+        }}
+        style={{ padding: "2px 8px" }}
+      >
+        Apply
+      </button>
+      {mismatch && (
+        <span data-testid="dim-env-inline-mismatch"
+              style={{ color: "#92400e", marginLeft: 8 }}>
+          ⚠ {mismatch}
+        </span>
+      )}
+    </div>
+  );
+}
diff --git a/vbgui/tests/DimEnvEditor.test.tsx b/vbgui/tests/DimEnvEditor.test.tsx
@@ -0,0 +1,65 @@
+import { describe, it, expect, vi } from "vitest";
+import { render, screen, fireEvent } from "@testing-library/react";
+import { DimEnvEditor } from "@/components/DimEnvEditor";
+
+describe("V7-F56b/F53 DimEnvEditor", () => {
+  it("populates inputs from the value prop", () => {
+    render(<DimEnvEditor
+      value={{ H: 128, nh: 2, head_dim: 64, B: 1, S: 8 }}
+      onApply={() => {}}
+    />);
+    expect((screen.getByTestId("dim-env-H") as HTMLInputElement).value).toBe("128");
+    expect((screen.getByTestId("dim-env-nh") as HTMLInputElement).value).toBe("2");
+    expect((screen.getByTestId("dim-env-head_dim") as HTMLInputElement).value).toBe("64");
+  });
+
+  it("Apply calls onApply with the parsed numeric draft", () => {
+    const onApply = vi.fn();
+    render(<DimEnvEditor
+      value={{ H: 128, nh: 2, head_dim: 64, B: 1, S: 8 }}
+      onApply={onApply}
+    />);
+    fireEvent.change(screen.getByTestId("dim-env-H"),
+      { target: { value: "256" } });
+    fireEvent.change(screen.getByTestId("dim-env-nh"),
+      { target: { value: "4" } });
+    fireEvent.click(screen.getByTestId("dim-env-apply"));
+    expect(onApply).toHaveBeenCalledTimes(1);
+    const next = onApply.mock.calls[0]?.[0];
+    expect(next.H).toBe(256);
+    expect(next.nh).toBe(4);
+    expect(next.head_dim).toBe(64);
+  });
+
+  it("renders an inline mismatch indicator when nh*head_dim != H", () => {
+    render(<DimEnvEditor
+      value={{ H: 128, nh: 3, head_dim: 50, B: 1, S: 8 }}
+      onApply={() => {}}
+    />);
+    const warn = screen.getByTestId("dim-env-inline-mismatch");
+    expect(warn).toBeDefined();
+    expect(warn.textContent).toContain("150");
+    expect(warn.textContent).toContain("128");
+  });
+
+  it("hides the mismatch indicator when nh*head_dim == H", () => {
+    render(<DimEnvEditor
+      value={{ H: 128, nh: 2, head_dim: 64, B: 1, S: 8 }}
+      onApply={() => {}}
+    />);
+    expect(screen.queryByTestId("dim-env-inline-mismatch")).toBeNull();
+  });
+
+  it("re-evaluates the mismatch live as the user types (before Apply)", () => {
+    render(<DimEnvEditor
+      value={{ H: 128, nh: 2, head_dim: 64, B: 1, S: 8 }}
+      onApply={() => {}}
+    />);
+    expect(screen.queryByTestId("dim-env-inline-mismatch")).toBeNull();
+    fireEvent.change(screen.getByTestId("dim-env-nh"),
+      { target: { value: "3" } });
+    // 3 * 64 = 192 ≠ 128 → mismatch surfaces without Apply.
+    const warn = screen.getByTestId("dim-env-inline-mismatch");
+    expect(warn.textContent).toContain("192");
+  });
+});
