pgw#1014: task rig:gauntlet — every mint shape in one command, plus GPU-REAL cycles on this box

Taskfile.yml

@@ -47,6 +47,39 @@ tasks:
     cmds:
       - nice -n 10 uv run python scripts/micro_mint_rig.py --vehicle micro {{.CLI_ARGS}}
 
+  rig:gauntlet:
+    desc: |
+      pgw#1014 — EVERY mint shape the fleet runs, in one command, with a table.
+      Each variant is a full production cycle (spawn, export+AOTI, seal,
+      publish, second-process adopt + parity) on a model that compiles in
+      under a minute.
+
+      Variants declare the outcome they EXPECT, because some exist to
+      demonstrate a REFUSAL — so the gauntlet reports agreement, and a flip in
+      either direction is news. Exit 0 = every variant matched.
+
+      Cardless by default. For the GPU lane (real sm_89, no synthetic probe):
+        eval "$(./scripts/rig_gpu_env.sh --export-only)"
+        task rig:gauntlet -- --device cuda --python "$GEN_WORKER_RIG_GPU_PYTHON"
+    cmds:
+      - nice -n 10 uv run python scripts/rig_gauntlet.py {{.CLI_ARGS}}
+
+  rig:gpu:
+    desc: |
+      pgw#983 — one rig cycle on the REAL card. Builds/refreshes the isolated
+      cu126 interpreter (see scripts/rig_gpu_env.sh: cu128 has no torch
+      2.13.0, cu126 does) and runs with --device cuda, so the cell carries a
+      MEASURED sm_89 instead of a supplied one.
+
+      The repo's own .venv and the cu130 pin are untouched; this is a second
+      interpreter, selected explicitly, never a default.
+    cmds:
+      - ./scripts/rig_gpu_env.sh > /dev/null
+      - |
+        eval "$(./scripts/rig_gpu_env.sh --export-only)"
+        nice -n 10 "$GEN_WORKER_RIG_GPU_PYTHON" scripts/micro_mint_rig.py \
+          --device cuda {{.CLI_ARGS}}
+
   rig:test:
     desc: |
       pgw#978 — the rig as a pytest row, including the full cycle (opt-in via

changelog.d/pgw1014.md

@@ -0,0 +1,44 @@
+### Added
+
+- **`task rig:gauntlet` (pgw#1014) — every mint shape the fleet runs, in one
+  command, with a table.** Each variant is a full production cycle (real child
+  spawn, `torch.export` + AOTInductor, seal, publish over the real wire,
+  second-process discover + arm + parity) on a model that compiles in under a
+  minute. Variants declare the outcome they EXPECT, because some exist to
+  demonstrate a REFUSAL; the gauntlet reports agreement, and a flip in either
+  direction is news. Exit 0 = every variant matched.
+  - `micro` — plain lane, 3 entries, container input + trailing plain input.
+  - `micro-lora` — `lora_bucket=64`, 5 entries (adapter fork), the
+    branch-bearing graph and the lifted-arm gates (pgw#1001).
+  - `micro-4d` — pgw#998's shape: a 4-D latent with BOTH spatial axes dynamic,
+    so every matmul's M extent is NONLINEAR in the traced symbols. This is
+    z-image's declaration, and it re-proves pgw#998's fix on every run instead
+    of trusting a changelog.
+  - `micro-lora-plain-parent` — EXPECTED RED: a bucket-bearing cell offered to
+    a plain-lane parent is rejected at discovery with `lane_mismatch`, before
+    any arm runs. That is A1 step 8's cold-second-pod path, and the reason
+    attempt 27 must dispatch the adopting pod with the cell's own bucket.
+
+- **`task rig:gpu` and `scripts/rig_gpu_env.sh` (pgw#983) — GPU-real cycles on
+  this box.** An isolated second interpreter (`.venv-cu126`, shared across
+  worktrees) with **torch 2.13.0+cu126**: the fleet's exact torch version, on
+  an index whose CUDA runs on this box's 12.8 driver. **cu128 was the approved
+  index and is wrong — it has no torch 2.13.0, only 2.11.0**, and `torch` is an
+  axis `verify_declared` checks strictly, so a cu128 cell would be further from
+  the fleet, not closer. The repo's `.venv` and the cu130 pin are untouched;
+  the GPU interpreter is selected explicitly and is never a default.
+
+  This retires the synthetic-`sm` caveat: cells minted through it carry a
+  MEASURED `sm_89` — the same compute capability as the fleet's L40S — instead
+  of a supplied one. Carve-out re-measured on the real card: 3.00 GiB mint +
+  1.00 GiB adopt = 4.00 GiB of a 7.63 GiB card, the cap reaches the child
+  (`vram_cap_bytes = 3.00 GiB`), `CardCensus` reads `basis=sampled`,
+  `GEN_WORKER_HOST_MOVE_GUARD` untouched, and measured peak is **0.011 GiB**.
+
+### Fixed
+
+- The rig's adopt child places the pipeline on the mint's device. A code-only
+  cell binds constants from RESIDENT weights, so a CPU-resident parent adopting
+  a CUDA cell fails inside AOTI itself
+  (`update_constant_buffer_func_ ... API call failed`) — reachable only on a
+  real card, and the first thing the GPU lane found.

examples/micro-diffusion/src/micro_diffusion/aot_declaration_4d.py

@@ -0,0 +1,75 @@
+"""The pgw#998 declaration: a 4-D latent grid, both spatial axes dynamic.
+
+This is the shape z-image declares (`H_lat` and `W_lat` on a 4-D latent under
+`dynamic-collapse`) and the shape pgw#998 measured as unlowerable across the
+mint's `torch.export.save`/`load` hand-off to its compile child:
+
+    InductorError: LoweringException: RuntimeError:
+      ('unexpected None!', 512*s18*s57)
+
+pgw#998's fix (the export handoff carries the ShapeEnv's symbol values) landed
+on master. This declaration exists so the gauntlet re-proves that on every run
+instead of trusting the changelog — and so that if it ever regresses, it does
+so on a 1.1 MB toy in 20 seconds rather than on a rented A100.
+
+Deliberately kept to ONE target and ONE fork-free class set: the point is the
+nonlinear EXTENT, not entry count, and a second target would only add wall
+clock to a variant whose whole job is a single lowering question.
+"""
+
+from __future__ import annotations
+
+from gen_worker import (
+    Compile,
+    Dim,
+    GraphClass,
+    Input,
+    register_export_declaration,
+)
+
+FAMILY = "micro-4d"
+
+#: The two latent grids, as GRID extents (not token counts).
+LATENT_ROWS = (32, 48)
+VAE_SCALE = 8
+PIXEL_ROWS = tuple((n * VAE_SCALE, n * VAE_SCALE) for n in LATENT_ROWS)
+COND_LEN = 16
+#: Batch is fixed here: the variable under test is the SPATIAL product, and a
+#: fork would just double the compile for no extra coverage of it.
+ARITY = 1
+
+
+def build_declaration() -> Compile:
+    return Compile(
+        family=FAMILY,
+        targets=("transformer",),
+        shapes=PIXEL_ROWS,
+        text_len=COND_LEN,
+        dims=(
+            Dim("N", carried_by=(("t", 0),)),
+            # BOTH spatial axes dynamic on ONE input — this is the whole
+            # point. Their product is the M extent of every matmul inside.
+            Dim("H_lat", carried_by=(("x", 1),), multiple_of=2),
+            Dim("W_lat", carried_by=(("x", 2),), multiple_of=2),
+        ),
+        classes=tuple(
+            GraphClass(dims={"N": ARITY, "H_lat": n, "W_lat": n})
+            for n in LATENT_ROWS),
+        inputs=(
+            Input("x", shape=(("config", "in_channels"), "H_lat", "W_lat"),
+                  repeat="N", dtype="float32"),
+            Input("t", shape=("N",), dtype="float32"),
+            Input("cond", shape=(COND_LEN, ("config", "cond_dim")),
+                  repeat="N", dtype="float32"),
+        ),
+        shape_strategy="dynamic-collapse",
+        warm_changes_key=False,
+    )
+
+
+DECLARATION = build_declaration()
+
+register_export_declaration(DECLARATION, replace=True)
+
+__all__ = ["ARITY", "COND_LEN", "DECLARATION", "FAMILY", "LATENT_ROWS",
+           "PIXEL_ROWS", "VAE_SCALE", "build_declaration"]

examples/micro-diffusion/src/micro_diffusion/main_4d.py

@@ -0,0 +1,85 @@
+"""The pgw#998 variant's worker function — same family shape, GRID latents.
+
+Separate from `main` because it is a different FAMILY (`micro-4d`) with a
+different traced signature; sharing a module would put two families' endpoints
+in one discovery namespace for no benefit. Everything else — catalog slot with
+no code default, declaration registered at import, `ctx.slots` dereferenced
+first — is deliberately identical to `main`, so a difference in outcome is a
+difference in the SHAPE under test and nothing else.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import msgspec
+import torch
+
+from gen_worker import Compile, RequestContext, Resources, Slot, endpoint
+from gen_worker.families import GenerationDefaults, family
+
+from .aot_declaration_4d import (  # noqa: F401 — registers at import
+    ARITY,
+    COND_LEN,
+    DECLARATION,
+    FAMILY,
+    LATENT_ROWS,
+    PIXEL_ROWS,
+)
+from .pipeline import MicroGridPipeline
+
+
+@family(FAMILY)
+class Micro4dDefaults(GenerationDefaults, frozen=True):
+    steps: int = 2
+
+
+class Micro4dIn(msgspec.Struct):
+    prompt: str = ""
+    model: str = ""
+
+
+class Micro4dOut(msgspec.Struct):
+    checkpoint: str = ""
+    shape: str = ""
+
+
+@endpoint(
+    models={"pipeline": Slot(MicroGridPipeline, selected_by="model")},
+    compile=Compile(
+        family=FAMILY, targets=("transformer",), shapes=PIXEL_ROWS,
+        text_len=COND_LEN),
+    resources=Resources(gpu=True),
+)
+class Generate4d:
+    def setup(self, pipeline: MicroGridPipeline) -> None:
+        self.pipe = pipeline
+
+    def generate_4d(
+        self, ctx: RequestContext[Micro4dDefaults], data: Micro4dIn,
+    ) -> Micro4dOut:
+        resolved = ctx.slots["pipeline"]
+        grid = LATENT_ROWS[0]
+        device = self.pipe.device
+        config = self.pipe.config
+        generator = ctx.generator(998)
+        x: List[torch.Tensor] = [
+            torch.randn(config.in_channels, grid, grid, generator=generator,
+                        device=device, dtype=torch.float32)
+            for _ in range(ARITY)
+        ]
+        t = torch.full((ARITY,), 100.0, device=device, dtype=torch.float32)
+        cond: List[torch.Tensor] = [
+            torch.randn(COND_LEN, config.cond_dim, generator=generator,
+                        device=device, dtype=torch.float32)
+            for _ in range(ARITY)
+        ]
+        with torch.no_grad():
+            out = self.pipe.transformer(x, t, cond)
+        return Micro4dOut(
+            checkpoint=str(resolved.ref.path),
+            shape=str(tuple(int(n) for n in out.shape)))
+
+
+__all__ = ["DECLARATION", "FAMILY", "Generate4d", "Micro4dDefaults",
+           "Micro4dIn", "Micro4dOut"]

examples/micro-diffusion/src/micro_diffusion/model.py

@@ -197,4 +197,36 @@ def forward(self, latent: torch.Tensor) -> torch.Tensor:
         return torch.tanh(self.proj_out(self.norm(h)))
 
 
-__all__ = ["MicroConfig", "MicroDecoder", "MicroDenoiser"]
+__all__ = ["MicroConfig", "MicroDecoder", "MicroDenoiser",
+           "MicroGridDenoiser"]
+
+
+class MicroGridDenoiser(MicroDenoiser):
+    """The pgw#998 shape: a 4-D latent GRID whose H and W are BOTH dynamic.
+
+    Same weights and same blocks as :class:`MicroDenoiser` — only the
+    interface differs. Each ``x`` element is ``(C, H, W)`` and is flattened to
+    tokens INSIDE the forward, so every matmul's M extent becomes the PRODUCT
+    ``H*W``: an extent that is NONLINEAR in the traced symbols.
+
+    That is z-image's declared shape (``H_lat`` and ``W_lat`` on a 4-D latent
+    under ``dynamic-collapse``), and it is what pgw#998 found unlowerable
+    across the mint's own `torch.export.save`/`load` hand-off to its compile
+    child. This class exists so the gauntlet keeps testing that seam on every
+    run rather than trusting a changelog.
+    """
+
+    def forward(  # type: ignore[override]
+        self,
+        x: List[torch.Tensor],
+        t: torch.Tensor,
+        cond: List[torch.Tensor],
+    ) -> torch.Tensor:
+        tokens = []
+        for grid in x:
+            channels, height, width = grid.shape
+            tokens.append(
+                grid.reshape(channels, height * width).transpose(0, 1))
+        out = super().forward(tokens, t, cond)
+        channels, height, width = x[0].shape
+        return out.transpose(1, 2).reshape(len(x), channels, height, width)

examples/micro-diffusion/src/micro_diffusion/pipeline.py

@@ -13,7 +13,12 @@
 
 import torch
 
-from .model import MicroConfig, MicroDecoder, MicroDenoiser
+from .model import (
+    MicroConfig,
+    MicroDecoder,
+    MicroDenoiser,
+    MicroGridDenoiser,
+)
 from .weights import SEED, load_config, load_state, materialize
 
 
@@ -128,4 +133,20 @@ def __call__(
         return self.unpatchify(cells, grid)
 
 
-__all__ = ["MicroConfig", "MicroPipeline"]
+__all__ = ["MicroConfig", "MicroGridPipeline", "MicroPipeline"]
+
+
+class MicroGridPipeline(MicroPipeline):
+    """The pgw#998 vehicle's pipeline: same weights, GRID-shaped transformer.
+
+    Only `.transformer` differs — it is a :class:`MicroGridDenoiser`, so the
+    traced call takes 4-D latents and every matmul's M extent is the product
+    of two dynamic symbols. See `aot_declaration_4d`.
+    """
+
+    @classmethod
+    def from_pretrained(cls, path: str, **_kw: Any) -> "MicroGridPipeline":
+        base = MicroPipeline.from_pretrained(path)
+        grid = MicroGridDenoiser(base.config)
+        grid.load_state_dict(base.transformer.state_dict(), strict=True)
+        return cls(grid.eval(), base.decoder, source=base.source)