diff --git a/README.md b/README.md
index f2827e8..53d449c 100644
--- a/README.md
+++ b/README.md
@@ -54,16 +54,18 @@ The model is trained from a random initialization until convergence, which is de
 1. If running the benchmark for the first time, or running with different fractal parameters (`n_categories`, `variance_threshold`) than previously, generate fractal classes and instances:  
     `scaffold generate_fractals -c ScaFFold/configs/benchmark_default.yml`
 
-    Note that the benchmark ships with an initial set of 50 fractal classes.
+    Fractal category libraries are generated deterministically from the
+    configured seed (under `fract_base_dir/var<...>/seed<...>/`) and reused by
+    later runs with the same seed.
 
 1. Once fractal generation completes, run the benchmark:  
     `torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml`
 
 ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution.
 
-`benchmark` creates a folder for the benchmark run(s) at `base_run_dir` set in the config file. For reproducibility, we store a copy of the benchmark run config yml. Within each run subfolder, `benchmark` creates a yml config for that specific run.
+Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml as `base_config.yaml` plus the fully merged `config.yaml` for that run.
 
-After each run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in `<base_run_dir/figures`.
+After the run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in `<base_run_dir/figures`.
 
 Parameters are set in a `.yml` config file and can be modified by the user. See
 [`ScaFFold/configs/benchmark_default.yml`](ScaFFold/configs/benchmark_default.yml)
@@ -143,7 +145,7 @@ The weights we use to scale IFS parameters look like the following:
 
 ### Dataset generation
 
-Finally,  we are ready to generate a dataset for training our model. Each sample in the dataset is composed of several fractal instances (default=3), randomly selected from any category, overlain with eachother in a 3D voxel grid. Each fractal instance in a sample is placed in a random, non-centered location in the voxel grid if the `scale` parameter set in the sweep config file is <1; otherwise, each fractal instance is centered on the center of the voxel grid. Below is an outline of the data generation process:
+Finally,  we are ready to generate a dataset for training our model. Each sample in the dataset is composed of several fractal instances (default=3), randomly selected from any category, overlain with eachother in a 3D voxel grid. Each fractal instance in a sample is placed in a random, non-centered location in the voxel grid if the `scale` parameter set in the benchmark config file is <1; otherwise, each fractal instance is centered on the center of the voxel grid. Below is an outline of the data generation process:
 ```
 For n  in n_volumes:
 
@@ -172,7 +174,7 @@ For n  in n_volumes:
 
 #### 1. Profiling with the PyTorch Profiler
 
-Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/).
+Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). The trace is written into the run directory. `1`, `true`, `on` and `yes` (any case) enable profiling; every other value, including `0`, `false`, `no` and `off`, leaves it disabled.
 
 #### 2. Profiling with Caliper & Adiak
 
diff --git a/ScaFFold/benchmark.py b/ScaFFold/benchmark.py
index 3b94150..a43d534 100644
--- a/ScaFFold/benchmark.py
+++ b/ScaFFold/benchmark.py
@@ -12,45 +12,18 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 
-import itertools
 import shutil
 from argparse import Namespace
 from pathlib import Path, PosixPath
 
-import yaml
 from mpi4py import MPI
 
 from ScaFFold import worker
-from ScaFFold.utils.config_utils import Config
 from ScaFFold.utils.distributed import get_world_rank
 from ScaFFold.utils.perf_measure import adiak_init, adiak_value
 from ScaFFold.utils.utils import setup_mpi_logger
 
 
-def expand_sweep_combinations(config_dict):
-    """
-    Expand list-valued sweep parameters into per-run scalar configs.
-
-    Scalar-typed config fields holding lists are treated as sweep dimensions;
-    the cross product of their values yields one config dict per combination.
-    An all-scalar config yields exactly one combination. Legitimately
-    list-typed fields (e.g. dc_num_shards) are never treated as sweeps.
-    """
-    sweep_keys = sorted(
-        k for k in Config._SCALAR_KEYS if isinstance(config_dict.get(k), list)
-    )
-    if not sweep_keys:
-        return [dict(config_dict)]
-
-    combinations = []
-    value_lists = [config_dict[k] for k in sweep_keys]
-    for values in itertools.product(*value_lists):
-        combo = dict(config_dict)
-        combo.update(dict(zip(sweep_keys, values)))
-        combinations.append(combo)
-    return combinations
-
-
 def main(kwargs_dict: dict = {}):
     args = Namespace(**kwargs_dict)
     log = setup_mpi_logger(__file__, args.verbose)
@@ -60,65 +33,42 @@ def main(kwargs_dict: dict = {}):
     rank = get_world_rank(required=True)
     log.debug("args found: %s", args)
 
-    run_dicts = None
-    # Now set up and start benchmark run(s)
+    run_dict = None
+    # Now set up and start the benchmark run. Each invocation runs exactly one
+    # benchmark run, in the run directory the CLI resolved.
+    benchmark_run_dir = args.benchmark_run_dir
     if args.restart:
-        # Resume the run in the directory the CLI resolved. The worker path
-        # reads config.run_dir/run_iter directly (a missing run_dir would crash
-        # BaseTrainer), so fill both from the resolved benchmark_run_dir here
-        # just as the fresh, single-combination path does below.
-        benchmark_run_dir = args.benchmark_run_dir
-        kdict = {k: v for k, v in vars(args).items() if k not in ["command"]}
-        kdict["run_dir"] = str(benchmark_run_dir)
-        kdict["run_iter"] = Path(f"{benchmark_run_dir}/run")
-        run_dicts = [kdict]
+        # Resume the run in that directory. The worker path reads
+        # config.run_dir/run_iter directly (a missing run_dir would crash
+        # BaseTrainer), so fill both in here just as the fresh path does below.
+        run_dict = {k: v for k, v in vars(args).items() if k not in ["command"]}
+        run_dict["run_dir"] = str(benchmark_run_dir)
+        run_dict["run_iter"] = Path(f"{benchmark_run_dir}/run")
     elif rank == 0:
-        # Get run dir
-        benchmark_run_dir = args.benchmark_run_dir
-
-        # Save copy of benchmark config yml to run dir
+        # Save a copy of the benchmark config yml in the run dir, under a fixed
+        # name of its own. Copying it in under its original name overwrites the
+        # merged config.yaml the CLI just wrote whenever the base config file
+        # happens to be called config.yaml -- and restart.sh points -c at
+        # $RUN_DIR/config.yaml, so the restart would reload the raw base config
+        # with none of the run's overrides or metadata.
         bench_config_path = Path(args.config)
-        shutil.copy(bench_config_path, benchmark_run_dir)
-
-        base_dict = {k: v for k, v in vars(args).items() if k not in ["command"]}
-        combinations = expand_sweep_combinations(base_dict)
+        shutil.copy(bench_config_path, Path(benchmark_run_dir) / "base_config.yaml")
 
-        # One run (and run directory) per parameter combination. The single
-        # all-scalar combination keeps the original flat layout.
-        run_dicts = []
-        if len(combinations) == 1:
-            kdict = combinations[0]
-            kdict["run_dir"] = str(benchmark_run_dir)
-            kdict["run_iter"] = Path(f"{benchmark_run_dir}/run")
-            run_dicts.append(kdict)
-        else:
-            for i, kdict in enumerate(combinations):
-                run_dir = Path(benchmark_run_dir) / f"param_set_{i}"
-                run_dir.mkdir(parents=True, exist_ok=True)
-                kdict["run_dir"] = str(run_dir)
-                kdict["run_iter"] = run_dir / "run"
-                with open(run_dir / "run_config.yaml", "w") as file:
-                    yaml.dump(
-                        {
-                            k: str(v) if isinstance(v, PosixPath) else v
-                            for k, v in kdict.items()
-                        },
-                        file,
-                    )
-                run_dicts.append(kdict)
+        run_dict = {k: v for k, v in vars(args).items() if k not in ["command"]}
+        run_dict["run_dir"] = str(benchmark_run_dir)
+        run_dict["run_iter"] = Path(f"{benchmark_run_dir}/run")
 
     comm.Barrier()
-    run_dicts = comm.bcast(run_dicts, root=0)
+    run_dict = comm.bcast(run_dict, root=0)
 
     adiak_init(comm)
-    for kdict in run_dicts:
-        # Add all config params as metadata
-        for key, value in kdict.items():
-            if isinstance(value, dict):
-                log.debug("Adiak: skipping key with dict value '%s'", key)
-                continue
-            if isinstance(value, PosixPath):
-                value = str(value)
-            adiak_value(key, value)
-
-        worker.main(kwargs_dict=kdict)
+    # Add all config params as metadata
+    for key, value in run_dict.items():
+        if isinstance(value, dict):
+            log.debug("Adiak: skipping key with dict value '%s'", key)
+            continue
+        if isinstance(value, PosixPath):
+            value = str(value)
+        adiak_value(key, value)
+
+    worker.main(kwargs_dict=run_dict)
diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py
index 2ab5cbb..105e87a 100644
--- a/ScaFFold/cli.py
+++ b/ScaFFold/cli.py
@@ -13,6 +13,7 @@
 # SPDX-License-Identifier: (Apache-2.0)
 
 import argparse
+import builtins
 import socket
 import sys
 from datetime import datetime
@@ -24,9 +25,49 @@
 from ScaFFold.utils import config_utils
 from ScaFFold.utils.collect_scheduler_info import collect_scheduler_metadata
 from ScaFFold.utils.create_restart_script import create_restart_script
+from ScaFFold.utils.distributed import get_world_size
 from ScaFFold.utils.utils import setup_mpi_logger
 
 
+def check_launcher_world_size(mpi_world_size):
+    """Verify that the MPI world spans the whole job.
+
+    ScaFFold uses two different sources of truth for the job shape: the CLI and
+    the benchmark driver make their job-wide decisions on ``MPI.COMM_WORLD``
+    rank 0, while the training path takes its rank and world size from the
+    launcher's environment (``get_world_rank`` / ``get_world_size``). Those
+    agree only when the job was started by an MPI-aware launcher.
+
+    Under a plain ``torchrun`` (or a bare ``python`` invocation of several
+    processes) mpi4py initializes as an independent singleton in every process,
+    so every process believes it is MPI rank 0: each one runs the rank-0 block,
+    atomically claims its *own* timestamped run directory, and the job then
+    diverges -- non-zero launcher ranks crash on a broadcast that never
+    happened while rank 0 blocks in the first collective until it times out.
+
+    Fail loudly here, before any run directory is created, instead of leaving
+    that mess behind. ``get_world_size`` falls back to the MPI communicator
+    when the environment reports nothing, so an unlauncher-ed single process
+    trivially agrees with itself.
+    """
+    env_world_size = get_world_size()
+    if env_world_size != mpi_world_size:
+        raise RuntimeError(
+            f"Launcher/MPI world size mismatch: MPI.COMM_WORLD reports "
+            f"{mpi_world_size} rank(s) but the launcher environment reports "
+            f"{env_world_size}. ScaFFold decides run directories and restart "
+            "state on MPI rank 0 and broadcasts them, so an MPI world that "
+            "does not span the job is unrecoverable: every process acts as "
+            "rank 0, each claims a separate run directory, and the job hangs "
+            "in the first collective. This is what a plain 'torchrun' (or "
+            "launching the processes directly) produces, because mpi4py then "
+            "initializes as a singleton in every process. Launch ScaFFold "
+            "with an MPI-aware launcher (torchrun-hpc, flux run, srun, "
+            "mpirun) so that MPI spans all ranks, or run a single process "
+            "with no launcher environment set."
+        )
+
+
 def _make_fresh_run_dir(base_run_dir, timestamp):
     """Create a fresh timestamped run directory without clobbering an existing one.
 
@@ -50,25 +91,105 @@ def _make_fresh_run_dir(base_run_dir, timestamp):
             candidate = base_run_dir / f"{timestamp}-{suffix}"
 
 
+def explicit_cli_keys(args, parsers):
+    """Return the names of the options actually given on the command line.
+
+    argparse does not record which options were supplied, so a value counts as
+    explicit when it differs from the default of the first parser in
+    ``parsers`` that defines one (subcommand parser first, then the top-level
+    parser). Only these may outrank a config-file setting; everything else in
+    the namespace is an argparse default, which is the weakest source.
+
+    The one ambiguity is a flag passed with exactly its default value: it is
+    indistinguishable from an absent flag, so a config-file entry outranks it.
+    Where the flag has no default (``None``) that is harmless -- passing a
+    value always makes it explicit -- but the two flags that do have one,
+    ``--datagen-batch-size`` (10000) and ``-v`` (0), lose the argument in that
+    one case: ``--datagen-batch-size 10000`` next to ``datagen_batch_size: 500``
+    in the config file yields 500. The alternative is to give every flag a
+    ``None`` default and re-derive the real defaults elsewhere, which buys a
+    narrow correctness win by scattering the defaults; the ambiguity is
+    documented instead.
+    """
+    explicit = set()
+    for name, value in vars(args).items():
+        default = None
+        for parser in parsers:
+            default = parser.get_default(name)
+            if default is not None:
+                break
+        if value != default:
+            explicit.add(name)
+    return explicit
+
+
+def missing_checkpoint_error(combined_config):
+    """Return the "nothing to resume from" message, or None if a restart can run.
+
+    Reports the first problem found rather than raising, so the caller can make
+    this a rank-0 decision and broadcast the verdict instead of letting every
+    rank stat the shared filesystem and possibly disagree.
+
+    The directory checked is the *resolved* benchmark run dir -- the one this
+    launch will actually train in -- and not the raw ``run_dir`` key. The two
+    can differ: a config.yaml dumped by a restarted run carries that run's
+    ``run_dir``, so reusing it as a base config had this check stat another
+    run's checkpoints, pass, and let the job die hours later with its dataset
+    already generated. ``resolve_run_dir`` keeps the two in agreement; this
+    prefers the resolved value so they cannot drift apart again.
+    """
+    run_dir = combined_config.get("benchmark_run_dir") or combined_config.get("run_dir")
+    if not run_dir:
+        return (
+            "Restart requested but no run directory was resolved. Pass "
+            "'--run-dir <path>' (or set run_dir in the config file)."
+        )
+    checkpoint_dir = Path(run_dir) / combined_config.get(
+        "checkpoint_dir", "checkpoints"
+    )
+    expected_checkpoints = (
+        checkpoint_dir / "checkpoint_last.pth",
+        checkpoint_dir / "checkpoint_best.pth",
+    )
+    if any(path.exists() for path in expected_checkpoints):
+        return None
+    expected = " or ".join(str(path) for path in expected_checkpoints)
+    return f"Restart requested but no checkpoint was found. Expected {expected}."
+
+
 def resolve_run_dir(args_dict, combined_config):
     """Decide the benchmark run directory and whether this launch resumes a run.
 
     The semantics are fixed and unambiguous:
 
-    * ``--run-dir DIR`` (with or without ``--restart``): resume in that exact
-      directory. ``train_from_scratch`` is forced off and ``restart`` on so the
-      downstream benchmark driver takes its restart path.
-    * ``--restart`` without ``--run-dir``: rejected with a clear error. The run
-      directory to resume must be named explicitly; the most recent directory
-      is never guessed.
-    * neither flag: create a fresh timestamped directory under ``base_run_dir``,
+    * a run directory (``--run-dir DIR``, or ``run_dir`` in the config file),
+      with or without a restart flag: resume in that exact directory.
+      ``train_from_scratch`` is forced off and ``restart`` on so the downstream
+      benchmark driver takes its restart path.
+    * a restart requested with no run directory anywhere: rejected with a clear
+      error, before any directory is created. The directory to resume must be
+      named explicitly; the most recent one is never guessed.
+    * neither: create a fresh timestamped directory under ``base_run_dir``,
       retrying with a numeric suffix on a same-second name collision.
 
-    ``combined_config['benchmark_run_dir']`` is set in every path so the driver
-    can always read it. Returns ``(benchmark_run_dir: Path, restarting: bool)``.
+    Both keys are read from the *merged* config, not from the command line
+    alone. A config file is a first-class source for them -- the generated
+    ``restart.sh`` replays a dumped ``config.yaml``, which carries both -- and
+    an absent ``--restart`` cannot outrank a file that sets it, because an
+    unset ``store_true`` flag is indistinguishable from its default. Reading
+    only the command line here let the two disagree: a config-file restart
+    created a *fresh* run directory and only then failed for want of a run dir,
+    and a stale ``run_dir`` inherited from a reused ``config.yaml`` was left in
+    the config for the restart pre-check to stat while training happened
+    somewhere else entirely.
+
+    The resolved answer is written back to ``benchmark_run_dir``, ``restart``
+    and ``run_dir``, so every later reader -- the pre-check, the dumped
+    ``config.yaml``, the benchmark driver -- sees exactly what was decided here.
+    Returns ``(benchmark_run_dir: Path, restarting: bool)``.
     """
-    restart_flag = bool(args_dict.get("restart"))
-    run_dir_arg = args_dict.get("run_dir")
+    restart_flag = bool(combined_config.get("restart") or args_dict.get("restart"))
+    run_dir_arg = combined_config.get("run_dir") or args_dict.get("run_dir")
 
     if run_dir_arg is not None:
         benchmark_run_dir = Path(run_dir_arg)
@@ -78,9 +199,10 @@ def resolve_run_dir(args_dict, combined_config):
         restarting = True
     elif restart_flag:
         raise ValueError(
-            "--restart requires --run-dir: pass the directory of the run to "
-            "resume (e.g. '--restart --run-dir <path>'). The most recent run "
-            "directory is not resolved automatically."
+            "A restart was requested (--restart, or restart: true in the "
+            "config file) but no run directory was given: pass the directory "
+            "of the run to resume (e.g. '--restart --run-dir <path>'). The "
+            "most recent run directory is not resolved automatically."
         )
     else:
         base_run_dir = Path(combined_config["base_run_dir"])
@@ -94,13 +216,154 @@ def resolve_run_dir(args_dict, combined_config):
         )
         restarting = False
 
+    # Write the resolution back, so nothing downstream can re-derive a
+    # different answer from the raw inputs. ``run_dir`` is cleared for a fresh
+    # run: left set, a value inherited from a reused config.yaml names a
+    # directory this run has nothing to do with.
     combined_config["benchmark_run_dir"] = str(benchmark_run_dir)
+    combined_config["restart"] = restarting
+    combined_config["run_dir"] = str(benchmark_run_dir) if restarting else None
     if restarting:
         combined_config["train_from_scratch"] = False
-        combined_config["restart"] = True
     return benchmark_run_dir, restarting
 
 
+def rebuild_error(type_name, message):
+    """Rebuild rank 0's configuration error on a peer that never saw it.
+
+    Only the type name and message cross the wire: an arbitrary exception
+    object may not survive a pickle round trip, and a failure to unpickle on
+    the receiving side would turn the error being reported back into the hang
+    it was reported to avoid. Builtin exception types are rebuilt as
+    themselves, so a caller's ``except ValueError`` still catches what rank 0
+    raised; anything else degrades to ``RuntimeError`` naming the original
+    type.
+    """
+    cls = getattr(builtins, type_name, None)
+    if isinstance(cls, type) and issubclass(cls, Exception):
+        return cls(message)
+    return RuntimeError(f"{type_name}: {message}")
+
+
+def build_run_config(args, parsers, log, world_size):
+    """Build the job-wide config, and lay down the benchmark's run directory.
+
+    Rank-0-only work: it reads and merges the config files, applies the
+    command-line overrides, validates the result, and -- for the benchmark
+    subcommand -- resolves the run directory, dumps the configs into it and
+    writes its restart script. Returns the merged config the caller broadcasts.
+
+    Raises whatever the validation or the filesystem raises; the caller runs
+    this inside a guard and broadcasts the outcome, since every peer is already
+    waiting in the barrier that follows.
+    """
+    log.debug("args = %s", args)
+
+    # --config may be a single path (generate_fractals) or a list of
+    # paths (benchmark, action="append"): base config plus overrides.
+    config_paths = args.config if isinstance(args.config, list) else [args.config]
+    merged_dict = config_utils.load_config_files(config_paths)
+    # Validate the merged result and derive dependent settings. Every run
+    # parameter must be single-valued; a list is rejected here by name.
+    bench_config = config_utils.Config(merged_dict)
+    bench_config_dict = vars(bench_config)
+    cli_args = vars(args)
+    # Downstream consumers expect a single config path (e.g. to copy it
+    # into the run dir); keep the base config there.
+    cli_args["config"] = config_paths[0]
+
+    # Combine configs, in increasing order of precedence:
+    # argparse default < config file < explicit command-line flag.
+    combined_config = bench_config_dict.copy()
+    # Config only keeps the keys it consumes; the auxiliary keys it accepts
+    # (verbose, datagen_batch_size, ...) never become attributes, so put
+    # the file's values back first. Without this they are absent below and
+    # the argparse default overwrites what the user wrote in the config.
+    for key, value in merged_dict.items():
+        combined_config.setdefault(key, value)
+
+    explicit_cli = explicit_cli_keys(args, parsers)
+    for key, value in cli_args.items():
+        if key == "command":
+            continue
+        if key not in combined_config:
+            combined_config[key] = value
+        elif key in explicit_cli and value is not None:
+            log.info(
+                "Overriding '%s=%s' with '%s=%s'",
+                key,
+                combined_config[key],
+                key,
+                value,
+            )
+            combined_config[key] = value
+    # The subcommand is always owned by the command line.
+    combined_config["command"] = cli_args["command"]
+
+    # Recalculate unet_layers to capture any CLI overrides. The overridden
+    # pair has to be re-validated: Config only saw the config-file values.
+    config_utils.validate_unet_dims(
+        combined_config["problem_scale"], combined_config["unet_bottleneck_dim"]
+    )
+    combined_config["unet_layers"] = (
+        combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"]
+    )
+    config_utils.require_positive_int("n_categories", combined_config["n_categories"])
+
+    # Resolve paths to absolute, matching Config() behavior
+    if "base_run_dir" in combined_config and combined_config["base_run_dir"]:
+        combined_config["base_run_dir"] = str(
+            Path(combined_config["base_run_dir"]).resolve()
+        )
+
+    if "dataset_dir" in combined_config and combined_config["dataset_dir"]:
+        combined_config["dataset_dir"] = str(
+            Path(combined_config["dataset_dir"]).resolve()
+        )
+
+    if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]:
+        combined_config["fract_base_dir"] = str(
+            Path(combined_config["fract_base_dir"]).resolve()
+        )
+
+    # Calculate these variables after override
+    combined_config["vol_size"] = pow(2, combined_config["problem_scale"])
+    combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256)
+
+    # The run directory, its config dumps and its restart script belong to
+    # the benchmark subcommand alone. Fractal generation writes nothing
+    # there, and the restart script it used to get replayed
+    # `generate_fractals --restart --run-dir ...` -- flags that subparser
+    # rejects, so the script could only ever exit 2.
+    if args.command == "benchmark":
+        # Resolve the run directory and whether this launch resumes a run.
+        # This sets combined_config["benchmark_run_dir"] on every path, and
+        # writes the resolved restart/run_dir back into the config.
+        benchmark_run_dir, restarting = resolve_run_dir(cli_args, combined_config)
+        if restarting:
+            log.info("Resuming in existing directory: %s", benchmark_run_dir)
+
+        # Add scheduler metadata and machine name to config.yaml
+        combined_config["scheduler_metadata"] = collect_scheduler_metadata()
+        combined_config["machine_name"] = socket.gethostname()
+
+        # Dump configs (Overwrite is okay/desired on restart to capture new job IDs)
+        overrides = {
+            k: v for k, v in cli_args.items() if v is not None and k != "command"
+        }
+        with open(benchmark_run_dir / "overrides.yaml", "w") as file:
+            yaml.dump(overrides, file)
+        with open(benchmark_run_dir / "config.yaml", "w") as file:
+            yaml.dump(combined_config, file)
+
+        # 4. Generate/Update the restart script in the directory. The
+        # communicator size is ground truth for the job scale; environment
+        # sniffing is only the fallback for callers that lack it.
+        create_restart_script(benchmark_run_dir, world_size=world_size)
+
+    return combined_config
+
+
 def main():
     """
     Command line interface for ScaFFold.
@@ -170,8 +433,8 @@ def main():
         help="Run the benchmark.",
         description=(
             "The default run method for ScaFFold. "
-            "Users may specify lists of run parameters in the config file. "
-            "This subcommand runs one instance of the benchmark for each parameter combination. "
+            "Each invocation runs exactly one instance of the benchmark, using "
+            "the single-valued run parameters given in the config file. "
             "Requires path to config file."
         ),
     )
@@ -301,6 +564,19 @@ def main():
     rank = comm.Get_rank()
     # Parse the command-line arguments.
     args = parser.parse_args()
+    # Every rank runs this identically, before any run directory is created, so
+    # a mis-launched job aborts uniformly instead of leaving per-rank run dirs
+    # behind and hanging. It runs *after* parsing so that the arguments argparse
+    # handles by itself -- ``--help``, a usage error -- still behave: asking
+    # what the flags are is not a job launch, and answering it with a launcher
+    # mismatch (which is exactly the environment someone debugging one is
+    # sitting in) helps nobody.
+    check_launcher_world_size(comm.Get_size())
+    subcommand_parsers = {
+        "benchmark": benchmark_parser,
+        "generate_fractals": generate_fractals_parser,
+    }
+    active_parser = subcommand_parsers[args.command]
     log = setup_mpi_logger(__file__, args.verbose)
     combined_config = None
 
@@ -318,109 +594,51 @@ def main():
             "resume (e.g. '--restart --run-dir <path>')."
         )
 
+    # Rank 0 builds the config for the whole job while every other rank waits
+    # in the barrier below. Everything in there can fail on user input (an
+    # unknown config key, an out-of-range bottleneck, a restart with no run
+    # dir) or on the filesystem, and a rank-0-only raise leaves the peers
+    # blocked in that barrier -- a hang instead of the error message the user
+    # needs. The outcome is therefore broadcast, exactly like the restart
+    # pre-check below, and every rank raises the same error together.
+    config_error = None
+    rank0_error = None
     if rank == 0:
-        log.debug("args = %s", args)
-
-        # --config may be a single path (generate_fractals) or a list of
-        # paths (benchmark, action="append"): base config plus overrides.
-        config_paths = args.config if isinstance(args.config, list) else [args.config]
-        merged_dict = config_utils.load_config_files(config_paths)
-        # Validate the merged result and derive dependent settings
-        # (list-valued sweep params are allowed here; the benchmark driver
-        # expands them per run).
-        bench_config = config_utils.Config(merged_dict, allow_sweeps=True)
-        bench_config_dict = vars(bench_config)
-        cli_args = vars(args)
-        # Downstream consumers expect a single config path (e.g. to copy it
-        # into the run dir); keep the base config there.
-        cli_args["config"] = config_paths[0]
-
-        # Combine configs: CLI args override config file values
-        combined_config = bench_config_dict.copy()
-        for key, value in cli_args.items():
-            if key not in combined_config:
-                combined_config[key] = value
-            elif value is not None and key != "command":
-                log.info(
-                    "Overriding '%s=%s' with '%s=%s'",
-                    key,
-                    combined_config[key],
-                    key,
-                    value,
-                )
-                combined_config[key] = value
-
-        # Recalculate unet_layers to capture any CLI overrides
-        combined_config["unet_layers"] = (
-            combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"]
-        )
-        config_utils.require_positive_int(
-            "n_categories", combined_config["n_categories"]
-        )
-
-        # Resolve paths to absolute, matching Config() behavior
-        if "base_run_dir" in combined_config and combined_config["base_run_dir"]:
-            combined_config["base_run_dir"] = str(
-                Path(combined_config["base_run_dir"]).resolve()
-            )
-
-        if "dataset_dir" in combined_config and combined_config["dataset_dir"]:
-            combined_config["dataset_dir"] = str(
-                Path(combined_config["dataset_dir"]).resolve()
-            )
-
-        if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]:
-            combined_config["fract_base_dir"] = str(
-                Path(combined_config["fract_base_dir"]).resolve()
+        try:
+            combined_config = build_run_config(
+                args, (active_parser, parser), log, comm.Get_size()
             )
-
-        # Calculate these variables after override
-        combined_config["vol_size"] = pow(2, combined_config["problem_scale"])
-        combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256)
-
-        # Resolve the run directory and whether this launch resumes a run.
-        # This sets combined_config["benchmark_run_dir"] on every path and,
-        # when resuming, forces train_from_scratch off / restart on.
-        benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config)
-        if restarting:
-            log.info("Resuming in existing directory: %s", benchmark_run_dir)
-
-        # Add scheduler metadata and machine name to config.yaml
-        combined_config["scheduler_metadata"] = collect_scheduler_metadata()
-        combined_config["machine_name"] = socket.gethostname()
-
-        # Dump configs (Overwrite is okay/desired on restart to capture new job IDs)
-        overrides = {
-            k: v for k, v in cli_args.items() if v is not None and k != "command"
-        }
-        with open(benchmark_run_dir / "overrides.yaml", "w") as file:
-            yaml.dump(overrides, file)
-        with open(benchmark_run_dir / "config.yaml", "w") as file:
-            yaml.dump(combined_config, file)
-
-        # 4. Generate/Update the restart script in the directory
-        create_restart_script(benchmark_run_dir)
+        except Exception as e:
+            combined_config = None
+            rank0_error = e
+            config_error = (type(e).__name__, str(e))
 
     comm.Barrier()
+    config_error = comm.bcast(config_error, root=0)
+    if config_error is not None:
+        # Rank 0 re-raises the original (keeping its traceback); the peers
+        # rebuild it from what crossed the wire.
+        if rank0_error is not None:
+            raise rank0_error
+        raise rebuild_error(*config_error)
     combined_config = comm.bcast(combined_config, root=0)
 
-    if combined_config.get("restart", False):
-        run_dir = combined_config.get("run_dir")
-        if not run_dir:
-            raise ValueError("--restart requires --run-dir")
-
-        checkpoint_dir = Path(run_dir) / combined_config.get(
-            "checkpoint_dir", "checkpoints"
-        )
-        expected_checkpoints = (
-            checkpoint_dir / "checkpoint_last.pth",
-            checkpoint_dir / "checkpoint_best.pth",
-        )
-        if not any(path.exists() for path in expected_checkpoints):
-            expected = " or ".join(str(path) for path in expected_checkpoints)
-            raise FileNotFoundError(
-                f"Restart requested but no checkpoint was found. Expected {expected}."
-            )
+    # Restart pre-check. Like every other decision here it is made once, on
+    # rank 0, and broadcast: the check reads the filesystem, and ranks can see
+    # different views of a shared filesystem (stale NFS/Lustre attribute
+    # caches). A rank that decided for itself would either abort alone --
+    # stranding its peers in benchmark.py's timeout-less barrier -- or keep
+    # running after rank 0 had already aborted.
+    # Only the benchmark subcommand has run directories or checkpoints, so it
+    # is the only one this applies to: fractal generation reading a benchmark's
+    # config.yaml must not be judged on that run's restart state.
+    restart_precheck_error = None
+    if args.command == "benchmark" and combined_config.get("restart", False):
+        if rank == 0:
+            restart_precheck_error = missing_checkpoint_error(combined_config)
+    restart_precheck_error = comm.bcast(restart_precheck_error, root=0)
+    if restart_precheck_error is not None:
+        raise FileNotFoundError(restart_precheck_error)
 
     if rank == 0:
         log.debug("combined_config = %s", combined_config)
diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml
index c41f4ed..29be50f 100644
--- a/ScaFFold/configs/benchmark_default.yml
+++ b/ScaFFold/configs/benchmark_default.yml
@@ -39,5 +39,6 @@ loss_freq: 1                       # Number of epochs between logging the overal
 normalize: 1                       # Cateogry search normalization parameter
 group_norm_groups: 8              # Number of groups used by GroupNorm in the UNet blocks.
 warmup_batches: 64                 # How many warmup batches per rank to run before training.
+activation_checkpointing: 0        # If 1, recompute UNet block activations during the backward pass instead of storing them: less memory, more compute.
 ce_weight_sample_fraction: 0.1     # Fraction of training masks to sample when estimating background vs foreground CE weights.
 dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse.
diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py
index 246666f..3cc1f81 100644
--- a/ScaFFold/datagen/category_search.py
+++ b/ScaFFold/datagen/category_search.py
@@ -26,6 +26,7 @@
 import numpy as np
 from mpi4py import MPI
 
+from ScaFFold.datagen import layout
 from ScaFFold.datagen.generate_fractal_points import generate_fractal_points
 from ScaFFold.datagen.rng import SEED_MASK, derive_seed, seed_numba
 from ScaFFold.utils.config_utils import Config
@@ -334,12 +335,41 @@ def save_valid_category(
     target = os.path.join(fracts_write_dir, "%06d.csv" % idx)
     if os.path.exists(target):
         raise FileExistsError(f"Refusing to overwrite existing category file: {target}")
-    np.savetxt(target, params, delimiter=",")
+    _savetxt_atomic(target, params)
     existing_indices.append(idx)
     existing_params.append(params)
     return idx
 
 
+def _savetxt_atomic(target: str, params: np.array) -> None:
+    """Write one category's parameters to ``target`` atomically.
+
+    A category CSV truncated by a killed job is poison: the six-digit name is
+    all the resume scan looks at, so the category counts as done forever, while
+    every consumer (instance generation, and the search's own resume) dies
+    parsing it. The file is therefore written to a temp name in the same
+    directory -- one that neither the resume glob (``NNNNNN.csv``) nor the
+    instance loader's ``*.csv`` filter can match -- flushed, fsynced, and only
+    then ``os.replace``d onto the final name.
+    """
+    directory, name = os.path.split(target)
+    tmp_path = os.path.join(directory, f".{name}.tmp{os.getpid()}")
+    try:
+        with open(tmp_path, "w") as handle:
+            np.savetxt(handle, params, delimiter=",")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(tmp_path, target)
+    except BaseException:
+        # A failed write must leave nothing behind: no temp file, and no
+        # partial file under the name resume would accept.
+        try:
+            os.remove(tmp_path)
+        except OSError:
+            pass
+        raise
+
+
 def _attempt_state_path(fracts_write_dir: str, rank: int) -> str:
     return os.path.join(fracts_write_dir, f".rng_attempt_rank{rank}")
 
@@ -373,6 +403,35 @@ def write_attempt_counter(fracts_write_dir: str, rank: int, attempt_index: int)
     os.replace(tmp, path)
 
 
+def _sweep_stale_temp_files(fracts_write_dir: str, log) -> None:
+    """Remove temp files stranded by killed writes in the category directory.
+
+    Both atomic writers here (``_savetxt_atomic`` and ``write_attempt_counter``)
+    unlink their temp file when the write raises, but a SIGKILL -- walltime, an
+    OOM, a node failure -- skips that Python-level cleanup and strands it. The
+    names carry the writer's pid, so they accumulate one per killed process and
+    nothing else ever removes them; ``instance.py`` sweeps its equivalents for
+    exactly this reason.
+
+    Called on rank 0 before any rank has written anything this run, and
+    best-effort: this is housekeeping, and it runs just before a Barrier the
+    peers are heading into, so it must not raise.
+    """
+    patterns = (
+        # .NNNNNN.csv.tmp<pid> -- a partially written category CSV.
+        f"{fracts_write_dir}/.*.csv.tmp*",
+        # .rng_attempt_rank<r>.tmp<pid> -- a partially written attempt counter.
+        f"{fracts_write_dir}/.rng_attempt_rank*.tmp*",
+    )
+    for pattern in patterns:
+        for stale in glob.glob(pattern):
+            try:
+                os.remove(stale)
+                log.info("Removed stale category-search temp file %s", stale)
+            except OSError as exc:
+                log.warning("Could not remove stale temp file %s: %s", stale, exc)
+
+
 def main(config: Config) -> None:
     """
     Generate fractal categories.
@@ -403,17 +462,22 @@ def main(config: Config) -> None:
 
     log.info("MPI size = %s", size)
 
-    # Setup directories
-    fracts_sub_dir = f"var{config.variance_threshold}"
-    fracts_write_dir = os.path.join(
-        config.fract_base_dir, fracts_sub_dir, "3DIFS_param"
-    )
+    # Setup directories. The library is keyed by seed (see
+    # ScaFFold.datagen.layout): categories are drawn from a seed-derived
+    # candidate stream, so a run under a different seed must never resume onto
+    # another seed's parameter files.
+    fracts_write_dir = layout.category_param_dir(config)
     if rank == 0:
         log.info("Writing fractals to %s", fracts_write_dir)
+        # A library in the pre-seed layout is invisible to everything below, so
+        # say why it is being ignored rather than appearing to regenerate work
+        # that is plainly still on disk.
+        layout.warn_if_legacy_library(config, log)
         if os.path.exists(fracts_write_dir) and config.datagen_from_scratch:
             log.info("Removing existing fractals directory")
             shutil.rmtree(fracts_write_dir)
         os.makedirs(fracts_write_dir, exist_ok=True)
+        _sweep_stale_temp_files(fracts_write_dir, log)
 
     # Wait until dir setup completes
     comm.Barrier()
@@ -423,18 +487,52 @@ def main(config: Config) -> None:
     # the ones a fresh run produced.
     attempt_index = read_attempt_counter(fracts_write_dir, rank)
 
-    # Parse existing category files (rank 0 owns saving/dedup). Free indices are
-    # derived from these parsed names -- filling holes, never overwriting.
-    existing_indices = parse_category_indices(fracts_write_dir)
+    # Parse existing category files on rank 0 alone and broadcast the result.
+    # Free indices are derived from these parsed names -- filling holes, never
+    # overwriting -- and, critically, the count derived below gates a loop that
+    # contains collectives. Scanning the shared filesystem independently per
+    # rank lets divergent views (stale metadata caches, a concurrent job, a
+    # partially visible directory) put one rank inside the loop while another is
+    # past it, so the two post mismatched collectives on COMM_WORLD and the job
+    # hangs. One scan, one broadcast, one shared verdict.
+    #
+    # Rank 0 also loads the parameters of the categories already on disk (only
+    # it writes, so only it needs them for the duplicate guard). That read is
+    # part of the same rank-0-only window: a category CSV that will not parse --
+    # ragged, or hand-edited -- would otherwise kill rank 0 *after* the peers
+    # had already taken the broadcast and moved on to the next collective. Scan
+    # and load are therefore one guarded decision, reported through one
+    # broadcast, exactly as ``get_dataset`` reports its selection.
     existing_params = []
+    interrupt = None
     if rank == 0:
-        for idx in existing_indices:
-            existing_params.append(
-                np.loadtxt(
-                    os.path.join(fracts_write_dir, "%06d.csv" % idx),
-                    delimiter=",",
+        try:
+            existing_indices = parse_category_indices(fracts_write_dir)
+            for idx in existing_indices:
+                existing_params.append(
+                    np.loadtxt(
+                        os.path.join(fracts_write_dir, "%06d.csv" % idx),
+                        delimiter=",",
+                    )
                 )
+            scan = ("ok", existing_indices)
+        except BaseException as e:
+            existing_params = []
+            scan = (
+                "error",
+                f"rank 0 failed to scan existing categories in "
+                f"{fracts_write_dir}: {type(e).__name__}: {e}",
             )
+            interrupt = e if isinstance(e, KeyboardInterrupt) else None
+    else:
+        scan = None
+    status, payload = comm.bcast(scan, root=0)
+    if status == "error":
+        # Rank 0 keeps an operator's interrupt; every rank aborts either way.
+        if interrupt is not None:
+            raise interrupt
+        raise RuntimeError(f"category search failed: {payload}")
+    existing_indices = payload
 
     # Calculate number of remaining fractal categories to generate
     existing_categories = len(existing_indices)
diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py
index c536b14..168f002 100644
--- a/ScaFFold/datagen/get_dataset.py
+++ b/ScaFFold/datagen/get_dataset.py
@@ -16,9 +16,11 @@
 
 import hashlib
 import json
+import os
 import shutil
 import subprocess
 import time
+import uuid
 from argparse import Namespace
 from pathlib import Path
 from typing import Any, Dict
@@ -30,14 +32,34 @@
 from ScaFFold.utils.utils import setup_mpi_logger
 
 META_FILENAME = "meta.yaml"
+# Datasets are generated into a staging directory carrying this prefix and only
+# renamed into their final ``<timestamp>__<commit>`` name once complete, so a
+# reader never observes a half-written dataset. The prefix is also what the
+# reuse scan skips and what the orphan cleanup collects.
+TMP_PREFIX = ".tmp_"
+# How long a staging directory must have shown no sign of life before it is
+# treated as orphaned (left by a killed/OOM'd job) and reclaimed. See
+# ``_cleanup_stale_staging_dirs`` for the safety argument behind the value.
+STALE_STAGING_AGE_SECONDS = 24 * 60 * 60
+# Bounds on the liveness probe that backs up the heartbeat: how many directory
+# levels below a staging dir it looks at, and how many directories it is willing
+# to visit before it gives up and calls the tree live. Directory mtimes change
+# when entries are created in them, so a couple of levels is enough to notice a
+# writer without stat-ing every volume file.
+_STALE_PROBE_MAX_DEPTH = 3
+_STALE_PROBE_MAX_DIRS = 10000
 # Bumped from 2 to 3 when instance point clouds moved from float64 to float32:
 # the storage layout is unchanged, but float32 voxel binning shifts a handful of
 # boundary voxels, so a float64-era dataset must not be reused as if it were
-# float32. This version stamps new datasets, gates reuse below, and feeds the
-# config_id hash, so an older dataset is neither matched nor scanned. The loader
-# in data_loading.py keeps its own (lower) minimum-layout version and still reads
-# v3 through the modern dense path.
-DATASET_FORMAT_VERSION = 3
+# float32. Bumped from 3 to 4 when the voxel centering offset was corrected from
+# (grid_size - 1 - span)/2 to (grid_size - span)/2: every volume and mask
+# generated before that was misregistered by half a voxel (with the first
+# half-voxel of each axis clipped onto plane 0), so those datasets must be
+# regenerated rather than reused. This version stamps new datasets, gates reuse
+# below, and feeds the config_id hash, so an older dataset is neither matched nor
+# scanned. The loader in data_loading.py keeps its own (lower) minimum-layout
+# version and still reads v4 through the modern dense path.
+DATASET_FORMAT_VERSION = 4
 INCLUDE_KEYS = [
     "dataset_format_version",
     "n_categories",
@@ -84,11 +106,26 @@ def _hash_volume_config(volume_config: Dict[str, Any]) -> str:
     return hashlib.sha256(s).hexdigest()[:12]
 
 
-def _git_commit_short(log) -> str:
+def _git_commit_short(log, source_dir: Path | None = None) -> str:
+    """Return the short commit of the ScaFFold checkout, or ``"no-commit-id"``.
+
+    The commit identifies *the code that generated a dataset*: it is stamped
+    into ``meta.yaml``, into the published directory name, and is what
+    ``dataset_reuse_enforce_commit_id`` compares against. It must therefore be
+    read from the ScaFFold source tree rather than from the process working
+    directory, which is wherever the job was launched (a site workflow repo, a
+    scratch directory, ...) and has nothing to do with this code.
+
+    ``source_dir`` overrides the directory git runs in; it defaults to this
+    module's own location and exists so the non-checkout case can be tested.
+    """
+    if source_dir is None:
+        source_dir = Path(__file__).resolve().parent
     try:
         return (
             subprocess.check_output(
                 ["git", "rev-parse", "--short", "HEAD"],
+                cwd=str(source_dir),
                 stderr=subprocess.DEVNULL,  # Don't show console output to user
             )
             .decode()
@@ -108,6 +145,142 @@ def _git_commit_short(log) -> str:
         return "no-commit-id"
 
 
+def _has_recent_write(path: Path, cutoff: float) -> bool:
+    """Return True if the top levels of ``path`` were written after ``cutoff``.
+
+    A directory's mtime changes whenever an entry is created in it, so the
+    directory mtimes near the top of a staging tree are a cheap proxy for "a
+    writer is active down there" -- no stat of the (possibly hundreds of
+    thousands of) volume files is needed. The walk therefore stats the staging
+    directory, its immediate children, and directories down to
+    ``_STALE_PROBE_MAX_DEPTH``, and stops at the first recent entry, which makes
+    the live case (the one that must not be misjudged) the cheap one.
+
+    A tree wide enough to exceed ``_STALE_PROBE_MAX_DIRS`` is reported as recent
+    rather than walked further: failing to reclaim disk is recoverable, deleting
+    a running job's dataset is not.
+    """
+    stack = [(path, 0)]
+    dirs_seen = 0
+    while stack:
+        current, depth = stack.pop()
+        if current.stat().st_mtime > cutoff:
+            return True
+        if depth >= _STALE_PROBE_MAX_DEPTH:
+            continue
+        with os.scandir(current) as entries:
+            for entry in entries:
+                if entry.is_dir(follow_symlinks=False):
+                    dirs_seen += 1
+                    if dirs_seen > _STALE_PROBE_MAX_DIRS:
+                        return True
+                    stack.append((Path(entry.path), depth + 1))
+                elif depth == 0 and entry.stat().st_mtime > cutoff:
+                    # Files directly in the staging dir (the heartbeat,
+                    # volumes_contents.csv, meta.yaml) are few and cheap.
+                    return True
+    return False
+
+
+def _staging_dir_is_live(path: Path, cutoff: float) -> bool:
+    """Return True if ``path`` shows any sign of a generation still running.
+
+    Two signals, in order of authority:
+
+    1. ``<staging>/.heartbeat``, refreshed by every writing rank every few
+       minutes for as long as volumes are being written (see
+       ``volumegen.StagingHeartbeat``). This is the reliable one, because it
+       does not depend on where in the tree the writers currently are.
+    2. a bounded-depth mtime probe, which covers staging directories written
+       before the heartbeat existed, or killed before the first beat.
+
+    Raises ``OSError`` if the directory cannot be examined; the caller treats
+    that as "cannot tell" and leaves the directory alone.
+    """
+    heartbeat = path / volumegen.STAGING_HEARTBEAT_NAME
+    try:
+        if heartbeat.stat().st_mtime > cutoff:
+            return True
+    except OSError:
+        pass  # No heartbeat: fall back to the mtime probe.
+    return _has_recent_write(path, cutoff)
+
+
+def _cleanup_stale_staging_dirs(
+    base: Path, log, max_age: float = STALE_STAGING_AGE_SECONDS
+) -> None:
+    """Reclaim orphaned ``.tmp_*`` staging directories under one config base.
+
+    A generation killed by a walltime limit, an OOM, or a node failure leaves
+    its whole staging tree behind, and nothing ever removed it: every retry
+    stacked another (potentially multi-terabyte) copy under the same config_id.
+
+    Safety policy. Only directories that (a) live directly under *this*
+    config_id base, (b) carry the ``.tmp_`` prefix this module owns, and (c)
+    show no sign of life for ``max_age`` are removed. Published datasets and
+    anything outside ``base`` are never touched. Failures are logged and
+    ignored: cleanup is opportunistic and must never break the decision it runs
+    inside.
+
+    "No sign of life" is the delicate part, because the staging directory of a
+    *running* job is visible here (unique staging names mean two live jobs never
+    share a directory, but they do share the base). Age alone is not enough:
+    generation is not bounded by a day -- at the larger scales it is measured in
+    days -- and after the first minutes it writes only at depth >= 2, so the top
+    of the tree stops changing while the job is perfectly healthy. Judging by
+    the top-level mtimes alone therefore let a concurrent same-config start
+    rmtree a live generation out from under its peers. ``_staging_dir_is_live``
+    is the answer: an explicit heartbeat maintained by the writers, backed by a
+    bounded-depth mtime probe for directories that predate it.
+    """
+    now = time.time()
+    cutoff = now - max_age
+    for path in base.iterdir():
+        if not path.name.startswith(TMP_PREFIX) or not path.is_dir():
+            continue
+        try:
+            if _staging_dir_is_live(path, cutoff):
+                continue
+        except OSError as exc:
+            log.warning("Could not examine staging dir %s: %s", path, exc)
+            continue
+
+        log.info(
+            "Removing orphaned dataset staging dir %s (no write in the last "
+            "%.1f hours, and no live generation heartbeat)",
+            path,
+            max_age / 3600.0,
+        )
+        shutil.rmtree(path, ignore_errors=True)
+
+
+def _write_meta_atomic(meta_path: Path, meta: Dict[str, Any]) -> None:
+    """Write ``meta`` to ``meta_path`` atomically.
+
+    ``meta.yaml`` is what the loader reads to decide how every sample in the
+    dataset is interpreted, so a partially written one is worse than none at
+    all: a truncated file parses as empty and silently reclassifies a modern
+    dataset as legacy v1. The document is therefore written to a temp file in
+    the same directory, flushed and fsynced, and only then ``os.replace``d onto
+    the final name -- an atomic rename within one filesystem.
+    """
+    tmp_path = meta_path.parent / f".{meta_path.name}.tmp{os.getpid()}"
+    try:
+        with open(tmp_path, "w") as handle:
+            handle.write(yaml.safe_dump(meta, sort_keys=True, default_flow_style=False))
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(tmp_path, meta_path)
+    except BaseException:
+        # A failed write must not leave a temp file behind, and the final name
+        # must keep whatever complete document was already there.
+        try:
+            os.remove(tmp_path)
+        except OSError:
+            pass
+        raise
+
+
 def _decide_reuse_or_generate(
     base: Path,
     config_id: str,
@@ -123,15 +296,48 @@ def _decide_reuse_or_generate(
     staging and final paths for a new generation. Making this decision in one
     place and broadcasting it prevents ranks from diverging when their views of
     the shared filesystem differ.
+
+    The scan is deliberately forgiving: a candidate whose metadata is missing,
+    unreadable, or malformed is warned about and skipped rather than allowed to
+    raise. This function runs inside a window where every peer is already
+    waiting in the decision broadcast, so a crash here is a job-wide hang; a
+    poison directory (exactly what a killed job leaves behind) must never be
+    able to cause one.
     """
+    # Rank 0 is the only rank that touches this base, so this is also the one
+    # safe place to reclaim staging dirs orphaned by earlier killed jobs.
+    _cleanup_stale_staging_dirs(base, log)
+
     candidates = sorted(
         (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True
     )
     for dataset_path in candidates:
+        # Staging dirs are not datasets: a job killed between the meta write and
+        # the rename leaves a complete meta.yaml inside one, and reusing it hands
+        # back a partially generated (and cleanup-eligible) directory.
+        if dataset_path.name.startswith(TMP_PREFIX):
+            continue
         meta_path = dataset_path / META_FILENAME
         if not meta_path.exists():
             continue
-        meta = yaml.safe_load(meta_path.read_text())
+        try:
+            meta = yaml.safe_load(meta_path.read_text())
+        except Exception as exc:
+            log.warning(
+                "Skipping dataset candidate %s: unreadable %s (%s: %s)",
+                dataset_path,
+                META_FILENAME,
+                type(exc).__name__,
+                exc,
+            )
+            continue
+        if not isinstance(meta, dict):
+            log.warning(
+                "Skipping dataset candidate %s: %s is empty or malformed",
+                dataset_path,
+                META_FILENAME,
+            )
+            continue
         if meta.get("config_id") != config_id:
             continue
         if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION:
@@ -145,11 +351,30 @@ def _decide_reuse_or_generate(
     log.info("No valid existing dataset found at %s. Generating new dataset.", base)
     ts = time.strftime("%Y%m%d-%H%M%S")
     dest = base / f"{ts}__{commit}"
-    tmp = base / f".tmp_{ts}"
+    # The staging name must be unique per job: a bare 1-second-granularity
+    # timestamp let two same-config jobs starting in the same second collide on
+    # ``mkdir(exist_ok=False)``, killing one of them mid-consensus. Adding the
+    # pid and a random suffix makes the name unique even across nodes.
+    tmp = base / f"{TMP_PREFIX}{ts}_{os.getpid()}_{uuid.uuid4().hex[:8]}"
     tmp.mkdir(parents=True, exist_ok=False)
     return ("generate", str(tmp), str(dest))
 
 
+def _reraisable(exc: BaseException) -> BaseException | None:
+    """Return ``exc`` when the rank that caught it should re-raise it verbatim.
+
+    The consensus guards below turn any failure into a sentinel so every rank
+    aborts together, and every rank then raises a ``RuntimeError`` carrying the
+    collected messages. That is the right report for a genuine error -- and for
+    ``SystemExit``, which must never escape ``get_dataset`` (a peer would treat
+    the silent unwind as success). A ``KeyboardInterrupt`` is different: it is
+    not an error but an operator abort, so the rank that received it re-raises
+    it after posting the sentinel, keeping the interrupt's own semantics (and
+    exit status) while its peers still learn to stop.
+    """
+    return exc if isinstance(exc, KeyboardInterrupt) else None
+
+
 def get_dataset(
     config: Namespace,
     require_commit: bool = False,  # default: ignore commit mismatches for reuse
@@ -188,14 +413,38 @@ def get_dataset(
     # same branch. Scanning the shared filesystem independently per rank lets
     # divergent views (stale metadata caches, a racing job's rename) strand some
     # ranks in the generation collectives while others return early.
+    # Everything rank 0 does here happens while the peers are already blocked in
+    # the broadcast below, so a rank-0 exception would strand the whole job.
+    # Any failure is therefore turned into an error sentinel that travels
+    # through the same broadcast and makes every rank raise the same error.
+    interrupt = None
     if rank == 0:
-        decision = _decide_reuse_or_generate(
-            base, config_id, commit, require_commit, log
-        )
+        try:
+            decision = _decide_reuse_or_generate(
+                base, config_id, commit, require_commit, log
+            )
+        except BaseException as e:
+            # BaseException, not (Exception, SystemExit): a KeyboardInterrupt
+            # delivered to rank 0 alone (Ctrl-C on the launching terminal, a
+            # site watchdog SIGINT) would otherwise skip the broadcast and hang
+            # every peer -- the exact failure this guard exists to prevent.
+            decision = (
+                "error",
+                f"rank 0 failed to select a dataset under {base}: "
+                f"{type(e).__name__}: {e}",
+            )
+            interrupt = _reraisable(e)
     else:
         decision = None
     decision = comm.bcast(decision, root=0)
 
+    if decision[0] == "error":
+        # Rank 0 keeps the abort signal it was actually given; the peers, which
+        # only ever saw the sentinel, report it as a generation failure.
+        if interrupt is not None:
+            raise interrupt
+        raise RuntimeError(f"dataset selection failed: {decision[1]}")
+
     if decision[0] == "reuse":
         return Path(decision[1])
 
@@ -208,13 +457,15 @@ def get_dataset(
     err = ""
 
     # A worker failure must not skip any collective below: catch everything
-    # (including SystemExit, which is a BaseException and would otherwise bypass
-    # the consensus) so every rank always reaches the allreduce and gather.
+    # (BaseException, so neither SystemExit nor a KeyboardInterrupt delivered to
+    # one rank can bypass the consensus) so every rank always reaches the
+    # allreduce and gather.
     try:
         volumegen.main(config)
-    except (Exception, SystemExit) as e:
+    except BaseException as e:
         ok = False
         err = f"volumegen attempt failed: rank {rank}: {type(e).__name__}: {e}"
+        interrupt = _reraisable(e)
 
     # Reach a global verdict, then have every rank participate in the error
     # gather so no rank is left in a mismatched collective on the failure path.
@@ -226,25 +477,45 @@ def get_dataset(
             shutil.rmtree(tmp, ignore_errors=True)
         # Every rank raises with the collected messages, so a non-root rank
         # never returns an unfinalized dataset path.
+        if interrupt is not None:
+            raise interrupt
         msgs = "; ".join(e for e in errs if e)
         raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}")
 
     # rank 0 writes metadata into the staging dir, then renames it into place so
-    # readers never observe a half-written dataset.
+    # readers never observe a half-written dataset. This is another rank-0-only
+    # window inside a collective sequence: the rename can fail (a racing job
+    # already published this name, quota, ...), so the outcome is broadcast
+    # rather than allowed to kill rank 0 while the peers wait for it.
+    finalize_err = ""
     if rank == 0:
-        meta = {
-            "config_id": config_id,
-            "dataset_format_version": DATASET_FORMAT_VERSION,
-            "config_subset": volume_config,
-            "include_keys": INCLUDE_KEYS,
-            "code_commit": commit,
-            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
-        }
-        (tmp / META_FILENAME).write_text(
-            yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)
-        )
-        tmp.rename(dest)
+        try:
+            meta = {
+                "config_id": config_id,
+                "dataset_format_version": DATASET_FORMAT_VERSION,
+                "config_subset": volume_config,
+                "include_keys": INCLUDE_KEYS,
+                "code_commit": commit,
+                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            }
+            # The liveness marker described this directory while it was being
+            # written; it has no meaning in a published dataset.
+            (tmp / volumegen.STAGING_HEARTBEAT_NAME).unlink(missing_ok=True)
+            _write_meta_atomic(tmp / META_FILENAME, meta)
+            tmp.rename(dest)
+        except BaseException as e:
+            finalize_err = (
+                f"rank 0 failed to finalize dataset at {dest}: {type(e).__name__}: {e}"
+            )
+            interrupt = _reraisable(e)
+
+    # This broadcast doubles as the synchronization the old Barrier provided: no
+    # rank returns before rank 0 has published the rename (or reported that it
+    # could not), so nobody observes the staging path or a missing dataset.
+    finalize_err = comm.bcast(finalize_err, root=0)
+    if finalize_err:
+        if interrupt is not None:
+            raise interrupt
+        raise RuntimeError(f"dataset generation failed: {finalize_err}")
 
-    # ensure the rename is visible everywhere before returning
-    comm.Barrier()
     return dest
diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py
index e28d104..ea84a91 100644
--- a/ScaFFold/datagen/instance.py
+++ b/ScaFFold/datagen/instance.py
@@ -28,6 +28,7 @@
 import numpy as np
 from mpi4py import MPI
 
+from ScaFFold.datagen import layout
 from ScaFFold.datagen.generate_fractal_points import generate_fractal_points
 from ScaFFold.datagen.rng import derive_seed, seed_numba
 from ScaFFold.utils.config_utils import Config
@@ -239,12 +240,12 @@ def main(config: Config):
 
     log.info("MPI size = %s", size)
 
-    # Setup directories
-    fracts_sub_dir = f"var{config.variance_threshold}"
-    fracts_read_dir = os.path.join(config.fract_base_dir, fracts_sub_dir, "3DIFS_param")
-    instance_write_dir = os.path.join(
-        config.fract_base_dir, fracts_sub_dir, "instances", f"np{config.point_num}"
-    )
+    # Setup directories. The library is keyed by seed (see
+    # ScaFFold.datagen.layout): every instance is generated from
+    # (seed, category, instance), so resuming onto another seed's files would
+    # silently mix data from two different seeds into one dataset.
+    fracts_read_dir = layout.category_param_dir(config)
+    instance_write_dir = layout.instance_dir(config)
     if rank == 0:
         log.info(
             "Generating instances for num_points=%s, writing to %s",
@@ -345,14 +346,26 @@ def main(config: Config):
 
     start_time = time.time()
 
+    # One-entry cache of the most recently parsed category CSV. The work list
+    # is built category-major and block-sliced, so every rank's items for a
+    # given category are contiguous: a single entry is enough to turn the
+    # per-item re-parse (145 identical reads of the same small file off the
+    # shared filesystem, per category) into one parse per category per rank.
+    # ``generate_instance_points`` copies before scaling, so sharing the parsed
+    # array across instances cannot leak weights from one item into the next.
+    cached_category = None
+    params = None
+
     for i, category_instance_pair in enumerate(instances_to_generate_for_this_rank):
         category, instance = category_instance_pair
-        category_IFS_params = IFS_param_csv_names[category]
-        params = np.genfromtxt(
-            f"{fracts_read_dir}/{category_IFS_params}",
-            dtype=DEFAULT_NP_DTYPE,
-            delimiter=",",
-        )
+        if category != cached_category:
+            category_IFS_params = IFS_param_csv_names[category]
+            params = np.genfromtxt(
+                f"{fracts_read_dir}/{category_IFS_params}",
+                dtype=DEFAULT_NP_DTYPE,
+                delimiter=",",
+            )
+            cached_category = category
         weights = weights_all[instance]
 
         # Generate a validated, weighted point cloud. Weighting can turn a
diff --git a/ScaFFold/datagen/layout.py b/ScaFFold/datagen/layout.py
new file mode 100644
index 0000000..eff02f8
--- /dev/null
+++ b/ScaFFold/datagen/layout.py
@@ -0,0 +1,102 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""On-disk layout of the fractal library.
+
+Every artifact in the library is a deterministic function of the seed: category
+IFS parameters come from a ``(seed, rank, attempt)`` candidate stream, and each
+instance point cloud is generated from ``(seed, category, instance)``. Resume,
+by contrast, is a pure file-existence test -- an instance is "already done" if
+its file is on disk.
+
+Those two facts only compose safely if the path itself carries the seed.
+Without it, a run under a new seed found the previous seed's files, generated
+nothing, and produced a dataset whose metadata advertised the new seed while
+its contents came from the old one. Keying the directory by seed makes the
+existence question seed-specific, so data from one seed can never be mistaken
+for another's:
+
+    <fract_base_dir>/var<variance_threshold>/seed<seed>/3DIFS_param/
+    <fract_base_dir>/var<variance_threshold>/seed<seed>/instances/np<point_num>/
+
+Libraries written under the older, seed-agnostic layout are simply not found
+and are regenerated in the new location.
+
+These helpers are the single definition of that layout; every producer and
+consumer (``category_search``, ``instance``, ``volumegen``) goes through them
+so the two sides cannot drift apart.
+"""
+
+from __future__ import annotations
+
+import os
+
+
+def library_root(config) -> str:
+    """Return the root of the fractal library for this config's seed."""
+    return os.path.join(
+        str(config.fract_base_dir),
+        f"var{config.variance_threshold}",
+        f"seed{int(config.seed)}",
+    )
+
+
+def category_param_dir(config) -> str:
+    """Return the directory holding this seed's category IFS parameter CSVs."""
+    return os.path.join(library_root(config), "3DIFS_param")
+
+
+def legacy_category_param_dir(config) -> str:
+    """Return where the category CSVs lived before the layout was seed-keyed."""
+    return os.path.join(
+        str(config.fract_base_dir),
+        f"var{config.variance_threshold}",
+        "3DIFS_param",
+    )
+
+
+def warn_if_legacy_library(config, log) -> bool:
+    """Warn when a library in the old, seed-agnostic layout is being ignored.
+
+    The relayout is deliberately silent about old data -- an existence check in
+    the seed-keyed location simply does not find it -- which from the outside
+    looks like a library that was there yesterday being regenerated for no
+    reason (at large scales, hours of work). One line naming both directories
+    turns that into an explained, expected event. Returns whether the old
+    layout was present, so callers can test the condition directly.
+    """
+    legacy = legacy_category_param_dir(config)
+    if not os.path.isdir(legacy):
+        return False
+    log.warning(
+        "Found a fractal library in the old, seed-agnostic layout at %s. "
+        "Libraries are now keyed by seed, so this one cannot be reused (a run "
+        "under a different seed would silently adopt another seed's data) and "
+        "the categories for seed %s will be generated at %s. Delete the old "
+        "directory once you no longer need it.",
+        legacy,
+        int(config.seed),
+        category_param_dir(config),
+    )
+    return True
+
+
+def instance_dir(config) -> str:
+    """Return the directory holding this seed's instance point clouds.
+
+    Instances are additionally keyed by point count, which is a property of the
+    cloud rather than of the category, so several point counts can coexist for
+    one seed.
+    """
+    return os.path.join(library_root(config), "instances", f"np{config.point_num}")
diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py
index cfd6d1b..997a73f 100644
--- a/ScaFFold/datagen/volumegen.py
+++ b/ScaFFold/datagen/volumegen.py
@@ -22,10 +22,60 @@
 import numpy as np
 from mpi4py import MPI
 
+from ScaFFold.datagen import layout
 from ScaFFold.utils.config_utils import Config
 from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE
 from ScaFFold.utils.utils import setup_mpi_logger
 
+# Liveness marker for the directory being generated into. Volume writing is the
+# long phase of a generation and it happens deep inside the tree
+# (``volumes/<split>/N.npy``), so the top of the staging directory can look
+# untouched for many hours while the job is perfectly healthy. Every writing
+# rank therefore refreshes this file periodically, and
+# ``get_dataset._staging_dir_is_live`` reads it instead of trying to infer
+# liveness from mtimes it cannot cheaply see. The name is owned here, next to
+# the writer; ``get_dataset`` (which already imports this module) reads it from
+# here so the two sides cannot drift.
+STAGING_HEARTBEAT_NAME = ".heartbeat"
+# Refresh interval. Small enough to be negligible against the staleness
+# threshold (a day), large enough that it is one utime per rank per few minutes
+# no matter how fast volumes are written.
+STAGING_HEARTBEAT_INTERVAL_SECONDS = 5 * 60
+
+
+class StagingHeartbeat:
+    """Periodically touch a staging directory's heartbeat file.
+
+    ``beat()`` is called from the volume loop and is a no-op until the interval
+    has elapsed, so it costs one comparison per volume. Every rank writing into
+    the directory beats the same file: the marker means "somebody is still
+    working here", and a last-writer-wins utime is exactly the semantics wanted.
+    Failures are swallowed -- a heartbeat that cannot be written must never take
+    down a generation that is otherwise fine (the cleanup's second signal, the
+    bounded mtime probe, still applies).
+    """
+
+    def __init__(
+        self, staging_dir, interval: float = STAGING_HEARTBEAT_INTERVAL_SECONDS
+    ) -> None:
+        self.path = os.path.join(str(staging_dir), STAGING_HEARTBEAT_NAME)
+        self.interval = interval
+        self._last_beat = float("-inf")
+
+    def beat(self, now: float | None = None) -> bool:
+        """Touch the marker if the interval has elapsed; return whether it did."""
+        now = time.time() if now is None else now
+        if now - self._last_beat < self.interval:
+            return False
+        self._last_beat = now
+        try:
+            with open(self.path, "a"):
+                pass
+            os.utime(self.path, (now, now))
+        except OSError:
+            return False
+        return True
+
 
 def load_np_ptcloud(path: str) -> np.ndarray:
     """
@@ -81,9 +131,15 @@ def points_to_voxel_indices(
     scaled = (points - mins) / voxel_size
 
     # 4) Center the occupied region: the largest axis fills the grid while the
-    #    shorter axes are offset so their span sits in the middle.
+    #    shorter axes are offset so their span sits in the middle. The free
+    #    space to split between the two margins is (grid_size - span) voxels,
+    #    measured in the same voxel units as ``scaled``; subtracting an extra 1
+    #    (as if the offset were an index rather than a length) shifted every
+    #    cloud half a voxel toward the origin, floored the first half-voxel of
+    #    each filled axis to -1, and let the clip below fold those points onto
+    #    plane 0.
     span = scaled.max(axis=0)
-    offset = (grid_size - 1 - span) / 2.0
+    offset = (grid_size - span) / 2.0
     idx = np.floor(scaled + offset).astype(int)
 
     # 5) Clip to valid range (guards float rounding at the boundaries).
@@ -222,6 +278,7 @@ def main(config: Dict):
     # the failure is then propagated to all ranks via an allreduce.
     ok = True
     err = ""
+    interrupt = None
 
     try:
         if start_idx >= end_idx:
@@ -240,11 +297,21 @@ def main(config: Dict):
             fractal_colors = np.random.rand(config.n_categories, 3)
 
             grid_size = resolve_grid_size(config)
-            fract_base_dir = str(config.fract_base_dir)
-
-            # Generation loop
+            # The instance library is keyed by seed (see ScaFFold.datagen
+            # .layout), so a volume can only ever be built from point clouds
+            # this run's seed produced. Resolved once, outside the loop.
+            instances_dir = layout.instance_dir(config)
+
+            # Generation loop. Every rank reports that this staging directory
+            # is still being written to, so a concurrent job's orphan cleanup
+            # can tell a live multi-hour generation from one killed a day ago
+            # (the volumes themselves land two levels down, where a cheap
+            # top-level mtime probe cannot see them).
+            heartbeat = StagingHeartbeat(dataset_dir)
+            heartbeat.beat()
             start_time = time.time()
             for i, curr_vol in enumerate(volumes_contents_subset):
+                heartbeat.beat()
                 if i % 10 == 0:
                     log.debug("Rank %s processing local volume %s", rank, i)
 
@@ -269,12 +336,7 @@ def main(config: Dict):
                     curr_instance = curr_vol[1 + 2 * curr_fract + 1]
                     fractal_color = fractal_colors[curr_category]
 
-                    instances_dir = (
-                        f"var{config.variance_threshold}/instances/np{config.point_num}"
-                    )
-
                     point_cloud_path = os.path.join(
-                        fract_base_dir,
                         instances_dir,
                         f"{curr_category:06d}",
                         f"{curr_category:06d}_{curr_instance:04d}.npy",
@@ -332,11 +394,14 @@ def main(config: Dict):
                     total_time,
                     len(volumes_contents_subset) / total_time,
                 )
-    except (Exception, SystemExit) as e:
+    except BaseException as e:
         # Capture the failure locally instead of letting it unwind past the
-        # collective below, which would desynchronize the ranks.
+        # collective below, which would desynchronize the ranks. BaseException,
+        # not (Exception, SystemExit): a KeyboardInterrupt delivered to one rank
+        # would otherwise skip the consensus and hang the others.
         ok = False
         err = f"rank {rank}: {type(e).__name__}: {e}"
+        interrupt = e if isinstance(e, KeyboardInterrupt) else None
 
     # Consensus on the generation status. This replaces a bare Barrier: every
     # rank always executes exactly this collective (regardless of success or
@@ -345,6 +410,10 @@ def main(config: Dict):
     all_ok = comm.allreduce(1 if ok else 0, op=MPI.MIN) == 1
     errs = comm.allgather(err)
     if not all_ok:
+        # The interrupted rank re-raises the operator's abort verbatim; the
+        # others report the gathered failure.
+        if interrupt is not None:
+            raise interrupt
         msgs = "; ".join(e for e in errs if e)
         raise RuntimeError(f"volume generation failed: {msgs or 'unknown error'}")
 
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv
deleted file mode 100644
index 824252c..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.113880658597081297e-01,1.525788435216699490e-01,-1.677513246653659085e-01,9.450327943163072675e-01,1.437558218378187647e-01,-7.781853097899118499e-01,-7.520371219608632529e-01,4.049719188683833515e-01,3.599506321979530910e-01,-9.209462080360120151e-01,9.797064915552313735e-01,2.046581090479335785e-01,5.952758981070249700e-01
--3.273838635851094025e-01,5.796182968623042608e-01,2.204038860996404559e-02,7.254736836236364006e-02,-3.624484794849551772e-01,3.976647887245436941e-01,-6.272009791693362590e-01,2.300328757717320372e-01,1.439593364137665699e-01,-1.758502593004573900e-01,7.577588977727744979e-01,-4.555583614631009137e-01,4.047241018929750300e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv
deleted file mode 100644
index fca2ed3..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-2.473679574324576524e-01,1.597856635526939684e-02,-1.261210131422938474e-01,2.860947005456315750e-01,8.355771386484407426e-01,-1.930246242945548030e-01,-4.801705885775757743e-01,1.904275728954856195e-01,7.806079367730189844e-01,7.953170851152218113e-01,-5.816986628081313171e-01,6.209542113667643193e-01,1.988928060480311955e-01
--7.041010339743891677e-02,-2.454306473738818717e-01,9.405874516413117448e-01,7.268593945019246050e-01,6.070031387640555387e-01,-9.102525915665182765e-02,6.339356801972018118e-01,-1.532413735258475462e-01,3.654495384649214529e-02,-4.532769270980656628e-01,1.964159205402833397e-01,-5.359870555644301593e-01,8.011071939519688323e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv
deleted file mode 100644
index c3e5385..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--9.413772084350502389e-02,5.254718465280536766e-01,-4.808322979877577286e-01,5.811542374432154823e-01,-4.454009637949891687e-01,-4.249149866229446904e-01,-3.000182472995254201e-01,4.022670562303001240e-01,-2.934854435005793682e-01,-6.605442110513501941e-01,-6.253032129726847632e-01,-5.719885961970654353e-01,4.224941885002395647e-01
--4.407723055180619021e-01,1.325088026302621014e-01,1.840254341100278079e-01,-5.508321197309193895e-01,8.464500652955120330e-01,5.041753392202226181e-01,5.651656783277005935e-01,-8.325090493066080732e-01,-1.381507388162375172e-01,5.571341416171868843e-01,-1.025781698153782617e-01,-8.838988798440101657e-01,5.775058114997604353e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv
deleted file mode 100644
index 18274fd..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--1.294466006455852192e-01,1.974941747383867074e-01,4.427934740152306148e-01,2.336055684661912935e-01,1.300312996679606758e-01,8.725504616000674396e-01,-1.902059536103457571e-01,-3.532649574684778582e-01,-8.430890597933748953e-01,-9.596612485558557726e-01,6.321097711736041180e-01,8.135413458612708038e-01,2.489526692860659640e-01
--1.399185307891788188e-01,-1.602045901301751840e-01,4.335897375838684287e-01,-9.384258169419115170e-01,-1.069737283930594085e-01,2.989241164297768982e-01,-6.355754104429354179e-01,2.746690226736152596e-01,2.732139814068013095e-01,9.765182437472603727e-01,-4.228379364440870702e-02,-5.170594086561124403e-01,7.510473307139340360e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv
deleted file mode 100644
index ac107cb..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--7.074987909298180444e-01,-1.415561527827760013e-01,6.453558418406579733e-02,9.717892764870739164e-01,4.692315433286387005e-02,-9.087912295403599572e-02,-7.306123549338172651e-01,5.199024492355210914e-01,3.242137268086431323e-01,-7.685133956354921470e-01,-2.890512330200536439e-01,-4.414459117103588515e-01,1.767437117859010087e-01
--5.783609543921015561e-01,-6.261288719190640784e-01,6.083280289156212106e-01,5.708229102420392387e-01,6.839066129793194282e-01,-3.730802596414783956e-01,4.320233943214566441e-01,-2.520497672481589735e-01,2.188189615840767654e-01,-5.142628557025761271e-01,9.167783772484443539e-01,-9.404699392000455127e-01,8.232562882140990190e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv
deleted file mode 100644
index 898a388..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--6.016222179722698904e-01,-8.568626391040801149e-01,4.339904630487174675e-01,8.566190336460910437e-02,5.955551394131912701e-01,-7.233426894750305536e-02,-3.439913328223502820e-01,-1.295598008765910247e-01,5.054214022018066466e-01,5.403179943497680160e-01,9.887300162494017108e-01,5.065235980037172681e-01,2.036756775286962806e-01
-4.255811510519635910e-02,-5.861700389763935259e-01,1.539188938105566784e-01,7.464882854730132689e-01,9.985987594821947866e-01,-3.852978585836623893e-01,5.848150331480403974e-01,-8.218568908737176049e-01,-4.786801609318307449e-01,-7.997867917133967275e-01,-3.226694643578802424e-01,-9.055189731557820032e-01,7.963243224713036916e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv
deleted file mode 100644
index 345f2a4..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-4.449557690511229957e-02,1.859345105226184458e-02,-5.492709805712474580e-01,5.002339770143411357e-01,6.022059536539468017e-01,-2.581455709557001210e-01,5.384026026512653829e-01,3.525173161797230392e-01,4.535819387452528773e-01,5.054013100483643051e-01,5.534237431467781132e-01,6.486536701245786407e-01,5.893263018191835512e-01
--4.935786839220561717e-01,-3.483405932567602559e-01,-4.920557235894853498e-01,-3.225857219554804090e-01,-5.369237708506735540e-01,-4.285587807484951828e-01,-2.187430959790792606e-02,8.181072307028691704e-02,4.255774677072714507e-01,-8.957260803321442921e-01,-6.285123937957712847e-01,-7.274955330956049959e-01,4.106736981808163378e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv
deleted file mode 100644
index fa04c8c..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-4.567482715684858530e-01,-2.162812897008825619e-01,8.688925426344686898e-01,9.403035945928239769e-01,-4.923377969407702892e-01,9.039413044894315519e-01,-7.007586459777630505e-01,9.958966466641376858e-01,-2.148840978447630334e-02,-8.820192395688601916e-01,-6.962417306340771272e-01,6.989108727983952551e-01,6.175732443919085268e-01
--1.311082673391426034e-01,2.651764139972054846e-01,2.416531271291553207e-01,2.658826889156997719e-01,5.321835857677645887e-01,8.850640913149898648e-01,-4.146324561488827776e-01,6.016104083259397051e-01,-6.006142415160782289e-01,6.902838974364782221e-01,5.080744400285828188e-01,-5.543689403825666773e-01,3.824267556080915287e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv
deleted file mode 100644
index b017040..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv
+++ /dev/null
@@ -1,3 +0,0 @@
-1.955211818468614027e-01,-1.860012909088504252e-02,1.679370276531577666e-01,-2.474658865690229081e-01,2.002748188638228122e-01,-2.296462069177727106e-01,4.272668474314074150e-01,5.518250090890584048e-01,-1.814099098142296640e-01,-7.759223213118924267e-01,2.396552032441381375e-01,-9.984320127889498853e-01,3.158600947042149998e-01
-5.269496943889822038e-01,5.090567464178141766e-01,3.638063575297016961e-01,-9.409697282110314198e-01,-8.838575707448308449e-01,-9.681971946347656122e-02,-6.464087288109372498e-01,-4.699168427008522109e-01,4.879093108230523335e-01,7.225552560446799610e-01,5.544701020084910059e-01,6.687273760328038552e-01,6.074177711483321751e-01
-3.702243567552572223e-01,2.665339551565091281e-01,-7.302570963874586152e-01,7.128226896410159164e-01,-1.412872752126193010e-01,-5.914873767070649713e-01,4.858707230431502655e-01,2.870662392522644879e-01,-8.632888349796041805e-01,8.908557359947888443e-02,-1.967254792508705830e-01,-6.141181327493672182e-01,7.672213414745278348e-02
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv
deleted file mode 100644
index f56b11b..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-2.948060069505098468e-01,6.887599776806021534e-01,7.611229734423119453e-02,4.932702124290413437e-01,-4.882699298136421451e-01,-3.811733224774245254e-01,7.278710433957895631e-01,6.852985410340497463e-01,-7.967518041925278904e-01,6.365195144826076845e-01,7.940485085563173673e-01,-7.048888688479171272e-02,6.538648345727335887e-01
-2.849341248173462571e-01,-5.989665908310131126e-01,-8.440934355375118159e-01,-5.676180605895972953e-01,2.243341626806083511e-02,-4.393691291202828086e-01,-6.561728033554823369e-01,-1.447073895277080080e-01,-3.027716266953806024e-01,-1.488449066054653436e-01,-7.989435645608391479e-01,6.936025509345438156e-01,3.461351654272663003e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv
deleted file mode 100644
index af8b628..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--2.352521358159322951e-01,-8.804830225066977434e-01,2.445692685925464627e-01,-3.804268259383307704e-01,7.002478757426469080e-01,4.949375238071500593e-01,6.340824702532006363e-01,-1.661009284036518707e-01,6.403120148905094844e-02,-4.544636202226781663e-03,-5.614632388117013484e-01,-8.225067911333789894e-01,7.543884465970274178e-01
--6.205195826406055826e-01,2.082688392648848197e-01,7.051714514820610624e-01,6.615150945575909436e-01,-5.430135870340266901e-01,-7.590298999219038389e-01,-2.484973171764268685e-01,4.507062852245327100e-01,-3.971212226677103274e-01,3.603919868974059249e-01,-6.455363862654992513e-01,6.976679002200263380e-01,2.456115534029726655e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv
deleted file mode 100644
index 60713bc..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--2.647713297675924338e-01,4.532537039449799909e-01,9.126597053670870707e-02,6.297942569420966752e-01,-7.970366903049754814e-01,-5.805081819440571778e-01,1.847828454797417752e-01,1.632357583131875955e-01,1.112052944356816120e-01,5.581182508420681199e-01,-1.235235401936412014e-01,5.525957004344514978e-01,3.917940489256083736e-01
--3.328722698643555855e-01,-6.568958305252938779e-01,9.548663670929014025e-02,2.135325385282085264e-01,-7.994872824743104456e-01,3.120617004916548254e-01,-8.149604985800276147e-01,3.155236614340959367e-01,-5.804561978693489888e-01,-3.465348059373873912e-01,8.779303525751596116e-01,-4.964138499509700431e-01,6.082059510743916819e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv
deleted file mode 100644
index d75a5f8..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.732047938669280640e-01,7.457174836228177561e-01,-1.134959963046562326e-01,-2.524155789085107404e-01,6.491113757559063835e-01,5.923873502246719269e-01,-4.163840681869901417e-01,-9.955960643681671662e-02,-4.399149502577506254e-01,-3.298124569267795181e-01,-4.480024355377016931e-01,4.831814457889009873e-01,7.022590685295611035e-01
-3.640511591326804908e-01,-8.138787463613117446e-01,6.945993218897195121e-01,1.561110571982835538e-01,1.420189997595979747e-01,8.921447905553927527e-01,8.983211135903812483e-02,-7.873694636782158085e-02,-1.924595656350303052e-01,3.444934594172788245e-01,1.245560132691707622e-01,-2.441271125304933509e-01,2.977409314704388410e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv
deleted file mode 100644
index bf897a7..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--9.477179726636597579e-02,1.411777244118233021e-01,-3.210166902811444345e-01,8.538383221528667022e-02,5.145040185660725296e-01,5.449889235641092178e-01,1.961008681633682471e-01,5.250846593606550705e-01,2.072609251252632845e-01,8.556312962216905404e-01,7.079983983577677886e-01,-5.970939704307753892e-01,1.986178895167198533e-01
-2.340421417737004184e-02,-2.331838675905175684e-01,8.789698166911528165e-01,4.337374001870073492e-01,-6.744896384287986102e-01,6.995573996808193140e-01,-9.371564747554361752e-02,7.122203806164681961e-01,-3.226197608587253463e-01,5.228286702908606642e-01,6.219275677253508494e-01,6.982539942183398907e-01,8.013821104832802300e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv
deleted file mode 100644
index 6fd2a78..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--5.079619751044526232e-01,-8.448208339950418200e-01,3.379361696842797524e-01,-7.424757815754379209e-01,-2.914860368679528246e-01,7.863175574507479393e-01,-3.913127807203062858e-01,-3.395110295018073376e-01,1.729246237110684259e-01,1.155790784528218929e-01,-1.778336752766045414e-01,-8.237750280597011532e-01,8.128837685573842009e-01
--5.053085469842428790e-01,3.413777808104254685e-01,5.072078013440961541e-02,-2.297771208541388166e-01,1.179085091368414773e-01,-5.353419781682577927e-02,3.033551256644697602e-01,-8.399081644307380135e-01,-2.580876839431038849e-01,-2.915499965121404191e-01,-8.274711651504977894e-01,-3.831424144957116251e-01,1.871162314426157436e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv
deleted file mode 100644
index 040ab5e..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-1.959234370454023821e-01,7.195879279286905295e-01,-8.847434444070787496e-01,-2.025052870905292846e-01,-6.099125369957605347e-01,4.364854998130400787e-01,2.605566328535515730e-01,6.345444901137731186e-02,2.845477814208994261e-01,3.682037052176181380e-01,-5.482600452804842206e-01,8.228596647488930493e-01,2.413865640087159981e-01
--3.105662416993948405e-02,1.541091773938909615e-01,9.832144790547923119e-01,2.617717287788325908e-01,2.964779787430120717e-01,2.867701916196363499e-01,-1.595092390377108593e-01,5.184743556024293820e-01,7.055853027879592787e-01,-2.253650882070545869e-02,3.223158737418712061e-01,-3.130584359547923246e-01,7.586134359912839464e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv
deleted file mode 100644
index 1262706..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--1.300116634675734240e-01,1.252286096526100678e-01,2.499311708856801761e-01,7.830409313847326302e-01,-3.485355694096761159e-01,-2.040759407213834642e-01,3.636277174597402073e-01,-9.857803444259161108e-01,4.953546011072127442e-01,-9.229964481663586184e-01,8.261855069165144894e-01,7.432769447647864514e-01,5.045106210869472196e-01
--9.419389020480599672e-01,-2.909144464059969515e-01,2.371749371359630487e-01,-9.140320234271737121e-02,2.578548696646267846e-01,-6.247044474540388581e-01,5.120075657145519710e-01,7.003880575966952016e-01,-7.351851386341190508e-01,-3.580478535291309328e-02,-3.983862417971921754e-01,-8.667134200938420019e-01,4.954893789130526693e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv
deleted file mode 100644
index 0a04245..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.510560689158552794e-01,-2.754404723500034624e-01,-2.684881793088436108e-01,-1.948162421580679204e-01,1.792154998944521793e-01,-3.924494838945669084e-01,-5.433151157278728327e-01,3.050183437518030338e-01,-5.297884663705332287e-01,-4.402434005450059917e-01,3.033450134513224761e-01,5.080005533136775497e-01,1.765071572885704987e-01
--3.296381408504345245e-02,8.213734474661573692e-01,2.884182165116890850e-01,-5.803866935447996589e-01,3.041849216693381930e-01,5.667920446859808781e-01,2.250925284276168448e-01,8.236273827714153395e-01,-4.114514278581935525e-01,2.460381915048910351e-01,-1.767376131099762659e-01,-3.058606496513562867e-01,8.234928427114294180e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv
deleted file mode 100644
index bca2acf..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--5.517066569583584101e-01,6.457030079067409556e-01,1.312452431616923931e-01,-9.654657530459793691e-01,1.001632917413284307e-01,4.808171457391481329e-01,2.887487385246092497e-01,9.356047048047881898e-01,1.042754743685525565e-01,-5.230136305881964986e-01,7.032043120523423507e-01,8.023204324713959501e-01,7.872166104195181813e-01
--5.753978519462528141e-01,7.086092336367590949e-01,8.862996973827064195e-01,-4.508828284143380216e-01,-1.498953015148070111e-01,-1.479318170320651493e-01,-3.642550894021263641e-01,1.915439725996275211e-01,7.097594434528575746e-02,4.751618882704935487e-01,3.779955724380190674e-02,2.814613891706765347e-01,2.127833895804817355e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv
deleted file mode 100644
index a21ea50..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-1.016398917604210972e-01,5.367241283642387728e-01,-6.818635487683848417e-01,-6.492151914242854094e-01,2.951107976835281033e-01,1.210983103608658240e-01,-1.914729327636166545e-01,2.084900726109191194e-01,-9.913674641194303305e-03,-1.318795603716438336e-01,-4.591725358362441778e-01,3.168885180251357347e-02,2.613304272286978147e-01
--2.518998374757852599e-01,-1.545791580965771850e-01,2.591599328760181287e-01,-9.227238644730679784e-01,-6.523143964998687760e-01,1.325715926486024099e-01,-2.634595670714392490e-01,-6.084513786591814188e-01,5.945406152909971098e-01,-7.413295649974527279e-01,-9.836113191580013737e-01,-6.613550934121019687e-01,7.386695727713020743e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv
deleted file mode 100644
index a24562e..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-6.921866875646716100e-01,5.108296431449499408e-03,-8.297196043425478784e-01,-8.782207198751741384e-01,4.881693383540661735e-01,7.602399262614862874e-01,7.929628517423199519e-01,2.747828187373047015e-01,-8.565292578863379358e-01,-8.028470318278093654e-01,-7.666729973480264082e-01,-3.618054675021058486e-01,6.474809588861711873e-01
--4.195656375439105190e-02,3.829378217234820081e-01,-4.738148839593532280e-01,-1.013350142767268647e-01,6.518052790254404982e-01,-2.989039824402659473e-01,3.810369849047898771e-01,-4.761398937049154956e-01,-8.736647575845664093e-01,8.671975808320215862e-01,4.518295413402169114e-01,3.044248270515332866e-01,3.525190411138288682e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv
deleted file mode 100644
index d861c0d..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-1.667868049973333378e-01,-9.628664509718032338e-01,6.014400922608953426e-01,2.972001483931350219e-01,4.941659040388075574e-01,-6.774253220554804500e-01,-1.993299704297657460e-01,6.394927405908403806e-01,-1.095498513905281968e-01,1.797181927692330650e-01,8.966277273316463070e-01,2.586404039871941229e-01,5.552709887608781036e-01
-8.904906062660500332e-01,3.456353981853357293e-01,-5.089078336971060157e-01,5.564474698994368307e-01,1.613474062509465679e-01,-4.536022710624829646e-01,7.605093590363174449e-01,-3.094273061226622268e-01,-6.937981682429052999e-01,8.031386423573720901e-01,-7.976443746091934628e-01,6.529445161208398130e-01,4.447290112391218964e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv
deleted file mode 100644
index d9cf9f6..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--5.460898717556510462e-01,5.858062690874137335e-01,2.194226187473129475e-01,9.188630804699871035e-02,-1.707137636423472493e-01,-3.178585699238225537e-01,3.633095509488837305e-01,-6.887805288505008949e-01,-1.468860672171268256e-02,-2.862987915279111562e-01,-7.723918312056159419e-01,8.126650586523440634e-01,4.458798622363638331e-01
-5.547486983713716402e-01,-5.675784088131794469e-01,-6.187415038981081139e-02,6.835868811760836827e-02,7.617313312413553916e-03,5.521100702462113929e-01,5.867352194847055280e-01,-8.129437845772602422e-01,-1.254568018619570680e-01,8.382074387221385425e-01,9.974232510580696154e-01,1.902631261807281593e-01,5.541201377636362224e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv
deleted file mode 100644
index c1d0e5a..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv
+++ /dev/null
@@ -1,3 +0,0 @@
-2.812878692303713013e-01,-5.380780165501937162e-01,-1.995746226435606285e-01,3.556609488625348536e-01,-3.570058976622436653e-02,3.226211706468127272e-01,7.604367324861782684e-02,-6.726022941928089249e-01,-1.244510164772598682e-01,-5.242674835142189238e-01,6.111367647898253708e-01,5.446881171705224567e-01,1.870729455165833777e-01
--1.804572792799459258e-01,-5.615234318506530098e-01,-5.678667666618428811e-01,-1.973882851907129421e-01,-9.342817705045243226e-01,-2.696256933877165807e-01,-4.754476132734983818e-01,8.379911380637130591e-01,-1.430754625396699620e-01,6.703727340713743210e-01,-7.950803881856465249e-01,1.646446229560936114e-01,5.812129149706145581e-01
--7.303265877923637017e-01,5.163190035467037919e-01,6.970165582377609859e-01,-2.277609099967679018e-01,5.773156478708756367e-01,7.154948998404351279e-01,-4.402130399012762485e-01,-5.815862404831599886e-01,-9.434039458871834594e-01,-3.147493553756923745e-01,-6.215697199671859074e-01,1.823798551890962738e-01,2.317141395128020642e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv
deleted file mode 100644
index 3eb8c33..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-2.166769833429884606e-01,-2.331594788641218052e-01,-8.006870701326018747e-01,-5.266093172746755258e-01,-8.322914134913979023e-01,7.446673507975531958e-01,-2.886575388105372397e-02,-2.990935453031380309e-01,3.936163556018033027e-01,-3.542536932837532238e-01,2.834018169147607402e-01,-9.586592464730825380e-01,4.259584301232757220e-01
-8.416839786286203218e-01,-5.562313049662526154e-02,-3.647897665777326548e-01,-6.475376571556650251e-01,-3.655643797640151238e-01,-3.196437920194410420e-01,5.442135586745493470e-01,5.508198826067678411e-03,-8.528986112361351957e-01,-6.533886080407793617e-01,-8.277524479526696677e-01,-6.784479736983972664e-01,5.740415698767242780e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv
deleted file mode 100644
index 4f7835f..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.328441505555097546e-01,-1.387237543865100786e-01,-8.385785443603737122e-01,-6.460892635786552596e-01,-6.206316056012477489e-01,-2.482763796893627806e-02,8.854727682161622759e-02,-3.627820436506015156e-01,-2.149069615655596621e-02,-8.638772932614158240e-01,-9.940866772762460002e-01,7.740573663063279319e-01,4.901549186583971096e-01
--5.533853946343669783e-02,6.783593497875153311e-01,-3.470947229869234540e-01,-8.081860602696091522e-01,-2.234790784515090500e-01,5.246894583705352666e-02,4.184989961414606885e-01,5.085439272585972059e-01,2.229705198326268345e-01,-2.452705816636990832e-02,-3.598764438670341015e-01,-4.548214145779025941e-01,5.098450813416028904e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv
deleted file mode 100644
index cf18c74..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-9.215997820884049840e-01,-8.886305956631088687e-01,7.918854213071213621e-02,8.549195867901333568e-01,-7.539338040902159310e-01,2.127998942252871117e-01,9.334583063583679063e-02,2.117336123070212572e-01,-3.905267778142720303e-01,4.922987444042437044e-01,3.538552884251142672e-01,-4.995810348035092385e-01,3.171922691986440168e-01
--2.710991105815863111e-01,6.963869144978018788e-02,-2.614443785391928898e-01,-1.745302413226517135e-01,2.582270480554498260e-01,-7.700432645988855018e-01,-9.312241172350603780e-01,5.326596688321914019e-01,6.785270353740011640e-01,-6.480593296020313865e-01,-5.862245859094883382e-01,-8.798502128158780522e-02,6.828077308013559277e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv
deleted file mode 100644
index 5645631..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.363859180192507736e-01,-2.005356922417877996e-02,4.105299683991425752e-03,4.520536430696797670e-01,-2.397748696508128496e-01,-9.149591317297329773e-01,-2.687516636255817826e-01,2.502854147579025579e-02,-5.658580150798908637e-01,-3.444933603162623204e-01,3.980110762078155062e-01,6.019507953845388837e-01,2.306736239861526538e-01
--9.924270668037067367e-01,1.880015741363572079e-01,7.482014409903954277e-01,-8.717841019029326510e-01,3.178806298249037265e-01,4.369798231563537527e-02,-4.363922340879544670e-01,-3.301960725800534568e-01,5.876506713161555595e-01,-6.176395243356269660e-02,-3.698941006853062596e-01,-6.461432233004753556e-01,7.693263760138474572e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv
deleted file mode 100644
index 3901f6d..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--9.488858097444059503e-01,4.079256280033416449e-01,-7.273191618538961123e-01,-2.164528040135893505e-02,1.801964698567355416e-01,2.110643357786530228e-01,-6.481561009612546442e-02,-2.415201336393544285e-01,8.209213192744986287e-01,1.342520688379920113e-01,6.654275517850556376e-01,4.674644534794492046e-02,2.727174515041750347e-01
-5.889383911906385105e-01,-4.716661362499952048e-01,8.398891628230311657e-01,9.444296451726641450e-01,6.852416895301041144e-01,-2.718291944059443299e-01,9.759285577328178363e-01,4.221736051403957024e-01,-5.860661824356676597e-01,6.789765242522334265e-01,-6.864599914320155261e-01,8.076484878207652596e-01,7.272825484958250764e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv
deleted file mode 100644
index 7f0e048..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.156004044367652916e-01,4.539226489835086475e-01,3.386522834986691599e-01,-6.897275360788317489e-01,-6.389720151715840846e-01,-3.773349166157233814e-01,1.813606744935740700e-01,-6.205184082106425247e-01,-3.558681169774229325e-01,3.970835926210039002e-01,-8.915932500008973971e-01,-7.113662644290583703e-01,2.064773053549207871e-01
--2.222304390465441593e-01,2.761428589612251461e-01,-4.556405012454298742e-01,1.006421622526911808e-01,1.163138470867763896e-01,4.306829086925703098e-01,7.112480073546407766e-01,8.255028592978328472e-01,-9.373835503173011396e-02,-7.714024572181317208e-01,9.021782591265152806e-01,-1.140412270850570398e-01,7.935226946450792962e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv
deleted file mode 100644
index 76e9192..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.054766017168615289e-01,3.285035685711270581e-01,3.339186103540698891e-01,8.138242613802961767e-01,-5.205905000696835483e-01,1.529228105404809579e-01,8.466677815415111219e-01,-7.061057292640156025e-01,-1.216890214920127722e-01,-1.716434030659630405e-01,-3.234074144154277519e-01,-5.208430114476063633e-01,5.154903134260301334e-01
-8.269501238092979989e-02,9.228403821803805585e-01,-9.202930276956799993e-01,-1.265069840482921926e-01,2.493516975241008016e-01,-3.998067562089782090e-01,-4.105033927816157391e-01,2.800337569818500683e-02,-5.976660197456356016e-01,4.214440523544955575e-01,4.705808164701692498e-01,4.943614828825371177e-01,4.845096865739698111e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv
deleted file mode 100644
index 297ec32..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-5.931682635522148583e-01,8.699251727039116755e-01,-8.264376444320893356e-01,6.777316299673374900e-03,1.259812425166662031e-02,7.526947568173611991e-01,-6.456449467182574509e-01,-1.867494885666358684e-01,-5.751692444575797758e-01,1.762422487332968579e-01,-4.008003102356285652e-01,-8.546965786953328870e-01,6.116071270069189936e-01
-9.061805329925218810e-01,5.725102286958729803e-01,6.505301386644584127e-01,-5.634820490409915283e-01,1.322069399068586115e-01,-5.320988916994693341e-01,-8.018914495696449762e-02,-1.672922951958961679e-01,-5.236940947370711807e-01,2.913933672990698387e-01,-7.670309932568115663e-01,2.266688501822122781e-01,3.883928729930810619e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv
deleted file mode 100644
index 7b3c85b..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--4.029727807509135840e-01,2.221234395356765923e-01,1.653311729326147894e-01,2.854384277332262521e-01,-5.299668975252573855e-01,2.277161429752445621e-01,-8.605748587143935424e-01,2.679358058329555092e-01,-1.753822281114920667e-01,-4.080365861113419701e-01,-7.990197638852070128e-02,6.451913142929026623e-01,2.646000967988433872e-01
--5.825951636626720553e-01,9.933789450712708913e-01,6.017632717440635215e-01,1.285660707793123692e-01,9.225675787278662110e-01,7.629667460674474100e-01,-1.269325586916481008e-01,-2.399724811514538647e-01,2.244383569074661633e-01,3.864334357836947120e-01,1.407009021060205978e-01,-4.470113212371948919e-01,7.353999032011566683e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv
deleted file mode 100644
index a6af9d6..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.343344504186707145e-01,-1.788377734180790490e-02,5.238121179021124618e-01,-1.801466014877619592e-01,2.444788960837480651e-01,5.831808836713883171e-01,-8.022784620348435425e-01,-4.295701503377884478e-01,-3.972178837057684930e-01,7.118887696693811939e-01,-8.639924687935303105e-01,4.474523280618651899e-01,4.561726115084100419e-01
--3.159104803788210791e-01,-8.391006409409682565e-01,-3.390316252365024319e-01,1.553039413054588813e-01,8.170194759350031255e-01,7.527728133817603862e-01,6.978584447967217663e-01,1.128938492491866619e-01,-8.250572721858402403e-01,-5.512768790771354066e-01,3.630337106170709038e-01,-2.222138950067416019e-01,5.438273884915899581e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv
deleted file mode 100644
index d18f1b5..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-5.697263627900581717e-01,9.724449544795059630e-01,-1.674757567980476036e-01,-1.372771123535918569e-01,-7.764922940367462445e-01,5.410911337820014655e-01,3.028425898014133200e-01,4.025386950980982537e-01,-1.476033977383572893e-01,3.294617484232353899e-01,-5.024261635732010234e-01,-4.464667963982544840e-01,6.101539033152794111e-01
--7.758966548598598134e-01,6.989084956917288594e-01,7.492372266276172699e-01,1.570862193960709252e-01,-2.601786362671802966e-01,2.727571671873869619e-01,-1.349546389088733811e-01,1.024561812582318598e-01,5.506075183810597018e-01,2.805526623632508265e-01,-9.364615058079310828e-01,5.565384469523662059e-01,3.898460966847205333e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv
deleted file mode 100644
index f61ede0..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--7.898235741040993130e-01,5.427817216628392227e-01,6.300274118070616769e-01,-5.211653741616515401e-01,-6.981622668479592342e-02,9.578781805862026655e-01,1.666377264587282081e-01,1.637757194655065085e-01,7.571393872230314237e-02,9.020579858164654574e-01,8.210889142502095783e-01,9.377974534116928496e-01,8.615315349321857052e-01
-2.679807526537674178e-01,-1.111544236701951238e-01,-5.024699421507621278e-01,7.457314846733458236e-01,-4.885664083857288453e-02,-7.673089298447215434e-01,-8.883324858657668521e-01,-2.464007149940794505e-01,6.129545902906914367e-01,2.592414698046310306e-01,-3.794121113471968787e-01,1.341081142088484945e-02,1.384684650678143225e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv
deleted file mode 100644
index bcb6f88..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--7.654273192322054165e-01,-4.094450069738897469e-02,-7.595383743431962653e-01,-8.557505909970739566e-01,9.064524921829164583e-01,1.517470713028805651e-01,9.075466163718142187e-01,-4.142547321400837923e-01,4.595383450523824465e-01,2.242844526762681756e-01,6.565226571017350743e-01,-6.406292457846078925e-01,2.164925796749257170e-01
-4.134224368111192316e-01,-1.279193807354295220e-01,1.395348624965011552e-01,-7.113928155897839556e-01,-2.037598375484022117e-01,-1.397999201487876153e-01,-8.100399318887028244e-01,-3.540527934775026253e-01,-8.142623919357250273e-01,-7.886095433794413356e-01,-3.674068192061692439e-01,-8.141558656567047247e-01,7.835074203250743663e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv
deleted file mode 100644
index d666294..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-6.234462065020718313e-01,1.300245200362042386e-01,-2.842597130938973038e-01,1.595438141453187075e-01,-4.713993691785938189e-01,1.486436548511567146e-01,-6.366613328949060069e-01,7.675794196907781419e-02,-9.937538509582277690e-02,5.037179007533050257e-01,-6.402726075739686440e-01,-3.221377519083947760e-01,4.313968476057338797e-01
--2.673181139834390763e-03,-2.286479913327021940e-02,6.059211548471905573e-01,-3.648636787800783043e-01,-4.414015900975392093e-01,2.960060241178898988e-01,-3.576760461529826518e-01,-9.644156958072134245e-01,-5.959032460738362680e-01,-7.214674024314973177e-01,5.674627493986834637e-01,3.261734992328386706e-01,5.686031523942661758e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv
deleted file mode 100644
index df4a712..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--7.420103884864390764e-01,-6.381524612883886505e-01,-1.613639618245801266e-01,7.199764450024330742e-01,-1.006210277087187244e-01,5.736088536664056825e-01,-2.078207707668398019e-01,3.313608111485661922e-01,-1.567612276534600113e-01,3.527694725218057936e-01,-1.767897433287017872e-01,-3.553205709558269199e-01,5.623580600962118092e-01
-1.803325284453796140e-01,1.231843219533304001e-01,5.739878271111997776e-03,-8.697875645645707365e-01,4.277980606243028117e-01,1.010399067634688564e-01,9.441666464685340987e-01,-2.209014237470598996e-01,3.359090300130838092e-01,5.940980517447476128e-01,2.482757629902709873e-01,-1.834052354989281763e-01,4.376419399037881908e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv
deleted file mode 100644
index a71f835..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-7.369773596082351830e-02,-2.482646733587001719e-01,5.057638234463792681e-01,-6.242771609481161388e-01,-7.429818940347492351e-01,-4.707267465509237248e-01,-4.626361468138358024e-01,-6.467146033806137062e-02,2.131345959536605772e-01,7.725903761561427885e-01,-4.240286243640691843e-01,-4.218771823875742122e-01,5.823322236392094453e-01
--6.554490572494988676e-01,-7.815506609888773770e-01,-7.332898854928275867e-01,6.612085576900739170e-02,2.178726193278017753e-01,5.569210460230202830e-01,-2.319012715747397202e-01,-9.022663389078762197e-01,-4.304867213238519064e-01,-5.467178901348748177e-01,5.657685803999605856e-01,7.731699809585836913e-01,4.176677763607904992e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv
deleted file mode 100644
index 2012dad..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-3.730533296626739048e-01,1.433183679154126366e-01,3.637273449949744997e-01,2.796358796157163429e-01,9.971687232829262726e-03,7.480311521939866370e-01,3.937355475477484212e-01,4.560000728934507919e-01,-4.984654853379433259e-01,-3.455550319734554954e-01,-9.835584381128927856e-01,-2.633451200187155727e-01,1.651205815987711323e-01
-8.112653261883815414e-01,-3.950454378848380355e-01,3.176817652040930806e-01,-5.239337818191858176e-03,4.250105481704846699e-01,3.604098546798490954e-01,5.490989594952853103e-01,-7.353484188252294995e-01,1.375370159687769878e-01,3.739887895489748537e-01,7.194665959008399447e-01,-2.052330554944428176e-01,8.348794184012288122e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv
deleted file mode 100644
index 9697871..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-1.985151568679008882e-01,-4.118089649714407052e-02,-1.173690980919479543e-01,-6.147768674282323431e-01,5.384205158174693029e-01,-8.092057756629023046e-02,2.779284151732872576e-01,2.862361042030712177e-01,3.840843723060081150e-01,9.898121724234101304e-01,9.713830508604506253e-01,7.822167263660317893e-01,7.111855449058056555e-01
-5.892885167675354641e-01,-1.963827275432317165e-03,-7.606511330012744043e-01,5.586488854728632880e-01,7.846991525938085132e-02,-6.184691865438018965e-01,2.366672318317677437e-01,-4.646883212217289838e-02,2.801530524023061464e-01,-6.722047100717012391e-01,-7.126605435777322306e-01,-6.314786757874539802e-01,2.888144550941944000e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv
deleted file mode 100644
index 9760c67..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--5.511430473636560468e-01,2.075233910672833471e-01,9.276713288167097726e-01,-1.822552515865676348e-01,4.964535495465318693e-01,4.198110192190658285e-01,-7.110023362885946607e-01,3.517890700491155265e-02,1.069361824193393318e-02,-7.884513208221513025e-01,-8.560972160329762826e-01,-9.047373606248456657e-01,8.266997873488395321e-01
--5.807999610356651132e-02,-1.039348906028856323e-01,3.141343076567044701e-01,-2.056624804791262751e-01,5.792891943323019710e-02,-9.509068945320611199e-01,-3.764789001983093186e-01,2.488096127203558439e-01,-1.853125052175255139e-01,8.301118240605835918e-01,-1.216569088549470656e-01,8.391547337468705514e-01,1.733002126511604402e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv
deleted file mode 100644
index a0119e8..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--5.633361199129609531e-01,-1.242868737984672567e-01,-8.018904487385585256e-01,-5.010751601456076010e-02,-5.054966895260570858e-02,4.270084989655040797e-01,5.613658733821358382e-02,1.067847504111867352e-01,9.715247219631861775e-02,-5.884270629564716248e-01,2.869301460815778526e-01,-8.916215668672899941e-01,1.127259639768470878e-01
-5.813484676405529239e-01,-7.608577159641982668e-02,-1.726841049851548515e-01,1.495963314716997061e-01,9.282798648258450136e-02,-9.208945702831410340e-01,-3.570382956573512345e-01,-3.065376051478874153e-01,-3.776733142035426649e-01,-9.376120443254172265e-01,-1.646485576867517953e-01,4.424834047857026942e-01,8.872740360231529122e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv
deleted file mode 100644
index 088d30f..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-1.245129250727008419e-01,-6.636582935874648648e-01,-7.095021996374384354e-01,-5.708048087285482186e-01,-2.613952915421429157e-01,4.360332099195141087e-01,4.386695896700476549e-01,2.591668816930605690e-01,2.710449407828230406e-01,-1.994249360309625629e-01,-9.884196752400828956e-01,-8.948557108504355817e-01,4.008892167512856930e-01
--7.977101440753309181e-01,-5.891213181154428824e-01,-5.394841756954262824e-01,6.845658522320718919e-01,5.996325324640972010e-01,-7.950515319880713250e-02,1.314403032344408917e-01,8.974954145234985692e-01,3.185168035203989056e-02,-8.571054947316427697e-01,8.942330986453337349e-01,4.800334580863008238e-01,5.991107832487143625e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv
deleted file mode 100644
index 379affb..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-4.156697549963126459e-01,1.938812896637998051e-02,3.762020023825551895e-02,-4.543409658210944002e-01,-4.160769028245498991e-01,-4.404768450621054932e-02,-1.226606756646233531e-01,2.246644739370420307e-01,-1.930408370461424994e-01,4.195049700021282746e-01,-7.252898605016289135e-01,-1.186045724865631978e-01,2.376483080486445632e-01
--5.048426707707873717e-01,2.678319927728600724e-01,-2.276997135948761741e-01,5.838793463819900165e-01,-5.592597689481533241e-01,5.818759639934625305e-01,-5.955380441213122822e-01,-6.677174966534085154e-01,2.195016551675315064e-01,-5.308958776483587716e-01,8.413556421996504220e-02,6.724915832844733377e-01,7.623516919513554368e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv
deleted file mode 100644
index 674f6e6..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--1.237880297177609279e-01,-5.389232504861669604e-01,8.351615604828057648e-01,-5.197299868564686509e-01,-6.477376740352802642e-02,6.235481432700507032e-01,8.648963276371262054e-01,-4.513602462196104614e-01,1.121472463841981515e-01,9.872965050461273151e-01,9.679965448742415823e-01,-1.147346194826741606e-01,3.912812329938487044e-01
--5.999286838715911507e-02,-2.097115124782031881e-01,3.409859382911564207e-01,-1.638644314851909201e-01,-5.392371544653988824e-01,1.453514884790074735e-01,-9.158051602213537201e-01,7.348874336021522513e-01,5.552381537384321053e-01,-2.804504387202546578e-01,7.872454148265923823e-01,2.773799831190322251e-01,6.087187670061512401e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv
deleted file mode 100644
index 4580b5a..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--4.781135907766911330e-01,-1.276623027908179164e-01,-5.244677607605363612e-02,7.793736542701723558e-02,-4.965349935613236898e-01,-7.450674808538539917e-01,2.215427033187817862e-01,-2.835553192890256646e-01,-5.879019088656525227e-01,-8.655409501489235158e-01,1.484038914439216317e-01,2.340787049439099210e-01,4.441034483063370786e-01
-8.918907925254648816e-02,9.583522210758066429e-02,3.380023039545712038e-01,-2.757160640620373027e-01,9.491265801938739699e-01,-2.949240243853041843e-01,2.368000271881480767e-01,-3.274535743006037336e-01,2.320714265769052709e-01,6.681697133141641931e-01,-8.151860272732713852e-02,-5.223786728288988268e-01,5.558965516936629214e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv
deleted file mode 100644
index 4b22744..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--2.430657731263985433e-01,-5.883176849573654721e-01,-1.051727327384854860e-02,-5.613318623862584289e-01,-3.654238252241008844e-01,-2.903699040052056812e-01,-3.376162129993560690e-01,-1.937379177653244522e-01,1.447762165448007732e-01,9.132031396935695877e-01,2.309821948964190241e-01,4.969710849235846606e-01,3.556727315444812021e-01
--5.700025402755233284e-01,3.962163276669257161e-01,4.106860990482503748e-01,-4.445226161030975121e-01,2.216226823069946672e-01,5.337581703209348660e-01,-5.759837441605704100e-01,-3.921802841679722373e-01,-5.148726436855111110e-01,-4.398568392874997457e-01,-8.504659419589648550e-01,6.585902602061193267e-01,6.443272684555187979e-01
diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv
deleted file mode 100644
index e66cf87..0000000
--- a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv
+++ /dev/null
@@ -1,2 +0,0 @@
--3.200427851155407399e-01,3.576359420223551577e-01,-2.513016740058249265e-01,-7.546934815443087086e-01,-2.786933152965342941e-01,2.327058657559735178e-01,7.113022836300710861e-01,6.680266543764468157e-01,2.330623308024013518e-01,-6.132416434718550580e-01,-1.253524191738015769e-01,4.190854187019641408e-01,6.428379861042887722e-01
--2.051650609723685292e-01,-4.086715839711074771e-02,1.337484780156206199e-02,-3.217966923911506072e-01,2.645280631801716353e-01,-5.515626691943973370e-01,-7.300788769231552067e-01,-9.944065418179377502e-01,-6.942904879158731113e-01,3.461188780593014158e-01,-2.042033665456883806e-01,-6.116735230775620646e-01,3.571620138957111168e-01
diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py
new file mode 100644
index 0000000..f6b134f
--- /dev/null
+++ b/ScaFFold/unet/group_norm.py
@@ -0,0 +1,869 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""GroupNorm with a Triton fast path and a ``torch.compile``d one behind it.
+
+Three kernels, tried in order, all of them producing the same numbers:
+
+1. **Native channels-last Triton** (:mod:`ScaFFold.unet.triton_group_norm`),
+   whenever that module's ``is_supported`` accepts the input.  Production runs
+   set ``PYTORCH_MIOPEN_SUGGEST_NHWC=1``, under which every convolution emits
+   ``channels_last_3d`` -- and every *stock* GroupNorm (eager or Inductor) reads
+   that layout through the logical NCDHW order, a strided gather, and returns a
+   contiguous tensor, breaking the layout chain at all 22 call sites of the
+   forward.  The Triton kernel is NDHWC in and NDHWC out and is 6.5x faster than
+   the compiled kernel on that input; with the ReLU fused it takes another 38%
+   off the forward.  See that module's docstring for the measurements.
+2. **``torch.compile``d ``F.group_norm``**, for inputs the Triton kernel does
+   not serve (contiguous NCDHW, non-5-D, unsupported dtypes) and as the landing
+   place if the Triton path ever raises.  ATen's own kernel launches one
+   workgroup per ``(batch, group)`` row -- 8 of them at this benchmark's
+   defaults -- so on a 228-CU MI300A it runs at a small fraction of achievable
+   bandwidth: measured 87 ms of a 187 ms step (47%) at scale 7, against 7% of a
+   184.7 ms step once Inductor tiles the reduction across the device.
+3. **Stock eager ``F.group_norm``**, which is what every rejection falls back
+   to and what defines the semantics the other two must match.
+
+``FastGroupNorm`` is a drop-in ``nn.GroupNorm``: same parameters, same names,
+same shapes, same numerics -- only the kernel differs, so checkpoints are
+interchangeable in both directions with any other GroupNorm-based build.  The
+one addition is the optional fused ``activation`` (see below), which adds no
+state either.
+
+All three return the input's memory format, so the rungs are interchangeable in
+everything a caller can observe (see :func:`_match_memory_format`).
+
+Routing rejections, in the order they are tested:
+
+* an explicit opt-out via ``SCAFFOLD_GROUPNORM_TRITON=0`` /
+  ``SCAFFOLD_GROUPNORM_COMPILE=0``,
+* a rung that has failed in this process, for every module that has not already
+  had a call served by it (see "Latches" below),
+* an active ``torch.func`` transform -- a ``vmap``/``grad``/``jvp`` layer is a
+  routing miss, not a kernel failure, and the stock kernel handles it,
+* non-CUDA tensors -- the CPU test suite pays neither compile latency nor the
+  Triton import,
+* tensor subclasses, whose ``__torch_dispatch__`` wrappers have unknown
+  semantics -- except DistConv's ``DCTensor``, which is unwrapped to its local
+  shard around both fast kernels (see ``FastGroupNorm``),
+* for the Triton kernel, anything its ``is_supported`` rejects (a layout, dtype,
+  degenerate shape or affine-parameter dtype it does not serve); for the
+  compiled one, an already-compiled enclosing region (the functional call
+  inlines instead).
+
+Latches
+=======
+Both fast rungs are optimizations, never correctness requirements: a broken
+Triton install must degrade a multi-node run, not kill it.  So a *kernel*
+failure is caught, logged once and retried on the next rung down.
+
+"A kernel failure" is an allowlist, not the absence of one.  The Triton rung is
+caught on ``triton_group_norm.TritonKernelError``, which that module raises for
+anything its launch region produces; the compiled rung on
+``torch._dynamo.exc.TorchDynamoException``, the root of every Dynamo and
+Inductor compile failure, *plus* ``FailOnRecompileLimitHit``, which despite the
+name derives from ``Exception`` and not from that root (see
+:func:`_compiled_kernel_failures`).  Everything else propagates -- saved-tensor pack
+hooks, ``torch.utils.checkpoint``'s recompute control flow, a user's offloading
+hook, ``torch.OutOfMemoryError``, an error from a shape the kernel mishandles
+badly enough to corrupt the graph.  The previous shape of this code caught
+``Exception`` and re-raised a denylist of framework mechanisms, which was wrong
+twice (``_StopRecomputationError``, then ``CheckpointError``): the set of things
+torch may raise through a forward is open, the set of ways a kernel can be
+broken is closed at its own boundary.  Both allowlisted exceptions are also
+raised strictly *before* their rung saves anything for backward (the Triton op
+saves in ``_setup_context``, after its launch region; a Dynamo/Inductor failure
+is a compile-time failure, before any execution), so the retry cannot double-fire
+saved-tensor hooks.
+
+A failure latches the rung off **for modules that have never had a call served
+by it**.  A module that has already run on a rung keeps it.  That is not a
+performance nicety: ``torch.utils.checkpoint``'s non-reentrant recompute
+compares the metadata of every tensor the recomputed forward saves against the
+originals, and the three rungs intrinsically save *different tensors* -- Triton
+saves ``(input, weight, bias, mean, rstd)``, the other two
+``(input, weight, mean, rstd, relu_output)``.  A latch that flipped between a
+block's forward and its recompute would therefore kill the step with
+``CheckpointError: Recomputed values ... have different metadata``, which is the
+exact opposite of the contract above (measured; matching the output memory
+format is *not* sufficient on its own).  Keeping a proven rung pins each
+module's choice for the life of the process, so forward and recompute always
+agree.
+
+The same reasoning bounds the *fallback* itself, which the latch alone does not:
+a proven module still has to answer the call its rung just failed, and answering
+it eagerly is exactly the flip the paragraph above forbids -- if that call is a
+checkpoint recompute.  So the fallback is declined in the one case where it
+would corrupt rather than degrade: a module proven on the rung, failing while an
+autograd graph task is in flight (:func:`_replaying_a_forward`), re-raises.
+Every other failure -- and in particular every *first* failure, which is what a
+broken Triton install, an unwritable Inductor cache or a missing compiler
+produce -- still degrades, which is where the "must not kill a multi-node run"
+contract actually lives.
+
+Note that a latch is process-local: under DDP one rank can end up running a
+different kernel from its peers.  All three kernels agree to fp32 rounding, not
+bitwise, so a rank that latches shifts that rank's gradients and therefore the
+all-reduced ones -- a real (measured) change to the job's trajectory, and a
+2.1x straggler besides.  That is the price of degrading instead of dying, but it
+is why the latch is as narrow as it is, and why ``torch.OutOfMemoryError`` --
+transient by nature, and no cheaper on any other rung -- does not latch anything
+at all.  :func:`set_triton_enabled` / :func:`set_compile_enabled` with ``True``
+clear the latch, which is the supported way to retry after a transient failure.
+
+Determinism: all three kernels are bitwise reproducible.  Two separate processes
+running the scale-7 UNet under ``more_determinism``
+(``use_deterministic_algorithms(True, warn_only=True)``, ``cudnn.benchmark=False``,
+fixed seeds) hash identically with the Triton path, the compiled path and the
+eager one alike, so no determinism gate is needed.  The Triton kernel's grid,
+split count and tile sizes are pure functions of the shape and it uses no float
+atomics, which is what buys that.
+
+Fused activation
+================
+Every GroupNorm in the UNet is immediately followed by a ReLU, and the Triton
+kernel can fold that into its forward store for free (it is store-bound) while
+removing a whole streaming pass -- 38% of the forward at the shapes that
+dominate.  ``FastGroupNorm(..., activation="relu")`` therefore *always* applies
+the ReLU: fused inside the Triton kernel where that path is taken, and as an
+explicit in-place ``F.relu`` on the compiled and eager paths.  ``DoubleConv``
+consequently holds an ``nn.Identity`` where its ``nn.ReLU`` used to be, so the
+positional keys of its ``nn.Sequential`` -- and therefore every checkpoint --
+are unchanged (neither module has parameters or buffers).  Correctness holds on
+every path; the fusion is purely an optimization inside the module.
+"""
+
+import logging
+import os
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+#: Opt-out (``0``/``false``/``off``/``no``) or explicit opt-in (``1``/``true``/
+#: ``on``/``yes``) for the compiled GroupNorm path.  Unset means "on wherever it
+#: is safe", which is what every production run wants.
+COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE"
+
+#: The same, for the native channels-last Triton kernel, which is tried first.
+#: Same spellings, same "unset means on wherever it is safe" default -- the
+#: whole point of the kernel is that production takes it.
+TRITON_ENV_VAR = "SCAFFOLD_GROUPNORM_TRITON"
+
+#: Activations this module can apply after normalizing.  Must stay a subset of
+#: ``triton_group_norm.SUPPORTED_ACTIVATIONS`` (pinned by a test); spelled out
+#: here rather than imported so that constructing a module -- or running the
+#: whole CPU suite -- never imports the kernel module.
+SUPPORTED_ACTIVATIONS = (None, "relu")
+
+#: Dynamo caches one entry per distinct guard set on the traced function.  A
+#: UNet presents one entry per distinct activation shape (5 at scale 7) times
+#: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock
+#: limit of 8, which would silently drop the whole model back to eager mid-run.
+#: ``activation_checkpointing`` on a ``DCTensor`` doubles that again: the
+#: recompute reaches this module with ``__torch_function__`` subclass handling
+#: *disabled* (DistConv's backward runs below it), which is part of Dynamo's
+#: ``GLOBAL_STATE`` guard, so the recomputed forward misses every entry the
+#: original forward built and compiles a second set beside it -- 20 for the same
+#: 5 shapes (measured).  The traced function is a single ``F.group_norm`` call,
+#: so the extra entries cost only their one-time compilation.
+_MIN_RECOMPILE_LIMIT = 64
+
+# Lazily built on the first eligible forward: importing ScaFFold must not drag
+# in Dynamo, and a run that never reaches the GPU must not pay for it.
+_compiled_group_norm = None
+
+# Set once if torch.compile raises; the eager path is then used everywhere.
+_compile_failed = False
+
+# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_COMPILE or
+# by set_compile_enabled().
+_compile_override = None
+
+# The triton_group_norm module, imported on the first CUDA forward.  Importing
+# it registers two dispatcher ops, and a CPU-only run must pay neither that nor
+# the `triton` import the module itself defers to its first launch.
+_triton_module = None
+
+# The ladder's two allowlists, resolved on first use of the rung they guard --
+# importing either provider (the kernel module, torch._dynamo) is exactly what
+# the lazy _get_* helpers exist to avoid paying for on a CPU-only run.
+_TRITON_KERNEL_FAILURES = None
+_COMPILED_KERNEL_FAILURES = None
+
+# Set once if the Triton kernel raises; the compiled path is used from then on.
+_triton_failed = False
+
+# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_TRITON or
+# by set_triton_enabled().
+_triton_override = None
+
+
+def _env_override(name):
+    """Read boolean env var ``name``; ``None`` when unset or unparsable."""
+    raw = os.environ.get(name)
+    if raw is None:
+        return None
+    value = raw.strip().lower()
+    if value in ("1", "true", "on", "yes"):
+        return True
+    if value in ("0", "false", "off", "no"):
+        return False
+    logger.warning(
+        f"Ignoring unrecognized {name}={raw!r}; "
+        "expected one of 1/0/true/false/on/off/yes/no"
+    )
+    return None
+
+
+_compile_override = _env_override(COMPILE_ENV_VAR)
+_triton_override = _env_override(TRITON_ENV_VAR)
+
+
+def set_compile_enabled(enabled):
+    """Force the compiled path on (``True``) or off (``False``).
+
+    ``None`` restores the default, which is the environment variable if set and
+    otherwise "compile wherever it is safe".  Forcing it on does not override
+    the device and tensor-subclass checks -- those are correctness conditions,
+    not preferences -- but it *does* clear a failure latch: an explicit "use
+    this rung" is the supported way to retry after a transient failure, and
+    leaving the latch set would make this function silently do nothing.
+    Returns the previous setting so callers (tests) can restore it.
+    """
+    global _compile_override, _compile_failed
+    previous = _compile_override
+    _compile_override = (
+        _env_override(COMPILE_ENV_VAR) if enabled is None else bool(enabled)
+    )
+    if _compile_override is True:
+        _compile_failed = False
+    return previous
+
+
+def set_triton_enabled(enabled):
+    """Force the Triton path on (``True``) or off (``False``).
+
+    The exact counterpart of :func:`set_compile_enabled`: ``None`` restores the
+    default (``SCAFFOLD_GROUPNORM_TRITON`` if set, otherwise "wherever
+    ``is_supported`` accepts"), forcing it on clears any failure latch but not
+    the device, subclass or ``is_supported`` checks -- those are correctness
+    conditions -- and the previous setting is returned so tests can restore it.
+
+    ``None`` deliberately does *not* clear the latch: it restores a preference,
+    it does not assert that the kernel works again.
+    """
+    global _triton_override, _triton_failed
+    previous = _triton_override
+    _triton_override = (
+        _env_override(TRITON_ENV_VAR) if enabled is None else bool(enabled)
+    )
+    if _triton_override is True:
+        _triton_failed = False
+    return previous
+
+
+def _group_norm(input, num_groups, weight, bias, eps):
+    """The function Dynamo traces: plain functional GroupNorm, nothing else."""
+    return F.group_norm(input, num_groups, weight, bias, eps)
+
+
+def _raise_recompile_limit():
+    """Lift Dynamo's *global* recompile cap to cover every UNet GN shape.
+
+    Only ever raises it, so a caller that deliberately set a larger limit keeps
+    theirs -- but note the converse: a limit deliberately set *smaller* than
+    ours is clobbered up to ``_MIN_RECOMPILE_LIMIT``. ``cache_size_limit`` is
+    the older spelling of ``recompile_limit``; set whichever exists.
+
+    This is the *portable* half of the mitigation and, on its own, not a
+    sufficient one: ``torch._dynamo.config`` stores user overrides in a
+    ``ContextVar`` (``torch/utils/_config_module.py``: "User overrides are
+    thread-local"), so an assignment made here is invisible to every other
+    thread, which keeps reading the stock default of 8.  That matters because
+    ``torch.utils.checkpoint``'s non-reentrant recompute runs inside the
+    backward pass, i.e. on the autograd engine's device worker thread, and a
+    recompute that has to compile -- which it does on a ``DCTensor``, see
+    ``_MIN_RECOMPILE_LIMIT`` -- would hit 8 there no matter what this function
+    wrote on the main thread.  :func:`_compile_group_norm` therefore also asks
+    ``torch.compile`` for a per-region limit, which Dynamo applies on whichever
+    thread is compiling.
+    """
+    config = torch._dynamo.config
+    for name in ("recompile_limit", "cache_size_limit"):
+        current = getattr(config, name, None)
+        if isinstance(current, int) and current < _MIN_RECOMPILE_LIMIT:
+            setattr(config, name, _MIN_RECOMPILE_LIMIT)
+
+
+def _compile_group_norm():
+    """``torch.compile`` :func:`_group_norm` with a thread-proof recompile cap.
+
+    ``recompile_limit=`` is the per-region spelling of the cap: Dynamo applies
+    it with ``config.patch()`` around the compile itself, on whatever thread
+    that compile happens on, which is the only spelling that survives the
+    autograd worker thread (see :func:`_raise_recompile_limit`).  Older torches
+    have no such keyword -- there the global assignment is all there is, and the
+    checkpoint-recompute case is simply out of reach.
+    """
+    try:
+        return torch.compile(
+            _group_norm,
+            dynamic=False,
+            fullgraph=True,
+            recompile_limit=_MIN_RECOMPILE_LIMIT,
+        )
+    except TypeError:
+        # A torch too old for the keyword: still compile, because the rung is
+        # worth far more than the one configuration the keyword rescues.
+        return torch.compile(_group_norm, dynamic=False, fullgraph=True)
+
+
+def _get_compiled_group_norm():
+    """Build (once) the compiled functional GroupNorm shared by every module.
+
+    One compiled callable for the whole model, not one per module: the shapes,
+    not the instances, are what Dynamo specializes on, and sharing keeps the
+    18 GroupNorms of a scale-7 UNet down to 5 compilations.  ``dynamic=False``
+    keeps the specialized kernels (this benchmark runs fixed shapes);
+    ``fullgraph=True`` turns anything Dynamo cannot handle into an exception we
+    catch, rather than a silent graph break that reintroduces the slow kernel.
+    """
+    global _compiled_group_norm
+    if _compiled_group_norm is None:
+        _raise_recompile_limit()
+        _compiled_group_norm = _compile_group_norm()
+    return _compiled_group_norm
+
+
+def _get_triton_module():
+    """Import (once) :mod:`ScaFFold.unet.triton_group_norm`.
+
+    Deferred rather than imported at the top of this file: that module registers
+    two dispatcher ops and builds an autograd formula at import time, and a run
+    that never reaches the GPU (the whole CPU unit suite) must not pay for it.
+    Only ever called after the input has been shown to be a CUDA tensor, which
+    is also what keeps ``import triton`` -- which that module defers again, to
+    its first kernel launch -- out of a CPU-only process entirely.
+    """
+    global _triton_module
+    if _triton_module is None:
+        from . import triton_group_norm
+
+        _triton_module = triton_group_norm
+    return _triton_module
+
+
+def _triton_kernel_failures():
+    """The ladder's allowlist for the Triton rung: exactly ``TritonKernelError``.
+
+    The kernel module raises it for every failure of its own launch region --
+    a missing or mismatched ``triton``, an unwritable JIT cache, a compile
+    error, a bad launch -- and for nothing else, so this catches "the kernel is
+    broken" without also catching the framework mechanisms that legitimately
+    raise through a forward.  See that class's docstring for what is
+    deliberately left untagged (``OutOfMemoryError``, contract violations).
+
+    Resolved separately from :func:`_get_triton_module` so the except clause is
+    still available when the thing that failed *is* the module lookup.  An
+    empty tuple (no kernel module at all) means "catch nothing": the ladder
+    then re-raises, which is right, because with no kernel module there is
+    nothing that could have failed inside one.
+    """
+    global _TRITON_KERNEL_FAILURES
+    if _TRITON_KERNEL_FAILURES is None:
+        try:
+            from .triton_group_norm import TritonKernelError
+
+            _TRITON_KERNEL_FAILURES = (TritonKernelError,)
+        except ImportError:  # pragma: no cover - the module is in-tree
+            _TRITON_KERNEL_FAILURES = ()
+    return _TRITON_KERNEL_FAILURES
+
+
+def _compiled_kernel_failures():
+    """The compiled rung's allowlist: every Dynamo and Inductor compile failure.
+
+    ``torch._dynamo.exc.TorchDynamoException`` is the root of ``Unsupported``
+    (``fullgraph=True`` met something untraceable), ``BackendCompilerFailed``
+    and its ``InductorError`` subclass (the backend, and therefore also an
+    unwritable Inductor cache or a broken C++/Triton toolchain), and
+    ``InternalTorchDynamoError``.
+
+    ``FailOnRecompileLimitHit`` -- raised when a frame needs more cache entries
+    than the recompile limit allows, which under ``fullgraph=True`` is a hard
+    error rather than a drop to eager -- is *not* under that root: it derives
+    straight from ``Exception`` (``torch/_dynamo/exc.py``), so catching only
+    ``TorchDynamoException`` lets it kill the run.  It is named separately
+    rather than assumed, and only added when it really is outside the root, so
+    a torch that later reparents it does not produce a duplicate entry.
+
+    All of these are raised while *compiling*, i.e. before the compiled callable
+    has executed or saved anything, which is what makes the fallback safe to
+    retry.
+
+    Resolved on demand and cached: importing ``torch._dynamo`` is precisely the
+    cost :func:`_get_compiled_group_norm` defers.  An empty tuple (a torch
+    without the module) means "catch nothing", which fails loudly rather than
+    silently swallowing.
+    """
+    global _COMPILED_KERNEL_FAILURES
+    if _COMPILED_KERNEL_FAILURES is None:
+        try:
+            import torch._dynamo.exc as dynamo_exc
+        except ImportError:  # pragma: no cover - torch always ships it
+            _COMPILED_KERNEL_FAILURES = ()
+        else:
+            failures = [dynamo_exc.TorchDynamoException]
+            limit_hit = getattr(dynamo_exc, "FailOnRecompileLimitHit", None)
+            if isinstance(limit_hit, type) and not issubclass(
+                limit_hit, dynamo_exc.TorchDynamoException
+            ):
+                failures.append(limit_hit)
+            _COMPILED_KERNEL_FAILURES = tuple(failures)
+    return _COMPILED_KERNEL_FAILURES
+
+
+def _replaying_a_forward():
+    """``True`` while this thread is executing inside an autograd graph task.
+
+    ``torch._C._current_graph_task_id()`` is ``-1`` outside a backward pass and
+    the running task's id inside one; it is the same signal
+    ``torch.utils.checkpoint`` keys its own recompute bookkeeping on
+    (``torch/utils/checkpoint.py``'s ``unpack_hook``).
+
+    A GroupNorm *forward* that runs while a backward is in flight is not a new
+    call: it is a checkpoint recompute (or a double backward) replaying a
+    forward that has already happened and whose saved tensors are already held.
+    That is the one place where quietly answering on a different rung than the
+    original forward used is not a fallback but a corruption -- the rungs save
+    different tensors for backward, so the recompute's saved set no longer
+    matches (measured: ``CheckpointError: Recomputed values ... have different
+    metadata``, and on one shape a GPU memory fault instead).  See
+    :meth:`FastGroupNorm.forward`.
+
+    ``is_compiling()`` first, for the same reason :func:`_warn_rung_failure`
+    checks it: the probe below is a ``torch._C`` builtin returning an ``int``,
+    which Dynamo cannot trace ("Unsupported torch.* op returned non-Tensor"),
+    so a caller who wraps this ``forward`` in ``torch.compile(fullgraph=True)``
+    would get a hard error where the fallback belongs.  Dynamo folds it to
+    ``True`` at trace time, leaving ``False`` here as a constant -- which is
+    also the right answer: tracing is not replaying, and the recompute this
+    guards against runs with Dynamo disabled anyway
+    (``torch.utils.checkpoint``'s ``_run_fn_with_dynamo_disabled``).
+    """
+    if torch.compiler.is_compiling():
+        return False
+    task_id = getattr(torch._C, "_current_graph_task_id", None)
+    if task_id is None:  # pragma: no cover - every supported torch has it
+        return False
+    return task_id() != -1
+
+
+#: ``True`` while a ``torch.func`` transform (``vmap``/``grad``/``jvp``) is on
+#: the stack.  Both fast rungs decline then: a functorch layer is a routing
+#: question, not a kernel defect, and the stock kernel handles every transform.
+#: Not merely a performance choice -- ``is_supported``'s
+#: ``is_contiguous(memory_format=...)`` raises outright under ``vmap``
+#: ("NYI: querying is_contiguous inside of vmap"), and the Triton op has no
+#: batching rule -- so without this the module is not the drop-in
+#: ``nn.GroupNorm`` it claims to be for any caller using ``torch.func``.
+_functorch_active = getattr(torch._C, "_are_functorch_transforms_active", lambda: False)
+
+# Set once if a predicate raised while deciding; see _use_triton.
+_predicate_warned = False
+
+
+def _use_triton(input, num_groups, weight, bias, activation, proven=False):
+    """Whether this particular input should take the native Triton kernel.
+
+    ``proven`` is the caller's "this module has already had a call served by
+    this rung", which keeps a proven module on it even after a *global* latch;
+    see the module docstring's "Latches".
+
+    Ordered so that the cheap local tests come first and the module import last:
+    a CPU tensor is rejected before ``_get_triton_module`` is ever called.
+    """
+    if _triton_override is False:
+        return False
+    if _triton_failed and not proven:
+        return False
+    if _functorch_active():
+        return False
+    # Same policy as _use_compiled: an unknown __torch_dispatch__ wrapper has
+    # unknown semantics and keeps the stock kernel.  is_supported() would accept
+    # one (it only asks isinstance), so this check is load-bearing here, not a
+    # copy for symmetry.  DistConv's DCTensor never reaches it -- forward()
+    # unwraps to the local shard first.
+    if type(input) is not torch.Tensor:
+        return False
+    if not input.is_cuda:
+        return False
+    # is_supported() is cheap and side-effect free: a handful of attribute reads
+    # and one stride check, no allocation, no launch, no triton import.  The
+    # broad catch is right *here* and nowhere else in this module: a predicate
+    # that cannot answer has a correct answer available ("no"), it has done no
+    # work anyone can observe, and the failure is a routing miss rather than a
+    # broken kernel -- so it must not latch the rung off, which is what letting
+    # it fall into the ladder's handler used to do.
+    try:
+        return _get_triton_module().is_supported(
+            input, num_groups, weight, bias, activation
+        )
+    except Exception as e:
+        _warn_once_about_the_predicate(e)
+        return False
+
+
+def _warn_rung_failure(what, error, fallback, env_var):
+    """Log a rung failure, without graph-breaking a compiled caller.
+
+    Dynamo cannot trace ``logging.Logger`` methods ("Unsupported: logging.Logger
+    method not supported for non-export cases"), so a bare ``logger.warning``
+    in the ladder's handler turns a *fallback* into a hard Dynamo error for
+    anyone who wraps this ``forward`` in ``torch.compile(fullgraph=True)`` --
+    the one caller for whom the fallback matters most, since the failure it is
+    reacting to is usually a compile failure.  ``is_compiling()`` is a Dynamo
+    intrinsic that folds to ``True`` at trace time, so the call below becomes
+    dead code inside a traced region and the fallback traces cleanly.  The
+    latch itself is a global assignment, which Dynamo does replay, so the
+    fallback is still recorded -- only this message is dropped, and only for a
+    caller that is compiling this module's forward (nothing in ScaFFold does).
+    """
+    if torch.compiler.is_compiling():
+        return
+    logger.warning(
+        f"{what} failed ({type(error).__name__}: {error}); falling back to the "
+        f"{fallback} for modules that have not already used it. "
+        f"Set {env_var}=0 to skip this attempt entirely."
+    )
+
+
+def _warn_once_about_the_predicate(error):
+    """Log the first ``is_supported`` failure; a repeat would log per call."""
+    global _predicate_warned
+    if _predicate_warned:
+        return
+    _predicate_warned = True
+    logger.warning(
+        f"Triton GroupNorm routing check failed ({type(error).__name__}: "
+        f"{error}); using the stock kernel for inputs like this one. This is a "
+        "routing miss, not a kernel failure, so nothing is latched off."
+    )
+
+
+def _use_compiled(input, proven=False):
+    """Whether this particular input should take the compiled path.
+
+    ``proven`` has the same meaning as in :func:`_use_triton`.
+    """
+    if _compile_override is False:
+        return False
+    if _compile_failed and not proven:
+        return False
+    # Dynamo cannot trace a functorch layer either, and under fullgraph=True
+    # that is an exception rather than a graph break.
+    if _functorch_active():
+        return False
+    # Tensor subclasses route their ops through __torch_dispatch__, which
+    # Dynamo cannot trace.  DistConv's DCTensor never reaches this check --
+    # forward() peeks at its local shard instead -- so anything rejected here
+    # is an unknown wrapper, and eager keeps its semantics exactly as they
+    # are today.
+    if type(input) is not torch.Tensor:
+        return False
+    # CPU GroupNorm is not the bottleneck and compiling it would put a
+    # multi-second C++ build in front of every unit test.
+    if not input.is_cuda:
+        return False
+    # Already inside a compiled region: let the functional call be inlined.
+    if torch.compiler.is_compiling():
+        return False
+    return True
+
+
+def _dctensor_ops(input):
+    """The ``distconv.distconv`` module when ``input`` is a DCTensor, else None.
+
+    Resolved through ``sys.modules`` instead of an import: a DCTensor can only
+    exist if DistConv is already imported, and this module must stay importable
+    (and the CPU suite runnable) without DistConv installed.
+    """
+    distconv = sys.modules.get("distconv.distconv")
+    if distconv is not None and isinstance(input, distconv.DCTensor):
+        return distconv
+    return None
+
+
+def _match_memory_format(out, reference):
+    """Give ``out`` ``reference``'s memory format, copying only if it differs.
+
+    ``F.group_norm`` -- eager or Inductor-compiled -- reads a
+    ``channels_last_3d`` input through the logical NCDHW order and returns a
+    *contiguous* tensor, which is the layout break this whole module exists to
+    avoid: with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` every convolution both sides
+    of it wants channels-last, so one fallback re-breaks the chain for the rest
+    of the network.  One relayout of the GroupNorm output is far cheaper than
+    the transposes the following convolutions would otherwise insert, and it
+    makes the three rungs agree on everything a caller can observe rather than
+    only on the values.
+
+    Free on the Triton rung (already channels-last) and on any contiguous input
+    (nothing to do); one copy on a fallback from a channels-last input, which is
+    the only case that reaches the copy at all.
+    """
+    if reference.dim() != 5:
+        # is_contiguous(memory_format=channels_last_3d) is only defined for 5-D.
+        return out
+    if _functorch_active():
+        # "NYI: querying is_contiguous inside of vmap for memory_format other
+        # than torch.contiguous_format" -- and a functorch transform has no
+        # layout chain to preserve anyway, since both fast rungs decline it.
+        return out
+    if not reference.is_contiguous(memory_format=torch.channels_last_3d):
+        return out
+    if out.is_contiguous(memory_format=torch.channels_last_3d):
+        return out
+    return out.contiguous(memory_format=torch.channels_last_3d)
+
+
+def _run_local(input, distconv, kernel):
+    """Run ``kernel`` on a plain tensor, DCTensor in -> DCTensor out.
+
+    ``distconv`` is ``None`` for a plain tensor, where this is just
+    ``kernel(input)``.  For a ``DCTensor`` the unwrap goes through DistConv's
+    ``_ToTensor``/``_FromTensor`` autograd pair (``DCTensor.from_shard`` is the
+    public spelling of the latter; there is no public unwrap yet -- upstream
+    ask) rather than a bare ``input._tensor`` read: DistConv's own dispatch may
+    read ``_tensor`` directly because it runs *below* autograd, while this runs
+    above it and a bare read would sever the graph back to the producing
+    convolution.
+    """
+    if distconv is None:
+        return kernel(input)
+    local = distconv._ToTensor.apply(input)
+    return distconv.DCTensor.from_shard(kernel(local), input._parallel_strategy)
+
+
+class FastGroupNorm(nn.GroupNorm):
+    """``nn.GroupNorm`` with a Triton GPU kernel and an optional fused ReLU.
+
+    Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no
+    buffers, so state dicts are interchangeable with plain ``nn.GroupNorm``
+    in both directions.  ``activation`` is a plain Python attribute, not a
+    submodule or a buffer, so setting it does not add a key either.
+
+    ``activation="relu"`` makes this module's forward *always* apply a ReLU --
+    fused into the Triton kernel's store where that path is taken, and as an
+    explicit in-place ``F.relu`` on the compiled and eager paths.  The
+    correctness of the model therefore does not depend on which kernel runs;
+    only the number of memory passes does.
+
+    DistConv's ``DCTensor`` gets the fast kernels too, by unwrapping to the
+    local shard in front of them rather than by letting the op dispatch through
+    the wrapper.  Both would work -- the Triton kernel is a real dispatcher op,
+    so ``DCTensor.__torch_dispatch__`` would intercept it, unwrap, run and
+    rewrap on its own -- but the explicit unwrap is what this module already
+    does for the compiled kernel, and it is better here for three reasons.
+    (1) It keeps the subclass policy in one place: ``is_supported`` accepts any
+    ``torch.Tensor`` *instance*, so relying on dispatch would silently extend
+    the fast path to every unknown wrapper subclass, which today keeps the
+    stock kernel.  (2) The eligibility predicates then examine the tensor the
+    kernel will actually touch -- its dtype, device, strides and shape -- rather
+    than a wrapper's mirrored metadata.  (3) The Triton and compiled paths share
+    one unwrap and one fallback ladder instead of needing two shapes of code,
+    and a Triton failure can be retried on the compiled kernel without a second
+    round trip through the wrapper.  Semantics are unchanged either way:
+    DistConv's generic ``__torch_dispatch__`` has no GroupNorm-specific
+    handling, so statistics are per-shard and no communication happens at any
+    shard count, exactly as before.
+    """
+
+    #: Class-level defaults, so that an instance restored from a *module*
+    #: pickle written before these attributes existed (``torch.save(model)``
+    #: rather than a state dict) still runs.  ``nn.Module.__setstate__``
+    #: replaces ``__dict__`` wholesale, so anything only ever set in
+    #: ``__init__`` is simply missing on such an instance.
+    activation = None
+
+    #: Per-module "a call has been served by this rung".  A global latch does
+    #: not demote a module that has one, which is what keeps a checkpointed
+    #: block's forward and its recompute on the same rung; see the module
+    #: docstring's "Latches".  Plain attributes, so they are not parameters,
+    #: buffers or state-dict keys.
+    _triton_ok = False
+    _compiled_ok = False
+
+    def __init__(
+        self,
+        num_groups,
+        num_channels,
+        eps=1e-5,
+        affine=True,
+        device=None,
+        dtype=None,
+        activation=None,
+    ):
+        if activation not in SUPPORTED_ACTIVATIONS:
+            raise ValueError(
+                f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}"
+            )
+        super().__init__(
+            num_groups, num_channels, eps=eps, affine=affine, device=device, dtype=dtype
+        )
+        self.activation = activation
+
+    def extra_repr(self):
+        base = super().extra_repr()
+        if self.activation is None:
+            return base
+        return f"{base}, activation={self.activation}"
+
+    def _activate(self, out):
+        """Apply the activation on the two paths that cannot fuse it.
+
+        In place, which is what the ``nn.ReLU(inplace=True)`` this module
+        absorbed did: ``out`` is a freshly allocated GroupNorm output with no
+        other consumer, and GroupNorm's backward reads its *input*, never its
+        output, so overwriting it is safe for autograd as well as for memory.
+
+        Validated *here* rather than only in ``__init__``: ``activation`` is a
+        plain attribute, so it can be assigned after construction, and the
+        Triton rung would then fuse an activation this method silently skipped
+        -- i.e. the network's function would depend on its input's memory
+        format.  This is also the guard that makes adding a third activation to
+        ``SUPPORTED_ACTIVATIONS`` a loud failure until it is implemented here.
+        """
+        activation = self.activation
+        if activation is None:
+            return out
+        if activation == "relu":
+            return F.relu(out, inplace=True)
+        raise ValueError(
+            f"activation must be one of {SUPPORTED_ACTIVATIONS}, got "
+            f"{activation!r}; this rung cannot apply it"
+        )
+
+    def _triton_forward(self, local):
+        """The native channels-last kernel, with the activation fused in."""
+        return _get_triton_module().triton_group_norm(
+            local, self.num_groups, self.weight, self.bias, self.eps, self.activation
+        )
+
+    def _compiled_forward(self, local):
+        return self._activate(
+            _match_memory_format(
+                _get_compiled_group_norm()(
+                    local, self.num_groups, self.weight, self.bias, self.eps
+                ),
+                local,
+            )
+        )
+
+    def _eager_forward(self, input):
+        # super().forward() is the stock kernel; deferring to it keeps the eager
+        # path identical to nn.GroupNorm's (plus the ReLU and the relayout) by
+        # construction.
+        return self._activate(_match_memory_format(super().forward(input), input))
+
+    def forward(self, input):
+        global _compile_failed, _triton_failed
+
+        distconv = _dctensor_ops(input)
+        # The eligibility checks look at the local shard for a DCTensor (the
+        # peek is a plain attribute read, no autograd involvement) and at the
+        # tensor itself otherwise.
+        local_view = input._tensor if distconv is not None else input
+
+        if _use_triton(
+            local_view,
+            self.num_groups,
+            self.weight,
+            self.bias,
+            self.activation,
+            proven=self._triton_ok,
+        ):
+            triton_failures = _triton_kernel_failures()
+            try:
+                out = _run_local(input, distconv, self._triton_forward)
+            except triton_failures as e:
+                # A broken or mismatched Triton install, an unwritable JIT cache
+                # or a shape the kernel mishandles must cost speed, not a
+                # multi-node run.  GroupNorm is pure and the kernel raises this
+                # only from its launch region -- before it has saved anything --
+                # so retrying the same call on the compiled kernel below is
+                # safe, and the compiled kernel, not eager, is the right landing
+                # place: it is still ~10x the stock one.
+                #
+                # Logged on the latch's False->True edge only.  A module that
+                # has already used the rung keeps trying it (that is what pins
+                # a checkpointed block to one rung), so a persistently broken
+                # kernel would otherwise warn once per call for the rest of the
+                # run; clearing the latch re-arms the message.
+                first = not _triton_failed
+                _triton_failed = True
+                if first:
+                    _warn_rung_failure(
+                        "Triton GroupNorm", e, "compiled kernel", TRITON_ENV_VAR
+                    )
+                # ... with one exception, shared with the compiled rung below
+                # and explained there: a module already proven on this rung must
+                # not be answered from a different one while a backward is
+                # replaying its forward.
+                if self._triton_ok and _replaying_a_forward():
+                    raise
+            else:
+                # Only written once: nn.Module.__setattr__ is not free, and
+                # after the first success this reads a class attribute.
+                if not self._triton_ok:
+                    self._triton_ok = True
+                return out
+
+        if not _use_compiled(local_view, proven=self._compiled_ok):
+            return self._eager_forward(input)
+        compile_failures = _compiled_kernel_failures()
+        try:
+            out = _run_local(input, distconv, self._compiled_forward)
+        except compile_failures as e:
+            # Compilation is an optimization, never a correctness requirement:
+            # a broken Inductor install, an unwritable cache directory or an
+            # untraceable input must degrade to the stock kernel, not kill a
+            # multi-node run.  Every exception caught here is a *compile*-time
+            # one, so nothing ran and retrying eagerly is safe.  Same
+            # once-per-latch-edge logging as the Triton rung above.
+            first = not _compile_failed
+            _compile_failed = True
+            if first:
+                _warn_rung_failure(
+                    "torch.compile of GroupNorm", e, "eager kernel", COMPILE_ENV_VAR
+                )
+            # The one call this rung must not answer eagerly: a module already
+            # proven on it, failing while a backward is in flight, is a
+            # checkpoint recompute of a forward that *did* run compiled.  The
+            # rungs save different tensors, so handing back the eager result
+            # makes the recomputed saved set disagree with the saved one and
+            # torch rejects the step -- a `CheckpointError`, or worse (both
+            # measured).  Degrading is for modules with nothing to contradict;
+            # here the honest answer is the original exception, which at least
+            # names the rung and the shape that could not be served.
+            if self._compiled_ok and _replaying_a_forward():
+                raise
+            return self._eager_forward(input)
+        else:
+            if not self._compiled_ok:
+                self._compiled_ok = True
+            return out
diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py
new file mode 100644
index 0000000..61c8105
--- /dev/null
+++ b/ScaFFold/unet/triton_group_norm.py
@@ -0,0 +1,1818 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Channels-last-native Triton GroupNorm (NDHWC in, NDHWC out).
+
+Why this exists
+===============
+With ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- which every production ScaFFold run
+sets -- every convolution in the UNet emits ``channels_last_3d`` activations,
+but *every* stock GroupNorm variant (eager or Inductor-compiled) consumes them
+through the *logical* NCDHW iteration order, which over a channels-last tensor
+is a strided gather, and then emits a **contiguous** tensor.  Measured on one
+MI300A at scale 8 that costs 6.4x on GroupNorm itself (443 ms/step of GN
+fwd+bwd against 69 ms/step here) *and* breaks the channels-last chain 22 times
+per forward, forcing the following convolution to convert back.
+
+A channels-last-3d contiguous ``(N, C, D, H, W)`` tensor is *physically* a dense
+``(N, S, C)`` array with ``S = D*H*W``; group ``g`` owns a contiguous run of
+``C/G`` channels *inside every voxel*.  The kernels below therefore let **one
+program handle all groups at once** for a chunk of voxels: they read a dense
+``(BLOCK_S, C)`` run, reshape the inner axis to ``(G, C/G)``, and get perfectly
+coalesced loads and stores with the group axis costing nothing.  Measured at
+95-98% of this device's streaming roofline at the two largest UNet shapes.
+
+Measured on one MI300A (228 CUs), fp32, ``num_groups=8``, median of 20, the six
+scale-8 UNet GroupNorm shapes, fwd / fwd+bwd in ms::
+
+    shape             this      compiled-CL   compiled-CONT   eager-CL fwd
+    [1,64,256^3]    4.23/11.39   19.51/77.01    4.77/11.85      151.2
+    [1,128,128^3]   1.15/ 3.00    7.23/24.64    1.22/ 3.07       36.5
+    [1,256,64^3]    0.35/ 0.97    2.27/ 7.42    0.34/ 0.84        9.1
+    [1,512,32^3]    0.14/ 0.57    0.19/ 1.03    0.11/ 0.33        1.7
+    [1,1024,16^3]   0.11/ 0.57    0.10/ 0.39    0.07/ 0.33        0.4
+    [1,2048,8^3]    0.11/ 0.57    0.08/ 0.40    0.07/ 0.33        0.1
+
+Over the 22 scale-8 call sites that is **442.8 -> 67.1 ms/step** of GroupNorm
+fwd+bwd against today's production path (compiled GroupNorm on channels-last
+input), i.e. **376 ms/step recovered**, and a dead heat with compiled GroupNorm
+on *contiguous* input (66.4 ms/step) while additionally not breaking the
+layout chain.  The three smallest shapes lose on host dispatch, not on GPU
+work -- see :func:`select_strategy`.
+
+The ``this`` column and the rollup were re-measured after the launch folding
+below (2+2 kernels instead of 3+4, retuned jointly); the same measurement of
+the unfused chain in the same process gives 4.26/11.59, 1.19/3.21, 0.35/1.01,
+0.14/0.63, 0.13/0.63, 0.13/0.62 and a 69.5 ms/step rollup, i.e. **-3.4% over
+the 22 sites** and -1.7% to -9.5% per shape.  The other three columns are from
+the earlier sweep and are unchanged.
+
+Public API
+==========
+``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5,
+activation=None)``
+    Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with
+    first-order autograd support.  Accepts *anything* ``F.group_norm`` accepts;
+    inputs the Triton kernel cannot serve fall back to ``F.group_norm``
+    internally (see "Layouts" below).
+
+``is_supported(input, num_groups, weight=None, bias=None, activation=None)``
+    Cheap, side-effect-free predicate: ``True`` exactly when the native Triton
+    kernel will run.  Callers that already have a good fallback (e.g. a
+    ``torch.compile``d GroupNorm) should test this and route rejects
+    themselves; ``triton_group_norm``'s own fallback is plain eager
+    ``F.group_norm``.
+
+Contract
+========
+For every input ``is_supported`` accepts, the result matches ``F.group_norm``
+to within fp32 reduction-order noise, with:
+
+* **dtype** -- output dtype is exactly ``F.group_norm``'s.  Verified
+  empirically on this build (torch 2.13.0+rocm7.2): without autocast the output
+  dtype is the input dtype (fp32/bf16/fp16); under ``torch.autocast("cuda",
+  ...)`` GroupNorm is an fp32-policy op, so the output is **fp32** for any
+  input dtype.  This module reproduces that rule (see ``_autocast_out_dtype``)
+  without materializing the fp32 copy of the input that autocast's cast would
+  create: the kernels read the input at its native width and accumulate in
+  fp32, which is bit-for-bit the same computation as upcasting first, but reads
+  half the bytes.  Gradients follow the same rule: ``d_input`` has the input's
+  dtype, ``d_weight``/``d_bias`` have the parameter's dtype.  Forward time at
+  ``[1,64,256^3]`` / ``[1,128,128^3]`` / ``[1,256,64^3]`` in ms: fp32
+  4.28/1.20/0.35, bf16 2.42/0.68/0.22, fp16 2.43/0.66/0.22, and bf16-in with
+  the fp32-out autocast contract 3.08/0.84/0.28 -- so honouring autocast's
+  fp32 output still buys 1.4x over fp32 end to end, because only the read side
+  narrows.
+* **statistics** -- always accumulated in fp32, never in the input dtype.
+* **memory format** -- the output has the *input's* memory format.  This is the
+  one deliberate difference from stock GroupNorm, which returns a contiguous
+  tensor for every input layout (measured, see the table in
+  ``review/gn-dctensor/triton/RESULTS.md``); preserving channels-last is the
+  entire point of the kernel.
+* **autograd** -- ``d_input``, ``d_weight``, ``d_bias``; ``weight=None`` and/or
+  ``bias=None`` supported.  **First order only**: the backward is itself a
+  custom op with no autograd formula of its own, so a second
+  ``torch.autograd.grad`` through this op raises ``RuntimeError: Trying to
+  backward through scaffold_gn.group_norm_backward.default but no autograd
+  formula was registered``.  Stock ``F.group_norm`` *does* support double
+  backward, so a gradient penalty or a Hessian-vector product must route
+  around this kernel (``is_supported`` says nothing about second derivatives;
+  it is documented there too).  It fails loudly rather than returning garbage.
+* **device** -- the kernels run on the *input's* device, whatever device is
+  current, matching ATen's ``DeviceGuard`` behaviour; see ``_device_guard``.
+* **determinism** -- bitwise reproducible run to run and process to process.
+  There are no float atomics anywhere, and the grid, split count and tile sizes
+  are pure functions of the shape (the tuning table is frozen in this file for
+  exactly that reason -- a *runtime* autotuner would break reproducibility by
+  changing the reduction order between runs).
+* **rejections** -- every shape/dtype/parameter combination ``F.group_norm``
+  raises on is one ``is_supported`` returns ``False`` for, including the
+  degenerate "1 value per channel" shape (``N*(C/G)*D*H*W == 1``), so a caller
+  that branches on ``is_supported`` never gets an answer where the op this
+  replaces would have raised.
+* **eps** -- one deliberate divergence, at a value no run uses: for a
+  *subnormal* fp32 ``eps`` (``< 1.18e-38``) on a zero-variance group the GPU
+  flushes ``var + eps`` to zero, so ``rstd`` is ``inf`` and ``y`` is ``NaN``
+  where ATen stays finite.  The boundary is exactly the normal/subnormal one
+  (``eps=1.2e-38`` gives ``rstd=9.1e18``, ``eps=1e-38`` gives ``inf``); at
+  ``eps == 0`` both implementations produce non-finite output identically.
+  Left as is rather than clamped because clamping would perturb every
+  ordinary call to defend a value nine orders of magnitude below the smallest
+  plausible one.
+
+Reduction strategy
+==================
+Group statistics span ``S * C/G`` elements (134M at the largest UNet shape), so
+one pass cannot produce them.  Split-K partial reductions land at a fixed
+scratch index and are combined by a fixed-order tree::
+
+    fwd:  stats_partial -> normalize    (2 kernels)
+    bwd:  bwd_partial   -> dx           (2 kernels)
+
+Traffic (``B = numel * itemsize``): 3B forward, 5B backward.
+
+Each pass is **two** launches, not the three and four an unfused split-K chain
+needs: the two finalize passes and the dweight/dbias row reduction are folded
+into the elementwise kernel that consumes them.  ``_normalize_kernel``
+re-derives ``mean``/``rstd`` from the split-K partials itself (and program 0
+stores them for the backward); ``_dx_kernel`` re-derives ``c1``/``c2`` the same
+way and its first ``ceil(C/BLOCK_C)`` programs also do the dweight/dbias
+reduction.  Folding plus the retuning below is worth 8.0-9.5% of fwd+bwd at the
+four smallest shapes, which are host-dispatch bound, 4.9-6.8% at the two middle
+ones and 1.7% at the largest.
+
+The catch, and the reason the tuning table was re-derived rather than inherited:
+**the fusion and the tiling are one problem, not two.**  A fused finalize is
+recomputed by every elementwise *program*, so its cost is
+``nprog_elem * nsplit`` triples of redundant (L2-resident) traffic.  Keeping the
+unfused table's ``nsplit_target=2048`` at ``[1,64,256^3]``, whose flat
+elementwise grid is 131072 programs, asks for 25.7 GB of redundant reads
+against a 4.3 GB tensor and costs **+34% of fwd+bwd** (+70% of the forward).
+Two things fix it, both of them in ``GNConfig``: the elementwise grid is capped
+at ``elem_progs`` programs which then stride over the tiles (so the redundancy
+is bounded by the *grid*, not by the tile count), and ``nsplit_target`` is
+retuned per shape against that cap.  With both, the same shape is 1-2% *faster*
+than the unfused chain.  Do not change one without re-running the other; the
+coordinate-descent tuner is ``review/gn-dctensor/kernel-opt/tune.py``.
+
+Why not one launch per pass
+---------------------------
+A device-scope software barrier (int32 atomics with volatile loads, no float
+atomics, so still bitwise deterministic) collapses each pass to a single
+launch and was measured at 2.3-2.6x on the four smallest shapes.  It is
+deliberately **not** used.  A grid barrier requires every workgroup to be
+co-resident, which caps the grid at the CU count (228 here); the kernel then
+tops out at 0.5-0.9 TB/s against split-K's 2.7, so it loses catastrophically
+the moment the shape is bandwidth-bound rather than dispatch-bound --
+**18.0 ms against 3.0 ms at [1,128,128^3]**, and it does not compile at all at
+``[1,64,256^3]``.  Serving both regimes therefore means shipping two kernel
+families plus a crossover rule, for a whole-model gain of ~3% (65.7 -> 63.6
+ms/step at scale 8); and under CUDA-graph capture, where launch count is free,
+the ten small-shape sites are already only 1.05 ms of a 64.1 ms/step total, so
+the gain is zero.  Hand-rolled inter-workgroup synchronisation is not a good
+trade for 3% in a benchmark whose value depends on being trustworthy and
+reproducible.  The measurements are in
+``review/gn-dctensor/triton-small/RESULTS.md``.
+
+Numerics: Welford, not ``E[x^2]-E[x]^2``
+========================================
+The prototype accumulated ``sum(x)`` and ``sum(x*x)`` and formed
+``var = E[x^2] - E[x]^2``.  That is split-friendly and cheap but cancels
+catastrophically once ``mean >> std``, because it subtracts two nearly equal
+large numbers to recover a small one.
+
+Here each tile instead produces ``(count, mean, M2)`` via a *corrected*
+two-pass over registers -- ``mean0 = sum(x)/n``, then ``corr = sum(x-mean0)/n``
+to recover the digits the first sum lost, then ``M2 = sum((x-mean0-corr)^2)``
+-- and tiles and splits are merged with Chan's parallel combine.  Every step is
+register-only (the tile is read from HBM exactly once either way) and
+atomic-free, so neither the traffic model nor determinism changes.
+
+Measured at ``[1,256,64^3]``, ``num_groups=8``, affine, relative error of the
+*output* against a float64 reference computed from the same fp32 samples:
+
+    x ~ N(mu, sigma)     this kernel   ATen fp32   E[x^2]-E[x]^2
+    mu=0,    sigma=1       1.6e-07      4.1e-07       1.8e-07
+    mu=10,   sigma=1       3.0e-07      6.9e-07       9.9e-06
+    mu=100,  sigma=1       8.0e-07      4.2e-06       5.6e-04
+    mu=1e3,  sigma=1e-2    1.1e-04      2.5e-03       2.3e+00
+
+At ``mu/sigma = 1e5`` the old formulation has lost the variance outright (the
+difference of the two ~1e6-sized fp32 terms is below one ulp, so ``rstd``
+saturates on ``eps`` and the output is meaningless), while this kernel is still
+good to 1.1e-04 -- and is 4-23x *more* accurate than ATen's own fp32 GroupNorm
+at every non-trivial mean.  The residual 1.1e-04 is the fp32 representation
+floor rather than an algorithm defect: a mean of 1e3 held in fp32 is quantized
+to ~6e-5, which is 6e-3 of a standard deviation here, and both kernels sit on
+that floor.
+
+Cost of the rewrite, isolated by timing the kernels alone against the
+prototype's: **+0.8%** on the forward at ``[1,64,256^3]`` (the shape that
+dominates the step), +3-8% at the middle shapes, and **0%** on the backward,
+which does not compute a variance.  A cheaper shifted-mean variant (two tile
+reductions instead of three, shift taken from a peeled first tile) would
+recover most of that; it was not worth the extra failure mode for ~4 ms/step.
+
+What the third pass (``corr``) is worth, separately
+---------------------------------------------------
+The accuracy above is mostly the *two-pass* structure; the ``corr`` term is a
+third reduction on top of it and deserves its own accounting.  Deleting it
+outright (keeping ``mean_t = mean0``, ``M2 = sum((x-mean0)^2)``) and comparing
+both against float64 on the same fp32 samples, relative error of ``rstd``,
+10 seeds each:
+
+    regime                                     with corr    without    ratio
+    one tile per group reduction (nsplit=1):
+      [2,64,8,4,4] G=4, mu/sigma=1e6            1.1e-07     1.6e-04    1472x
+      [2,64,8,4,4] G=2, mu/sigma=1e6            8.4e-08     4.9e-05     580x
+      [2,64,8,4,4] G=1, mu/sigma=1e6            9.8e-08     2.1e-05     216x
+      [2,64,8,4,4] G=4, mu/sigma=1e5            9.9e-08     8.4e-06      85x
+    many tiles and splits (the production configs):
+      [1,512,32^3]  G=8, mu/sigma=1e5           1.2e-05     4.4e-05     3.8x
+      [1,1024,16^3] G=8, mu/sigma=1e5           8.6e-06     2.5e-05     2.9x
+      [1,256,24^3]  G=8, mu/sigma=1e5           3.0e-05     3.7e-05     1.3x
+      [1,256,24^3]  G=8, mu/sigma=1e7           5.8e-04     2.3e-06    0.004x
+
+So it is decisively load-bearing exactly where the tile mean is formed from
+many large values -- up to 1472x on ``rstd`` -- and worth a steady 1.3-4x in
+the multi-split configs the tuning table actually picks, at ``mu/sigma = 1e5``.
+Past ``mu/sigma ~ 1e6`` with many splits it can go the *other* way (last row):
+there the true spread between tile means is smaller than one ulp of the means
+themselves, so Chan's between-tile term is computed from quantization noise
+either way and the uncorrected version's inflated ``M2`` partly cancels it.
+That regime is past fp32's floor for this computation (a mean of 1e6 held in
+fp32 quantizes to 0.06, i.e. 6% of a standard deviation at sigma=1) and no
+production input is near it.
+
+The *output* error is nearly unmoved by any of this -- at most ~1.4x in either
+direction -- which is why the term looks free to delete if you only measure
+``y``: the output is dominated by the fp32 representation of the mean, which
+``corr`` cannot improve (``mean0 + corr`` rounds straight back to ``mean0``
+once the mean is large).  It is ``rstd`` that carries the benefit.
+
+Price, measured the same way (median of 20 forwards, correction removed
+outright rather than zeroed): **+2.3%** of the forward at ``[1,64,256^3]``,
++3.0% at ``[1,128,128^3]``, +4.7% at ``[1,256,64^3]``, and nothing measurable
+(-1.2% to +0.5%, i.e. noise) at the three launch-bound shapes.  At the shape
+that dominates the step that is +0.10 ms of a 11.6 ms fwd+bwd, i.e. +0.9%.
+Kept: a 1.3-1472x accuracy factor on the statistic the whole rewrite exists to
+protect is worth ~1% of GroupNorm time.  The load-bearing case is pinned by
+``test_welford_correction_recovers_rstd_in_a_single_tile_reduction`` in
+``tests/test_triton_group_norm_edge.py``, so deleting the term now fails the
+suite instead of passing it silently.
+
+Layouts
+=======
+* ``channels_last_3d`` 5-D input -> **native Triton kernel**, channels-last
+  output.  This is the fast path and the only one ``is_supported`` accepts.
+* Plain contiguous NCDHW (and every other layout/rank) -> ``triton_group_norm``
+  falls back to ``F.group_norm``, which returns a contiguous tensor, so the
+  input's memory format is still preserved.  ``is_supported`` returns ``False``
+  so that callers keep their own (probably compiled) fallback rather than
+  silently dropping to the eager kernel.
+
+  This is a deliberate scope decision, not an oversight.  A native NCDHW kernel
+  would need a *different* tiling -- with C outermost the fast axis is spatial,
+  so a program must own one group and stream S, rather than owning all groups
+  and streaming voxels -- i.e. a second family of four kernels.  The payoff is
+  small: on contiguous input Inductor's compiled GroupNorm already reaches
+  89-92% of this device's measured streaming roofline (RESULTS.md 4) -- and the
+  table above confirms it, 66.4 ms/step against this kernel's 67.1 -- so a
+  native NCDHW kernel could win ~10% there, against the 6.4x it wins on
+  channels-last input.  If a mixed-layout model ever makes that 10% matter, the
+  place to add it is the strategy hook below.
+
+Addressing
+==========
+``[2, 64, 256^3]`` is *exactly* 2^31 elements, so int32 linear offsets block
+batch>1 at the largest UNet shape and every shape above it.  The kernels widen
+only the **scalar tile base** to int64 (``INT64`` is a ``tl.constexpr``, so
+shapes that fit still emit pure 32-bit code); the vector offsets inside a tile
+span at most ``BLOCK_S*C + C`` elements and stay int32 either way.  That is why
+the wide path is free: forcing int64 on every scale-8 shape moves fwd+bwd by
+-0.8% to +0.8% and forward by -3.7% to +4% (a 0.02 ms swing on the two smallest,
+launch-bound shapes) -- noise in both directions.  The switch is kept anyway
+because it costs one constexpr and documents where the boundary is; correctness
+above 2^31 elements is covered by a test at ``[2, 64, 256, 256, 257]``
+(2_155_872_256 elements, 2.5e-05 relative error on *both* batch items, i.e. the
+same reduction noise a 134M-element fp32 reduction has anywhere).
+
+Fused activation
+================
+``activation="relu"`` folds the ReLU into the forward store.  In a store-bound
+kernel that is free (one compare and one select) and it removes an entire 2B
+streaming pass.  Measured against ``F.relu(triton_group_norm(x))``: 39% off the forward
+and 35% off fwd+bwd at ``[1,64,256^3]`` (6.83 -> 4.20 ms and 17.82 -> 11.61
+ms), 38%/35% at ``[1,128,128^3]``, 34%/30% at ``[1,256,64^3]``, tapering to
+21%/9% at ``[1,512,32^3]`` and below, where the call is host bound and there is
+less streaming pass to remove.
+
+The backward gates the incoming gradient on the sign of the **pre-activation**
+value, which it *recomputes* from the saved ``(x, mean, rstd, weight, bias)``
+using the identical expression the forward used.  Recomputation costs two FLOPs
+on values already in registers and is bit-exact -- same inputs, same operation
+order, same fp32 rounding -- so the sign always agrees with the forward's.  The
+alternative, testing ``y > 0`` on the saved output, would need the output kept
+alive *in addition to* ``x`` (which the GroupNorm backward needs regardless),
+and in bf16/fp16 it would also mis-gate any element whose positive
+pre-activation rounded to zero on the store.
+
+Both the store and the gate are spelled as the *complement* of the usual test
+(``tl.where(y <= 0, 0, y)``, ``tl.where(pre <= 0, 0, dy)``) rather than as
+``tl.maximum(y, 0)`` / ``tl.where(pre > 0, dy, 0)``.  The two are identical on
+every finite value but not on NaN: ``tl.maximum`` returns the *non*-NaN operand
+and ``NaN > 0`` is False, so both of the usual spellings silently map a NaN to
+0.0, while ``F.relu`` propagates it and ``threshold_backward(grad, result, 0)``
+-- ReLU's real backward -- passes its gradient (``NaN <= 0`` is False too).
+Matching ``F.relu`` here is not pedantry: a diverging run whose forward comes
+back finite because the fused activation ate the NaN passes straight through
+ScaFFold's non-finite-loss abort and checkpoints a broken model.  ``+-Inf`` and
+``-0.0`` are bit-identical under either spelling (``-0.0`` flushes to ``+0.0``,
+as ``F.relu`` does).  Cost: nil, measured -- see ``FastGroupNorm``'s tests and
+``review/gn-dctensor/wiring-fixes``.
+
+Composition
+===========
+Registered as real dispatcher ops (``scaffold_gn::group_norm`` /
+``scaffold_gn::group_norm_backward``) via ``torch.library.custom_op``, with a
+fake/meta kernel and ``register_autograd``.  Consequences:
+
+* ``torch.compile(..., fullgraph=True)`` traces through without a graph break.
+* Tensor subclasses that dispatch via ``__torch_dispatch__`` -- notably
+  DistConv's ``DCTensor`` -- intercept the op, unwrap to the local shard, run
+  it, and rewrap, so a DCTensor goes in and a DCTensor comes out with the graph
+  intact.  As with the rest of DistConv today, statistics are per-shard.
+
+Where the host time goes, and what is left
+==========================================
+Composing has a price, and at the launch-bound shapes it is now the *dominant*
+cost.  Peeling the layers at ``[1,2048,8^3]``, steady-state wall clock per
+fwd+bwd (median of 200, min of 7 rounds; GPU work is 0.030 ms)::
+
+    kernels + this file's Python (_forward/_backward called directly)  0.145 ms
+    + torch.library dispatcher (both custom ops)                      +0.054 ms
+    + autograd (register_autograd node, save_for_backward, ctx)       +0.338 ms
+    = triton_group_norm(x).backward(dy)                                0.537 ms
+
+    for scale: an *empty* python torch.autograd.Function, fwd+bwd      0.065 ms
+
+So **63% of the call is the autograd layer** and 10% is the dispatcher -- both
+of them the cost of being a real dispatcher op that ``torch.compile`` and
+``DCTensor`` can see, which is the whole point of registering it that way.  Of
+the 0.145 ms this file is responsible for, 0.030 ms is GPU and the remaining
+0.115 ms is four launches plus the allocations, plan lookup and argument
+binding around them: launching every kernel twice measures the marginal cost of
+a whole invocation site at **35 us**, so the four of them are ~0.14 ms of host
+work that the GPU work does not cover.
+
+Two things were considered for that 0.14 ms and rejected:
+
+* **Bypassing ``JITFunction.run`` for a cached ``CompiledKernel`` handle**
+  (8.68 us -> 3.94 us per launch on this node) would recover ~19 us, i.e. 3.3%
+  of the call.  It buys that by asserting that Triton's specialization key --
+  including 16-byte pointer alignment -- is a pure function of the shape.  It is
+  not: this module accepts channels-last *views with a storage offset* and
+  non-contiguous affine parameters, both of which the test suite exercises, and
+  a stale specialization there is a wrong answer rather than a crash.  3.3% on
+  the host-bound shapes only is not worth a silent-miscompute failure mode.
+* **Caching the scratch buffers** across calls saves ~1.7 us per allocation,
+  ~20 us here, and makes the buffers shared mutable state across call sites --
+  correct on one stream, wrong on two, and this module has no way to know.
+
+What that leaves: the kernels are at 95-98% of the streaming roofline at the
+two largest shapes and the fused chain is 1.7-6.8% faster than the unfused one
+there, so there is no meaningful GPU headroom left.  At the launch-bound
+shapes the remaining 0.4 ms is torch's own plumbing, and the two ways to remove
+it are both outside this file: CUDA-graph capture of the training step (which
+takes the ten small scale-8 sites to ~1.05 ms/step of GPU time in total), or a
+C++ autograd node.
+
+Triton is imported lazily, on the first call that actually reaches the kernel,
+so importing this module (or running the CPU test suite) costs nothing.
+"""
+
+import contextlib
+import functools
+import importlib.util
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+__all__ = [
+    "triton_group_norm",
+    "is_supported",
+    "select_strategy",
+    "GNConfig",
+    "default_config",
+    "SUPPORTED_ACTIVATIONS",
+    "TritonKernelError",
+]
+
+
+class TritonKernelError(RuntimeError):
+    """A failure of the Triton kernels themselves, with the original as ``__cause__``.
+
+    Raised in place of whatever ``_forward``/``_backward`` raised -- a missing or
+    mismatched ``triton``, an unwritable JIT cache, a compile error, a launch
+    failure, an API change between Triton releases.  It exists so that a caller
+    with a fallback (``ScaFFold.unet.group_norm``'s ladder) can catch *exactly*
+    "the kernel is broken" and nothing else, instead of catching ``Exception``
+    and trying to enumerate every framework mechanism that legitimately raises
+    through a forward -- saved-tensor pack hooks, ``torch.utils.checkpoint``'s
+    recompute control flow, functorch, a user's offloading hook.
+
+    Two things are deliberately *not* tagged and therefore propagate unchanged:
+
+    * ``torch.OutOfMemoryError``, which is a resource condition rather than a
+      defect (every fallback allocates an output of the same size, so retrying
+      one is a second, differently-shaped OOM at a call site the caller did not
+      ask about), and
+    * the ``ValueError``s ``_validate`` raises, which are contract violations by
+      the caller.  ``is_supported`` accepts exactly what ``_validate`` accepts,
+      so a caller that branches on it can never see one; if one escapes, that
+      is a bug in this module and must be loud.
+
+    The tagged region contains no autograd-observable work -- allocations and
+    kernel launches only, with ``save_for_backward`` happening in
+    ``_setup_context`` strictly *after* ``_forward`` returns -- so an exception
+    that carries this type is guaranteed to have been raised before the op saved
+    anything.  That is what makes retrying the call on another kernel safe.
+    """
+
+
+#: The activations that may be fused into the forward store.
+SUPPORTED_ACTIVATIONS = (None, "relu")
+
+#: Input dtypes the kernels read directly (statistics are always fp32).
+SUPPORTED_DTYPES = (torch.float32, torch.bfloat16, torch.float16)
+
+#: Largest linear element index representable in int32.
+_INT32_MAX = 2**31 - 1
+
+
+# --------------------------------------------------------------------------- #
+# tiling configuration
+# --------------------------------------------------------------------------- #
+class GNConfig:
+    """Tiling knobs.  A pure function of the shape => bitwise determinism.
+
+    ``stats_tile``/``elem_tile`` are *element* budgets per program (the spatial
+    block is ``tile // channels_per_voxel``, rounded down to a power of two);
+    ``nsplit_target`` is the total number of split-K partials wanted across the
+    batch, so the per-sample split count is ``nsplit_target // N``;
+    ``elem_progs`` caps the elementwise grid, each program then striding over
+    ``ceil(nblk_elem / elem_progs)`` tiles (0 = one program per tile).
+
+    These are **not** independent knobs, and in particular they stopped being
+    independent when the finalize passes were folded into the elementwise
+    kernels: each elementwise *program* now re-reads all ``nsplit`` split-K
+    partials, so the redundant traffic is ``min(nblk_elem, elem_progs) *
+    nsplit`` triples.  Raising ``nsplit_target`` for stats-kernel occupancy and
+    lowering ``elem_progs`` for redundancy pull against each other and were
+    tuned together; see the module docstring.
+    """
+
+    __slots__ = (
+        "stats_tile",
+        "stats_warps",
+        "nsplit_target",
+        "elem_tile",
+        "elem_warps",
+        "elem_progs",
+    )
+
+    def __init__(
+        self,
+        stats_tile=8192,
+        stats_warps=4,
+        nsplit_target=2048,
+        elem_tile=8192,
+        elem_warps=4,
+        elem_progs=2048,
+    ):
+        self.stats_tile = stats_tile
+        self.stats_warps = stats_warps
+        self.nsplit_target = nsplit_target
+        self.elem_tile = elem_tile
+        self.elem_warps = elem_warps
+        self.elem_progs = elem_progs
+
+    def key(self):
+        return (
+            self.stats_tile,
+            self.stats_warps,
+            self.nsplit_target,
+            self.elem_tile,
+            self.elem_warps,
+            self.elem_progs,
+        )
+
+    def __eq__(self, other):
+        return isinstance(other, GNConfig) and self.key() == other.key()
+
+    def __hash__(self):
+        return hash(self.key())
+
+    def __repr__(self):
+        return (
+            "GNConfig(stats_tile=%d, stats_warps=%d, nsplit_target=%d, "
+            "elem_tile=%d, elem_warps=%d, elem_progs=%d)" % self.key()
+        )
+
+
+#: Frozen tuning table, produced by coordinate descent on fwd+bwd time on one
+#: MI300A (228 CUs) at fp32 with ``num_groups=8``, keyed by the
+#: ``(num_channels, cube-root spatial extent)`` of the scale-8 ScaFFold UNet
+#: GroupNorm sites plus the ``[1,4096,4^3]`` tail.  Frozen -- never autotuned at
+#: run time -- because the split count fixes the reduction order and therefore
+#: the bits of the result.
+#:
+#: Re-derived for the fused (2+2 launch) kernels: every candidate was timed
+#: **interleaved against the incumbent** in one process, because this node moves
+#: the host-bound shapes by +-35% over the minutes a sweep takes.  The table is
+#: keyed on ``(C, edge)`` and not on ``N``: ``nsplit_target`` is a target for
+#: the split count *summed over the batch* (the per-sample count is
+#: ``nsplit_target // N``), so the same entry serves ``N > 1`` with the same
+#: total number of stats programs.  Verified at ``[2,1024,16^3]``,
+#: ``[4,2048,8^3]`` and ``[2,256,64^3]``.
+_TUNED = {
+    (64, 256): GNConfig(16384, 4, 2048, 8192, 4, 2048),
+    (128, 128): GNConfig(16384, 4, 2048, 16384, 4, 912),
+    (256, 64): GNConfig(16384, 4, 512, 16384, 8, 912),
+    (512, 32): GNConfig(32768, 8, 4096, 16384, 4, 0),
+    (1024, 16): GNConfig(65536, 4, 8192, 16384, 4, 0),
+    (2048, 8): GNConfig(16384, 8, 1024, 8192, 8, 228),
+    (4096, 4): GNConfig(16384, 4, 32, 4096, 8, 0),
+}
+
+_DEFAULT_CONFIG = GNConfig()
+
+
+def default_config(num_channels: int, spatial: int) -> GNConfig:
+    """Tiling for ``num_channels`` channels and ``spatial = D*H*W`` voxels."""
+    edge = round(spatial ** (1.0 / 3.0))
+    if edge**3 != spatial:
+        edge = None
+    return _TUNED.get((num_channels, edge), _DEFAULT_CONFIG)
+
+
+# --------------------------------------------------------------------------- #
+# small-shape dispatch hook
+# --------------------------------------------------------------------------- #
+#: Every strategy name ``select_strategy`` may return.  Only ``"split_k"`` is
+#: implemented; anything else raises rather than silently doing the wrong
+#: thing.
+STRATEGIES = ("split_k",)
+
+#: Spatial extent (``D*H*W``) below which the split-K chain is host-dispatch
+#: bound rather than bandwidth bound.  Measured on MI300A: at ``[1,2048,8^3]``
+#: the kernels do 0.030 ms of GPU work behind ~0.58 ms of
+#: Python/autograd/launch cost, and 0.086 ms of that is the *empty*
+#: ``torch.autograd.Function`` wrapper -- i.e. 68% of the remaining call is
+#: torch's plumbing, not this file's.  Purely informational --
+#: ``select_strategy`` does not use it.
+SMALL_SPATIAL_THRESHOLD = 4096
+
+
+def select_strategy(n: int, num_channels: int, spatial: int, num_groups: int) -> str:
+    """### SMALL-SHAPE DISPATCH HOOK ### -- the single point where a different
+    kernel strategy is chosen for a shape.
+
+    Returns a name from :data:`STRATEGIES`.  Today it always returns
+    ``"split_k"``: two forward and two backward kernels with split-K partial
+    reductions and the finalize passes fused into their consumers.  That is
+    bandwidth-optimal for the large shapes and, after the fusion, within
+    ~0.09 ms of the floor a Python ``autograd.Function`` can reach at the small
+    ones -- so there is much less left here than there looks.  Below roughly
+    ``SMALL_SPATIAL_THRESHOLD`` voxels the call is host bound, but the host cost
+    is now dominated by autograd and the dispatcher rather than by launches:
+    see "Why not one launch per pass" in the module docstring for the one
+    strategy that *would* cut it further and why it is not here.
+
+    If a second strategy ever lands, add its name to :data:`STRATEGIES`, return
+    it from here on a rule that is a **pure function of the shape**
+    (determinism depends on it), and branch on it in ``_dispatch`` -- which is
+    the only caller, sits in front of the memoized tiling plan, and is itself
+    called by both ``_forward`` and ``_backward``.  Nothing else in this file
+    needs to change.
+    """
+    return "split_k"
+
+
+# --------------------------------------------------------------------------- #
+# planning helpers
+# --------------------------------------------------------------------------- #
+def _prev_pow2(x: int) -> int:
+    p = 1
+    while p * 2 <= x:
+        p *= 2
+    return p
+
+
+def _next_pow2(x: int) -> int:
+    p = 1
+    while p < x:
+        p *= 2
+    return p
+
+
+def _cdiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+class _Plan:
+    """Everything the launcher needs, derived only from the shape + config."""
+
+    __slots__ = (
+        "n",
+        "channels",
+        "spatial",
+        "groups",
+        "group_channels",
+        "groups_p2",
+        "group_channels_p2",
+        "masked_c",
+        "int64",
+        "block_s_stats",
+        "nsplit",
+        "chunk",
+        "block_s_elem",
+        "nblk_elem",
+        "nprog_elem",
+        "elements_per_group",
+        "dwdb_rows",
+        "dwdb_block_c",
+        "dwdb_block_r",
+        "dwdb_progs",
+        "grid_dx",
+        "zero_dx",
+        "cfg",
+    )
+
+    def __init__(self, n, channels, spatial, groups, cfg, numel):
+        self.n = n
+        self.channels = channels
+        self.spatial = spatial
+        self.groups = groups
+        self.group_channels = channels // groups
+        self.groups_p2 = _next_pow2(groups)
+        self.group_channels_p2 = _next_pow2(self.group_channels)
+        # Only power-of-two group/channel counts tile the (G, C/G) axes exactly;
+        # anything else is rounded up and masked, which is correct but reads a
+        # few lanes it throws away.
+        self.masked_c = (
+            self.groups_p2 != groups or self.group_channels_p2 != self.group_channels
+        )
+        # int64 addressing is needed once a linear element index can exceed
+        # INT32_MAX.  [2,64,256^3] is exactly 2^31 elements, so this is not
+        # hypothetical at scale.  Only the *scalar* tile base is widened (see
+        # the kernels), which measurement showed to be free.
+        self.int64 = numel + channels > _INT32_MAX
+        self.cfg = cfg
+
+        voxel = self.groups_p2 * self.group_channels_p2
+        self.block_s_stats = max(1, _prev_pow2(cfg.stats_tile // max(1, voxel)))
+        ntiles = max(1, spatial // self.block_s_stats)
+        self.nsplit = _prev_pow2(
+            max(1, min(ntiles, max(1, cfg.nsplit_target // max(1, n))))
+        )
+        self.chunk = _cdiv(spatial, self.nsplit)
+        self.block_s_elem = max(1, _prev_pow2(cfg.elem_tile // max(1, voxel)))
+        self.nblk_elem = _cdiv(spatial, self.block_s_elem)
+        # Grid cap for the two elementwise kernels; each program then strides
+        # over its share of the tiles.  Bounds the cost of the fused finalize,
+        # which every *program* pays once.
+        self.nprog_elem = (
+            self.nblk_elem
+            if cfg.elem_progs <= 0
+            else min(self.nblk_elem, cfg.elem_progs)
+        )
+        self.elements_per_group = float(spatial * self.group_channels)
+        # Everything the fused dweight/dbias reduction in _dx_kernel needs.
+        # Precomputed rather than derived per call: the launch-bound shapes pay
+        # every Python statement in _backward, and _next_pow2 is a loop.
+        self.dwdb_rows = n * self.nsplit
+        self.dwdb_block_c = min(256, max(64, _next_pow2(channels)))
+        self.dwdb_block_r = 32 if self.dwdb_rows >= 32 else 1
+        self.dwdb_progs = _cdiv(channels, self.dwdb_block_c)
+        # Programs past nprog_elem run no elementwise loop iterations; they
+        # exist only when there are more dweight/dbias blocks than tiles.
+        self.grid_dx = max(self.nprog_elem, self.dwdb_progs)
+        self.zero_dx = self.group_channels * spatial == 1
+
+
+@functools.lru_cache(maxsize=256)
+def _plan(n, channels, spatial, groups, numel) -> _Plan:
+    """Memoized: a UNet presents a handful of shapes, and the three smallest
+    GroupNorm sites are host-dispatch bound, so rebuilding the plan (two
+    power-of-two loops and a dict lookup) on every call is measurable there.
+    Memoization cannot affect results -- the plan is a pure function of its
+    arguments, which is also what makes the kernels bitwise reproducible."""
+    return _Plan(n, channels, spatial, groups, default_config(channels, spatial), numel)
+
+
+def _dispatch(n, channels, spatial, groups, numel) -> _Plan:
+    """Consult the strategy hook, then build (or reuse) the tiling plan."""
+    strategy = select_strategy(n, channels, spatial, groups)
+    if strategy != "split_k":
+        raise NotImplementedError(
+            f"kernel strategy {strategy!r} selected by select_strategy() is not "
+            f"implemented; known strategies are {STRATEGIES}"
+        )
+    return _plan(n, channels, spatial, groups, numel)
+
+
+# --------------------------------------------------------------------------- #
+# Triton kernels (built lazily -- importing this module must not import triton)
+# --------------------------------------------------------------------------- #
+triton = None
+tl = None
+_welford_combine = None
+_stats_partial_kernel = None
+_normalize_kernel = None
+_bwd_partial_kernel = None
+_dx_kernel = None
+
+
+_TRITON_AVAILABLE = None
+
+
+def triton_available() -> bool:
+    """Whether Triton is importable, without importing it.
+
+    ``find_spec`` on a top-level name only touches the finders, so this stays
+    side-effect free and is safe to call from :func:`is_supported`.  The answer
+    is memoized in a plain global rather than an ``lru_cache`` because Dynamo
+    warns (loudly, once per process) when it traces through a cache wrapper.
+    """
+    global _TRITON_AVAILABLE
+    if _TRITON_AVAILABLE is None:
+        try:
+            _TRITON_AVAILABLE = importlib.util.find_spec("triton") is not None
+        except (ImportError, ValueError):
+            _TRITON_AVAILABLE = False
+    return _TRITON_AVAILABLE
+
+
+def _build_kernels():
+    """Import Triton and install the JIT kernels into this module's globals.
+
+    The kernels are defined inside a function purely so that ``import triton``
+    is deferred to the first GPU call; they are written into ``globals()`` so
+    Triton's name resolution (which reads ``fn.__globals__``) sees them.
+    """
+    global triton, tl
+    import triton as _triton
+    import triton.language as _tl
+
+    triton = _triton
+    tl = _tl
+
+    # ---------------------------------------------------------------- stats --
+    @_triton.jit
+    def _welford_combine(cnt_a, mean_a, m2_a, cnt_b, mean_b, m2_b):
+        """Chan's parallel merge of two (count, mean, M2) triples.
+
+        Exact for empty partials on either side (``cnt == 0`` leaves the other
+        triple untouched), which matters because the last split of a shape
+        whose spatial extent is not a multiple of the chunk size can be empty.
+        """
+        cnt = cnt_a + cnt_b
+        denom = tl.where(cnt == 0.0, 1.0, cnt)
+        delta = mean_b - mean_a
+        mean = mean_a + delta * (cnt_b / denom)
+        m2 = m2_a + m2_b + delta * delta * (cnt_a * cnt_b / denom)
+        return cnt, mean, m2
+
+    @_triton.jit
+    def _stats_partial_kernel(
+        X,
+        PCNT,
+        PMEAN,
+        PM2,
+        S,
+        CHUNK,
+        C: tl.constexpr,
+        G: tl.constexpr,
+        CG: tl.constexpr,
+        GP: tl.constexpr,
+        CGP: tl.constexpr,
+        NSPLIT: tl.constexpr,
+        BLOCK_S: tl.constexpr,
+        MASKED_C: tl.constexpr,
+        INT64: tl.constexpr,
+    ):
+        """One program per ``(split, n)``: Welford partials for every group.
+
+        Reads a dense ``(BLOCK_S, C)`` run of memory per step -- perfectly
+        coalesced -- and produces the statistics of all G groups at once, the
+        group axis being the inner channel axis reshaped to ``(G, C/G)``.
+        """
+        sp = tl.program_id(0)
+        n = tl.program_id(1)
+
+        offs_g = tl.arange(0, GP)
+        offs_j = tl.arange(0, CGP)
+        offs_s = tl.arange(0, BLOCK_S)
+        inner = offs_g[None, :, None] * CG + offs_j[None, None, :]
+        cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)
+        off = offs_s[:, None, None] * C + inner
+
+        s_begin = sp * CHUNK
+        s_end = tl.minimum(s_begin + CHUNK, S)
+
+        cnt = tl.zeros((GP,), dtype=tl.float32)
+        mean = tl.zeros((GP,), dtype=tl.float32)
+        m2 = tl.zeros((GP,), dtype=tl.float32)
+
+        for s0 in range(s_begin, s_end, BLOCK_S):
+            # Only the scalar tile base is ever widened to int64; the vector
+            # offsets stay int32 because they span at most BLOCK_S*C+C
+            # elements.  That keeps the wide arithmetic off the hot path.
+            if INT64:
+                base = (n.to(tl.int64) * S + s0) * C
+            else:
+                base = (n * S + s0) * C
+            nvalid = tl.minimum(BLOCK_S, s_end - s0)
+            m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP))
+            if MASKED_C:
+                m = m & cmask
+            # `other` is not load-bearing for any *bounded* value: `cnt_t`
+            # counts only the valid lanes, so `corr` below evaluates to
+            # `sum_valid(x)/cnt_t - mean0` and `mean_t = mean0 + corr` is the
+            # true mean whatever the masked lanes contributed, while `d`/`dd`
+            # are re-masked before they reach `m2_t`.  (Bounded: `other=1e30`
+            # would swamp `mean0` and the correction with it, and `other=inf`
+            # or `nan` would poison it outright.)  0.0 is kept because it is
+            # the value that survives all three of those, not because the
+            # cancellation is something to rely on.
+            x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32)
+
+            # Corrected two-pass within the tile: the first mean loses digits
+            # to the magnitude of the data, `corr` puts them back, and the
+            # centred squares are then accurate to fp32 roundoff.  Everything
+            # here is register traffic; the tile is read from HBM exactly once.
+            # `corr` is worth 1.3x to 1472x on `rstd` once the data's mean
+            # dominates its spread -- see "What the third pass is worth" in the
+            # module docstring for the measurements, for the one regime where
+            # it goes the other way, and for why the *output* error barely
+            # moves even where `rstd` improves by three orders of magnitude.
+            cnt_t = (nvalid * CG).to(tl.float32)
+            mean0 = tl.sum(tl.sum(x, 2), 0) / cnt_t
+            d = tl.where(m, x - mean0[None, :, None], 0.0)
+            corr = tl.sum(tl.sum(d, 2), 0) / cnt_t
+            dd = tl.where(m, d - corr[None, :, None], 0.0)
+            m2_t = tl.sum(tl.sum(dd * dd, 2), 0)
+            mean_t = mean0 + corr
+
+            new_cnt = cnt + cnt_t
+            delta = mean_t - mean
+            mean = mean + delta * (cnt_t / new_cnt)
+            m2 = m2 + m2_t + delta * delta * (cnt * cnt_t / new_cnt)
+            cnt = new_cnt
+
+        o = (n * NSPLIT + sp) * G + offs_g
+        gm = offs_g < G
+        tl.store(PCNT + o, cnt, mask=gm)
+        tl.store(PMEAN + o, mean, mask=gm)
+        tl.store(PM2 + o, m2, mask=gm)
+
+    # ------------------------------------------------------------ normalize --
+    @_triton.jit
+    def _normalize_kernel(
+        X,
+        Y,
+        PCNT,
+        PMEAN,
+        PM2,
+        MEAN,
+        RSTD,
+        W,
+        B,
+        S,
+        M,
+        eps,
+        C: tl.constexpr,
+        G: tl.constexpr,
+        CG: tl.constexpr,
+        GP: tl.constexpr,
+        CGP: tl.constexpr,
+        NSPLIT: tl.constexpr,
+        BLOCK_S: tl.constexpr,
+        NBLK: tl.constexpr,
+        NPROG: tl.constexpr,
+        RELU: tl.constexpr,
+        HAS_W: tl.constexpr,
+        HAS_B: tl.constexpr,
+        MASKED_C: tl.constexpr,
+        INT64: tl.constexpr,
+    ):
+        """Finalize the split-K statistics, then normalize NBLK/NPROG tiles.
+
+        The finalize is *recomputed by every program* rather than round-tripped
+        through its own kernel launch: merging NSPLIT Welford triples is a few
+        KB of L2-resident traffic and a tree reduction over a ``(NSPLIT, GP)``
+        tile, which is cheaper than the ~9 us launch it replaces.  What it is
+        *not* cheap enough for is being paid once per tile at the largest
+        shapes, where the flat grid is 10^5 programs: the grid is therefore
+        capped at ``NPROG`` and each program strides over its share of the
+        ``NBLK`` tiles, so the redundant read costs ``NPROG * NSPLIT`` and not
+        ``NBLK * NSPLIT``.  See :class:`GNConfig` -- ``nsplit_target``,
+        ``elem_tile`` and ``elem_progs`` are one joint tuning problem, not
+        three independent knobs.
+
+        Every program reads the same partials with the same tile shape, so they
+        all get bit-identical ``mean``/``rstd``; program 0 stores them for the
+        backward.  The loop carries nothing across iterations, so the striding
+        cannot affect the result.
+        """
+        pid = tl.program_id(0)
+        n = tl.program_id(1)
+
+        offs_g = tl.arange(0, GP)
+        offs_j = tl.arange(0, CGP)
+        offs_s = tl.arange(0, BLOCK_S)
+        inner = offs_g[None, :, None] * CG + offs_j[None, None, :]
+        off = offs_s[:, None, None] * C + inner
+        wb = offs_g[:, None] * CG + offs_j[None, :]
+        wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG)
+
+        gm = offs_g < G
+        offs_p = tl.arange(0, NSPLIT)
+        pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :]
+        pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP))
+        # Padded group lanes load cnt == 0, which _welford_combine treats as the
+        # identity, so they merge to (0, 0, 0) and are masked off on the store.
+        # Reduced over axis 0 -- the *slowest* axis -- deliberately: reducing
+        # the fastest axis of a 2-D tile makes Triton stage the whole tile
+        # through LDS, which for a (G, NSPLIT) tile is 64 KB per array.
+        cnt_p = tl.load(PCNT + pidx, mask=pm, other=0.0)
+        _cnt, mu, m2 = tl.reduce(
+            (
+                cnt_p,
+                tl.load(PMEAN + pidx, mask=pm, other=0.0),
+                tl.load(PM2 + pidx, mask=pm, other=0.0),
+            ),
+            0,
+            _welford_combine,
+        )
+        # `_cnt` equals M by construction; M is passed in so the divisor is the
+        # exact element count rather than a float accumulated from partials.
+        rs = 1.0 / tl.sqrt(m2 / M + eps)
+        if pid == 0:
+            tl.store(MEAN + n * G + offs_g, mu, mask=gm)
+            tl.store(RSTD + n * G + offs_g, rs, mask=gm)
+        mean = mu[None, :, None]
+        rstd = rs[None, :, None]
+        if HAS_W:
+            w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+        else:
+            w = tl.full((1, GP, CGP), 1.0, tl.float32)
+        if HAS_B:
+            b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+        else:
+            b = tl.zeros((1, GP, CGP), dtype=tl.float32)
+
+        for blk in tl.range(pid, NBLK, NPROG):
+            s0 = blk * BLOCK_S
+            m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP))
+            if MASKED_C:
+                m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG))
+            if INT64:
+                base = (n.to(tl.int64) * S + s0) * C
+            else:
+                base = (n * S + s0) * C
+            x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32)
+            xhat = (x - mean) * rstd
+            y = xhat * w + b
+            if RELU:
+                # `tl.maximum(y, 0.0)` and `tl.where(y > 0, y, 0.0)` both map NaN
+                # to 0.0 (the first returns the non-NaN operand, the second
+                # because `NaN > 0` is False), while `F.relu` propagates it.
+                # Testing the *complement* keeps NaN on the pass-through side:
+                # `NaN <= 0` is also False, so NaN falls to `y`.  Bit-identical
+                # to `F.relu` on NaN, +-Inf and -0.0 (which both flush to +0.0),
+                # for one comparison and one select -- see the module docstring.
+                y = tl.where(y <= 0.0, 0.0, y)
+            tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m)
+
+    # ------------------------------------------------------------- backward --
+    @_triton.jit
+    def _bwd_partial_kernel(
+        X,
+        DY,
+        MEAN,
+        RSTD,
+        W,
+        B,
+        PS1,
+        PS2,
+        PDW,
+        PDB,
+        S,
+        CHUNK,
+        C: tl.constexpr,
+        G: tl.constexpr,
+        CG: tl.constexpr,
+        GP: tl.constexpr,
+        CGP: tl.constexpr,
+        NSPLIT: tl.constexpr,
+        BLOCK_S: tl.constexpr,
+        RELU: tl.constexpr,
+        HAS_W: tl.constexpr,
+        HAS_B: tl.constexpr,
+        MASKED_C: tl.constexpr,
+        INT64: tl.constexpr,
+    ):
+        """Partials for the two per-``(n, g)`` reductions used by dx, and for
+        the per-channel dweight / dbias reductions."""
+        sp = tl.program_id(0)
+        n = tl.program_id(1)
+
+        offs_g = tl.arange(0, GP)
+        offs_j = tl.arange(0, CGP)
+        offs_s = tl.arange(0, BLOCK_S)
+        inner = offs_g[None, :, None] * CG + offs_j[None, None, :]
+        cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)
+        off = offs_s[:, None, None] * C + inner
+        wb = offs_g[:, None] * CG + offs_j[None, :]
+        wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG)
+
+        s_begin = sp * CHUNK
+        s_end = tl.minimum(s_begin + CHUNK, S)
+
+        gm = offs_g < G
+        mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None]
+        rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None]
+        if HAS_W:
+            w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+        else:
+            w = tl.full((1, GP, CGP), 1.0, tl.float32)
+        if HAS_B:
+            b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+        else:
+            b = tl.zeros((1, GP, CGP), dtype=tl.float32)
+
+        acc1 = tl.zeros((GP,), dtype=tl.float32)
+        acc2 = tl.zeros((GP,), dtype=tl.float32)
+        accdw = tl.zeros((GP, CGP), dtype=tl.float32)
+        accdb = tl.zeros((GP, CGP), dtype=tl.float32)
+
+        for s0 in range(s_begin, s_end, BLOCK_S):
+            if INT64:
+                base = (n.to(tl.int64) * S + s0) * C
+            else:
+                base = (n * S + s0) * C
+            nvalid = tl.minimum(BLOCK_S, s_end - s0)
+            m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP))
+            if MASKED_C:
+                m = m & cmask
+            x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32)
+            dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32)
+            xhat = (x - mean) * rstd
+            if RELU:
+                # Identical expression (and therefore identical rounding) to
+                # the forward's pre-activation, so the sign test agrees with
+                # the forward bit for bit.  Masked lanes carry dy == 0, so
+                # gating cannot resurrect them.  Spelled as the *complement*
+                # (`pre <= 0` zeroes) rather than `pre > 0` passes, so that a
+                # NaN pre-activation passes the gradient through: that is what
+                # `threshold_backward(grad, result, 0)` -- ReLU's real backward
+                # -- does, since `NaN <= 0` is False.  See the forward store.
+                dy = tl.where(xhat * w + b <= 0.0, 0.0, dy)
+            dyw = dy * w
+            acc1 += tl.sum(tl.sum(dyw, 2), 0)
+            acc2 += tl.sum(tl.sum(dyw * xhat, 2), 0)
+            accdw += tl.sum(dy * xhat, 0)
+            accdb += tl.sum(dy, 0)
+
+        o = (n * NSPLIT + sp) * G + offs_g
+        tl.store(PS1 + o, acc1, mask=gm)
+        tl.store(PS2 + o, acc2, mask=gm)
+        row = (n * NSPLIT + sp) * C + wb
+        tl.store(PDW + row, accdw, mask=wbm)
+        tl.store(PDB + row, accdb, mask=wbm)
+
+    @_triton.jit
+    def _dx_kernel(
+        X,
+        DY,
+        DX,
+        MEAN,
+        RSTD,
+        W,
+        B,
+        PS1,
+        PS2,
+        PDW,
+        PDB,
+        DW,
+        DB,
+        ROWS,
+        S,
+        M,
+        C: tl.constexpr,
+        G: tl.constexpr,
+        CG: tl.constexpr,
+        GP: tl.constexpr,
+        CGP: tl.constexpr,
+        NSPLIT: tl.constexpr,
+        BLOCK_S: tl.constexpr,
+        NBLK: tl.constexpr,
+        NPROG: tl.constexpr,
+        NDW: tl.constexpr,
+        BLOCK_C: tl.constexpr,
+        BLOCK_R: tl.constexpr,
+        RELU: tl.constexpr,
+        HAS_W: tl.constexpr,
+        HAS_B: tl.constexpr,
+        MASKED_C: tl.constexpr,
+        INT64: tl.constexpr,
+        ZERO_DX: tl.constexpr,
+    ):
+        """The whole backward tail: dweight/dbias, the c1/c2 finalize, and dx.
+
+        Two reductions that used to be their own launches ride along here.  The
+        per-channel dweight/dbias row reduction is done by the first ``NDW``
+        programs of ``n == 0`` (a single pass over an ``(n*nsplit, C)`` scratch,
+        i.e. a few hundred KB); the per-``(n, g)`` c1/c2 finalize is recomputed
+        redundantly by every program, once, before the tile loop -- exactly as
+        in ``_normalize_kernel``, and capped the same way.  The grid is
+        ``(max(NPROG, NDW), n)``; programs past ``NPROG`` exist only to cover
+        the dweight/dbias rows and run no loop iterations.
+        """
+        pid = tl.program_id(0)
+        n = tl.program_id(1)
+
+        # ---- dweight / dbias: rows of the split-K scratch, once per channel --
+        if n == 0:
+            if pid < NDW:
+                offs_c = pid * BLOCK_C + tl.arange(0, BLOCK_C)
+                mc = offs_c < C
+                accw = tl.zeros((BLOCK_C,), dtype=tl.float32)
+                accb = tl.zeros((BLOCK_C,), dtype=tl.float32)
+                for r0 in range(0, ROWS, BLOCK_R):
+                    offs_r = r0 + tl.arange(0, BLOCK_R)
+                    rm = (offs_r[:, None] < ROWS) & mc[None, :]
+                    roff = offs_r[:, None] * C + offs_c[None, :]
+                    accw += tl.sum(tl.load(PDW + roff, mask=rm, other=0.0), 0)
+                    accb += tl.sum(tl.load(PDB + roff, mask=rm, other=0.0), 0)
+                tl.store(DW + offs_c, accw, mask=mc)
+                tl.store(DB + offs_c, accb, mask=mc)
+
+        offs_g = tl.arange(0, GP)
+        offs_j = tl.arange(0, CGP)
+        offs_s = tl.arange(0, BLOCK_S)
+        inner = offs_g[None, :, None] * CG + offs_j[None, None, :]
+        off = offs_s[:, None, None] * C + inner
+        wb = offs_g[:, None] * CG + offs_j[None, :]
+        wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG)
+
+        if ZERO_DX:
+            # One element per group: mean == x and var == 0 identically, so
+            # xhat is the constant 0 and y does not depend on x at all -- the
+            # exact d_input is zero everywhere.  The expression below would
+            # instead return rstd * (dy*w - c1), and since the compiler
+            # contracts that to fma(dy, w, -c1) while c1 was accumulated from
+            # the *rounded* product, what survives is the product's rounding
+            # error amplified by rstd = 1/sqrt(eps) ~ 316 (2.2e-05 at
+            # eps=1e-5).  Answering with the exact zero costs one constexpr.
+            zero = tl.zeros((BLOCK_S, GP, CGP), dtype=tl.float32)
+            for blk in tl.range(pid, NBLK, NPROG):
+                s0 = blk * BLOCK_S
+                m = tl.broadcast_to(
+                    (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)
+                )
+                if MASKED_C:
+                    m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG))
+                if INT64:
+                    base = (n.to(tl.int64) * S + s0) * C
+                else:
+                    base = (n * S + s0) * C
+                tl.store(DX + base + off, zero.to(DX.dtype.element_ty), mask=m)
+        else:
+            gm = offs_g < G
+            offs_p = tl.arange(0, NSPLIT)
+            pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :]
+            pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP))
+            c1 = (tl.sum(tl.load(PS1 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None]
+            c2 = (tl.sum(tl.load(PS2 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None]
+
+            mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None]
+            rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None]
+            if HAS_W:
+                w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+            else:
+                w = tl.full((1, GP, CGP), 1.0, tl.float32)
+            if HAS_B:
+                b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :]
+            else:
+                b = tl.zeros((1, GP, CGP), dtype=tl.float32)
+
+            for blk in tl.range(pid, NBLK, NPROG):
+                s0 = blk * BLOCK_S
+                m = tl.broadcast_to(
+                    (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)
+                )
+                if MASKED_C:
+                    m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG))
+                if INT64:
+                    base = (n.to(tl.int64) * S + s0) * C
+                else:
+                    base = (n * S + s0) * C
+                x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32)
+                dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32)
+                xhat = (x - mean) * rstd
+                if RELU:
+                    # Same complement spelling as _bwd_partial_kernel: a NaN
+                    # pre-activation must pass the gradient, exactly as
+                    # `threshold_backward(grad, result, 0)` does.
+                    dy = tl.where(xhat * w + b <= 0.0, 0.0, dy)
+                dyw = dy * w
+                dx = rstd * (dyw - c1 - xhat * c2)
+                tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m)
+
+    globals().update(
+        _welford_combine=_welford_combine,
+        _stats_partial_kernel=_stats_partial_kernel,
+        _normalize_kernel=_normalize_kernel,
+        _bwd_partial_kernel=_bwd_partial_kernel,
+        _dx_kernel=_dx_kernel,
+    )
+
+
+def _ensure_kernels():
+    if _stats_partial_kernel is None:
+        _build_kernels()
+
+
+# --------------------------------------------------------------------------- #
+# python drivers
+# --------------------------------------------------------------------------- #
+_CL_FORMAT = torch.channels_last_3d
+
+#: Reused so the common (already-current device) path allocates nothing.
+_NO_GUARD = contextlib.nullcontext()
+
+
+def _device_guard(device: torch.device):
+    """Make ``device`` current for the kernel launches inside the ``with``.
+
+    A Triton launch goes to whatever device is *current*, not to the device the
+    argument tensors live on, so without this a tensor on ``cuda:1`` while
+    ``cuda:0`` is current makes the kernel dereference another device's pointers
+    and the process dies with ``Memory access fault by GPU node-N``.  ATen ops
+    (including ``F.group_norm``) carry a ``DeviceGuard`` and handle the same
+    call, so this is required for the drop-in contract, not a nicety.
+
+    The ``current_device()`` test is not about correctness but about *cost*.
+    Measured on this node (median of 200k calls, torch 2.13.0+rocm7.2):
+    ``with torch.cuda.device(t.device)`` is **1.55 us** of host time per call,
+    ``with torch.cuda._DeviceGuard(t.device.index)`` **0.61 us**, and this
+    helper **0.51 us** when the tensor is already on the current device --
+    which it is on every ScaFFold call, since ScaFFold pins one device per
+    rank.  Two of those (forward + backward) against the 0.65 ms fwd+bwd of the
+    two smallest scale-8 shapes, which are host-dispatch bound, is 0.16%
+    instead of 0.48%.
+    """
+    if device.index == torch.cuda.current_device():
+        return _NO_GUARD
+    return torch.cuda.device(device)
+
+
+def _shape_of(input: torch.Tensor):
+    n, channels = input.shape[0], input.shape[1]
+    spatial = 1
+    for d in input.shape[2:]:
+        spatial *= d
+    return n, channels, spatial
+
+
+def _tag_kernel_failures(fn):
+    """Re-raise anything ``fn`` raises as :class:`TritonKernelError`.
+
+    Applied to the two functions that do nothing but import Triton, allocate
+    scratch and launch kernels.  The region is *closed*: it runs no
+    autograd-observable op, so a blanket ``except Exception`` here cannot
+    swallow framework control flow the way one at the call site would -- there
+    is no pack hook, no recompute stop and no functorch layer inside it.  That
+    closure is what lets the caller's fallback ladder use a one-element
+    allowlist instead of an ever-growing denylist.
+
+    ``torch.OutOfMemoryError`` is passed through untagged; see
+    :class:`TritonKernelError`.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        try:
+            return fn(*args, **kwargs)
+        except torch.OutOfMemoryError:
+            raise
+        except TritonKernelError:
+            raise
+        except Exception as e:
+            raise TritonKernelError(
+                f"{fn.__name__} failed ({type(e).__name__}: {e})"
+            ) from e
+
+    return wrapper
+
+
+@_tag_kernel_failures
+def _forward(input, num_groups, weight, bias, eps, activation, out_dtype):
+    _ensure_kernels()
+    n, channels, spatial = _shape_of(input)
+    plan = _dispatch(n, channels, spatial, num_groups, input.numel())
+    groups = num_groups
+    device = input.device
+
+    with _device_guard(device):
+        pcnt = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32)
+        pmean = torch.empty_like(pcnt)
+        pm2 = torch.empty_like(pcnt)
+        mean = torch.empty((n, groups), device=device, dtype=torch.float32)
+        rstd = torch.empty_like(mean)
+        out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT)
+
+        _stats_partial_kernel[(plan.nsplit, n)](
+            input,
+            pcnt,
+            pmean,
+            pm2,
+            spatial,
+            plan.chunk,
+            C=channels,
+            G=groups,
+            CG=plan.group_channels,
+            GP=plan.groups_p2,
+            CGP=plan.group_channels_p2,
+            NSPLIT=plan.nsplit,
+            BLOCK_S=plan.block_s_stats,
+            MASKED_C=plan.masked_c,
+            INT64=plan.int64,
+            num_warps=plan.cfg.stats_warps,
+        )
+        _normalize_kernel[(plan.nprog_elem, n)](
+            input,
+            out,
+            pcnt,
+            pmean,
+            pm2,
+            mean,
+            rstd,
+            weight,
+            bias,
+            spatial,
+            plan.elements_per_group,
+            eps,
+            C=channels,
+            G=groups,
+            CG=plan.group_channels,
+            GP=plan.groups_p2,
+            CGP=plan.group_channels_p2,
+            NSPLIT=plan.nsplit,
+            BLOCK_S=plan.block_s_elem,
+            NBLK=plan.nblk_elem,
+            NPROG=plan.nprog_elem,
+            RELU=activation == "relu",
+            HAS_W=weight is not None,
+            HAS_B=bias is not None,
+            MASKED_C=plan.masked_c,
+            INT64=plan.int64,
+            num_warps=plan.cfg.elem_warps,
+        )
+    return out, mean, rstd
+
+
+@_tag_kernel_failures
+def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation):
+    _ensure_kernels()
+    n, channels, spatial = _shape_of(input)
+    plan = _dispatch(n, channels, spatial, num_groups, input.numel())
+    groups = num_groups
+    device = input.device
+
+    with _device_guard(device):
+        ps1 = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32)
+        ps2 = torch.empty_like(ps1)
+        pdw = torch.empty(
+            n * plan.nsplit * channels, device=device, dtype=torch.float32
+        )
+        pdb = torch.empty_like(pdw)
+
+        _bwd_partial_kernel[(plan.nsplit, n)](
+            input,
+            grad_out,
+            mean,
+            rstd,
+            weight,
+            bias,
+            ps1,
+            ps2,
+            pdw,
+            pdb,
+            spatial,
+            plan.chunk,
+            C=channels,
+            G=groups,
+            CG=plan.group_channels,
+            GP=plan.groups_p2,
+            CGP=plan.group_channels_p2,
+            NSPLIT=plan.nsplit,
+            BLOCK_S=plan.block_s_stats,
+            RELU=activation == "relu",
+            HAS_W=weight is not None,
+            HAS_B=bias is not None,
+            MASKED_C=plan.masked_c,
+            INT64=plan.int64,
+            num_warps=plan.cfg.stats_warps,
+        )
+
+        d_weight = torch.empty(channels, device=device, dtype=torch.float32)
+        d_bias = torch.empty_like(d_weight)
+        d_input = torch.empty_like(input, memory_format=_CL_FORMAT)
+        _dx_kernel[(plan.grid_dx, n)](
+            input,
+            grad_out,
+            d_input,
+            mean,
+            rstd,
+            weight,
+            bias,
+            ps1,
+            ps2,
+            pdw,
+            pdb,
+            d_weight,
+            d_bias,
+            plan.dwdb_rows,
+            spatial,
+            plan.elements_per_group,
+            C=channels,
+            G=groups,
+            CG=plan.group_channels,
+            GP=plan.groups_p2,
+            CGP=plan.group_channels_p2,
+            NSPLIT=plan.nsplit,
+            BLOCK_S=plan.block_s_elem,
+            NBLK=plan.nblk_elem,
+            NPROG=plan.nprog_elem,
+            NDW=plan.dwdb_progs,
+            BLOCK_C=plan.dwdb_block_c,
+            BLOCK_R=plan.dwdb_block_r,
+            RELU=activation == "relu",
+            HAS_W=weight is not None,
+            HAS_B=bias is not None,
+            MASKED_C=plan.masked_c,
+            INT64=plan.int64,
+            ZERO_DX=plan.zero_dx,
+            num_warps=plan.cfg.elem_warps,
+        )
+    return d_input, d_weight, d_bias
+
+
+# --------------------------------------------------------------------------- #
+# torch.library registration
+# --------------------------------------------------------------------------- #
+def _one_value_per_channel(input, num_groups: int) -> bool:
+    """Whether ``F.group_norm`` would reject this shape as degenerate.
+
+    ``F.group_norm`` runs ``_verify_batch_size([N*C//G, G, *spatial])``, which
+    raises ``ValueError("Expected more than 1 value per channel when
+    training")`` exactly when ``N * (C/G) * D*H*W == 1``.  All three factors are
+    positive, so that holds iff ``N == 1``, ``C == num_groups`` and the spatial
+    extent is 1 -- i.e. iff ``numel == C == num_groups``, which is the cheap
+    form used here (``numel`` is wanted by the caller anyway).
+
+    Rejected rather than served: the kernel *can* compute it (it returns
+    ``bias``, since every group has zero variance), but a caller that branches
+    on :func:`is_supported` would then get a result where the op this replaces
+    raises, which is a worse failure than being slower.
+    """
+    channels = input.shape[1]
+    return channels == num_groups and input.numel() == channels
+
+
+def _validate(input, num_groups, weight, bias, activation):
+    if activation not in SUPPORTED_ACTIVATIONS:
+        raise ValueError(
+            f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}"
+        )
+    if input.dim() != 5:
+        raise ValueError(f"expected a 5-D NCDHW tensor, got {tuple(input.shape)}")
+    if num_groups <= 0 or input.shape[1] % num_groups != 0:
+        raise ValueError(
+            f"num_channels={input.shape[1]} is not divisible by num_groups={num_groups}"
+        )
+    if _one_value_per_channel(input, num_groups):
+        # Same rejection, and the same exception type, as F.group_norm's
+        # _verify_batch_size; see _one_value_per_channel.
+        raise ValueError(
+            f"Expected more than 1 value per channel when training, got input "
+            f"size {tuple(input.shape)} with num_groups={num_groups}"
+        )
+    if input.dtype not in SUPPORTED_DTYPES:
+        raise ValueError(f"unsupported input dtype {input.dtype}")
+    if not input.is_contiguous(memory_format=_CL_FORMAT):
+        # Required, not converted: the fake kernel promises the *input's*
+        # memory format for the output, so silently converting here would make
+        # the traced and eager results disagree on strides.  The public
+        # ``triton_group_norm`` routes non-channels-last input to
+        # ``F.group_norm`` before it ever reaches this op.
+        raise ValueError(
+            "input must be channels_last_3d-contiguous; use triton_group_norm() "
+            "which falls back to F.group_norm for other layouts"
+        )
+    for name, t in (("weight", weight), ("bias", bias)):
+        if t is not None and t.numel() != input.shape[1]:
+            raise ValueError(
+                f"{name} has {t.numel()} elements, expected {input.shape[1]}"
+            )
+
+
+@torch.library.custom_op(
+    "scaffold_gn::group_norm", mutates_args=(), device_types="cuda"
+)
+def _group_norm_op(
+    input: torch.Tensor,
+    num_groups: int,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    eps: float,
+    activation: Optional[str],
+    out_dtype: Optional[torch.dtype],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Channels-last GroupNorm forward: returns ``(output, mean, rstd)``.
+
+    ``mean``/``rstd`` are ``(N, num_groups)`` fp32 tensors kept for the
+    backward; they are marked non-differentiable in ``_setup_context``
+    (nothing produces a gradient for them), so they come back with
+    ``requires_grad=False`` and differentiating through them raises rather than
+    returning zeros.  Callers should treat them as opaque.
+    """
+    _validate(input, num_groups, weight, bias, activation)
+    weight = None if weight is None else weight.contiguous()
+    bias = None if bias is None else bias.contiguous()
+    out, mean, rstd = _forward(
+        input, num_groups, weight, bias, eps, activation, out_dtype or input.dtype
+    )
+    return out, mean, rstd
+
+
+@_group_norm_op.register_fake
+def _(input, num_groups, weight, bias, eps, activation, out_dtype):
+    # empty_like preserves the input's memory format, which is the contract.
+    out = torch.empty_like(input, dtype=out_dtype or input.dtype)
+    mean = input.new_empty((input.shape[0], num_groups), dtype=torch.float32)
+    rstd = input.new_empty((input.shape[0], num_groups), dtype=torch.float32)
+    return out, mean, rstd
+
+
+@torch.library.custom_op(
+    "scaffold_gn::group_norm_backward", mutates_args=(), device_types="cuda"
+)
+def _group_norm_backward_op(
+    grad_out: torch.Tensor,
+    input: torch.Tensor,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    mean: torch.Tensor,
+    rstd: torch.Tensor,
+    num_groups: int,
+    activation: Optional[str],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Returns ``(d_input, d_weight, d_bias)``.
+
+    ``d_weight``/``d_bias`` are zero-element tensors when the corresponding
+    parameter is ``None``.  ``d_input`` always has the input's dtype and
+    channels-last memory format.
+    """
+    if not grad_out.is_contiguous(memory_format=_CL_FORMAT):
+        grad_out = grad_out.contiguous(memory_format=_CL_FORMAT)
+    if not input.is_contiguous(memory_format=_CL_FORMAT):
+        input = input.contiguous(memory_format=_CL_FORMAT)
+    weight = None if weight is None else weight.contiguous()
+    bias = None if bias is None else bias.contiguous()
+    d_input, d_weight, d_bias = _backward(
+        grad_out, input, weight, bias, mean, rstd, num_groups, activation
+    )
+    # No `d_input.to(input.dtype)`: `_backward` allocates it with
+    # `empty_like(input)` and `_dx_kernel` stores through `DX.dtype.element_ty`,
+    # so it already *is* the input's dtype.
+    if weight is None:
+        d_weight = d_weight.new_empty(0)
+    else:
+        d_weight = d_weight.to(weight.dtype)
+    if bias is None:
+        d_bias = d_bias.new_empty(0)
+    else:
+        d_bias = d_bias.to(bias.dtype)
+    return d_input, d_weight, d_bias
+
+
+@_group_norm_backward_op.register_fake
+def _(grad_out, input, weight, bias, mean, rstd, num_groups, activation):
+    channels = input.shape[1]
+    # channels_last_3d, *not* the input's own format: the real op relayouts a
+    # non-channels-last `input` and always returns a channels-last `d_input`,
+    # so promising `empty_like(input)` here would hand torch.compile the wrong
+    # strides for any contiguous NCDHW input -- silently, since eager never
+    # consults the fake kernel.
+    d_input = torch.empty_like(input, memory_format=_CL_FORMAT)
+    d_weight = input.new_empty(
+        channels if weight is not None else 0,
+        dtype=weight.dtype if weight is not None else torch.float32,
+    )
+    d_bias = input.new_empty(
+        channels if bias is not None else 0,
+        dtype=bias.dtype if bias is not None else torch.float32,
+    )
+    return d_input, d_weight, d_bias
+
+
+def _setup_context(ctx, inputs, output):
+    input, num_groups, weight, bias, eps, activation, out_dtype = inputs
+    _out, mean, rstd = output
+    # Outputs 1 and 2 are backward state, not results: nothing produces a
+    # gradient for them.  Without this they come back requiring grad, and
+    # differentiating through them *succeeds* -- autograd materializes an
+    # all-zero cotangent for the unused `out` and runs the whole backward to
+    # return zeros, which is a plausible wrong answer rather than an error.
+    ctx.mark_non_differentiable(mean, rstd)
+    ctx.save_for_backward(input, weight, bias, mean, rstd)
+    ctx.num_groups = num_groups
+    ctx.activation = activation
+    ctx.needs = (
+        ctx.needs_input_grad[0],
+        ctx.needs_input_grad[2],
+        ctx.needs_input_grad[3],
+    )
+
+
+def _autograd_backward(ctx, grad_out, grad_mean, grad_rstd):
+    input, weight, bias, mean, rstd = ctx.saved_tensors
+    need_x, need_w, need_b = ctx.needs
+    if not (need_x or need_w or need_b):
+        return None, None, None, None, None, None, None
+    d_input, d_weight, d_bias = torch.ops.scaffold_gn.group_norm_backward(
+        grad_out, input, weight, bias, mean, rstd, ctx.num_groups, ctx.activation
+    )
+    return (
+        d_input if need_x else None,
+        None,  # num_groups
+        d_weight if need_w else None,
+        d_bias if need_b else None,
+        None,  # eps
+        None,  # activation
+        None,  # out_dtype
+    )
+
+
+torch.library.register_autograd(
+    "scaffold_gn::group_norm", _autograd_backward, setup_context=_setup_context
+)
+
+
+# --------------------------------------------------------------------------- #
+# public API
+# --------------------------------------------------------------------------- #
+def _autocast_active(input: torch.Tensor) -> bool:
+    """Whether autocast is enabled for this tensor's device type."""
+    try:
+        return bool(torch.is_autocast_enabled(input.device.type))
+    except (RuntimeError, TypeError):  # device type autocast does not know
+        return False
+
+
+def _autocast_out_dtype(input: torch.Tensor) -> Optional[torch.dtype]:
+    """``F.group_norm``'s output dtype for this input, or None for "unchanged".
+
+    ``at::group_norm`` carries autocast's ``fp32`` cast policy, so under an
+    enabled autocast region it upcasts its input and returns fp32 whatever came
+    in.  Verified empirically on torch 2.13.0+rocm7.2 for fp32/bf16/fp16 input
+    and both bf16 and fp16 autocast dtypes.
+    """
+    if input.dtype is not torch.float32 and _autocast_active(input):
+        return torch.float32
+    return None
+
+
+def is_supported(
+    input,
+    num_groups: int,
+    weight=None,
+    bias=None,
+    activation: Optional[str] = None,
+) -> bool:
+    """Whether the native channels-last Triton kernel can serve this call.
+
+    Cheap (a handful of attribute reads and one stride check) and side-effect
+    free -- in particular it does not import Triton, allocate, or launch.
+    ``False`` means "use ``F.group_norm``": the fast path needs a 5-D CUDA
+    tensor that is ``channels_last_3d``-contiguous, an fp32/bf16/fp16 dtype, a
+    channel count divisible by ``num_groups``, and affine parameters whose
+    dtype ``F.group_norm`` would itself accept for this input (equal to the
+    input's, or fp32 under autocast, which is what autocast would produce).
+    Shapes ``F.group_norm`` itself rejects are rejected here too, so that
+    branching on this predicate can never turn a stock ``ValueError`` into an
+    answer (see :func:`_one_value_per_channel`).
+
+    ``True`` promises the *first* derivative only: the backward is itself a
+    custom op with no autograd formula, so a second ``torch.autograd.grad``
+    raises where stock ``F.group_norm`` would succeed.  Callers that need a
+    gradient penalty or a Hessian-vector product must not take this path.
+
+    Note that this is a capability predicate, not a layout classifier: for
+    shapes whose spatial *and* channel extents make the contiguous and
+    channels-last-3d stride patterns coincide (e.g. ``(N, C, 1, 1, 1)``), a
+    plain contiguous tensor is accepted, correctly -- it is the same bytes.
+    """
+    if activation not in SUPPORTED_ACTIVATIONS:
+        return False
+    if not isinstance(input, torch.Tensor):
+        return False
+    if input.device.type != "cuda" or not triton_available():
+        return False
+    if input.dim() != 5 or input.dtype not in SUPPORTED_DTYPES:
+        return False
+    if not isinstance(num_groups, int) or num_groups <= 0:
+        return False
+    channels = input.shape[1]
+    if channels % num_groups != 0 or input.numel() == 0:
+        return False
+    if _one_value_per_channel(input, num_groups):
+        return False
+    if not input.is_contiguous(memory_format=_CL_FORMAT):
+        return False
+    autocast = None
+    for t in (weight, bias):
+        if t is None:
+            continue
+        if not isinstance(t, torch.Tensor):
+            return False
+        if t.dim() != 1 or t.numel() != channels:
+            return False
+        if t.device != input.device:
+            return False
+        if t.dtype is not input.dtype:
+            if t.dtype is not torch.float32:
+                return False
+            if autocast is None:
+                autocast = _autocast_active(input)
+            if not autocast:
+                # F.group_norm would raise "expected scalar type ..." here;
+                # reject so the caller reproduces that behaviour exactly.
+                return False
+    return True
+
+
+def triton_group_norm(
+    input,
+    num_groups: int,
+    weight=None,
+    bias=None,
+    eps: float = 1e-5,
+    activation: Optional[str] = None,
+):
+    """GroupNorm with an optionally fused activation, channels-last native.
+
+    A drop-in replacement for ``F.group_norm(input, num_groups, weight, bias,
+    eps)`` (followed by ``F.relu`` when ``activation="relu"``).  Inputs that
+    :func:`is_supported` rejects are served by ``F.group_norm`` itself, which
+    keeps this function total but means such calls get the *eager* kernel --
+    callers with a faster fallback should branch on :func:`is_supported`
+    themselves.
+
+    The output has the input's memory format and ``F.group_norm``'s dtype; see
+    the module docstring for the full contract.
+    """
+    if activation not in SUPPORTED_ACTIVATIONS:
+        raise ValueError(
+            f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}"
+        )
+    if not is_supported(input, num_groups, weight, bias, activation):
+        out = F.group_norm(input, num_groups, weight, bias, eps)
+        return F.relu(out) if activation == "relu" else out
+    out, _mean, _rstd = torch.ops.scaffold_gn.group_norm(
+        input,
+        num_groups,
+        weight,
+        bias,
+        float(eps),
+        activation,
+        _autocast_out_dtype(input),
+    )
+    return out
diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py
index 681e44b..5df1cee 100644
--- a/ScaFFold/unet/unet_parts.py
+++ b/ScaFFold/unet/unet_parts.py
@@ -20,22 +20,90 @@
 
 from ScaFFold.utils.perf_measure import annotate
 
+from .group_norm import FastGroupNorm
+
 _doubleconv_annotate = annotate(fmt="DoubleConv.{}")
 _down_annotate = annotate(fmt="Down.{}")
 _up_annotate = annotate(fmt="Up.{}")
 _outconv_annotate = annotate(fmt="OutConv.{}")
 
 
-def _group_norm(num_groups, num_channels):
+def _group_norm(num_groups, num_channels, activation=None):
     if num_channels % num_groups != 0:
         raise ValueError(
             f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}"
         )
-    return nn.GroupNorm(num_groups, num_channels)
+    # FastGroupNorm is nn.GroupNorm plus a Triton/compiled GPU kernel; it holds
+    # the same parameters under the same names, and `activation` is a plain
+    # attribute rather than a submodule, so checkpoints are unaffected.
+    return FastGroupNorm(num_groups, num_channels, activation=activation)
+
+
+def _consumer_dtype(*tensors):
+    """The dtype the convolution consuming a concatenation will actually see.
+
+    Inside an enabled autocast region the answer is autocast's dtype, because
+    ``aten::convolution`` carries the ``lower_precision_fp`` cast policy and
+    casts whatever it is handed.  Producing that dtype from the concatenation
+    is *bitwise identical* to producing ATen's promoted dtype and letting the
+    convolution narrow it -- the promoted tensor holds exact widenings of both
+    sources, so narrowing before or after the copy rounds the same values once
+    -- while writing and reading back half the bytes.
+
+    Outside autocast the answer is ``torch.cat``'s ordinary promotion, so eval,
+    ``inference_mode`` and pure-fp32 runs are unchanged.
+    """
+    dtype = tensors[0].dtype
+    for tensor in tensors[1:]:
+        dtype = torch.promote_types(dtype, tensor.dtype)
+    device_type = tensors[0].device.type
+    try:
+        if not torch.is_autocast_enabled(device_type):
+            return dtype
+        autocast_dtype = torch.get_autocast_dtype(device_type)
+    except (RuntimeError, TypeError):  # a device type autocast does not know
+        return dtype
+    # Only ever narrow: if autocast's dtype is the wider of the two, keep the
+    # promotion ATen would have done.
+    if torch.promote_types(autocast_dtype, dtype) is autocast_dtype:
+        return dtype
+    return autocast_dtype
+
+
+def _skip_concat(skip, upsampled):
+    """``torch.cat([skip, upsampled], dim=1)`` at the consumer's dtype.
+
+    Under ``torch.autocast`` the two halves do not share a dtype: the skip
+    comes from a GroupNorm, an fp32-policy op, while the upsampled half comes
+    from a ``ConvTranspose3d`` and is bf16.  ``torch.cat`` carries the
+    ``promote`` policy, so it widens the bf16 half to fp32, concatenates at
+    fp32, and the following convolution narrows the whole double-width result
+    straight back down -- three full-resolution passes to deliver one.  Casting
+    the inputs first collapses that to one, and the convolution reads the same
+    bits either way (see :func:`_consumer_dtype`).
+    """
+    dtype = _consumer_dtype(skip, upsampled)
+    return torch.cat([skip.to(dtype), upsampled.to(dtype)], dim=1)
 
 
 class DoubleConv(nn.Module):
-    """(convolution => GroupNorm => ReLU) * 2"""
+    """(convolution => GroupNorm => ReLU) * 2
+
+    The ReLU lives *inside* the GroupNorm (``activation="relu"``), because the
+    Triton GroupNorm kernel folds it into its forward store for free and thereby
+    removes an entire streaming pass -- 38% of the forward at the shapes that
+    dominate the step.  ``FastGroupNorm`` applies the ReLU on every path,
+    including eager, so the network's function is unchanged; only the number of
+    memory passes differs.
+
+    The ``nn.ReLU`` slots are held open by ``nn.Identity`` rather than removed:
+    ``nn.Sequential`` names its children by position, so deleting them would
+    renumber the two convolutions and the second GroupNorm and invalidate every
+    existing checkpoint.  Neither ``nn.ReLU`` nor ``nn.Identity`` has parameters
+    or buffers, so with the placeholders in place the state dict is byte
+    identical to the pre-fusion model's (pinned by
+    ``tests/test_groupnorm.py::test_state_dict_matches_plain_groupnorm_model``).
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=None):
         super().__init__()
@@ -43,11 +111,11 @@ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=No
             mid_channels = out_channels
         self.double_conv = nn.Sequential(
             nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, mid_channels),
-            nn.ReLU(inplace=True),
+            _group_norm(group_norm_groups, mid_channels, activation="relu"),
+            nn.Identity(),
             nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, out_channels),
-            nn.ReLU(inplace=True),
+            _group_norm(group_norm_groups, out_channels, activation="relu"),
+            nn.Identity(),
         )
 
     @_doubleconv_annotate
@@ -71,7 +139,23 @@ def forward(self, x):
 
 
 class Up(nn.Module):
-    """Upscaling then double conv"""
+    """Upscaling then double conv
+
+    The skip concatenation goes through :func:`_skip_concat` rather than
+    ``torch.cat`` directly, so that it emits the dtype the following
+    convolution will use instead of ``torch.cat``'s promoted one.  The tensor
+    that convolution reads is bitwise unchanged either way; it is written and
+    read back at half the width.  This rests on ``self.conv`` beginning with a
+    convolution, which the constructor below guarantees on either branch.
+
+    Measured at scale 7: 1.09 ms of a 92.8 ms step, and 0.50 GiB of peak
+    memory.  A channels-last-native Triton concatenation kernel was built and
+    measured too -- ``cat``'s *backward* is a narrowed view that consumers force
+    contiguous, at 51-63% of this device's streaming roofline against the
+    kernel's 90-103% -- but it was worth a further 0.08 ms of the step, which
+    did not justify a second hand-written kernel in a benchmark other people
+    have to trust.  ``review/skip-path/RESULTS.md`` has the numbers.
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True):
         super().__init__()
@@ -114,7 +198,9 @@ def forward(self, x1, x2):
         # if you have padding issues, see
         # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
         # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
-        x = torch.cat([x2, x1], dim=1)
+        # torch.cat([x2, x1], dim=1) with the dtype and the layout the
+        # convolution below actually wants; see the class docstring.
+        x = _skip_concat(x2, x1)
         return self.conv(x)
 
 
diff --git a/ScaFFold/utils/checkpoint_validators.py b/ScaFFold/utils/checkpoint_validators.py
deleted file mode 100644
index c2212b6..0000000
--- a/ScaFFold/utils/checkpoint_validators.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
-# Produced at the Lawrence Livermore National Laboratory.
-# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
-# the CONTRIBUTORS file. See the top-level LICENSE file for details.
-#
-# LLNL-CODE-697807.
-# All rights reserved.
-#
-# This file is part of LBANN: Livermore Big Artificial Neural Network
-# Toolkit. For details, see http://software.llnl.gov/LBANN or
-# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
-#
-# SPDX-License-Identifier: (Apache-2.0)
-
-import os
-
-# import pdb
-
-if hasattr(os, "sched_getaffinity"):
-    _orig_affinity = os.sched_getaffinity(0)
-else:
-    _orig_affinity = None
-
-import torch
-
-
-def compare_state_dicts(quantity: str, dict1, dict2):
-    def _compare(dict1, dict2, prefix=""):
-        equal = True
-        if dict1.keys() != dict2.keys():
-            missing_in_dict2 = dict1.keys() - dict2.keys()
-            missing_in_dict1 = dict2.keys() - dict1.keys()
-            if missing_in_dict2:
-                print(
-                    f"train.py: {quantity} missing in dict2: {', '.join(missing_in_dict2)} at {prefix}"
-                )
-                equal = False
-            if missing_in_dict1:
-                print(
-                    f"train.py: {quantity} missing in dict1: {', '.join(missing_in_dict1)} at {prefix}"
-                )
-                equal = False
-
-        for key in dict1.keys() & dict2.keys():
-            full_key = f"{prefix}.{key}" if prefix else key
-            if isinstance(dict1[key], torch.Tensor) and isinstance(
-                dict2[key], torch.Tensor
-            ):
-                if not torch.equal(dict1[key], dict2[key]):
-                    print(
-                        f"train.py: {quantity} tensor discrepancy at {full_key}: dict1[{key}] != dict2[{key}]"
-                    )
-                    equal = False
-            elif isinstance(dict1[key], dict) and isinstance(dict2[key], dict):
-                if not _compare(dict1[key], dict2[key], prefix=full_key):
-                    equal = False
-            else:
-                if dict1[key] != dict2[key]:
-                    print(
-                        f"train.py: {quantity} value discrepancy at {full_key}: dict1[{key}]={dict1[key]}, dict2[{key}]={dict2[key]}"
-                    )
-                    equal = False
-        return equal
-
-    return _compare(dict1, dict2)
-
-
-def compare_state_dicts2(*dicts):
-    keys = dicts[0].keys()
-    for key in keys:
-        values = [d[key] for d in dicts]
-        tensor_comparisons = [
-            (
-                torch.equal(values[i], values[i + 1])
-                if torch.is_tensor(values[i])
-                else values[i] == values[i + 1]
-            )
-            for i in range(len(values) - 1)
-        ]
-        if not all(tensor_comparisons):
-            return False
-    return True
-
-
-def compare_tensors(tensor1, tensor2):
-    return torch.all(torch.eq(tensor1, tensor2))
-
-
-def compare_items(item1, item2):
-    if isinstance(item1, torch.Tensor) and isinstance(item2, torch.Tensor):
-        return compare_tensors(item1, item2)
-    elif isinstance(item1, dict) and isinstance(item2, dict):
-        return compare_dicts3(item1, item2)
-    else:
-        return item1 == item2
-
-
-def compare_dicts3(dict1, dict2):
-    if dict1.keys() != dict2.keys():
-        return False
-    for key in dict1.keys():
-        if not compare_items(dict1[key], dict2[key]):
-            return False
-    return True
-
-
-#
-# Usage in `train.py` below:
-#
-
-# For debugging, write the saved optimizer state to file to compare to loaded state on restart
-# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-# optim_saved_path = f"{dir_checkpoint}/restarts/optim_saved_epoch{epoch}_{timestamp}.txt"
-# with open(optim_saved_path, "w") as optim_f:
-#     optim_f.write(str(state_dict['optimizer_state_dict']))
-
-#
-# For debugging purposes, load that checkpoint into new model, optimizer, etc and compare to active
-#
-# prev_checkpoint = torch.load(checkpoint_path)
-# newmodel = UNet(n_channels=3, n_classes=n_classes, trilinear=False, layers=unet_layers)
-# newmodel = newmodel.to(memory_format=torch.channels_last_3d)
-# newmodel.to(device=device)
-# newmodel = torch.nn.parallel.DistributedDataParallel(newmodel, device_ids=[get_cuda_device()], output_device=get_cuda_device())
-# newmodel.module.load_state_dict(prev_checkpoint['model_state_dict'])
-# newoptimizer = optim.RMSprop(newmodel.parameters(),
-#                 lr=learning_rate, weight_decay=weight_decay, momentum=momentum, foreach=True)
-# if optimizer_name == "ADAM":
-#     print(f"train.py(w{rank}|l{local_rank}): using ADAM optimizer .........")
-#     newoptimizer = optim.Adam(newmodel.parameters(), lr=learning_rate)
-# elif optimizer_name == "SGD":
-#     print(f"train.py(w{rank}|l{local_rank}): using SGD optimizer .........")
-#     newoptimizer = optim.SGD(newmodel.parameters(), lr=learning_rate, momentum=0.9)
-# newoptimizer.load_state_dict(prev_checkpoint['optimizer_state_dict'])
-# newscheduler = optim.lr_scheduler.ReduceLROnPlateau(newoptimizer, 'max', patience=25)
-# newscheduler.load_state_dict(prev_checkpoint['scheduler_state_dict'])
-
-# # Compare model state dicts
-# model_compare = compare_state_dicts("model", model.state_dict(), newmodel.state_dict())
-# optimizer_compare = compare_state_dicts("optimizer", optimizer.state_dict(), newoptimizer.state_dict())
-# scheduler_compare = compare_state_dicts("scheduler", scheduler.state_dict(), newscheduler.state_dict())
-# print(f"train.py: model_compare={model_compare}, optimizer_compare={optimizer_compare}, scheduler_compare={scheduler_compare}")
-# print(f"train.py: all equal? {all(compare_dicts3(sd, optimizer.state_dict()) for sd in [state_dict['optimizer_state_dict'], prev_checkpoint['optimizer_state_dict'], newoptimizer.state_dict()])}")
diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py
index 823f154..00dfc14 100644
--- a/ScaFFold/utils/checkpointing.py
+++ b/ScaFFold/utils/checkpointing.py
@@ -15,6 +15,7 @@
 import math
 import os
 import random
+import shutil
 import traceback
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
@@ -25,6 +26,17 @@
 import torch.distributed as dist
 
 
+class CheckpointSaveError(RuntimeError):
+    """A rank-0 checkpoint operation failed (a write, a cleanup, a load).
+
+    Raised identically on every rank. Only rank 0 touches the run directory,
+    but its outcome is broadcast, so the peers report the real disk error
+    instead of the unmatched-collective symptom (an opaque gloo transport
+    error, or an NCCL watchdog timeout minutes later) that a rank-0-only raise
+    produces.
+    """
+
+
 class CheckpointManager:
     """
     Checkpoint Manager for DDP/Single-Process.
@@ -90,6 +102,10 @@ def __init__(
         # Async handling
         self.executor = None
         self.future = None
+        # The exception behind the most recently reported save failure, kept
+        # only so rank 0 can chain it (and its traceback) onto the
+        # CheckpointSaveError every rank raises.
+        self._save_error_exc: Optional[BaseException] = None
         if self.async_save and self.world_rank == 0:
             # We only need 1 worker for serializing writes
             self.executor = ThreadPoolExecutor(max_workers=1)
@@ -97,50 +113,211 @@ def __init__(
         # Ensure base directory exists (Rank 0 only)
         if self.world_rank == 0:
             self.base_dir.mkdir(parents=True, exist_ok=True)
+            try:
+                self._sweep_orphaned_tmp_files()
+            except Exception as e:
+                # Construction is rank-0-only work outside any collective, so a
+                # raise here aborts rank 0 alone and leaves the peers waiting in
+                # the manager's first collective (``cleanup``). The sweep only
+                # reclaims space, so a directory that cannot be listed (a stale
+                # NFS/Lustre handle, a permissions oddity) degrades to a warning
+                # rather than taking the job down asymmetrically.
+                self._log(
+                    f"Could not sweep orphaned checkpoint temp files in "
+                    f"{self.base_dir}: {type(e).__name__}: {e}"
+                )
 
     def cleanup(self, train_from_scratch: bool) -> None:
-        """Clear existing checkpoints if training from scratch."""
-        # Ensure any pending async saves are finished before deleting
-        self.wait_for_save()
+        """Clear existing checkpoints if training from scratch.
+
+        Rank-symmetric, like every other collective point here: any pending
+        async write is drained, the rank-0 deletion is fenced, and whichever
+        failed is broadcast, so a failure raises on all ranks together (see
+        ``save_checkpoint``). The broadcast payload is ``(phase, description)``
+        so the peers -- which saw neither the write nor the deletion -- report
+        the same operation rank 0 did.
+        """
+        # Ensure any pending async save is finished before deleting.
+        error = self._drain_pending_save()
+        failure = None if error is None else ("save", error)
+
+        if train_from_scratch:
+            # Drop the cached state that described the run being deleted. Both
+            # fields are seeded from disk (or a previous save), so keeping them
+            # would let a deleted run's best gate this run's is_best decisions
+            # -- the fresh run would then never write a best checkpoint until
+            # it beat a score no file backs any more, leaving it with no
+            # best-checkpoint fallback.
+            self.best_val_loss = math.inf
+            self.last_saved_epoch = None
+
+            if self.world_rank == 0:
+                # Rank 0 alone touches the filesystem here, while every peer is
+                # already committed to the broadcast below. Individual unlinks
+                # are tolerated inside, but the glob/stat around them can still
+                # raise on a shared filesystem (ESTALE, EACCES), and raising in
+                # this window strands the peers in an unmatched collective --
+                # the R05 hazard. Report the failure through the broadcast, like
+                # a failed write.
+                try:
+                    self._remove_checkpoint_files()
+                except Exception as e:
+                    if failure is None:
+                        self._save_error_exc = e
+                        failure = ("cleanup", f"{type(e).__name__}: {e}")
+                    else:
+                        # A drained write already failed; that outcome is the
+                        # one being reported, so this is only logged.
+                        self._log(f"Clearing existing checkpoints also failed: {e}")
+
+        failure = self._broadcast_obj(failure)
+        self._barrier()
+        if failure is not None:
+            phase, description = failure
+            self._raise_save_error(description, phase=phase)
+
+    def _remove_checkpoint_files(self) -> None:
+        """Delete this run's checkpoint files and debris (rank 0 only).
+
+        Besides the two canonical files, a run directory can hold
+        ``checkpoint_*.pth.tmp.<pid>`` (an interrupted write whose Python-level
+        cleanup never ran) and ``checkpoint_*.pth.corrupt`` (a checkpoint
+        quarantined on resume). Both are full-checkpoint-sized and nothing else
+        removes them, so a "from scratch" cleanup that left them would claim to
+        have cleared the checkpoints while keeping their bytes on disk.
+        """
+        debris = sorted(
+            set(self.base_dir.glob("checkpoint_*.pth.tmp.*"))
+            | set(self.base_dir.glob("checkpoint_*.pth.corrupt"))
+        )
+        for p in (self.last_ckpt_path, self.best_ckpt_path, *debris):
+            if p.exists():
+                try:
+                    p.unlink()
+                    self._log(f"Removed existing checkpoint: {p}")
+                except Exception as e:
+                    self._log(f"Failed to remove {p}: {e}")
+
+    def _sweep_orphaned_tmp_files(self) -> None:
+        """Delete temp files stranded by checkpoint writes that were killed.
+
+        ``_atomic_save`` unlinks its ``<name>.tmp.<pid>`` file when the write
+        raises, but a SIGKILL (walltime, node failure) skips that Python-level
+        cleanup and strands a full-checkpoint-sized file. These accumulate one
+        per killed pid: the kill/restart cycle that produces them always takes
+        the *resume* path, so ``cleanup(train_from_scratch=True)`` never gets a
+        chance to clear them.
+
+        Sweeping at construction is safe because run directories are per-run
+        and not shared between concurrently running jobs (F55), so any temp
+        file here belongs to a dead process -- except one from this pid, which
+        another manager in this process could still be writing.
+
+        ``*.corrupt`` files are deliberately left alone here:
+        ``_quarantine_corrupt`` renames onto a fixed name, so at most two can
+        ever exist (they cannot accumulate) and they are the only evidence of
+        what a resume discarded. The from-scratch cleanup removes them.
+        """
+        own_suffix = f".tmp.{os.getpid()}"
+        for path in sorted(self.base_dir.glob("checkpoint_*.pth.tmp.*")):
+            if path.name.endswith(own_suffix):
+                continue
+            try:
+                path.unlink()
+                self._log(f"Removed orphaned checkpoint temp file: {path}")
+            except OSError as e:
+                self._log(f"Failed to remove {path}: {e}")
+
+    def wait_for_save(self):
+        """Block until the background save (if any) is complete.
+
+        The writer's exception is re-raised rather than logged and dropped: a
+        save that could not complete has to surface, otherwise the manager
+        state (and the process exit code) claims a checkpoint that is not on
+        disk. The future is consumed exactly once, so the failure is reported
+        at exactly one point.
 
-        if not train_from_scratch:
-            self._barrier()
+        Callers that are inside a collective region must not let this
+        propagate directly -- use ``_drain_pending_save`` instead, per the
+        collective invariant documented on ``save_checkpoint``.
+        """
+        if self.future is None:
             return
+        # Clear the handle *before* blocking on it so a failed write is
+        # reported once and does not re-raise at some later, arbitrary point.
+        future, self.future = self.future, None
+        if not future.done():
+            self._log("Waiting for background checkpoint save to complete...")
+        future.result()  # Blocks and re-raises whatever the writer raised
+
+    def _drain_pending_save(self) -> Optional[str]:
+        """Consume the in-flight async save, reporting failure without raising.
+
+        Returns a description of the writer's failure, or ``None``. Only rank 0
+        ever has a pending write, so raising here directly would leave the
+        peers blocked in the next collective; callers broadcast this result and
+        raise on every rank together.
+        """
+        try:
+            self.wait_for_save()
+        except Exception as e:
+            self._save_error_exc = e
+            return f"{type(e).__name__}: {e}"
+        return None
 
-        if self.world_rank == 0:
-            for p in (self.last_ckpt_path, self.best_ckpt_path):
-                if p.exists():
-                    try:
-                        p.unlink()
-                        self._log(f"Removed existing checkpoint: {p}")
-                    except Exception as e:
-                        self._log(f"Failed to remove {p}: {e}")
-        self._barrier()
+    def _raise_save_error(self, description: str, phase: str = "save") -> None:
+        """Raise a broadcast rank-0 failure on this rank.
 
-    def wait_for_save(self):
-        """Blocks until the background save (if any) is complete."""
-        if self.future is not None:
-            # check if running
-            if not self.future.done():
-                self._log("Waiting for background checkpoint save to complete...")
-            try:
-                self.future.result()  # Blocks and raises exceptions if any occurred
-            except Exception as e:
-                self._log(f"Background save failed with error: {e}")
-            self.future = None
+        ``phase`` names the operation that failed (``save``, ``cleanup``,
+        ``load``) so the message describes what actually went wrong. Rank 0
+        chains the original exception so its traceback survives; the peers never
+        saw it and raise the same message on its own.
+        """
+        cause, self._save_error_exc = self._save_error_exc, None
+        raise CheckpointSaveError(
+            f"Checkpoint {phase} failed on rank 0: {description}"
+        ) from cause
+
+    def finalize_saves(self) -> None:
+        """Consume the outcome of the run's last save before the run ends.
+
+        Nothing touches the manager after the training loop, so an
+        asynchronous write that failed there would never be observed: the
+        process would exit successfully having written no checkpoint (or left
+        a stale one), the benchmark would report success, and a later
+        ``--restart`` would resume from the wrong epoch or fail its pre-check.
+        Rank-symmetric, like ``save_checkpoint``.
+        """
+        error = self._drain_pending_save()
+        error = self._broadcast_obj(error)
+        self._barrier()
+        if error is not None:
+            self._raise_save_error(error)
 
     def snapshot_training_state(self) -> Dict[str, Any]:
-        """Capture mutable in-memory training state without writing a checkpoint."""
+        """Capture mutable in-memory training state without writing a checkpoint.
+
+        The copy is taken on the *host*. Cloning device-to-device instead would
+        hold roughly 3x parameter bytes of accelerator memory (the weights plus
+        Adam's two moment buffers) for as long as the snapshot lives -- which is
+        the whole of warmup, precisely where the run establishes its peak
+        memory. ``restore_training_state`` puts the state back on the model's
+        own device, so the detour is invisible to callers.
+        """
         model_ref = self.model.module if hasattr(self.model, "module") else self.model
         return {
-            "model_state_dict": self._clone_state_dict(model_ref.state_dict()),
-            "optimizer_state_dict": self._clone_state_dict(self.optimizer.state_dict())
+            "model_state_dict": self._transfer_dict_to_cpu(model_ref.state_dict()),
+            "optimizer_state_dict": self._transfer_dict_to_cpu(
+                self.optimizer.state_dict()
+            )
             if self.optimizer
             else None,
-            "scheduler_state_dict": self._clone_state_dict(self.scheduler.state_dict())
+            "scheduler_state_dict": self._transfer_dict_to_cpu(
+                self.scheduler.state_dict()
+            )
             if self.scheduler
             else None,
-            "grad_scaler_state_dict": self._clone_state_dict(
+            "grad_scaler_state_dict": self._transfer_dict_to_cpu(
                 self.grad_scaler.state_dict()
             )
             if self.grad_scaler
@@ -150,7 +327,13 @@ def snapshot_training_state(self) -> Dict[str, Any]:
         }
 
     def restore_training_state(self, snapshot: Dict[str, Any]) -> None:
-        """Restore an in-memory training snapshot."""
+        """Restore an in-memory training snapshot.
+
+        The snapshot is host-resident (see ``snapshot_training_state``); the
+        ``load_state_dict`` calls below copy into the live parameters and move
+        optimizer state onto each parameter's device, so the restored state
+        ends up exactly where it started.
+        """
         model_ref = self.model.module if hasattr(self.model, "module") else self.model
         model_ref.load_state_dict(snapshot["model_state_dict"])
 
@@ -175,7 +358,10 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int:
         With ``require_checkpoint`` (an explicit ``--restart``), a missing
         checkpoint raises instead of silently starting over.
         """
-        self.wait_for_save()  # Safety: don't load while writing
+        # Safety: don't load while writing. A failure from that write is
+        # folded into the decision broadcast below rather than raised here, so
+        # rank 0 never abandons its peers inside the broadcast.
+        error = self._drain_pending_save()
 
         # 1. Rank 0 is the sole reader: it selects the newest readable
         # checkpoint and deserializes it once, then broadcasts the loaded
@@ -184,10 +370,28 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int:
         # read of one (multi-GB) file from the shared filesystem on restart --
         # a restart I/O storm that serializes on the parallel FS. Peer ranks
         # therefore never open the checkpoint files at all.
-        result = self._select_and_load() if self.world_rank == 0 else None
+        result = None
+        if self.world_rank == 0:
+            if error is not None:
+                result = ("save_failed", error)
+            else:
+                # The selection itself stats, deserializes and renames files
+                # while the peers are already blocked in the broadcast below, so
+                # anything it raises (a stale handle on ``exists``, a rename
+                # denied, an unpickling MemoryError) has to become a decision
+                # rather than a rank-0-only death.
+                try:
+                    result = self._select_and_load()
+                except Exception as e:
+                    self._save_error_exc = e
+                    result = ("load_failed", f"{type(e).__name__}: {e}")
         status, payload = self._broadcast_obj(result)
 
         # 2. Every rank acts on the same decision rank 0 reached.
+        if status == "save_failed":
+            self._raise_save_error(payload)
+        if status == "load_failed":
+            self._raise_save_error(payload, phase="load")
         if status == "empty":
             if require_checkpoint:
                 # An explicit restart must resume real state; silently
@@ -273,6 +477,11 @@ def _select_and_load(self):
         * ``("empty", None)``           -- no checkpoint files exist;
         * ``("ok", checkpoint_dict)``   -- a candidate deserialized cleanly;
         * ``("unreadable", [paths])``   -- every candidate was corrupt.
+
+        The filesystem calls around those decisions (``exists``, the quarantine
+        rename) can still fail outright; the caller runs this inside a guard
+        that turns such a failure into a broadcast ``("load_failed", ...)``
+        decision, because raising here would strand the peers.
         """
         candidates = []
         if self.last_ckpt_path.exists():
@@ -301,22 +510,55 @@ def save_checkpoint(
         """
         Save checkpoint.
         If async_save is True, this returns immediately after CPU transfer.
+
+        Collective invariant: every rank posts exactly one ``_broadcast_obj``
+        followed by exactly one ``_barrier`` on every path through this method,
+        failures included. Rank 0 is the sole writer, but it never raises
+        before those collectives: what gets broadcast is the write's *outcome*
+        -- ``is_best`` on success, or an error description on failure,
+        including a *previous* asynchronous write whose failure is surfaced
+        here, at the next collective point. Every rank then raises the same
+        ``CheckpointSaveError`` together. Raising on rank 0 before the
+        broadcast would strand the peers in an unmatched collective, where a
+        plain disk error resurfaces as a gloo transport error or an NCCL
+        watchdog timeout that hides the real cause.
         """
-        is_best = False
+        # Non-zero ranks contribute nothing; their placeholder is overwritten
+        # by rank 0's outcome in the broadcast below (and is a harmless no-op
+        # in the degenerate non-distributed case).
+        outcome = ("ok", False)
         if self.world_rank == 0:
+            outcome = self._rank0_save(epoch, val_loss_avg, extras)
+
+        status, payload = self._broadcast_obj(outcome)
+
+        # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues.
+        # Even in async mode, we must wait for the CPU transfer to finish.
+        self._barrier()
+
+        if status == "error":
+            self._raise_save_error(payload)
+        return payload
+
+    def _rank0_save(self, epoch, val_loss_avg, extras):
+        """Perform rank 0's write and REPORT its outcome; never raises.
+
+        Returns ``("ok", is_best)`` or ``("error", description)``. The caller
+        broadcasts that outcome so every rank fails together -- see the
+        collective invariant on ``save_checkpoint``.
+        """
+        try:
+            # 1. Wait for previous async save to prevent OOM or race. If that
+            # write failed, this is where it surfaces.
+            if self.async_save:
+                self.wait_for_save()
+
             # Decide is_best from the cached best loss (single source of truth),
             # not by re-reading checkpoint_best.pth from disk. The cache is
             # seeded once at construction and updated below, so the decision
             # never races the background writer that may still be replacing the
             # best checkpoint in async mode.
-            if val_loss_avg < self.best_val_loss:
-                is_best = True
-                self.best_val_loss = val_loss_avg
-
-        if self.world_rank == 0:
-            # 1. Wait for previous async save to prevent OOM or race
-            if self.async_save:
-                self.wait_for_save()
+            is_best = val_loss_avg < self.best_val_loss
 
             model_to_save = (
                 self.model.module if hasattr(self.model, "module") else self.model
@@ -341,10 +583,6 @@ def save_checkpoint(
             if extras:
                 state_dict.update(extras)
 
-            # Record the epoch being written so callers can tell whether the
-            # last completed epoch has already been checkpointed.
-            self.last_saved_epoch = epoch
-
             # 2. Save Trigger
             if self.async_save:
                 # We must clone tensors to CPU now, because training will resume
@@ -371,13 +609,19 @@ def save_checkpoint(
                     self.log,
                 )
 
-        # Broadcast result (for logging elsewhere)
-        is_best = self._broadcast_obj(is_best)
-
-        # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues.
-        # Even in async mode, we must wait for the CPU transfer to finish.
-        self._barrier()
-        return is_best
+            # Only now claim the save: the bytes are on disk (sync) or handed
+            # to the writer (async). Recording the epoch lets callers tell
+            # whether the last completed epoch has already been checkpointed.
+            # An async write that fails later aborts the run at the next
+            # collective point, so this optimistic state is never observed by
+            # a run that keeps going.
+            if is_best:
+                self.best_val_loss = val_loss_avg
+            self.last_saved_epoch = epoch
+            return ("ok", is_best)
+        except Exception as e:
+            self._save_error_exc = e
+            return ("error", f"{type(e).__name__}: {e}")
 
     @staticmethod
     def _atomic_save(state_dict, path):
@@ -405,6 +649,35 @@ def _atomic_save(state_dict, path):
                 pass
             raise
 
+    @staticmethod
+    def _atomic_copy(src, dst):
+        """Copy an already-committed checkpoint onto another name atomically.
+
+        Same discipline as ``_atomic_save`` -- copy into a temp file in the
+        same directory, fsync it, then ``os.replace`` -- so ``dst`` is never
+        observed half-written and a crash mid-copy cannot damage the previous
+        good file there.
+        """
+        src = Path(src)
+        dst = Path(dst)
+        tmp_path = dst.with_name(f"{dst.name}.tmp.{os.getpid()}")
+        try:
+            with open(src, "rb") as fsrc, open(tmp_path, "wb") as fdst:
+                # A large buffer keeps the syscall count low on parallel
+                # filesystems (the default 64 KiB means ~1k read/write pairs
+                # per 64 MiB checkpoint).
+                shutil.copyfileobj(fsrc, fdst, length=16 * 1024 * 1024)
+                fdst.flush()
+                os.fsync(fdst.fileno())
+            os.replace(tmp_path, dst)
+        except Exception:
+            try:
+                if tmp_path.exists():
+                    tmp_path.unlink()
+            except OSError:
+                pass
+            raise
+
     @classmethod
     def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log):
         """Worker function to perform actual disk I/O.
@@ -417,10 +690,24 @@ def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log):
         try:
             # Save 'last' atomically.
             cls._atomic_save(state_dict, last_path)
-            # Save 'best' atomically (re-serialize rather than copy a file that
-            # a concurrent writer might still be replacing).
+            # 'best' is byte-identical to the 'last' just committed, so copy
+            # that file instead of pickling the same state a second time. The
+            # saving is the serialization CPU (pickle + zip of the full state
+            # dict); the filesystem traffic is roughly a wash -- the copy
+            # writes the same bytes and adds a read (measured ~1.02x faster on
+            # Lustre). Trade-off: 'best' is now a byte copy of 'last', so a
+            # silently corrupted 'last' write would propagate into 'best'
+            # rather than being an independent serialization.
+            #
+            # There is no concurrent writer to race: checkpoint writes are
+            # serialized through a single writer -- the caller's thread in sync
+            # mode, or the one-worker ThreadPoolExecutor in async mode, whose
+            # previous write ``_rank0_save`` drains before submitting the next
+            # -- and only rank 0 ever writes. So this very thread performed the
+            # ``os.replace`` onto ``last_path`` a moment ago and nothing else
+            # can be replacing it now.
             if is_best:
-                cls._atomic_save(state_dict, best_path)
+                cls._atomic_copy(last_path, best_path)
         except Exception:
             if log is not None:
                 log.error("Saving checkpoint failed:\n%s", traceback.format_exc())
@@ -434,8 +721,8 @@ def _transfer_dict_to_cpu(self, obj):
 
         ``Tensor.cpu()`` is a no-op for tensors already on CPU (it returns the
         same object), so CPU-resident state must be cloned explicitly;
-        otherwise the async writer would serialize tensors the training loop
-        keeps mutating in place.
+        otherwise the async writer -- or the warmup snapshot -- would keep an
+        alias of tensors the training loop mutates in place.
         """
         if torch.is_tensor(obj):
             t = obj.detach()
@@ -449,19 +736,6 @@ def _transfer_dict_to_cpu(self, obj):
         else:
             return obj
 
-    def _clone_state_dict(self, obj):
-        """Recursively clone tensors so in-memory snapshots are isolated."""
-        if torch.is_tensor(obj):
-            return obj.detach().clone()
-        elif isinstance(obj, dict):
-            return {k: self._clone_state_dict(v) for k, v in obj.items()}
-        elif isinstance(obj, list):
-            return [self._clone_state_dict(v) for v in obj]
-        elif isinstance(obj, tuple):
-            return tuple(self._clone_state_dict(v) for v in obj)
-        else:
-            return obj
-
     def _quarantine_corrupt(self, path):
         """Rename an unreadable checkpoint aside so it is not retried on the
         next restart (a persistent corrupt 'last' would otherwise break every
diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py
index fd7c701..3d7ebe3 100644
--- a/ScaFFold/utils/config_utils.py
+++ b/ScaFFold/utils/config_utils.py
@@ -18,8 +18,6 @@
 
 import yaml
 
-import ScaFFold.paths
-
 
 def require_positive_int(name: str, value: int) -> int:
     if not isinstance(value, int) or isinstance(value, bool) or value < 1:
@@ -27,6 +25,53 @@ def require_positive_int(name: str, value: int) -> int:
     return value
 
 
+def require_flag(name: str, value) -> bool:
+    """Validate an on/off config toggle written as 0/1 (or a YAML boolean).
+
+    ``bool(value)`` would quietly accept ``2``, ``-1`` or ``"no"`` (all true),
+    so a mistyped toggle would enable the feature it was meant to disable.
+    """
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int) and value in (0, 1):
+        return bool(value)
+    raise ValueError(f"{name} must be 0, 1, or a boolean; got {value!r}")
+
+
+def validate_unet_dims(problem_scale, unet_bottleneck_dim) -> int:
+    """Check that ``problem_scale``/``unet_bottleneck_dim`` describe a real U-Net.
+
+    The U-Net has ``unet_layers = problem_scale - unet_bottleneck_dim``
+    down/up levels over a ``2**problem_scale`` volume, so the bottleneck
+    exponent must satisfy ``0 <= unet_bottleneck_dim <= problem_scale - 1``:
+    a larger value asks for a bottleneck no smaller than the input (zero or
+    negative layers) and a negative one asks for more pooling levels than the
+    volume has. Both are only discovered later as an opaque
+    ``max_pool3d`` size error -- in production, after the whole dataset has
+    been generated -- so reject them here, at config time, naming the two keys
+    that have to change.
+
+    Returns the validated bottleneck dimension.
+    """
+    if isinstance(unet_bottleneck_dim, bool) or not isinstance(
+        unet_bottleneck_dim, int
+    ):
+        raise ValueError(
+            f"unet_bottleneck_dim must be an integer; got {unet_bottleneck_dim!r}"
+        )
+    unet_layers = problem_scale - unet_bottleneck_dim
+    if unet_bottleneck_dim < 0 or unet_layers < 1:
+        raise ValueError(
+            f"unet_bottleneck_dim={unet_bottleneck_dim} is out of range for "
+            f"problem_scale={problem_scale}: it must satisfy "
+            f"0 <= unet_bottleneck_dim <= problem_scale - 1 "
+            f"(i.e. <= {problem_scale - 1}) so that the U-Net has at least one "
+            f"layer, but unet_layers = problem_scale - unet_bottleneck_dim = "
+            f"{unet_layers}. Raise problem_scale or lower unet_bottleneck_dim."
+        )
+    return unet_bottleneck_dim
+
+
 class Config:
     """
     A class for storing configuration settings for a specific run.
@@ -42,7 +87,6 @@ class Config:
             "dataset_dir",
             "fract_base_dir",
             "job_name",
-            "library_root",
             "n_categories",
             "problem_scale",
             "unet_bottleneck_dim",
@@ -71,6 +115,7 @@ class Config:
             "normalize",
             "group_norm_groups",
             "warmup_batches",
+            "activation_checkpointing",
             "ce_weight_sample_fraction",
             "dataset_reuse_enforce_commit_id",
             "target_dice",
@@ -106,8 +151,8 @@ class Config:
         }
     )
 
-    # Fields that hold scalars; a list here indicates an unexpanded parameter
-    # sweep leaking into a single run.
+    # Fields that hold scalars. Parameter sweeps are not supported, so a list
+    # here is a user error and is rejected by name (see _validate_keys).
     _SCALAR_KEYS = frozenset(
         {
             "n_categories",
@@ -129,6 +174,7 @@ class Config:
             "loss_freq",
             "group_norm_groups",
             "warmup_batches",
+            "activation_checkpointing",
             "ce_weight_sample_fraction",
             "target_dice",
             "checkpoint_interval",
@@ -136,7 +182,7 @@ class Config:
     )
 
     @classmethod
-    def _validate_keys(cls, config_dict, strict, allow_sweeps):
+    def _validate_keys(cls, config_dict, strict):
         allowed = cls.KNOWN_KEYS | cls.AUX_KEYS
         unknown = sorted(set(config_dict) - allowed)
         if unknown:
@@ -144,21 +190,22 @@ def _validate_keys(cls, config_dict, strict, allow_sweeps):
             if strict:
                 raise ValueError(message)
             print(f"WARNING: {message}")
-        if allow_sweeps:
-            return
         lists = sorted(
             k for k in cls._SCALAR_KEYS if isinstance(config_dict.get(k), list)
         )
         if lists:
+            details = "; ".join(
+                f"{k}: parameter sweeps are no longer supported; "
+                f"got list {config_dict[k]}"
+                for k in lists
+            )
             raise ValueError(
-                f"Config key(s) {', '.join(lists)} hold list values but a "
-                "single run needs scalars. Parameter sweeps must be expanded "
-                "before constructing a run config."
+                f"{details}. Each of these keys must hold a single value; "
+                "launch one benchmark run per parameter setting."
             )
 
-    def __init__(self, config_dict, strict=True, allow_sweeps=False):
-        self._validate_keys(config_dict, strict, allow_sweeps)
-        self.library_root = str(ScaFFold.paths.scaffold_root).rstrip("/") + "/ScaFFold/"
+    def __init__(self, config_dict, strict=True):
+        self._validate_keys(config_dict, strict)
         self.base_run_dir = str(Path(config_dict["base_run_dir"]).resolve())
         self.dataset_dir = str(
             Path(config_dict.get("dataset_dir", "datasets/")).resolve()
@@ -180,7 +227,9 @@ def __init__(self, config_dict, strict=True, allow_sweeps=False):
                 "WARNING: problem_scale found to be non-integer. Truncating to nearest int."
             )
             self.problem_scale = math.floor(self.problem_scale)
-        self.unet_bottleneck_dim = config_dict["unet_bottleneck_dim"]
+        self.unet_bottleneck_dim = validate_unet_dims(
+            self.problem_scale, config_dict["unet_bottleneck_dim"]
+        )
         self.unet_layers = self.problem_scale - self.unet_bottleneck_dim
         self.n_fracts_per_vol = config_dict["n_fracts_per_vol"]
         self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"]
@@ -214,6 +263,10 @@ def __init__(self, config_dict, strict=True, allow_sweeps=False):
         self.normalize = config_dict["normalize"]
         self.group_norm_groups = config_dict.get("group_norm_groups", 8)
         self.warmup_batches = config_dict.get("warmup_batches")
+        self.activation_checkpointing = require_flag(
+            "activation_checkpointing",
+            config_dict.get("activation_checkpointing", 0),
+        )
         self.ce_weight_sample_fraction = config_dict.get(
             "ce_weight_sample_fraction", 0.1
         )
@@ -279,7 +332,11 @@ def load_config_files(file_paths):
 
 def load_config(file_path: str, config_type: str):
     """
-    Load run config from yaml file
+    Load a config from a yaml file.
+
+    ``config_type`` is either ``"benchmark"`` (a benchmark config, before a run
+    directory has been resolved) or ``"run"`` (a config for one run, which also
+    carries ``run_dir``/``run_iter``). Both require single-valued parameters.
 
     Returns:
         Config: A Config instance with settings loaded from the yaml file
@@ -290,13 +347,12 @@ def load_config(file_path: str, config_type: str):
     with open(file_path, "r") as file:
         config_dict = yaml.safe_load(file)
 
-    if config_type == "sweep":
-        # Sweep configs may hold list-valued parameters that the benchmark
-        # driver expands into one run per combination.
-        return Config(config_dict, allow_sweeps=True)
+    if config_type == "benchmark":
+        return Config(config_dict)
     elif config_type == "run":
         return RunConfig(config_dict)
     else:
         raise ValueError(
-            f"Invalid config type specified: {config_type}. Must be either 'sweep' or 'run'"
+            f"Invalid config type specified: {config_type}. "
+            "Must be either 'benchmark' or 'run'"
         )
diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py
index b43c0e3..1d5caf1 100644
--- a/ScaFFold/utils/create_restart_script.py
+++ b/ScaFFold/utils/create_restart_script.py
@@ -15,6 +15,7 @@
 # restart_script.py
 from __future__ import annotations
 
+import math
 import os
 import shlex
 import stat
@@ -23,10 +24,27 @@
 from pathlib import Path
 from typing import List, Union
 
+from ScaFFold.utils.distributed import get_local_size
+
 # Profiling toggles that must be reproduced on restart -- but only when they
 # were active in the generating run. Names mirror ScaFFold.utils.perf_measure.
 _PROFILING_ENV_VARS = ("PROFILE_TORCH", "CALI_CONFIG")
 
+# Launcher variables carrying the total rank count, in the same priority order
+# as ScaFFold.utils.distributed.get_world_size. The rank side and the restart
+# generator must recognize the same set, or a job launched under a launcher
+# only one of them knows about (e.g. Cray PALS) gets a restart script for the
+# wrong number of ranks.
+_WORLD_SIZE_ENV_VARS = (
+    "WORLD_SIZE",
+    "MV2_COMM_WORLD_SIZE",
+    "OMPI_COMM_WORLD_SIZE",
+    "PMI_SIZE",
+    "PALS_NRANKS",
+    "SLURM_NTASKS",
+    "FLUX_JOB_SIZE",
+)
+
 
 def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]:
     """
@@ -78,14 +96,26 @@ def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]:
     return new_args
 
 
+def _substitute_placeholder(tok: str, var_subs: dict[str, str]) -> str:
+    """Return ``tok`` with a placeholder replaced by its Bash expansion.
+
+    A placeholder may be a whole token (``--config __CFG__``) or the value half
+    of a combined token (``--config=__CFG__``); argparse accepts both spellings
+    on the command line, so the rewriter can emit either. Anything else is
+    shell-quoted verbatim.
+    """
+    if tok in var_subs:
+        return var_subs[tok]  # e.g., "$RUN_DIR/config.yaml"
+    flag, sep, value = tok.partition("=")
+    if sep and value in var_subs:
+        # --config=__CFG__ -> --config="$RUN_DIR/config.yaml"
+        return shlex.quote(flag + sep) + var_subs[value]
+    return shlex.quote(tok)
+
+
 def _bash_array(var_name: str, argv: List[str], var_subs: dict[str, str]) -> str:
     """Render a Bash array declaration VAR=( ... ), safely quoted, with simple placeholder substitution."""
-    parts = []
-    for tok in argv:
-        if tok in var_subs:
-            parts.append(var_subs[tok])  # e.g., "$RUN_DIR/config.yaml"
-        else:
-            parts.append(shlex.quote(tok))
+    parts = [_substitute_placeholder(tok, var_subs) for tok in argv]
     return f"{var_name}=( " + " ".join(parts) + " )"
 
 
@@ -231,8 +261,11 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]:
     Reads, in priority order:
       1. Flux (FLUX_JOB_SIZE is total tasks, FLUX_JOB_NNODES is node count),
       2. Slurm (SLURM_NTASKS / SLURM_NPROCS total tasks, SLURM_*NODES nodes),
-      3. generic launcher hints for total rank count: torchrun's WORLD_SIZE,
-         Open MPI's OMPI_COMM_WORLD_SIZE, and PMI's PMI_SIZE.
+      3. generic launcher hints for the total rank count, in the same order
+         and covering the same variables as
+         ``ScaFFold.utils.distributed.get_world_size``: keeping the two in
+         sync is what stops a restart script from relaunching the job at the
+         wrong scale.
 
     ``nodes`` is None when the environment does not report a node count.
     ``world_size`` is the best available total-rank estimate (>= 1).
@@ -248,7 +281,7 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]:
         total_tasks = int(env.get("SLURM_NTASKS") or env.get("SLURM_NPROCS") or 1)
     else:
         # No scheduler: fall back to generic launcher hints for the rank count.
-        for key in ("WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE"):
+        for key in _WORLD_SIZE_ENV_VARS:
             val = env.get(key)
             if val:
                 total_tasks = int(val)
@@ -269,6 +302,11 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) ->
     torchrun, Open MPI, PMI). The multi-rank torchrun-hpc template is emitted
     whenever the resulting world size is greater than one; the local
     single-process template is used only for a world size of one.
+
+    The node count comes from the scheduler when it reports one (Flux, Slurm);
+    otherwise it is derived from the per-node rank count
+    ``ScaFFold.utils.distributed.get_local_size`` reads, so a PALS or torchrun
+    job is not relaunched with every rank crammed onto one node.
     """
     run_dir = Path(run_dir)
     run_dir.mkdir(parents=True, exist_ok=True)
@@ -309,8 +347,24 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) ->
     if use_torchrun:
         # Calculate tasks per node for torchrun (-n arg).
         if nodes is None:
-            nodes = 1
-        tasks_per_node = max(1, total_tasks // nodes)
+            # No scheduler reported a node count: this is a PALS or plain
+            # torchrun launch. Assuming one node put the job's whole rank count
+            # on a single node (an 8-rank job across 2 nodes came back as
+            # NODES=1 TASKS_PER_NODE=8), which either oversubscribes one node or
+            # is rejected outright. The rank side's ``get_local_size`` reads the
+            # same launchers' per-node variables, so ask it how many ranks share
+            # this node and derive the node count from that. ``required=True``
+            # distinguishes "one rank per node" from "nothing reported a
+            # per-node count", where the historical single-node assumption is
+            # still the best guess available.
+            try:
+                local_size = get_local_size(required=True)
+            except RuntimeError:
+                local_size = total_tasks
+            tasks_per_node = max(1, min(local_size, total_tasks))
+            nodes = math.ceil(total_tasks / tasks_per_node)
+        else:
+            tasks_per_node = max(1, total_tasks // nodes)
 
         script = _render_torchrun_hpc_restart(
             py_array_decl, nodes, tasks_per_node, env_setup
diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py
index 564c380..8a75b1d 100644
--- a/ScaFFold/utils/data_loading.py
+++ b/ScaFFold/utils/data_loading.py
@@ -165,16 +165,36 @@ def __init__(
         customlog(f"Dataset format version: {self.dataset_format_version}")
 
         # Masks are handed off in a signed 16-bit carrier (widened to long on
-        # the compute device), so every class id must fit that range. Legacy
-        # masks are remapped to 0..len(mask_values)-1; optimized masks store
-        # dense ids that stay within the same bound.
-        max_class_id = len(self.mask_values) - 1
+        # the compute device), so the largest class id the carrier will hold
+        # must fit that range.
+        max_class_id = self._max_class_id()
         if max_class_id > np.iinfo(np.int16).max:
             raise ValueError(
-                f"{len(self.mask_values)} classes exceed the int16 mask carrier "
-                f"limit ({np.iinfo(np.int16).max})"
+                f"Mask class id {max_class_id} (from {len(self.mask_values)} "
+                f"classes) exceeds the int16 mask carrier limit "
+                f"({np.iinfo(np.int16).max}); it would wrap negative"
             )
 
+    def _max_class_id(self):
+        """Return the largest class id ``_to_mask_carrier`` will have to carry.
+
+        The bound differs by format, and using the wrong one is unsafe in one
+        direction and needlessly strict in the other. v2+ masks ship *raw*
+        ``category + 1`` ids, and the per-split table lists only the categories
+        present in that split -- so a sparse split can declare two classes while
+        holding an id in the tens of thousands, which the class *count* check
+        happily waved through. Legacy masks, by contrast, are remapped to
+        ``0..len(mask_values)-1``, so the count is exactly right there and their
+        (arbitrarily large) raw values are irrelevant.
+        """
+        if self.dataset_format_version < DATASET_FORMAT_VERSION:
+            return len(self.mask_values) - 1
+
+        ids = np.asarray(self.mask_values)
+        if ids.size == 0:
+            return 0
+        return int(ids.max())
+
     def _load_mask_values(self, data_dir):
         """Return the label-remap table for this split.
 
@@ -277,20 +297,41 @@ def _load_numpy_array(path, mmap_mode=None):
         return np.load(path, allow_pickle=False, mmap_mode=mmap_mode)
 
     def _load_dataset_format_version(self):
+        """Determine which on-disk layout this dataset uses.
+
+        Only a *missing* ``meta.yaml`` means legacy v1: those datasets predate
+        the metadata file. A metadata file that exists but cannot be read or
+        does not carry a usable version is a damaged modern dataset, and
+        falling back to the legacy loader there silently transposes
+        channels-first volumes and remaps already-dense labels -- corrupt
+        training data with no error. Such a dataset is rejected instead, with a
+        message naming the file so it can be repaired or regenerated.
+        """
         meta_path = self.dataset_root / META_FILENAME
         if not meta_path.exists():
             return LEGACY_DATASET_FORMAT_VERSION
 
         try:
             with open(meta_path, "r") as meta_file:
-                meta = yaml.safe_load(meta_file) or {}
+                meta = yaml.safe_load(meta_file)
         except Exception as exc:
-            customlog(
-                f"Failed to read dataset metadata from {meta_path}: {exc}. Falling back to legacy loader."
-            )
-            return LEGACY_DATASET_FORMAT_VERSION
+            raise ValueError(
+                f"Dataset metadata {meta_path} exists but could not be read "
+                f"({type(exc).__name__}: {exc}). A dataset carrying a "
+                f"{META_FILENAME} is not a legacy dataset; refusing to guess its "
+                "layout. Repair the file or regenerate the dataset."
+            ) from exc
 
-        return int(meta.get("dataset_format_version", LEGACY_DATASET_FORMAT_VERSION))
+        version = meta.get("dataset_format_version") if isinstance(meta, dict) else None
+        try:
+            return int(version)
+        except (TypeError, ValueError):
+            raise ValueError(
+                f"Dataset metadata {meta_path} is missing a usable "
+                f"'dataset_format_version' (got {version!r}). A dataset carrying "
+                f"a {META_FILENAME} is not a legacy dataset; refusing to guess "
+                "its layout. Repair the file or regenerate the dataset."
+            ) from None
 
     @staticmethod
     def _prepare_legacy_image(img):
diff --git a/ScaFFold/utils/distributed.py b/ScaFFold/utils/distributed.py
index c46a5bd..5cb0700 100644
--- a/ScaFFold/utils/distributed.py
+++ b/ScaFFold/utils/distributed.py
@@ -12,6 +12,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 
+import logging
 import os
 import os.path
 import socket
@@ -21,6 +22,52 @@
 import torch
 import torch.distributed
 
+logger = logging.getLogger(__name__)
+
+
+def _env_int(name: str) -> Optional[int]:
+    """Return the launcher variable ``name`` as an int, or None if unusable.
+
+    Launcher variables are not always what they claim to be. A site wrapper
+    that exports ``WORLD_SIZE=`` (empty, e.g. from an unset shell variable) or
+    a placeholder like ``auto`` is common enough, and a bare ``int()`` turned
+    it into a ``ValueError`` raised from the first of these helpers anything
+    called -- killing every ``scaffold`` invocation, ``--help`` included, with
+    a traceback that named neither the variable nor a remedy.
+
+    An unusable value is treated as absent so the next source in the priority
+    order is consulted (ultimately the MPI communicator, or the documented
+    default). A non-empty value that is not an integer is warned about, because
+    unlike an empty one it looks deliberate and the fallback may not be what
+    its author intended. ``create_restart_script._sniff_launch_shape`` skips
+    empty values for the same reason.
+    """
+    value = os.environ.get(name)
+    if value is None:
+        return None
+    value = value.strip()
+    if not value:
+        return None
+    try:
+        return int(value)
+    except ValueError:
+        logger.warning(
+            "Ignoring launcher variable %s=%r: not an integer. Falling back to "
+            "the next source for the job shape.",
+            name,
+            value,
+        )
+        return None
+
+
+def _first_env_int(names) -> Optional[int]:
+    """Return the first usable integer among ``names``, in priority order."""
+    for name in names:
+        value = _env_int(name)
+        if value is not None:
+            return value
+    return None
+
 
 def get_num_gpus() -> int:
     """Return the number of GPUs on this node."""
@@ -40,39 +87,74 @@ def _mpi_comm_world():
         return None
 
 
+_LOCAL_RANK_VARS = (
+    "LOCAL_RANK",
+    "MV2_COMM_WORLD_LOCAL_RANK",
+    "OMPI_COMM_WORLD_LOCAL_RANK",
+    "PMI_LOCAL_RANK",
+    "PALS_LOCAL_RANKID",
+    "SLURM_LOCALID",
+    "FLUX_TASK_LOCAL_ID",
+)
+
+_LOCAL_SIZE_VARS = (
+    "LOCAL_WORLD_SIZE",
+    "MV2_COMM_WORLD_LOCAL_SIZE",
+    "OMPI_COMM_WORLD_LOCAL_SIZE",
+    "PMI_LOCAL_SIZE",
+    "PALS_LOCAL_SIZE",
+)
+
+_WORLD_RANK_VARS = (
+    "RANK",
+    "MV2_COMM_WORLD_RANK",
+    "OMPI_COMM_WORLD_RANK",
+    "PMI_RANK",
+    "PALS_RANKID",
+    "SLURM_PROCID",
+    "FLUX_TASK_RANK",
+)
+
+_WORLD_SIZE_VARS = (
+    "WORLD_SIZE",
+    "MV2_COMM_WORLD_SIZE",
+    "OMPI_COMM_WORLD_SIZE",
+    "PMI_SIZE",
+    "PALS_NRANKS",
+    "SLURM_NTASKS",
+    "FLUX_JOB_SIZE",
+)
+
+
 def get_local_rank(required: bool = False) -> int:
     """Return the local MPI rank."""
-    if "LOCAL_RANK" in os.environ:
-        return int(os.environ["LOCAL_RANK"])
-    if "MV2_COMM_WORLD_LOCAL_RANK" in os.environ:
-        return int(os.environ["MV2_COMM_WORLD_LOCAL_RANK"])
-    if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
-        return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
-    if "PMI_LOCAL_RANK" in os.environ:
-        return int(os.environ["PMI_LOCAL_RANK"])
-    if "PALS_LOCAL_RANKID" in os.environ:
-        return int(os.environ["PALS_LOCAL_RANKID"])
-    if "SLURM_LOCALID" in os.environ:
-        return int(os.environ["SLURM_LOCALID"])
-    if "FLUX_TASK_LOCAL_ID" in os.environ:
-        return int(os.environ["FLUX_TASK_LOCAL_ID"])
+    value = _first_env_int(_LOCAL_RANK_VARS)
+    if value is not None:
+        return value
     if required:
         raise RuntimeError("Could not get local rank")
     return 0
 
 
 def get_local_size(required: bool = False) -> int:
-    """Return the number of local MPI ranks."""
-    if "MV2_COMM_WORLD_LOCAL_SIZE" in os.environ:
-        return int(os.environ["MV2_COMM_WORLD_LOCAL_SIZE"])
-    if "OMPI_COMM_WORLD_LOCAL_SIZE" in os.environ:
-        return int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"])
-    if "SLURM_NNODES" in os.environ and "SLURM_NTASKS" in os.environ:
-        return int(os.environ["SLURM_NTASKS"]) // int(os.environ["SLURM_NNODES"])
-    # Flux does not have an env variable for this, so we assume an
-    # even distribution.
-    if "FLUX_JOB_SIZE" in os.environ and "FLUX_JOB_NNODES" in os.environ:
-        return int(os.environ["FLUX_JOB_SIZE"]) // int(os.environ["FLUX_JOB_NNODES"])
+    """Return the number of local MPI ranks.
+
+    Recognizes the same launchers as ``get_local_rank``: a variable honored
+    there but not here silently yields 1, which makes per-node logic (e.g. the
+    profiler's one-rank-per-node gate) treat every rank as node-local.
+    """
+    value = _first_env_int(_LOCAL_SIZE_VARS)
+    if value is not None:
+        return value
+    # Slurm and Flux report only totals; assume an even distribution. A zero
+    # node count is as unusable as a non-numeric one, so it falls through
+    # rather than raising ZeroDivisionError.
+    ntasks, nnodes = _env_int("SLURM_NTASKS"), _env_int("SLURM_NNODES")
+    if ntasks is not None and nnodes:
+        return ntasks // nnodes
+    job_size, job_nodes = _env_int("FLUX_JOB_SIZE"), _env_int("FLUX_JOB_NNODES")
+    if job_size is not None and job_nodes:
+        return job_size // job_nodes
     if required:
         raise RuntimeError("Could not get local size")
     return 1
@@ -80,20 +162,9 @@ def get_local_size(required: bool = False) -> int:
 
 def get_world_rank(required: bool = False) -> int:
     """Return the global MPI rank."""
-    if "RANK" in os.environ:
-        return int(os.environ["RANK"])
-    if "MV2_COMM_WORLD_RANK" in os.environ:
-        return int(os.environ["MV2_COMM_WORLD_RANK"])
-    if "OMPI_COMM_WORLD_RANK" in os.environ:
-        return int(os.environ["OMPI_COMM_WORLD_RANK"])
-    if "PMI_RANK" in os.environ:
-        return int(os.environ["PMI_RANK"])
-    if "PALS_RANKID" in os.environ:
-        return int(os.environ["PALS_RANKID"])
-    if "SLURM_PROCID" in os.environ:
-        return int(os.environ["SLURM_PROCID"])
-    if "FLUX_TASK_RANK" in os.environ:
-        return int(os.environ["FLUX_TASK_RANK"])
+    value = _first_env_int(_WORLD_RANK_VARS)
+    if value is not None:
+        return value
     comm = _mpi_comm_world()
     if comm is not None:
         return comm.Get_rank()
@@ -104,20 +175,9 @@ def get_world_rank(required: bool = False) -> int:
 
 def get_world_size(required: bool = False) -> int:
     """Return the number of MPI ranks."""
-    if "WORLD_SIZE" in os.environ:
-        return int(os.environ["WORLD_SIZE"])
-    if "MV2_COMM_WORLD_SIZE" in os.environ:
-        return int(os.environ["MV2_COMM_WORLD_SIZE"])
-    if "OMPI_COMM_WORLD_SIZE" in os.environ:
-        return int(os.environ["OMPI_COMM_WORLD_SIZE"])
-    if "PMI_SIZE" in os.environ:
-        return int(os.environ["PMI_SIZE"])
-    if "PALS_NRANKS" in os.environ:
-        return int(os.environ["PALS_NRANKS"])
-    if "SLURM_NTASKS" in os.environ:
-        return int(os.environ["SLURM_NTASKS"])
-    if "FLUX_JOB_SIZE" in os.environ:
-        return int(os.environ["FLUX_JOB_SIZE"])
+    value = _first_env_int(_WORLD_SIZE_VARS)
+    if value is not None:
+        return value
     comm = _mpi_comm_world()
     if comm is not None:
         return comm.Get_size()
diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py
index c69b4b6..9d87a2b 100644
--- a/ScaFFold/utils/perf_measure.py
+++ b/ScaFFold/utils/perf_measure.py
@@ -12,12 +12,31 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 
+import logging
 import os
 from contextlib import nullcontext
 
 CALI_PERF_ENV_VAR = "CALI_CONFIG"
 TORCH_PERF_ENV_VAR = "PROFILE_TORCH"
 
+# This module is imported before (and independently of) the run's MPI logger,
+# so it keeps its own. Everything it has to say is about the user's profiling
+# request not being honored as written, which belongs on a diagnostic channel
+# that a caller can filter or capture -- not on stdout, where it lands in the
+# middle of whatever the run is printing.
+logger = logging.getLogger(__name__)
+
+
+def _profiler_env_flag(name):
+    """Return True only for an affirmative value of the environment variable.
+
+    Every profiler toggle -- the master switch and its sub-options alike --
+    goes through this, so "0"/"false"/"no"/"off"/"" all mean off and there is
+    no spelling that means the opposite of what it says.
+    """
+    return os.environ.get(name, "").lower() in ("1", "true", "on", "yes")
+
+
 _CALI_PERF_ENABLED = False
 TORCH_PERF_ENABLED = False
 if CALI_PERF_ENV_VAR in os.environ:
@@ -28,23 +47,24 @@
 
         _CALI_PERF_ENABLED = True
     except Exception as e:
-        print("User requested Caliper annotations, but could not import Caliper")
-        print(f"Exception: {e}")
+        logger.warning(
+            "User requested Caliper annotations, but could not import Caliper: %s: %s",
+            type(e).__name__,
+            e,
+        )
 
 # The torch profiler is gated purely on its own environment variable: Caliper
 # and the torch profiler may both be enabled at once.
-if (
-    TORCH_PERF_ENV_VAR in os.environ
-    and os.environ.get(TORCH_PERF_ENV_VAR).lower() != "off"
-):
+if _profiler_env_flag(TORCH_PERF_ENV_VAR):
     try:
         from torch.profiler import ProfilerActivity
         from torch.profiler import profile as torchprofile
 
         TORCH_PERF_ENABLED = True
     except Exception:
-        print(
-            "User requested PyTorch profiling, but could not import the PyTorch profiler"
+        logger.warning(
+            "User requested PyTorch profiling, but could not import the "
+            "PyTorch profiler"
         )
 
 
@@ -100,10 +120,6 @@ def _profiler_env_int(name, default):
         return default
 
 
-def _profiler_env_flag(name):
-    return os.environ.get(name, "").lower() in ("1", "true", "on", "yes")
-
-
 def get_torch_context(ranks_per_node, rank):
     if TORCH_PERF_ENABLED:
         TORCH_PERF_LOCAL = TORCH_PERF_ENABLED and (rank % ranks_per_node == 0)
@@ -118,7 +134,22 @@ def get_torch_context(ranks_per_node, rank):
         # trace. The window (skip `wait`, prime `warmup`, capture `active`, once)
         # is tunable via the environment. Callers must drive it with
         # ``prof.step()`` once per training step for the schedule to advance.
+        # The context is entered around checkpoint cleanup and the warmup
+        # batches, but prof.step() only advances once per *training* batch, so
+        # everything before the first training batch lands in step 0. The
+        # window must therefore skip at least one step: with wait=0 that whole
+        # prologue -- warmup_batches forward+backward passes per rank -- is
+        # buffered in host memory as a single unbounded step, which is the very
+        # thing the bounded window exists to prevent.
         wait = _profiler_env_int("PROFILE_TORCH_WAIT", 1)
+        if wait < 1:
+            logger.warning(
+                "PROFILE_TORCH_WAIT must be at least 1: the profiler window "
+                "opens before the warmup batches, whose work would otherwise "
+                "accumulate in host memory as one unbounded step. Using "
+                "PROFILE_TORCH_WAIT=1."
+            )
+            wait = 1
         warmup = _profiler_env_int("PROFILE_TORCH_WARMUP", 1)
         active = _profiler_env_int("PROFILE_TORCH_ACTIVE", 3) or 1
 
diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py
index 108df71..b284708 100644
--- a/ScaFFold/utils/trainer.py
+++ b/ScaFFold/utils/trainer.py
@@ -117,6 +117,10 @@ def __init__(self, model, config, device, log):
         self.global_step = 0
         self.total_optimizer_steps = 0
         self.start_epoch = -1
+        # Validation dice already achieved by the epoch we resume from (0 for a
+        # fresh run). The training loop's exit condition is a threshold on this
+        # score, so it has to survive a restart: see cleanup_or_resume.
+        self.start_val_dice = 0.0
         self.ps = getattr(self.config, "_parallel_strategy", None)
         self.spatial_mesh = None  # Spatial mesh for use w/ DistConv
         self.data_num_replicas = self.world_size
@@ -387,6 +391,7 @@ def cleanup_or_resume(self):
                         pass
 
             self.start_epoch = 1
+            self.start_val_dice = 0.0
         else:
             # Load checkpoint via manager. An explicit restart must find a
             # checkpoint; a plain non-scratch launch may simply start fresh.
@@ -399,6 +404,16 @@ def cleanup_or_resume(self):
             if "train_mask_values" in restored:
                 self.train_set.mask_values = restored["train_mask_values"]
 
+            # Resume the convergence state, not just the weights. The loop runs
+            # while the validation dice is below target_dice, so a run that had
+            # already converged must re-enter the loop with the score it
+            # achieved -- starting from 0 would unconditionally re-train (and
+            # re-log, and re-checkpoint) one full epoch before the condition is
+            # re-tested, inflating the epoch count and the FOM's total train
+            # time. Checkpoints written before this key existed simply fall
+            # back to 0, i.e. the old behaviour.
+            self.start_val_dice = restored.get("val_dice", 0.0)
+
             # Continue the optimizer-step counts from where the checkpoint
             # left off; otherwise a resumed run restarts them at 0 and
             # undercounts all pre-resume work in the reported step totals.
@@ -698,6 +713,79 @@ def _sync_gather_minibatch_timer(self, minibatch_events):
         minibatch_time_s = statistics.median(minibatch_times.cpu().tolist())
         return minibatch_time_s
 
+    def _warmup_ragged_batches(self, batch):
+        """Warm the narrower final batch each loader ends its epoch with.
+
+        Warmup runs only the *leading* batches of the train loader, which are
+        all ``local_batch_size`` wide, and neither loader drops its last batch.
+        So when a rank's sample count is not a multiple of the batch size, the
+        final batch of every epoch presents a set of convolution problems
+        nothing has ever run: with ``cudnn.benchmark`` on, its algorithm search
+        (MIOpen find) then happens inside the *first timed epoch* -- a one-off
+        stall measured at 95 s that lands in ``epoch_duration``, i.e. straight
+        in the FOM denominator, while staying invisible to the per-minibatch
+        timer (which excludes partial batches).
+
+        One extra iteration per distinct ragged size fixes that, cut from a
+        batch warmup already fetched so no additional I/O is needed. The
+        validation shapes are covered by the same (training) step: the forward
+        convolutions are what validation shares, and this runs inside warmup's
+        snapshot/restore envelope, so the extra step cannot affect training.
+        ``local_batch_size = 1`` never has a ragged batch and does no extra work.
+
+        The set of sizes is agreed across ranks first. The training shards are
+        padded to equal length, so their remainder is already identical
+        everywhere, but validation is sharded *unpadded* on purpose (F-series:
+        an unbiased SUM-reduced metric), so per-rank counts -- and their
+        remainders -- differ. A rank that decided on its own would run a step
+        its peers did not, and the collectives inside that step (the gradient
+        all-reduce, the sharded loss reductions) would deadlock.
+        """
+        local_batch_size = self.config.local_batch_size
+
+        # Agree on the ragged sizes FIRST. The all_gather below is a
+        # collective: every rank must post it even when its own warmup fetched
+        # no batch at all (``batch is None``), or its peers block in the
+        # gather. Batch availability is rank-invariant in practice (padded
+        # training shards give every rank the same loader length), but the
+        # collective pattern must not depend on local loader state, so the
+        # no-batch early return comes after the gather.
+        ragged_sizes = {len(self.train_sampler) % local_batch_size}
+        local_val_ragged = torch.tensor(
+            [len(self.val_sampler) % local_batch_size], device=self.device
+        )
+        gathered_val_ragged = [
+            torch.empty_like(local_val_ragged) for _ in range(self.world_size)
+        ]
+        torch.distributed.all_gather(gathered_val_ragged, local_val_ragged)
+        ragged_sizes.update(int(size.item()) for size in gathered_val_ragged)
+
+        if batch is None:
+            return
+        available = batch["image"].shape[0]
+
+        # The batches already run are all ``available`` wide, and ``available``
+        # is itself rank-invariant (the padded training shards give every rank
+        # the same leading batch size), so this loop is identical on all ranks.
+        warmed = {0, available}
+        for ragged in sorted(ragged_sizes):
+            if ragged in warmed:
+                continue
+            if ragged > available:
+                self.log.debug(
+                    f"  warmup: cannot build a {ragged}-sample batch from a "
+                    f"{available}-sample batch; skipping that shape"
+                )
+                continue
+            warmed.add(ragged)
+            self.log.debug(
+                f"  warmup: running the ragged batch shape ({ragged} samples)"
+            )
+            self._run_training_batch(
+                {key: value[:ragged] for key, value in batch.items()},
+                log_prefix=f"warmup ragged ({ragged}): ",
+            )
+
     def warmup(self):
         """Run warmup iterations before the main training loop."""
         warmup_batches = self.config.warmup_batches
@@ -720,10 +808,12 @@ def warmup(self):
         self.optimizer.zero_grad(set_to_none=True)
 
         try:
+            last_batch = None
             for batch_idx, batch in enumerate(self.train_loader):
                 if batch_idx >= max_batches:
                     break
 
+                last_batch = batch
                 self._run_training_batch(
                     batch,
                     log_prefix="warmup: ",
@@ -734,6 +824,8 @@ def warmup(self):
                     f"  warmup: batch {batch_idx} completed in {batch_t_end - start_warmup} seconds"
                 )
 
+            self._warmup_ragged_batches(last_batch)
+
             self.val_loader.sampler.set_epoch(0)
 
             if max_val_batches > 0:
@@ -766,7 +858,10 @@ def train(self, profiler=None):
         """
 
         epoch = self.start_epoch
-        dice_score_train = 0
+        # Seeded from the resumed checkpoint (0 for a fresh run) so an already
+        # converged run exits the loop immediately instead of training an
+        # extra epoch to rediscover that it converged.
+        dice_score_train = self.start_val_dice
         epoch_minibatch_times_s = []
         # Track the last epoch checkpointed inside the loop so the final-save
         # decision below is identical on every rank. The in-loop checkpoint
@@ -776,6 +871,11 @@ def train(self, profiler=None):
         # other ranks would make them disagree about whether to call
         # save_checkpoint on exit, deadlocking its internal collective.
         last_checkpoint_epoch = None
+        # Whether this invocation completed at least one NEW epoch. A resume
+        # whose checkpoint already covers config.epochs (or an --epochs lowered
+        # below the checkpointed epoch) leaves the loop at the max-epoch check
+        # before any epoch body runs, so none of the per-epoch metrics exist.
+        completed_new_epoch = False
         with open(self.outfile_path, "a", newline="") as outfile:
             start = time.time()
             while dice_score_train < self.config.target_dice:
@@ -947,6 +1047,25 @@ def train(self, profiler=None):
                 # Reduced sample-weighted total and per-sample mean val loss.
                 val_loss_epoch = val_info[1].item()
                 val_loss_avg = val_loss_epoch / global_val_samples
+
+                # Divergence check. The dice score below is computed from a
+                # hard argmax, so it stays finite even for an all-NaN model
+                # (argmax of NaN logits is 0 and the one-hots are finite) --
+                # the loss is the only value that actually goes non-finite. Bail
+                # out before the CSV row and the checkpoint: continuing would
+                # keep overwriting checkpoint_last.pth with NaN weights (which
+                # then poison the next restart) while the loop's dice threshold
+                # can never be met. Both values come out of the data-parallel
+                # reductions above, so they are identical on every rank and
+                # every rank raises here together, leaving no unmatched
+                # collective behind.
+                if not (math.isfinite(overall_loss) and math.isfinite(val_loss_avg)):
+                    raise ValueError(
+                        f"Non-finite loss at epoch {epoch} "
+                        f"(train_loss={overall_loss}, val_loss={val_loss_avg}): "
+                        "training diverged, aborting before checkpointing."
+                    )
+
                 if not self.config.disable_scheduler:
                     self.scheduler.step()
                 else:
@@ -1022,6 +1141,9 @@ def train(self, profiler=None):
                         "train_mask_values": self.train_set.mask_values,
                         "global_step": self.global_step,
                         "total_optimizer_steps": self.total_optimizer_steps,
+                        # The convergence state: what a resume must restore to
+                        # know this epoch already met (or missed) target_dice.
+                        "val_dice": val_score,
                     }
                     self.checkpoint_manager.save_checkpoint(epoch, val_loss_avg, extras)
                     last_checkpoint_epoch = epoch
@@ -1030,6 +1152,7 @@ def train(self, profiler=None):
 
                 dice_score_train = val_score
                 epoch += 1
+                completed_new_epoch = True
 
                 # This check must exist otherwise the condition dice_score_train < self.config.target_dice will evaluate to False and incorrectly exit the training
                 if math.isnan(dice_score_train):
@@ -1039,12 +1162,31 @@ def train(self, profiler=None):
 
         completed_epochs = epoch - 1
 
+        if not completed_new_epoch:
+            # The loop exited without running a single epoch: the epoch budget
+            # was already exhausted (a resume whose checkpoint covers every
+            # epoch asked for, or a fresh run configured with none) or the
+            # starting state already met target_dice. There is nothing new to
+            # save -- any checkpoint on disk already records epoch
+            # `completed_epochs`, and none of the per-epoch metrics the final
+            # save would write were ever computed -- so skip it and return
+            # normally; the caller's post-processing still has whatever CSV is
+            # there.
+            self.log.warning(
+                "No new epoch was trained (start epoch %s, 'epochs' %s, "
+                "starting val dice %s vs target_dice %s): this run had no epoch "
+                "left to run, and no checkpoint was written.",
+                self.start_epoch,
+                self.config.epochs,
+                self.start_val_dice,
+                self.config.target_dice,
+            )
         # Save a final checkpoint when the run exits (convergence or max epochs)
         # at an epoch that was not a checkpoint interval, so the converged
         # weights that produced the reported metrics are not lost. Skipped when
         # checkpointing is disabled, when no epoch completed, or when the last
         # completed epoch was already checkpointed inside the loop.
-        if (
+        elif (
             self.config.checkpoint_interval > 0
             and completed_epochs >= 1
             and last_checkpoint_epoch != completed_epochs
@@ -1053,11 +1195,20 @@ def train(self, profiler=None):
                 "train_mask_values": self.train_set.mask_values,
                 "global_step": self.global_step,
                 "total_optimizer_steps": self.total_optimizer_steps,
+                "val_dice": val_score,
             }
             self.checkpoint_manager.save_checkpoint(
                 completed_epochs, val_loss_avg, extras
             )
 
+        # Nothing downstream of the training loop touches the checkpoint
+        # manager, so this is the last chance to observe the outcome of the
+        # run's final (possibly asynchronous) write. Without it a failed final
+        # save would let the process exit successfully with no checkpoint at
+        # all, and the next --restart would resume from a stale epoch or fail
+        # its pre-check.
+        self.checkpoint_manager.finalize_saves()
+
         if epoch_minibatch_times_s:
             minibatch_time_s = statistics.median(epoch_minibatch_times_s)
             adiak_value("minibatch_time_s", minibatch_time_s)
diff --git a/ScaFFold/utils/utils.py b/ScaFFold/utils/utils.py
index 5d9e334..12c196b 100644
--- a/ScaFFold/utils/utils.py
+++ b/ScaFFold/utils/utils.py
@@ -113,12 +113,24 @@ def setup_mpi_logger(
 
 
 def mem_stats():
+    """Return this rank's GPU memory counters.
+
+    On a host with no visible GPU (a CPU/gloo run, or a job launched with the
+    devices masked off) there are no counters to read: report that instead of
+    raising, so a diagnostic call cannot take down a run that is otherwise
+    perfectly able to proceed.
+    """
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    if not torch.cuda.is_available():
+        return {"rank": rank, "device": "cpu", "cuda_available": False}
+
     dev = torch.cuda.current_device()
     free, total = torch.cuda.mem_get_info()  # device-level (driver) view
     stats = torch.cuda.memory_stats(dev)  # allocator internals
     return {
-        "rank": dist.get_rank() if dist.is_initialized() else 0,
+        "rank": rank,
         "device": dev,
+        "cuda_available": True,
         "allocated": torch.cuda.memory_allocated(
             dev
         ),  # bytes currently used by tensors
@@ -135,6 +147,15 @@ def mem_stats():
 def gather_and_print_mem(log, tag=""):
     if log.getEffectiveLevel() > 10:  # 10 -> DEBUG
         return
+    if not torch.cuda.is_available():
+        # Uniform across ranks, so returning here skips the all_gather on every
+        # rank rather than desynchronizing them.
+        log.debug(
+            "=== %s === no CUDA device visible on this rank; "
+            "GPU memory statistics unavailable",
+            tag,
+        )
+        return
     stats = mem_stats()
     if dist.is_initialized():
         world = dist.get_world_size()
diff --git a/ScaFFold/viz/standard_viz.py b/ScaFFold/viz/standard_viz.py
index 50d5beb..d9f2dda 100644
--- a/ScaFFold/viz/standard_viz.py
+++ b/ScaFFold/viz/standard_viz.py
@@ -27,6 +27,10 @@ def main(config: RunConfig):
     figures_path = Path(config.run_dir) / "figures"
     figures_path.mkdir(parents=True, exist_ok=True)
 
+    # pyplot keeps a strong reference to every figure until it is closed, and a
+    # sweep calls this once per combination in the same process, so an unclosed
+    # figure is retained (canvas included) for the rest of the run.
+    figures = []
     try:
         epochs = []
         train_loss = []
@@ -53,7 +57,7 @@ def main(config: RunConfig):
         legend_loc = (0, -0.17)
 
         # Plot training loss
-        plt.figure()
+        figures.append(plt.figure())
         plt.plot(epochs, train_loss, label="Train Loss", linewidth=line_thickness)
         plt.xlabel("Epoch", fontsize=fontsize)
         plt.ylabel("Train loss", fontsize=fontsize)
@@ -63,9 +67,10 @@ def main(config: RunConfig):
         plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight")
+        plt.close(figures[-1])
 
         # Plot validation dice
-        plt.figure()
+        figures.append(plt.figure())
         plt.plot(epochs, val_dice, label="Val Dice Score", linewidth=line_thickness)
         plt.xlabel("Epoch", fontsize=fontsize)
         plt.ylabel("Val dice score", fontsize=fontsize)
@@ -74,10 +79,11 @@ def main(config: RunConfig):
         plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight")
+        plt.close(figures[-1])
 
         # Plot validation loss if available
         if val_loss:
-            plt.figure()
+            figures.append(plt.figure())
             plt.plot(epochs, val_loss, label="Val Loss", linewidth=line_thickness)
             plt.xlabel("Epoch", fontsize=fontsize)
             plt.ylabel("Val loss", fontsize=fontsize)
@@ -86,5 +92,12 @@ def main(config: RunConfig):
             plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
             plt.grid(True, axis="y")
             plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight")
+            plt.close(figures[-1])
     except Exception as e:
         logger.error(f"Failed to generate figures: {e}")
+    finally:
+        # Errors here are logged and swallowed, so the close has to be
+        # unconditional: a failure between figure() and savefig() would
+        # otherwise leak exactly the figure nobody goes looking for.
+        for figure in figures:
+            plt.close(figure)
diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py
index 459a8cf..40458ba 100644
--- a/ScaFFold/worker.py
+++ b/ScaFFold/worker.py
@@ -17,6 +17,7 @@
 import socket
 import time
 from argparse import Namespace
+from pathlib import Path
 
 import numpy as np
 import psutil
@@ -96,6 +97,49 @@ def wrap_model_ddp(model, device, ps):
     )
 
 
+def export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node):
+    """Write this rank's chrome trace, reporting failures instead of raising.
+
+    A profiling rank reaches this while the others are already heading for the
+    barrier that precedes rank-0 post-processing, so an exception here does not
+    just lose a trace: it kills this rank and leaves every other rank blocked
+    in that barrier until the collective times out. Failures are real (a run
+    with zero training batches never starts the profiler, and export_chrome_trace
+    then raises; the trace can also fill the filesystem), so log them and let
+    the job finish.
+
+    The trace is written into ``config.run_dir`` so it lands with the rest of
+    the run's artifacts instead of wherever the job happened to be launched
+    from.
+
+    Returns the path written, or None if the trace could not be written.
+    """
+    # Round up: with a partly-filled last node, flooring would report one node
+    # too few (and a ranks_per_node larger than the job would report none).
+    nodes = max(1, math.ceil(world_size / max(1, ranks_per_node)))
+    tracename = (
+        f"torch-{socket.gethostname()}-r{rank}"
+        f"-N{nodes}-n{world_size}"
+        f"-ps{config.problem_scale}-e{config.epochs}"
+        f"-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json"
+    )
+    tracepath = Path(getattr(config, "run_dir", None) or os.getcwd()) / tracename
+    try:
+        prof.export_chrome_trace(str(tracepath))
+    except Exception as e:
+        log.error(
+            "Could not write PyTorch trace '%s': %s: %s. Continuing so the "
+            "run can finish; a run with zero profiled steps never starts the "
+            "profiler and has no trace to export.",
+            tracepath,
+            type(e).__name__,
+            e,
+        )
+        return None
+    log.info("Wrote PyTorch trace '%s'", tracepath)
+    return tracepath
+
+
 @annotate()
 def main(kwargs_dict: dict = {}):
     #
@@ -183,6 +227,15 @@ def main(kwargs_dict: dict = {}):
     )
 
     model = model.to(device, memory_format=torch.channels_last_3d)
+    if config.activation_checkpointing:
+        # Has to happen before the DDP wrap: afterwards the model is only
+        # reachable as ``model.module``, and the wrapper does not forward the
+        # method.
+        log.info(
+            "activation_checkpointing TRUE -- recomputing block activations in "
+            "the backward pass instead of storing them"
+        )
+        model.use_checkpointing()
     # Wrap with DistConvDDP that corrects gradient scaling for dc submesh
     model = wrap_model_ddp(model, device, ps)
     # Store ps for use in the training loop
@@ -261,10 +314,7 @@ def main(kwargs_dict: dict = {}):
         trainer.train(profiler=prof if TORCH_PERF_LOCAL else None)
         end_code_region("train")
     if TORCH_PERF_LOCAL:
-        hostname = socket.gethostname()
-        tracename = f"torch-{hostname}-r{rank}-N{world_size // ranks_per_node}-n{world_size}-ps{config.problem_scale}-e{config.epochs}-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json"
-        prof.export_chrome_trace(tracename)
-        log.info("Wrote PyTorch trace '%s'", tracename)
+        export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node)
 
     # Results are final here; synchronize before rank-0 post-processing so a
     # post-processing failure on rank 0 cannot strand the other ranks in a
diff --git a/pyproject.toml b/pyproject.toml
index c0e6c3e..80b5aa7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,6 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools]
 package-data = { "ScaFFold" = [
     "package_data/weights_ins145.csv",
-    "fractals/var0.15/3DIFS_param/*",
     "configs/*",
 ] }
 include-package-data = true
diff --git a/tests/conftest.py b/tests/conftest.py
index 376afbc..e35f739 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -467,6 +467,7 @@ def make(
         n_train: int = 4,
         n_val: int = 2,
         n: int = 16,
+        log_level: int = logging.INFO,
         config_overrides: Optional[dict] = None,
     ) -> "PyTorchTrainer":
         dataset_root = tiny_dataset(
@@ -495,9 +496,10 @@ def make(
         device = torch.device("cpu")
 
         log = logging.getLogger(f"tiny_trainer.{id(config)}")
-        # INFO (20) > DEBUG (10) => gather_and_print_mem short-circuits and
-        # never touches CUDA / torch.distributed.
-        log.setLevel(logging.INFO)
+        # At the INFO default (20 > DEBUG's 10) gather_and_print_mem
+        # short-circuits and never touches CUDA / torch.distributed; pass
+        # log_level=logging.DEBUG to exercise the memory diagnostics.
+        log.setLevel(log_level)
 
         return PyTorchTrainer(model, config, device, log)
 
diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py
index 224fefa..f07bd3a 100644
--- a/tests/datagen/test_artifacts.py
+++ b/tests/datagen/test_artifacts.py
@@ -37,6 +37,7 @@
 import pytest
 
 from ScaFFold.datagen import instance as inst
+from ScaFFold.datagen import layout
 from ScaFFold.datagen import mask_detection as md
 from ScaFFold.datagen.volumegen import (
     load_np_ptcloud,
@@ -74,13 +75,15 @@ def _seed_category(fract_base: Path, *, point_num: int, keep: range) -> Path:
 
     Pre-seeding all but instance 0 means a ``main`` run only has to generate the
     single missing instance, keeping the test fast. Returns the instance dir.
+    The library lives under the seed-keyed layout, so the paths are derived from
+    the same config the run under test uses.
     """
-    vt = 0.15
-    param_dir = fract_base / f"var{vt}" / "3DIFS_param"
+    config = _make_config(fract_base, point_num=point_num)
+    param_dir = Path(layout.category_param_dir(config))
     param_dir.mkdir(parents=True)
     np.savetxt(param_dir / "000000.csv", _contractive_params(), delimiter=",")
 
-    inst_dir = fract_base / f"var{vt}" / "instances" / f"np{point_num}" / "000000"
+    inst_dir = Path(layout.instance_dir(config)) / "000000"
     inst_dir.mkdir(parents=True)
     rng = np.random.default_rng(0)
     for i in keep:
@@ -231,6 +234,63 @@ def test_resume_rejects_truncated(tmp_path):
     assert np.isfinite(np.load(victim)).all()
 
 
+# ---------------------------------------------------------------------------
+# R36: a category's IFS parameters are parsed once, not once per instance
+# ---------------------------------------------------------------------------
+
+
+def test_category_params_parsed_once_per_category(tmp_path, monkeypatch):
+    """``main`` parses each category CSV once per rank, not once per work item.
+
+    The parse used to sit inside the per-item loop, so a full generation read
+    and re-parsed the same small CSV 145 times per category off the shared
+    filesystem. Work items for a category are contiguous in the block
+    partition, so a one-entry cache collapses that to one parse per category.
+    """
+    fract_base = tmp_path / "fractals"
+    point_num = 60
+    n_categories = 2
+    missing_per_category = 3
+
+    config = _make_config(fract_base, point_num=point_num)
+    config.n_categories = n_categories
+
+    param_dir = Path(layout.category_param_dir(config))
+    param_dir.mkdir(parents=True)
+    instance_root = Path(layout.instance_dir(config))
+    rng = np.random.default_rng(0)
+    for category in range(n_categories):
+        np.savetxt(
+            param_dir / f"{category:06d}.csv", _contractive_params(), delimiter=","
+        )
+        # Pre-seed all but a few instances so the run stays fast; the ones left
+        # missing are what the loop (and the parse) actually iterates over.
+        inst_dir = instance_root / f"{category:06d}"
+        inst_dir.mkdir(parents=True)
+        for i in range(missing_per_category, 145):
+            np.save(inst_dir / f"{category:06d}_{i:04d}.npy", rng.random((10, 3)))
+
+    parses = []
+    real_genfromtxt = np.genfromtxt
+
+    def counting_genfromtxt(fname, *args, **kwargs):
+        parses.append(Path(str(fname)).name)
+        return real_genfromtxt(fname, *args, **kwargs)
+
+    monkeypatch.setattr(inst.np, "genfromtxt", counting_genfromtxt)
+    inst.main(config)
+
+    category_parses = [name for name in parses if name[0].isdigit()]
+    # Every missing instance was generated ...
+    for category in range(n_categories):
+        for i in range(missing_per_category):
+            assert (
+                instance_root / f"{category:06d}" / f"{category:06d}_{i:04d}.npy"
+            ).exists()
+    # ... from n_categories parses, not one per (category, instance) item.
+    assert sorted(category_parses) == ["000000.csv", "000001.csv"]
+
+
 # ---------------------------------------------------------------------------
 # F62: mask scanner requires exactly one file per id
 # ---------------------------------------------------------------------------
@@ -309,20 +369,22 @@ def test_scale_config_rejected():
 # ---------------------------------------------------------------------------
 
 
-def _dense_reference_indices(points: np.ndarray, grid_size: int, eps=1e-6):
-    """The pre-refactor index computation, kept verbatim as a reference.
+def _reference_indices(points: np.ndarray, grid_size: int, eps=1e-6, *, clip=True):
+    """The index computation of ``points_to_voxel_indices``, spelled out.
 
-    This is the exact math the old dense ``points_to_voxelgrid`` ran before
-    scattering ``True`` into a full ``grid_size**3`` boolean array; the scatter
-    API must reproduce the identical occupied-voxel set and painted values.
+    Two tests need to see inside the function: the scatter-vs-dense equivalence
+    check (which needs the per-point indices the dense grid was built from) and
+    the centering check (which needs the indices *before* ``np.clip`` hides
+    out-of-range bins). Every test using this asserts the replica reproduces the
+    real function's output, so it cannot silently drift from it.
     """
     mins = points.min(axis=0)
     maxs = points.max(axis=0)
     voxel_size = (float((maxs - mins).max()) + eps) / grid_size
     scaled = (points - mins) / voxel_size
-    offset = (grid_size - 1 - scaled.max(axis=0)) / 2.0
+    offset = (grid_size - scaled.max(axis=0)) / 2.0
     idx = np.floor(scaled + offset).astype(int)
-    return np.clip(idx, 0, grid_size - 1)
+    return np.clip(idx, 0, grid_size - 1) if clip else idx
 
 
 def test_voxel_indices_match_dense_grid():
@@ -334,7 +396,7 @@ def test_voxel_indices_match_dense_grid():
     idx = points_to_voxel_indices(points, grid_size)
 
     # Reference dense grid built the old way, from the reference index math.
-    ref_idx = _dense_reference_indices(points, grid_size)
+    ref_idx = _reference_indices(points, grid_size)
     reference = np.zeros((grid_size,) * 3, dtype=bool)
     reference[ref_idx[:, 0], ref_idx[:, 1], ref_idx[:, 2]] = True
 
@@ -418,3 +480,80 @@ def test_dataset_version_bumped_past_float64_era():
     from ScaFFold.datagen import get_dataset as gd
 
     assert gd.DATASET_FORMAT_VERSION > 2
+
+
+# ---------------------------------------------------------------------------
+# R33: voxel centering is a whole voxel, not half of one
+# ---------------------------------------------------------------------------
+
+
+def _assert_replica_tracks_real(grid_size: int = 16) -> None:
+    """Pin ``_reference_indices`` to the real function on a sparse cloud.
+
+    A sparse cloud is essential here: a dense one occupies every voxel under
+    any offset, so the comparison would pass vacuously.
+    """
+    sparse = np.random.default_rng(7).random((300, 3)).astype(np.float32)
+    assert np.array_equal(
+        np.unique(_reference_indices(sparse, grid_size), axis=0),
+        points_to_voxel_indices(sparse, grid_size),
+    ), "the replicated index arithmetic no longer matches points_to_voxel_indices"
+
+
+def test_voxelization_never_bins_outside_the_grid():
+    """No point lands outside ``[0, grid_size)`` before the safety clip.
+
+    The centering offset positions a span of ``span`` voxels inside a grid of
+    ``grid_size`` voxels, so the free space to split between the two margins is
+    ``grid_size - span``. Using ``grid_size - 1 - span`` shifted every cloud
+    half a voxel toward the origin: points in the first half-voxel of each
+    filled axis floored to -1, and ``np.clip`` quietly folded them into bin 0.
+    """
+    grid_size = 16
+    _assert_replica_tracks_real(grid_size)
+    rng = np.random.default_rng(1234)
+    points = rng.random((200_000, 3)).astype(np.float32)
+
+    pre_clip = _reference_indices(points, grid_size, clip=False)
+    assert pre_clip.min() >= 0, (
+        f"{int((pre_clip < 0).any(axis=1).sum())} of {len(points)} points floored "
+        "below bin 0 and were clipped back in"
+    )
+    assert pre_clip.max() <= grid_size - 1
+
+
+def test_voxelization_density_is_uniform_at_the_boundaries():
+    """A uniform cloud fills the boundary planes like the interior ones.
+
+    The half-voxel shift piled the clipped points onto plane 0 (1.5x the
+    interior density) and starved the far plane (0.5x), a systematic
+    misregistration in every generated volume and mask.
+    """
+    grid_size = 16
+    _assert_replica_tracks_real(grid_size)
+    rng = np.random.default_rng(1234)
+    points = rng.random((200_000, 3)).astype(np.float32)
+
+    idx = _reference_indices(points, grid_size)
+    counts = np.bincount(idx[:, 0], minlength=grid_size)
+    interior = counts[2:-2].mean()
+    assert 0.9 <= counts[0] / interior <= 1.1, (
+        f"boundary plane 0 holds {counts[0] / interior:.2f}x the interior density"
+    )
+    assert 0.9 <= counts[-1] / interior <= 1.1, (
+        f"boundary plane {grid_size - 1} holds {counts[-1] / interior:.2f}x the "
+        "interior density"
+    )
+
+
+def test_dataset_version_bumped_past_half_voxel_era():
+    """The reuse marker advanced past 3: pre-fix datasets are misregistered.
+
+    Correcting the offset changes the voxel contents of every generated volume
+    and mask, so a dataset built before the fix must not be handed to a run
+    after it. Bumping the version both stops the reuse scan from matching those
+    directories and changes the config_id they hash to.
+    """
+    from ScaFFold.datagen import get_dataset as gd
+
+    assert gd.DATASET_FORMAT_VERSION > 3
diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py
index 8608884..3ff19a2 100644
--- a/tests/datagen/test_category_search.py
+++ b/tests/datagen/test_category_search.py
@@ -12,8 +12,16 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 
-"""Tests for category-search round sizing."""
+"""Tests for category-search round sizing and its work-scan consensus."""
 
+from argparse import Namespace
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from ScaFFold.datagen import category_search as cs
+from ScaFFold.datagen import layout
 from ScaFFold.datagen.category_search import compute_round_attempts
 
 
@@ -53,3 +61,315 @@ def test_at_least_one_attempt_per_rank_when_work_remains():
     # A positive remaining count always yields at least one attempt per rank so
     # the loop makes progress and can keep learning the acceptance rate.
     assert compute_round_attempts(1, 1024, 10000, 0.99) >= 1
+
+
+# ---------------------------------------------------------------------------
+# R30: the initial work scan is made once on rank 0 and broadcast.
+#
+# ``categories_remaining`` gates a while loop that contains collectives, so it
+# must be identical on every rank. Deriving it from a per-rank filesystem scan
+# lets divergent views (a stale metadata cache, a racing job, a partially
+# visible directory) put one rank inside the loop issuing ``bcast`` while
+# another is past it issuing ``reduce`` -- mismatched collectives on
+# COMM_WORLD, i.e. a hang. This mirrors the fix already applied in
+# ``instance.py``: rank 0 scans, everyone else consumes the broadcast.
+# ---------------------------------------------------------------------------
+
+
+class CategorySearchComm:
+    """Single-process stand-in for COMM_WORLD recording the collective order."""
+
+    def __init__(self, rank=0, size=1, bcast_returns=None):
+        self.rank = rank
+        self.size = size
+        self.calls = []
+        self.bcast_payloads = []
+        self._bcast_returns = list(bcast_returns or [])
+
+    def Get_rank(self):
+        return self.rank
+
+    def Get_size(self):
+        return self.size
+
+    def Barrier(self):
+        self.calls.append("Barrier")
+
+    def bcast(self, obj, root=0):
+        self.calls.append("bcast")
+        self.bcast_payloads.append(obj)
+        if self.rank == root:
+            return obj
+        return self._bcast_returns.pop(0)
+
+    def gather(self, obj, root=0):
+        self.calls.append("gather")
+        return [obj] if self.rank == root else None
+
+    def reduce(self, value, op=None, root=0):
+        self.calls.append("reduce")
+        return value if self.rank == root else None
+
+
+class FakeMPI:
+    """Namespace mimicking ``mpi4py.MPI`` for one ``CategorySearchComm``."""
+
+    def __init__(self, comm):
+        import mpi4py.MPI as real_mpi
+
+        self.COMM_WORLD = comm
+        self.SUM = real_mpi.SUM
+
+
+def _cs_config(fract_base: Path) -> Namespace:
+    return Namespace(
+        fract_base_dir=str(fract_base),
+        n_categories=1,
+        seed=42,
+        variance_threshold=0.15,
+        point_num=60,
+        normalize=1,
+        datagen_from_scratch=False,
+        datagen_batch_size=4,
+        verbose=0,
+    )
+
+
+def _seed_one_category(config: Namespace) -> None:
+    """Write the single category CSV this config asks for."""
+    param_dir = Path(layout.category_param_dir(config))
+    param_dir.mkdir(parents=True, exist_ok=True)
+    params = np.zeros((2, 13), dtype=np.float64)
+    params[:, 0] = params[:, 4] = params[:, 8] = 0.5
+    params[1, 9] = params[1, 10] = params[1, 11] = 0.5
+    params[0, 12] = 0.5
+    np.savetxt(param_dir / "000000.csv", params, delimiter=",")
+
+
+def test_work_scan_is_root_only_and_broadcast(tmp_path, monkeypatch):
+    """A non-root rank never scans; it consumes root's broadcast index list."""
+    config = _cs_config(tmp_path / "fractals")
+    comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])])
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    def no_scan(*_args, **_kwargs):
+        raise AssertionError("non-root rank must not scan the filesystem")
+
+    monkeypatch.setattr(cs, "parse_category_indices", no_scan)
+
+    cs.main(config)
+
+    # Root said category 0 already exists, so this rank has nothing to do and
+    # went straight to the post-loop reductions.
+    assert comm.calls == ["Barrier", "bcast", "reduce", "reduce", "reduce", "reduce"]
+    # It contributed nothing to the scan broadcast (it is not the scanner).
+    assert comm.bcast_payloads[0] is None
+
+
+def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch):
+    """Ranks disagreeing about the directory still issue identical collectives.
+
+    Rank 0 sees the finished category; rank 1's own view is empty. Before the
+    fix rank 1 entered the work loop (``bcast``) while rank 0 was already past
+    it (``reduce``). With the scan broadcast, rank 1's view is irrelevant.
+    """
+    # Rank 0: the category is on disk, so its scan finds it.
+    root_config = _cs_config(tmp_path / "root_view")
+    _seed_one_category(root_config)
+    root_comm = CategorySearchComm(rank=0, size=2)
+    monkeypatch.setattr(cs, "MPI", FakeMPI(root_comm))
+    cs.main(root_config)
+
+    # Rank 1: an empty directory (a divergent view), but root broadcast [0].
+    peer_config = _cs_config(tmp_path / "peer_view")
+    peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])])
+    monkeypatch.setattr(cs, "MPI", FakeMPI(peer_comm))
+    cs.main(peer_config)
+
+    assert root_comm.bcast_payloads[0] == ("ok", [0])
+    assert peer_comm.calls == root_comm.calls, (
+        "ranks with divergent filesystem views issued different collectives: "
+        f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# VB-4: temp files stranded by killed writes are swept, as in instance.py.
+#
+# Both atomic writers name their temp file after the writing pid, so a job
+# killed mid-write leaves one behind per killed process, forever: nothing in
+# this module ever looked at them again.
+# ---------------------------------------------------------------------------
+
+
+def test_stale_temp_files_are_swept(tmp_path, monkeypatch):
+    """Category and attempt-counter temps from dead pids are removed."""
+    config = _cs_config(tmp_path / "fractals")
+    param_dir = Path(layout.category_param_dir(config))
+    param_dir.mkdir(parents=True, exist_ok=True)
+    _seed_one_category(config)  # n_categories=1, so the search has no work
+
+    stale_csv = param_dir / ".000001.csv.tmp999999"
+    stale_csv.write_text("0.5,0.5\n")
+    stale_counter = param_dir / ".rng_attempt_rank3.tmp999999"
+    stale_counter.write_text("17")
+
+    comm = CategorySearchComm(rank=0, size=1)
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    cs.main(config)
+
+    assert not stale_csv.exists(), "a stranded category temp file was kept"
+    assert not stale_counter.exists(), "a stranded attempt-counter temp was kept"
+    # The real artifact is untouched: only the temp names are swept.
+    assert (param_dir / "000000.csv").exists()
+
+
+# ---------------------------------------------------------------------------
+# VB-6: a library in the old, seed-agnostic layout is explained, not ignored.
+# ---------------------------------------------------------------------------
+
+
+def test_old_layout_library_is_reported(tmp_path, monkeypatch, caplog):
+    """A pre-relayout library produces one warning naming both directories."""
+    config = _cs_config(tmp_path / "fractals")
+    legacy = Path(layout.legacy_category_param_dir(config))
+    legacy.mkdir(parents=True)
+    (legacy / "000000.csv").write_text("")
+    _seed_one_category(config)  # nothing to generate under the new layout
+
+    comm = CategorySearchComm(rank=0, size=1)
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    with caplog.at_level("WARNING"):
+        cs.main(config)
+
+    messages = " ".join(record.getMessage() for record in caplog.records)
+    assert str(legacy) in messages
+    assert layout.category_param_dir(config) in messages
+
+
+def test_no_warning_without_an_old_layout(tmp_path, monkeypatch, caplog):
+    """The warning does not fire for a fresh library (control)."""
+    config = _cs_config(tmp_path / "fractals")
+    _seed_one_category(config)
+
+    comm = CategorySearchComm(rank=0, size=1)
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    with caplog.at_level("WARNING"):
+        cs.main(config)
+
+    assert "seed-agnostic" not in " ".join(
+        record.getMessage() for record in caplog.records
+    )
+
+
+# ---------------------------------------------------------------------------
+# VB-2: the whole rank-0 scan window is fenced, not just the index parse.
+#
+# Rank 0 also reads the parameters of every category already on disk, right
+# after the scan broadcast. A CSV that will not parse therefore killed rank 0
+# while its peers had already consumed the broadcast and moved on -- the same
+# stranding the scan broadcast was introduced to prevent.
+# ---------------------------------------------------------------------------
+
+
+def test_unparseable_existing_category_is_broadcast_not_raised_on_root(
+    tmp_path, monkeypatch
+):
+    """A ragged category CSV becomes a broadcast error, not a rank-0-only death."""
+    config = _cs_config(tmp_path / "fractals")
+    param_dir = Path(layout.category_param_dir(config))
+    param_dir.mkdir(parents=True, exist_ok=True)
+    # Six-digit name, so the scan accepts it; ragged rows, so loadtxt raises.
+    (param_dir / "000000.csv").write_text("0.5,0.5,0.5\n0.5,0.5\n")
+
+    comm = CategorySearchComm(rank=0, size=2)
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    with pytest.raises(RuntimeError) as excinfo:
+        cs.main(config)
+
+    # Rank 0 stopped at the scan broadcast, and what it broadcast is the error
+    # sentinel its peers need in order to abort with it.
+    assert comm.calls == ["Barrier", "bcast"]
+    assert comm.bcast_payloads[-1][0] == "error"
+    assert "000000.csv" in str(excinfo.value) or "3DIFS_param" in str(excinfo.value)
+
+
+def test_peer_raises_on_broadcast_scan_error(tmp_path, monkeypatch):
+    """A peer receiving the scan sentinel raises instead of entering the loop."""
+    config = _cs_config(tmp_path / "fractals")
+    sentinel = ("error", "rank 0 failed to scan existing categories: ValueError: boom")
+    comm = CategorySearchComm(rank=1, size=2, bcast_returns=[sentinel])
+    monkeypatch.setattr(cs, "MPI", FakeMPI(comm))
+
+    with pytest.raises(RuntimeError, match="boom"):
+        cs.main(config)
+
+    # It never reached the work loop's collectives.
+    assert comm.calls == ["Barrier", "bcast"]
+
+
+# ---------------------------------------------------------------------------
+# R31: category CSVs appear complete or not at all.
+#
+# A category file truncated by a killed job is still counted as "done" by the
+# resume scan, so nothing ever regenerates it: instance generation then dies
+# parsing it, and the category search's own resume dies re-loading it. The
+# pipeline cannot self-heal -- the file has to be deleted by hand.
+# ---------------------------------------------------------------------------
+
+
+def _params() -> np.ndarray:
+    params = np.zeros((2, 13), dtype=np.float64)
+    params[:, 0] = params[:, 4] = params[:, 8] = 0.5
+    params[0, 12] = 0.5
+    return params
+
+
+def test_category_csv_write_is_atomic(tmp_path, monkeypatch):
+    """A killed mid-write leaves no category file under a name resume accepts."""
+    param_dir = tmp_path / "3DIFS_param"
+    param_dir.mkdir()
+
+    # One complete category, saved normally.
+    indices, saved = [], []
+    first = _params()
+    assert cs.save_valid_category(str(param_dir), first, indices, saved) == 0
+    assert np.loadtxt(param_dir / "000000.csv", delimiter=",").shape == (2, 13)
+
+    # The next save is interrupted after some bytes have been written.
+    observed = {}
+
+    def partial_then_raise(fname, arr, *args, **kwargs):
+        observed["listing"] = sorted(p.name for p in param_dir.iterdir())
+        handle = fname if hasattr(fname, "write") else open(fname, "w")
+        handle.write("0.5,0.5,0.5\n")
+        handle.flush()
+        if handle is not fname:
+            handle.close()
+        raise OSError("simulated SIGKILL mid-write")
+
+    monkeypatch.setattr(cs.np, "savetxt", partial_then_raise)
+
+    second = _params()
+    second[0, 0] = 0.25
+    with pytest.raises(OSError):
+        cs.save_valid_category(str(param_dir), second, indices, saved)
+
+    # No truncated category is visible: the resume scan still sees exactly the
+    # one complete category, and every file it names parses.
+    assert cs.parse_category_indices(str(param_dir)) == [0]
+    assert not (param_dir / "000001.csv").exists()
+    for idx in cs.parse_category_indices(str(param_dir)):
+        assert np.loadtxt(param_dir / f"{idx:06d}.csv", delimiter=",").shape == (2, 13)
+
+    # Mid-write, the partial data lived under a name neither the resume scan
+    # nor the instance loader (which takes every ``*.csv``) would pick up.
+    partial_names = [n for n in observed["listing"] if n != "000000.csv"]
+    assert all(not name.endswith(".csv") for name in partial_names), partial_names
+
+    # And nothing was left behind afterwards.
+    assert sorted(p.name for p in param_dir.iterdir()) == ["000000.csv"]
diff --git a/tests/datagen/test_library_layout.py b/tests/datagen/test_library_layout.py
new file mode 100644
index 0000000..d09a524
--- /dev/null
+++ b/tests/datagen/test_library_layout.py
@@ -0,0 +1,248 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""The fractal library is keyed by the seed that produced it (R29).
+
+Categories and instances are *derived from* ``config.seed``: the IFS parameters
+come from a seed-keyed candidate stream and every instance point cloud is
+seeded from ``(seed, category, instance)``. Resume, however, is a pure
+file-existence check, so a library laid out only by variance threshold and
+point count let a run under one seed silently adopt another seed's data --
+and then publish a dataset whose metadata claimed the new seed. Two datasets
+with identical provenance and different content.
+
+The fix puts the seed in the directory path, so the question "does this file
+exist" is asked in a seed-specific place and can only ever be answered with
+data that seed produced:
+
+    <fract_base_dir>/var<variance_threshold>/seed<seed>/3DIFS_param/
+    <fract_base_dir>/var<variance_threshold>/seed<seed>/instances/np<point_num>/
+
+Everything here runs single-process at tiny scale (one category, 60-point
+clouds) so the real generators run in well under a second.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from argparse import Namespace
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from ScaFFold.datagen import category_search as cs
+from ScaFFold.datagen import instance as inst
+from ScaFFold.datagen import layout, volumegen
+
+VT = 0.15
+POINT_NUM = 60
+
+
+def _contractive_params() -> np.ndarray:
+    """A 2-map IFS whose orbit stays bounded, so generation is fast and finite."""
+    params = np.zeros((2, 13), dtype=np.float64)
+    params[:, 0] = params[:, 4] = params[:, 8] = 0.5
+    params[1, 9] = params[1, 10] = params[1, 11] = 0.5
+    params[0, 12] = 0.5
+    return params
+
+
+def _param_dir(fract_base: Path, seed: int) -> Path:
+    """The category directory the new layout mandates, spelled out literally."""
+    return fract_base / f"var{VT}" / f"seed{seed}" / "3DIFS_param"
+
+
+def _instance_dir(fract_base: Path, seed: int) -> Path:
+    """The instance directory the new layout mandates, spelled out literally."""
+    return fract_base / f"var{VT}" / f"seed{seed}" / "instances" / f"np{POINT_NUM}"
+
+
+def _seed_params(fract_base: Path, seed: int, n_categories: int = 1) -> Path:
+    param_dir = _param_dir(fract_base, seed)
+    param_dir.mkdir(parents=True, exist_ok=True)
+    for category in range(n_categories):
+        np.savetxt(
+            param_dir / f"{category:06d}.csv", _contractive_params(), delimiter=","
+        )
+    return param_dir
+
+
+def _inst_config(fract_base: Path, seed: int) -> Namespace:
+    return Namespace(
+        fract_base_dir=str(fract_base),
+        n_categories=1,
+        seed=seed,
+        variance_threshold=VT,
+        point_num=POINT_NUM,
+        datagen_from_scratch=False,
+        verbose=0,
+    )
+
+
+def _cs_config(fract_base: Path, seed: int) -> Namespace:
+    return Namespace(
+        fract_base_dir=str(fract_base),
+        n_categories=1,
+        seed=seed,
+        variance_threshold=VT,
+        point_num=POINT_NUM,
+        normalize=1,
+        datagen_from_scratch=False,
+        datagen_batch_size=4,
+        verbose=0,
+    )
+
+
+def _volumegen_config(dataset_dir: Path, fract_base: Path, seed: int) -> Namespace:
+    return Namespace(
+        dataset_dir=str(dataset_dir),
+        fract_base_dir=str(fract_base),
+        n_categories=1,
+        n_instances_used_per_fractal=1,
+        n_fracts_per_vol=1,
+        seed=seed,
+        variance_threshold=VT,
+        val_split=0,
+        vol_size=8,
+        point_num=POINT_NUM,
+        scale=1,
+        verbose=0,
+    )
+
+
+def _library_digest(instance_dir: Path) -> tuple[int, str]:
+    """(file count, content digest) for one instance directory."""
+    digest = hashlib.sha256()
+    files = sorted(instance_dir.rglob("*.npy"))
+    for path in files:
+        digest.update(path.name.encode())
+        digest.update(path.read_bytes())
+    return len(files), digest.hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# The layout helper is the single definition of the seed-keyed paths.
+# ---------------------------------------------------------------------------
+
+
+def test_layout_helpers_key_every_path_by_seed(tmp_path):
+    """Both library directories carry the seed, and differ across seeds."""
+    fract_base = tmp_path / "fractals"
+
+    for seed in (7, 999):
+        config = _inst_config(fract_base, seed)
+        assert Path(layout.category_param_dir(config)) == _param_dir(fract_base, seed)
+        assert Path(layout.instance_dir(config)) == _instance_dir(fract_base, seed)
+
+    assert layout.category_param_dir(_inst_config(fract_base, 7)) != (
+        layout.category_param_dir(_inst_config(fract_base, 999))
+    )
+    assert layout.instance_dir(_inst_config(fract_base, 7)) != (
+        layout.instance_dir(_inst_config(fract_base, 999))
+    )
+
+
+# ---------------------------------------------------------------------------
+# Producers write under the seed; consumers read from under the seed.
+# ---------------------------------------------------------------------------
+
+
+def test_category_search_writes_under_the_seed_dir(tmp_path):
+    """A generated category CSV lands in this seed's parameter directory."""
+    fract_base = tmp_path / "fractals"
+
+    cs.main(_cs_config(fract_base, seed=42))
+
+    assert (_param_dir(fract_base, 42) / "000000.csv").exists()
+    # Nothing was written to a seed-agnostic location.
+    assert not (fract_base / f"var{VT}" / "3DIFS_param").exists()
+
+
+def test_instances_are_written_under_the_seed_dir(tmp_path):
+    """Instance point clouds land under this seed's instance directory."""
+    fract_base = tmp_path / "fractals"
+    _seed_params(fract_base, seed=7)
+
+    inst.main(_inst_config(fract_base, seed=7))
+
+    count, _digest = _library_digest(_instance_dir(fract_base, 7))
+    assert count == 145
+    assert not (fract_base / f"var{VT}" / "instances").exists()
+
+
+def test_same_seed_resume_generates_nothing_new(tmp_path):
+    """Re-running under the same seed reuses the library byte-for-byte.
+
+    The resume path must stay cheap: the whole point of the library is that a
+    second run under the same configuration regenerates nothing.
+    """
+    fract_base = tmp_path / "fractals"
+    _seed_params(fract_base, seed=7)
+    inst.main(_inst_config(fract_base, seed=7))
+
+    instance_dir = _instance_dir(fract_base, 7)
+    count_before, digest_before = _library_digest(instance_dir)
+    mtimes_before = {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))}
+
+    inst.main(_inst_config(fract_base, seed=7))
+
+    count_after, digest_after = _library_digest(instance_dir)
+    assert (count_after, digest_after) == (count_before, digest_before)
+    # No file was rewritten (0 new instances generated).
+    assert {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))} == (
+        mtimes_before
+    )
+
+
+def test_different_seed_cannot_reuse_another_seeds_instances(tmp_path):
+    """A second seed generates its own library instead of adopting the first.
+
+    Before the fix this run reported "Generated 0 instances" and left the
+    first seed's bytes in place, so the dataset built on top of it carried the
+    wrong seed's data under the new seed's provenance.
+    """
+    fract_base = tmp_path / "fractals"
+    _seed_params(fract_base, seed=7)
+    _seed_params(fract_base, seed=999)
+
+    inst.main(_inst_config(fract_base, seed=7))
+    count_7, digest_7 = _library_digest(_instance_dir(fract_base, 7))
+
+    inst.main(_inst_config(fract_base, seed=999))
+    count_999, digest_999 = _library_digest(_instance_dir(fract_base, 999))
+
+    # Both libraries are complete and independent...
+    assert count_7 == count_999 == 145
+    assert digest_999 != digest_7, "seed 999 reused seed 7's instances"
+    # ...and the first seed's data was left untouched.
+    assert _library_digest(_instance_dir(fract_base, 7)) == (count_7, digest_7)
+
+
+def test_volumegen_reads_instances_for_its_own_seed(tmp_path):
+    """volumegen resolves point clouds under the seed it was configured with."""
+    fract_base = tmp_path / "fractals"
+    _seed_params(fract_base, seed=7)
+    inst.main(_inst_config(fract_base, seed=7))
+
+    # Seed 7: the instances are where volumegen looks, so generation succeeds.
+    volumegen.main(_volumegen_config(tmp_path / "ds7", fract_base, seed=7))
+    assert list((tmp_path / "ds7" / "volumes").rglob("*.npy"))
+
+    # Seed 999: a different library entirely. Nothing has been generated for
+    # it, so volumegen must report the missing file rather than quietly
+    # rasterizing seed 7's clouds.
+    with pytest.raises(RuntimeError) as excinfo:
+        volumegen.main(_volumegen_config(tmp_path / "ds999", fract_base, seed=999))
+    assert "seed999" in str(excinfo.value)
diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py
index 6a67e89..5185eb3 100644
--- a/tests/datagen/test_mpi_consensus.py
+++ b/tests/datagen/test_mpi_consensus.py
@@ -38,7 +38,10 @@
 
 from __future__ import annotations
 
+import logging
+import os
 import re
+import time
 from argparse import Namespace
 from math import ceil
 from pathlib import Path
@@ -50,7 +53,7 @@
 
 from ScaFFold.datagen import get_dataset as gd
 from ScaFFold.datagen import instance as inst
-from ScaFFold.datagen import volumegen
+from ScaFFold.datagen import layout, volumegen
 
 RANK_SCRIPTS = Path(__file__).resolve().parents[1] / "helpers" / "rank_scripts"
 
@@ -340,6 +343,482 @@ def test_generation_success_finalizes_and_returns(tmp_path, monkeypatch):
     assert leftover == []
 
 
+# ---------------------------------------------------------------------------
+# R27: rank 0 must never die between the collectives its peers have entered.
+# Every rank-0-only step of the consensus (the reuse/generate decision and the
+# final meta-write + rename) is wrapped so a failure travels to the peers as a
+# broadcast sentinel instead of stranding them in ``bcast``/``Barrier``.
+# ---------------------------------------------------------------------------
+
+
+def _base_dir_for(config: Namespace) -> Path:
+    """The ``<dataset_dir>/<config_id>`` directory ``get_dataset`` scans."""
+    config_dict = vars(config).copy()
+    config_dict["dataset_format_version"] = gd.DATASET_FORMAT_VERSION
+    volume_config = gd._get_required_keys_dict(config_dict, gd.INCLUDE_KEYS)
+    return Path(config.dataset_dir) / gd._hash_volume_config(volume_config)
+
+
+def test_decision_failure_is_broadcast_not_raised_before_bcast(tmp_path, monkeypatch):
+    """A rank-0 decision failure reaches peers through the broadcast.
+
+    Any exception inside the rank-0-only decision (an unreadable base dir, a
+    staging ``mkdir`` hitting ENOSPC, ...) must be converted into an error
+    sentinel that is broadcast, so peers already waiting in ``bcast`` learn
+    about it and raise the same error. Before the fix rank 0 raised *before*
+    reaching the broadcast, leaving every peer blocked forever.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=2)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+
+    def explode(*_args, **_kwargs):
+        raise OSError("No space left on device")
+
+    monkeypatch.setattr(gd, "_decide_reuse_or_generate", explode)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        gd.get_dataset(config)
+
+    # Rank 0 reached the broadcast before raising, and the payload is the
+    # error sentinel every peer will see.
+    assert comm.calls == ["bcast"]
+    assert comm.bcast_payloads[0][0] == "error"
+    assert "No space left on device" in str(excinfo.value)
+
+
+def test_non_root_raises_on_broadcast_decision_error(tmp_path, monkeypatch):
+    """A peer receiving the error sentinel raises instead of generating."""
+    config = _reuse_config(tmp_path / "datasets")
+    sentinel = ("error", "rank 0: OSError: No space left on device")
+    comm = FakeComm(rank=1, size=2, bcast_returns=[sentinel])
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+
+    with pytest.raises(RuntimeError) as excinfo:
+        gd.get_dataset(config)
+
+    assert "No space left on device" in str(excinfo.value)
+    # The peer stopped at the decision broadcast: no generation collectives.
+    assert comm.calls == ["bcast"]
+
+
+def test_reuse_scan_skips_staging_dirs(tmp_path, monkeypatch):
+    """A complete ``meta.yaml`` stranded in a ``.tmp_*`` dir is never reused.
+
+    A job killed between the meta write and the rename leaves a fully valid
+    meta inside its staging dir. Treating that as a publishable dataset hands
+    back a half-generated directory (and one that cleanup may delete).
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    base = _base_dir_for(config)
+    config_id = base.name
+    stranded = base / ".tmp_20260101-000000_1234"
+    stranded.mkdir(parents=True)
+    (stranded / gd.META_FILENAME).write_text(
+        yaml.safe_dump(
+            {
+                "config_id": config_id,
+                "dataset_format_version": gd.DATASET_FORMAT_VERSION,
+            }
+        )
+    )
+
+    comm = FakeComm(rank=0, size=1)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+    log = logging.getLogger("test_reuse_scan_skips_staging_dirs")
+
+    decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log)
+
+    assert decision[0] == "generate", f"staging dir was reused: {decision}"
+
+
+def test_reuse_scan_tolerates_corrupt_meta(tmp_path, monkeypatch):
+    """A corrupt/unreadable candidate meta is skipped, not fatal.
+
+    A 0-byte ``meta.yaml`` (``yaml.safe_load`` -> ``None``) or an unparseable
+    one used to raise inside the rank-0-only scan. The scan must warn, skip the
+    directory, and keep looking -- here finding the good dataset next to it.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    base = _base_dir_for(config)
+    config_id = base.name
+    base.mkdir(parents=True)
+
+    # Sorted-descending scan order visits these two poison dirs first.
+    (base / "20260301-000000__zzz").mkdir()
+    (base / "20260301-000000__zzz" / gd.META_FILENAME).write_text("")
+    (base / "20260201-000000__yyy").mkdir()
+    (base / "20260201-000000__yyy" / gd.META_FILENAME).write_text("{[not yaml")
+
+    good = _write_reusable_dataset(base, config_id)
+    log = logging.getLogger("test_reuse_scan_tolerates_corrupt_meta")
+
+    decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log)
+
+    assert decision[0] == "reuse"
+    assert Path(decision[1]) == good
+
+
+def test_staging_dir_names_are_collision_proof(tmp_path, monkeypatch):
+    """Two decisions in the same second stage into different directories.
+
+    The old name was ``.tmp_%Y%m%d-%H%M%S`` with ``mkdir(exist_ok=False)``, so
+    two same-config jobs starting in the same second raced to a
+    ``FileExistsError`` on one of them -- inside the unguarded rank-0 window.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    base = _base_dir_for(config)
+    base.mkdir(parents=True)
+    log = logging.getLogger("test_staging_dir_names_are_collision_proof")
+
+    # Pin the clock so both decisions share a timestamp: only a non-time
+    # component can keep the names apart.
+    monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000")
+
+    first = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log)
+    second = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log)
+
+    assert first[0] == "generate" and second[0] == "generate"
+    assert first[1] != second[1], "same-second staging dirs collided"
+    assert Path(first[1]).is_dir() and Path(second[1]).is_dir()
+
+
+def test_finalize_failure_is_broadcast_not_left_to_barrier(tmp_path, monkeypatch):
+    """A rank-0 rename failure is broadcast; peers raise instead of hanging.
+
+    The rename happens *after* the generation consensus, so a failure there
+    (e.g. a racing job already created the destination) used to kill rank 0
+    while every peer sat in the final ``Barrier``. The fix carries the failure
+    through one more collective and raises everywhere.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=2, allreduce_result=1)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+    monkeypatch.setattr(volumegen, "main", lambda _config: None)
+
+    # Pin the clock so the destination name is predictable, then have a
+    # "racing job" occupy it with a non-empty directory: the rename fails with
+    # ENOTEMPTY exactly as it did in the field.
+    monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000")
+    base = _base_dir_for(config)
+    dest = base / "20260101-000000__abc123"
+    dest.mkdir(parents=True)
+    (dest / "placeholder").write_text("created by a racing job")
+
+    with pytest.raises(RuntimeError) as excinfo:
+        gd.get_dataset(config)
+
+    message = str(excinfo.value)
+    assert "20260101-000000__abc123" in message
+    # The failure travelled through a collective *after* the generation
+    # consensus, so peers learn about it rather than waiting in the barrier.
+    assert "allgather" in comm.calls
+    assert comm.calls.index("allgather") < len(comm.calls) - 1
+
+
+def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch):
+    """A peer receiving the finalize error raises rather than returning dest."""
+    config = _reuse_config(tmp_path / "datasets")
+    dest = tmp_path / "datasets" / "cid" / "20260101-000000__abc123"
+    comm, _tmp, _dest = _generate_decision_comm(
+        rank=1, size=2, dest=dest, allreduce_result=1
+    )
+    # Second bcast: root's finalize verdict (a failure message).
+    comm._bcast_returns.append("rank 0 failed to finalize: OSError: boom")
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+    monkeypatch.setattr(volumegen, "main", lambda _config: None)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        gd.get_dataset(config)
+
+    assert "boom" in str(excinfo.value)
+    assert not dest.exists()
+
+
+# ---------------------------------------------------------------------------
+# VB-3: the consensus guards catch BaseException, not (Exception, SystemExit).
+#
+# ``KeyboardInterrupt`` is neither, so an interrupt delivered to rank 0 alone
+# (Ctrl-C on the launching terminal, a site watchdog's SIGINT) unwound straight
+# past the broadcast and hung every peer -- the failure mode the guard exists to
+# prevent, arriving through the one exception class it did not cover.
+# ---------------------------------------------------------------------------
+
+
+def test_rank0_interrupt_still_posts_the_decision_sentinel(tmp_path, monkeypatch):
+    """A KeyboardInterrupt on rank 0 reaches the peers as an error sentinel."""
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=2)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+
+    def interrupted(*_args, **_kwargs):
+        raise KeyboardInterrupt
+
+    monkeypatch.setattr(gd, "_decide_reuse_or_generate", interrupted)
+
+    # Rank 0 keeps the operator's abort ...
+    with pytest.raises(KeyboardInterrupt):
+        gd.get_dataset(config)
+
+    # ... but only after telling the peers to stop.
+    assert comm.calls == ["bcast"]
+    assert comm.bcast_payloads[0][0] == "error"
+    assert "KeyboardInterrupt" in comm.bcast_payloads[0][1]
+
+
+def test_interrupt_during_generation_reaches_the_consensus(tmp_path, monkeypatch):
+    """An interrupt inside volumegen still drives the allreduce and allgather."""
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=2, allreduce_result=0, allgather_peers=[""])
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+
+    def interrupted(_config):
+        raise KeyboardInterrupt
+
+    monkeypatch.setattr(volumegen, "main", interrupted)
+
+    with pytest.raises(KeyboardInterrupt):
+        gd.get_dataset(config)
+
+    assert "allreduce" in comm.calls and "allgather" in comm.calls
+    assert "KeyboardInterrupt" in comm.allgather_payloads[0]
+
+
+# ---------------------------------------------------------------------------
+# R37: orphaned staging dirs are reclaimed instead of accumulating forever.
+# ---------------------------------------------------------------------------
+
+
+def _age_tree(path: Path, seconds: float) -> None:
+    """Backdate ``path`` and everything under it by ``seconds``."""
+    stamp = time.time() - seconds
+    for entry in sorted(path.rglob("*"), reverse=True):
+        os.utime(entry, (stamp, stamp))
+    os.utime(path, (stamp, stamp))
+
+
+def test_stale_staging_dirs_are_cleaned(tmp_path, monkeypatch):
+    """A long-orphaned ``.tmp_*`` dir is reclaimed; a live one is not.
+
+    Every killed generation leaves a full staging tree behind (potentially
+    terabytes) that nothing ever reclaims. Cleanup is age-gated so a *running*
+    job's staging dir -- which by construction is far younger than the
+    threshold -- is never deleted out from under it.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    base = _base_dir_for(config)
+    base.mkdir(parents=True)
+
+    orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef"
+    (orphan / "volumes" / "training").mkdir(parents=True)
+    (orphan / "volumes" / "training" / "0.npy").write_bytes(b"stale payload")
+    _age_tree(orphan, 10 * gd.STALE_STAGING_AGE_SECONDS)
+
+    live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe"
+    (live / "volumes").mkdir(parents=True)
+
+    log = logging.getLogger("test_stale_staging_dirs_are_cleaned")
+    gd._decide_reuse_or_generate(base, base.name, "abc123", False, log)
+
+    assert not orphan.exists(), "orphaned staging dir was not reclaimed"
+    assert live.exists(), "a concurrent job's staging dir was deleted"
+
+
+def test_cleanup_never_touches_published_datasets(tmp_path):
+    """Only ``.tmp_*`` dirs under this config_id base are ever removed.
+
+    An old published dataset is precisely what reuse is for, and other configs'
+    (or other users') directories are none of this job's business.
+    """
+    config = _reuse_config(tmp_path / "datasets")
+    base = _base_dir_for(config)
+    base.mkdir(parents=True)
+
+    published = _write_reusable_dataset(base, "some-other-config")
+    _age_tree(published, 10 * gd.STALE_STAGING_AGE_SECONDS)
+
+    # A staging dir belonging to a different config_id base entirely.
+    other_base = base.parent / "0123456789ab"
+    other_orphan = other_base / f"{gd.TMP_PREFIX}20200101-000000_333_f00d"
+    other_orphan.mkdir(parents=True)
+    _age_tree(other_orphan, 10 * gd.STALE_STAGING_AGE_SECONDS)
+
+    log = logging.getLogger("test_cleanup_never_touches_published_datasets")
+    gd._decide_reuse_or_generate(base, base.name, "abc123", False, log)
+
+    assert published.exists(), "an old published dataset was deleted"
+    assert other_orphan.exists(), "cleanup escaped this job's config_id base"
+
+
+# ---------------------------------------------------------------------------
+# VB-1: a staging dir is aged by its DEEPEST recent write, not its top level.
+#
+# Generation is not bounded by the staleness threshold -- at the larger scales
+# it runs for days -- and after the first minutes it writes only at depth >= 2
+# (``volumes/<split>/N.npy``). Judging liveness from the staging dir and its
+# immediate children alone therefore reported a healthy multi-day generation as
+# "untouched for 48 hours", and a concurrent same-config start rmtree'd it out
+# from under its peers (which then died on FileNotFoundError).
+# ---------------------------------------------------------------------------
+
+
+def _cleanup(base: Path, name: str) -> None:
+    gd._cleanup_stale_staging_dirs(base, logging.getLogger(name))
+
+
+def test_live_generation_survives_a_deep_write(tmp_path):
+    """A >24h-old staging dir with a fresh deep write is NOT reclaimed."""
+    base = tmp_path / "cid"
+    base.mkdir(parents=True)
+
+    live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe"
+    split = live / "volumes" / "training"
+    split.mkdir(parents=True)
+    _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS)
+
+    # The job is alive and still laying down volumes: the write lands two levels
+    # down, so only that directory's mtime is current.
+    (split / "0.npy").write_bytes(b"payload from a running job")
+
+    _cleanup(base, "test_live_generation_survives_a_deep_write")
+
+    assert live.exists(), "a live generation's staging dir was reclaimed"
+
+
+def test_live_generation_survives_on_its_heartbeat_alone(tmp_path):
+    """A fresh heartbeat keeps a staging dir whose whole tree looks ancient.
+
+    The mtime probe is bounded, so a generation writing deeper than it looks
+    (or on a filesystem with coarse directory mtimes) still has to be safe. The
+    heartbeat is the signal that does not depend on the shape of the tree.
+    """
+    base = tmp_path / "cid"
+    base.mkdir(parents=True)
+
+    live = base / f"{gd.TMP_PREFIX}20260101-000000_333_f00d"
+    (live / "volumes" / "training" / "deep" / "deeper").mkdir(parents=True)
+    heartbeat = live / volumegen.STAGING_HEARTBEAT_NAME
+    heartbeat.write_text("")
+    _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS)
+    now = time.time()
+    os.utime(heartbeat, (now, now))
+
+    _cleanup(base, "test_live_generation_survives_on_its_heartbeat_alone")
+
+    assert live.exists(), "a heartbeating generation's staging dir was reclaimed"
+
+
+def test_dead_staging_dir_with_stale_heartbeat_is_reclaimed(tmp_path):
+    """The heartbeat must not turn cleanup into a no-op (R37 still holds)."""
+    base = tmp_path / "cid"
+    base.mkdir(parents=True)
+
+    orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef"
+    split = orphan / "volumes" / "training"
+    split.mkdir(parents=True)
+    (split / "0.npy").write_bytes(b"stale payload")
+    (orphan / volumegen.STAGING_HEARTBEAT_NAME).write_text("")
+    _age_tree(orphan, 2 * gd.STALE_STAGING_AGE_SECONDS)
+
+    _cleanup(base, "test_dead_staging_dir_with_stale_heartbeat_is_reclaimed")
+
+    assert not orphan.exists(), "an orphaned staging dir was not reclaimed"
+
+
+def test_generation_writes_and_then_drops_the_heartbeat(tmp_path, monkeypatch):
+    """volumegen marks the staging dir live; publishing removes the marker."""
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=1, allreduce_result=1)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+
+    beating = {}
+
+    def fake_volumegen(cfg):
+        # Stand in for the write loop: report the staging dir as live.
+        volumegen.StagingHeartbeat(cfg.dataset_dir).beat()
+        beating["path"] = Path(cfg.dataset_dir) / volumegen.STAGING_HEARTBEAT_NAME
+        beating["existed_during_generation"] = beating["path"].exists()
+
+    monkeypatch.setattr(volumegen, "main", fake_volumegen)
+
+    published = Path(gd.get_dataset(config))
+
+    assert beating["existed_during_generation"], "no heartbeat during generation"
+    assert not (published / volumegen.STAGING_HEARTBEAT_NAME).exists(), (
+        "the staging heartbeat was published with the dataset"
+    )
+
+
+def test_heartbeat_respects_its_interval(tmp_path):
+    """``beat`` is a no-op until the interval elapses, then refreshes."""
+    staging = tmp_path / "staging"
+    staging.mkdir()
+    heartbeat = volumegen.StagingHeartbeat(staging, interval=60)
+
+    assert heartbeat.beat(now=1000.0) is True
+    first = Path(heartbeat.path).stat().st_mtime
+    assert heartbeat.beat(now=1030.0) is False  # inside the interval
+    assert Path(heartbeat.path).stat().st_mtime == first
+    assert heartbeat.beat(now=1090.0) is True
+    assert Path(heartbeat.path).stat().st_mtime > first
+
+
+# ---------------------------------------------------------------------------
+# R28: meta.yaml is published atomically, so no reader ever sees a partial one.
+# ---------------------------------------------------------------------------
+
+
+def test_meta_write_is_atomic(tmp_path, monkeypatch):
+    """An interrupted meta write leaves the previous file intact and no temp.
+
+    ``meta.yaml`` is the file that decides how every sample is interpreted (a
+    truncated one reclassifies the dataset as legacy v1), so it must appear at
+    its final name complete or not at all.
+    """
+    target = tmp_path / gd.META_FILENAME
+    gd._write_meta_atomic(target, {"dataset_format_version": gd.DATASET_FORMAT_VERSION})
+    good_bytes = target.read_bytes()
+
+    # Interrupt the write after bytes have reached the temp file but before the
+    # rename -- the shape of a kill mid-write.
+    def boom(_fd):
+        raise OSError("simulated SIGKILL mid-write")
+
+    monkeypatch.setattr(gd.os, "fsync", boom)
+
+    with pytest.raises(OSError):
+        gd._write_meta_atomic(target, {"dataset_format_version": 99})
+
+    # The final name still holds the complete previous file, byte-for-byte, and
+    # no temp file is left behind for the reuse scan to trip over.
+    assert target.read_bytes() == good_bytes
+    assert [p.name for p in tmp_path.iterdir()] == [gd.META_FILENAME]
+
+
+def test_published_dataset_has_no_partial_meta(tmp_path, monkeypatch):
+    """A successful generation publishes a parseable meta and no temp files."""
+    config = _reuse_config(tmp_path / "datasets")
+    comm = FakeComm(rank=0, size=1, allreduce_result=1)
+    monkeypatch.setattr(gd, "MPI", FakeMPI(comm))
+    monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123")
+    monkeypatch.setattr(volumegen, "main", lambda _config: None)
+
+    result = Path(gd.get_dataset(config))
+
+    meta = yaml.safe_load((result / gd.META_FILENAME).read_text())
+    assert meta["dataset_format_version"] == gd.DATASET_FORMAT_VERSION
+    # Nothing hidden alongside it (a temp meta would be a dotted sibling).
+    assert [p.name for p in result.iterdir() if p.name.startswith(".")] == []
+
+
 # ---------------------------------------------------------------------------
 # A missing instance file raises FileNotFoundError (a catchable Exception)
 # rather than calling sys.exit(1) (a BaseException that bypasses consensus), and
@@ -370,15 +849,10 @@ def _seed_one_instance(fract_base: Path, config: Namespace, *, present: bool) ->
 
     volumegen selects instance indices with ``random.sample(range(145), ...)``
     seeded by ``config.seed``; to be robust we populate every one of the 145
-    instance slots for category 0 when ``present`` is True.
+    instance slots for category 0 when ``present`` is True. The library path is
+    seed-keyed, so it is derived from the same config the run under test uses.
     """
-    inst_dir = (
-        fract_base
-        / f"var{config.variance_threshold}"
-        / "instances"
-        / f"np{config.point_num}"
-        / "000000"
-    )
+    inst_dir = Path(layout.instance_dir(config)) / "000000"
     inst_dir.mkdir(parents=True, exist_ok=True)
     if present:
         rng = np.random.default_rng(0)
@@ -489,7 +963,7 @@ def _instance_config(fract_base: Path) -> Namespace:
 
 def _seed_ifs_params(fract_base: Path, config: Namespace, n_categories: int) -> None:
     """Write a contractive IFS param CSV per category so generation stays fast."""
-    param_dir = fract_base / f"var{config.variance_threshold}" / "3DIFS_param"
+    param_dir = Path(layout.category_param_dir(config))
     param_dir.mkdir(parents=True, exist_ok=True)
     params = np.zeros((2, 13), dtype=np.float64)
     params[:, 0] = params[:, 4] = params[:, 8] = 0.5
@@ -555,14 +1029,7 @@ def no_scan(*_args, **_kwargs):
     assert rc == 0
 
     # Rank 1 received the broadcast list and generated its share (pair [1, 0]).
-    generated = (
-        fract_base
-        / f"var{config.variance_threshold}"
-        / "instances"
-        / f"np{config.point_num}"
-        / "000001"
-        / "000001_0000.npy"
-    )
+    generated = Path(layout.instance_dir(config)) / "000001" / "000001_0000.npy"
     assert generated.exists()
 
 
diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py
new file mode 100644
index 0000000..d907c39
--- /dev/null
+++ b/tests/datagen/test_provenance.py
@@ -0,0 +1,116 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Dataset provenance: the commit stamped on a dataset is ScaFFold's (R35).
+
+``meta.yaml``'s ``code_commit``, the published ``<timestamp>__<commit>``
+directory name, and the ``dataset_reuse_enforce_commit_id`` gate all key off
+one string. Reading it from the *launch* directory made it a property of
+wherever the job happened to start -- a site workflow repo, a scratch
+directory -- instead of the code that generated the data. Reuse was then gated
+on an unrelated repo's churn while real ScaFFold changes went undetected.
+"""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+from pathlib import Path
+
+from ScaFFold.datagen import get_dataset as gd
+
+LOG = logging.getLogger("test_provenance")
+
+# The ScaFFold source tree: the checkout whose commit must be stamped.
+PACKAGE_DIR = Path(gd.__file__).resolve().parent
+
+
+def _head_of(repo: Path) -> str:
+    return (
+        subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=repo)
+        .decode()
+        .strip()
+    )
+
+
+def _make_repo(path: Path) -> str:
+    """Create a throwaway git repo with one commit; return its short HEAD."""
+    path.mkdir(parents=True, exist_ok=True)
+    subprocess.run(["git", "init", "-q"], cwd=path, check=True)
+    (path / "README").write_text("an unrelated project\n")
+    subprocess.run(["git", "add", "README"], cwd=path, check=True)
+    subprocess.run(
+        [
+            "git",
+            "-c",
+            "user.name=test",
+            "-c",
+            "user.email=test@example.invalid",
+            "-c",
+            "commit.gpgsign=false",
+            "commit",
+            "-qm",
+            "init",
+        ],
+        cwd=path,
+        check=True,
+    )
+    return _head_of(path)
+
+
+def test_commit_is_read_from_the_scaffold_tree_not_the_cwd(tmp_path, monkeypatch):
+    """Running from an unrelated repo still stamps ScaFFold's commit."""
+    expected = _head_of(PACKAGE_DIR)
+
+    other = tmp_path / "workflow-repo"
+    other_head = _make_repo(other)
+    assert other_head != expected, "the throwaway repo must differ from ScaFFold"
+
+    monkeypatch.chdir(other)
+    assert gd._git_commit_short(LOG) == expected
+
+
+def test_commit_survives_a_non_repo_working_directory(tmp_path, monkeypatch):
+    """A scratch launch directory does not degrade provenance to no-commit-id."""
+    expected = _head_of(PACKAGE_DIR)
+
+    scratch = tmp_path / "scratch-cwd"
+    scratch.mkdir()
+    monkeypatch.chdir(scratch)
+
+    assert gd._git_commit_short(LOG) == expected
+
+
+def test_non_repo_install_reports_no_commit_id(tmp_path, monkeypatch):
+    """An installed (non-git) ScaFFold still degrades gracefully.
+
+    Provenance is best-effort: when the source tree is not a checkout there is
+    no commit to record, and reuse simply is not gated on one.
+
+    "Not a checkout" has to be made true of the *whole path*, not just the leaf:
+    git walks upwards until it finds a repository, so with ``--basetemp`` inside
+    a ScaFFold checkout this directory inherits that checkout's HEAD and the
+    test fails on where it was run rather than on what it tests. The ceiling
+    stops the walk at ``tmp_path``.
+    """
+    monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path))
+    not_a_repo = tmp_path / "site-packages" / "ScaFFold" / "datagen"
+    not_a_repo.mkdir(parents=True)
+
+    assert gd._git_commit_short(LOG, source_dir=not_a_repo) == "no-commit-id"
+
+
+def test_missing_source_dir_reports_no_commit_id(tmp_path):
+    """A source directory that does not exist is handled, not raised."""
+    assert gd._git_commit_short(LOG, source_dir=tmp_path / "gone") == "no-commit-id"
diff --git a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py
index 3799874..bc7e7a4 100644
--- a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py
+++ b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py
@@ -45,7 +45,7 @@
 from mpi4py import MPI
 
 import ScaFFold.datagen.get_dataset as gd
-from ScaFFold.datagen import volumegen
+from ScaFFold.datagen import layout, volumegen
 
 VT = 0.15
 PN = 64
@@ -71,14 +71,10 @@ def _config(dataset_dir: Path, fract_base: Path) -> Namespace:
 
 
 def _instance_path(fract_base: Path, cat: int, inst: int) -> Path:
-    return (
-        fract_base
-        / f"var{VT}"
-        / "instances"
-        / f"np{PN}"
-        / f"{cat:06d}"
-        / f"{cat:06d}_{inst:04d}.npy"
-    )
+    # The instance library is keyed by seed; derive the path from the same
+    # config the run under test uses.
+    inst_root = Path(layout.instance_dir(_config(Path("unused"), fract_base)))
+    return inst_root / f"{cat:06d}" / f"{cat:06d}_{inst:04d}.npy"
 
 
 def _seed_instances(fract_base: Path) -> None:
diff --git a/tests/helpers/rank_scripts/datagen_instance_partition.py b/tests/helpers/rank_scripts/datagen_instance_partition.py
index 10df10f..37d0c98 100644
--- a/tests/helpers/rank_scripts/datagen_instance_partition.py
+++ b/tests/helpers/rank_scripts/datagen_instance_partition.py
@@ -44,6 +44,7 @@
 from mpi4py import MPI
 
 import ScaFFold.datagen.instance as inst
+from ScaFFold.datagen import layout
 
 VT = 0.15
 PN = 64
@@ -62,7 +63,7 @@ def _config(fract_base: Path) -> Namespace:
 
 
 def _seed_ifs_params(fract_base: Path) -> None:
-    param_dir = fract_base / f"var{VT}" / "3DIFS_param"
+    param_dir = Path(layout.category_param_dir(_config(fract_base)))
     param_dir.mkdir(parents=True, exist_ok=True)
     params = np.zeros((2, 13), dtype=np.float64)
     params[:, 0] = params[:, 4] = params[:, 8] = 0.5
@@ -82,7 +83,7 @@ def main() -> None:
         _seed_ifs_params(fract_base)
     comm.Barrier()
 
-    inst_root = fract_base / f"var{VT}" / "instances" / f"np{PN}"
+    inst_root = Path(layout.instance_dir(_config(fract_base)))
 
     # Give each rank a divergent view of pre-existing instances. Only rank 0's
     # view should matter after the fix (its list is broadcast); rank 1's phantom
diff --git a/tests/helpers/rank_scripts/groupnorm_shards_2rank.py b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py
new file mode 100644
index 0000000..c347365
--- /dev/null
+++ b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py
@@ -0,0 +1,123 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Two-rank check that FastGroupNorm's DCTensor route is shard-count agnostic.
+
+Run under ``torchrun --nproc_per_node=2`` with the gloo backend (see
+``tests/test_groupnorm.py``).  Every other DCTensor test uses
+``num_shards=(1, 1, 1)``, where sharding is a no-op and the local shard is the
+whole tensor; this one shards a spatial dim across two ranks so the claim the
+fast path actually rests on -- GroupNorm statistics are per-shard, and the
+route does not add communication or change which elements are reduced together
+-- is exercised where it can fail.
+
+Each rank prints one ``RESULT ...`` line plus ``DONE``; the parent asserts.
+"""
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+sys.path.insert(0, os.environ.get("SCAFFOLD_ROOT", "/usr/WS1/dryden1/ScaFFold"))
+
+from ScaFFold.unet import group_norm as gn_mod  # noqa: E402
+from ScaFFold.unet.group_norm import FastGroupNorm  # noqa: E402
+
+GROUPS = 8
+CHANNELS = 16
+SIZE = 8  # dim 2 is split into two shards of 4
+
+
+def run():
+    dist.init_process_group(backend="gloo")
+    rank = dist.get_rank()
+
+    import distconv
+
+    ps = distconv.ParallelStrategy(num_shards=(2,), shard_dim=(2,), device_type="cpu")
+
+    # The same global volume on both ranks; each takes its own slab.
+    generator = torch.Generator().manual_seed(41)
+    volume = torch.randn(1, CHANNELS, SIZE, SIZE, SIZE, generator=generator)
+    half = SIZE // 2
+    local = volume.narrow(2, rank * half, half).contiguous()
+
+    norm = FastGroupNorm(GROUPS, CHANNELS)
+    param_generator = torch.Generator().manual_seed(97)
+    with torch.no_grad():
+        norm.weight.normal_(1.0, 0.1, generator=param_generator)
+        norm.bias.normal_(0.0, 0.1, generator=param_generator)
+
+    def forward(compiled):
+        """Run the wrapped GroupNorm with the compiled route on or off.
+
+        The compiled route is forced on CPU by standing in the stock functional
+        kernel for the compiled callable: what is under test here is the
+        unwrap/rewrap plumbing at shard counts > 1, not Inductor.
+        """
+        original_use, original_get = (
+            gn_mod._use_compiled,
+            gn_mod._get_compiled_group_norm,
+        )
+        if compiled:
+            gn_mod._use_compiled = lambda t, **kw: type(t) is torch.Tensor
+            gn_mod._get_compiled_group_norm = lambda: F.group_norm
+        else:
+            gn_mod._use_compiled = lambda t, **kw: False
+        try:
+            norm.zero_grad(set_to_none=True)
+            x = local.clone().requires_grad_(True)
+            out = norm(distconv.DCTensor.from_shard(x, ps))
+            assert isinstance(out, distconv.DCTensor), type(out)
+            distconv.distconv._ToTensor.apply(out).pow(2).sum().backward()
+            weight_grad = norm.weight.grad
+            if isinstance(weight_grad, distconv.DCTensor):
+                weight_grad = weight_grad._tensor
+            return out._tensor.detach().clone(), x.grad.clone(), weight_grad.clone()
+        finally:
+            gn_mod._use_compiled = original_use
+            gn_mod._get_compiled_group_norm = original_get
+
+    eager = forward(compiled=False)
+    compiled = forward(compiled=True)
+    identical = all(torch.equal(a, b) for a, b in zip(eager, compiled))
+
+    # Per-shard statistics: this rank's output normalizes its own slab only.
+    with torch.no_grad():
+        per_shard = F.group_norm(local, GROUPS, norm.weight, norm.bias, norm.eps)
+        global_slice = F.group_norm(
+            volume, GROUPS, norm.weight, norm.bias, norm.eps
+        ).narrow(2, rank * half, half)
+
+    # No spaces in any field: torchrun interleaves the ranks' stdout and a
+    # line can arrive without its trailing newline, so the parent's regex has
+    # to be able to tell two RESULT lines apart when they run together.
+    shape = "x".join(str(dim) for dim in compiled[0].shape)
+    print(
+        f"RESULT rank={rank} shape={shape} "
+        f"identical={identical} "
+        f"per_shard={torch.equal(compiled[0], per_shard)} "
+        f"global={torch.allclose(compiled[0], global_slice, atol=1e-6)}",
+        flush=True,
+    )
+    print("DONE", flush=True)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    run()
diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py
index 4e1737d..5d7afa7 100644
--- a/tests/test_checkpointing.py
+++ b/tests/test_checkpointing.py
@@ -26,8 +26,12 @@
 
 from __future__ import annotations
 
+import math
+import os
 import re
+import textwrap
 import time
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import pytest
@@ -35,7 +39,7 @@
 import torch.distributed as dist
 
 import ScaFFold.utils.trainer as trainer_mod
-from ScaFFold.utils.checkpointing import CheckpointManager
+from ScaFFold.utils.checkpointing import CheckpointManager, CheckpointSaveError
 from ScaFFold.utils.trainer import PyTorchTrainer
 from tests.helpers import mpi_runner
 
@@ -186,6 +190,185 @@ def always_raise(obj, f, *args, **kwargs):
         mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)
 
 
+# ---------------------------------------------------------------------------
+# R04 -- async save failures are reported, and the run's LAST save is consumed
+# ---------------------------------------------------------------------------
+
+
+def _failing_torch_save(obj, f, *args, **kwargs):
+    raise RuntimeError("writer boom")
+
+
+def test_async_save_failure_surfaces_at_next_save(tmp_path, monkeypatch):
+    """A background write that failed is reported at the next save, not dropped.
+
+    In async mode ``save_checkpoint`` returns as soon as the CPU snapshot is
+    handed to the writer thread, so the failure can only be observed later.
+    Swallowing it leaves ``last_saved_epoch``/``best_val_loss`` claiming a
+    checkpoint that does not exist on disk.
+    """
+    mgr, _ = _make_manager(tmp_path, async_save=True)
+    monkeypatch.setattr(torch, "save", _failing_torch_save)
+
+    mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)  # write fails in background
+
+    with pytest.raises(RuntimeError, match="writer boom"):
+        mgr.save_checkpoint(epoch=2, val_loss_avg=0.4)
+
+    assert not mgr.last_ckpt_path.exists()
+
+
+def test_async_save_failure_surfaces_at_finalize(tmp_path, monkeypatch):
+    """The run's final save has its outcome consumed before the run ends."""
+    mgr, _ = _make_manager(tmp_path, async_save=True)
+    monkeypatch.setattr(torch, "save", _failing_torch_save)
+
+    mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)
+
+    with pytest.raises(RuntimeError, match="writer boom"):
+        mgr.finalize_saves()
+
+    assert not mgr.last_ckpt_path.exists()
+
+
+def test_final_async_save_failure_fails_the_run(tiny_trainer, monkeypatch):
+    """``train()`` must not return successfully after a failed final save.
+
+    Nothing downstream of the training loop touches the checkpoint manager, so
+    without an explicit wait the child process exits 0 with no checkpoint at
+    all: the benchmark reports success and a later ``--restart`` resumes from a
+    stale epoch or fails its pre-check.
+    """
+    trainer = tiny_trainer(
+        config_overrides={
+            "checkpoint_interval": 1,
+            "epochs": 1,
+            "target_dice": 0.95,
+        }
+    )
+    # The config has no async knob at this scale; drive the manager directly.
+    mgr = trainer.checkpoint_manager
+    mgr.async_save = True
+    mgr.executor = ThreadPoolExecutor(max_workers=1)
+
+    monkeypatch.setattr(
+        trainer,
+        "_run_training_batch",
+        lambda batch, **kw: (1, torch.tensor(0.3), torch.tensor(0.5)),
+    )
+    monkeypatch.setattr(
+        trainer_mod, "evaluate", lambda *a, **k: (0.5 * 2, 0.4 * 2, 0.4, 2, 2)
+    )
+    monkeypatch.setattr(torch, "save", _failing_torch_save)
+
+    trainer.cleanup_or_resume()
+    try:
+        with pytest.raises(RuntimeError, match="writer boom"):
+            trainer.train()
+    finally:
+        mgr.executor.shutdown(wait=True)
+
+    assert not mgr.last_ckpt_path.exists()
+
+
+# ---------------------------------------------------------------------------
+# R05 -- a rank-0 write failure fails every rank with the same error
+# ---------------------------------------------------------------------------
+
+# Two-rank script: rank 0's torch.save fails. Both ranks must come out of
+# save_checkpoint with the SAME real error rather than rank 0 raising the disk
+# error while its peers die (gloo) or stall (NCCL) in an unmatched collective.
+# Kept inline rather than in tests/helpers/rank_scripts/ because it is only
+# meaningful together with the assertions below.
+SAVE_FAIL_RANK_SCRIPT = textwrap.dedent(
+    '''\
+    """Two-rank save-failure rank script (gloo, CPU)."""
+
+    import os
+    import sys
+
+    import torch
+    import torch.distributed as dist
+
+    from ScaFFold.utils.checkpointing import CheckpointManager
+
+    rank = int(os.environ["RANK"])
+    dist.init_process_group(backend="gloo")
+
+    torch.manual_seed(0)
+    mgr = CheckpointManager(
+        model=torch.nn.Linear(64, 64),
+        base_dir=os.environ["CKPT_DIR"],
+        world_rank=rank,
+        dist_enabled=True,
+    )
+
+    if rank == 0:
+        def _boom(obj, f, *args, **kwargs):
+            raise RuntimeError("injected disk failure on rank 0")
+
+        torch.save = _boom
+
+    # Markers are delimited (trailing '.', '<<...>>') because two ranks writing
+    # the same pipe can interleave without a newline between them.
+    try:
+        mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)
+    except BaseException as exc:  # noqa: BLE001 - the point is what we caught
+        print(f"RANK {rank} RAISED {type(exc).__name__}.", flush=True)
+        message = str(exc).replace(chr(10), " ")
+        print(f"RANK {rank} MESSAGE <<{message}>>", flush=True)
+    else:
+        print(f"RANK {rank} NO_RAISE.", flush=True)
+
+    print(f"RANK {rank} DONE.", flush=True)
+    sys.stdout.flush()
+    try:
+        dist.destroy_process_group()
+    except Exception:
+        pass
+    '''
+)
+
+
+@_requires_gloo
+def test_rank0_save_failure_fails_all_ranks(tmp_path):
+    """Under 2 gloo ranks, a rank-0 disk error reaches the peer as itself.
+
+    Rank 0 is the only writer, so raising its error before the broadcast the
+    peers are already waiting in leaves them in an unmatched collective: gloo
+    reports an opaque "Connection closed by peer" and NCCL stalls until the
+    watchdog timeout, in both cases hiding the disk error that actually
+    happened. The write's OUTCOME must be broadcast instead, so both ranks
+    raise the same error together.
+    """
+    ckpt_dir = tmp_path / "checkpoints"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    script = tmp_path / "save_fail_2rank.py"
+    script.write_text(SAVE_FAIL_RANK_SCRIPT)
+
+    rc, out, err = mpi_runner.torchrun_gloo(
+        str(script), n=2, timeout=90, env={"CKPT_DIR": str(ckpt_dir)}
+    )
+
+    done = set(re.findall(r"RANK (\d+) DONE\.", out))
+    assert rc == 0 and {"0", "1"} <= done, (
+        f"expected both ranks to fail cleanly and finish, rc={rc}\n"
+        f"stdout:\n{out}\nstderr:\n{err[-3000:]}"
+    )
+
+    raised = dict(re.findall(r"RANK (\d+) RAISED (\w+)\.", out))
+    messages = dict(re.findall(r"RANK (\d+) MESSAGE <<(.*?)>>", out))
+    assert set(raised) == {"0", "1"}, f"both ranks must raise\nstdout:\n{out}"
+    assert raised["0"] == raised["1"], (
+        f"ranks raised different error types: {raised}\nstdout:\n{out}"
+    )
+    for rank in ("0", "1"):
+        assert "injected disk failure on rank 0" in messages.get(rank, ""), (
+            f"rank {rank} did not see the real disk error: "
+            f"{messages.get(rank)!r}\nstdout:\n{out}"
+        )
+
+
 # ---------------------------------------------------------------------------
 # F41 -- race-free best decision (cached best loss, no per-save probe)
 # ---------------------------------------------------------------------------
@@ -244,6 +427,264 @@ def spy_load(path, *args, **kwargs):
     assert final_best["val_loss_avg"] == pytest.approx(0.2)
 
 
+# ---------------------------------------------------------------------------
+# R02 -- a from-scratch cleanup drops the deleted run's best, not just its files
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_from_scratch_resets_best(tmp_path):
+    """``cleanup(train_from_scratch=True)`` resets the cached best-loss state.
+
+    The manager seeds ``best_val_loss`` from ``checkpoint_best.pth`` at
+    construction so a resumed run does not call its first epoch "best". When
+    the same directory is then wiped for a fresh run, that cached score
+    outlives the file it came from: every ``is_best`` decision of the new run
+    is gated by a deleted run's score, so no ``checkpoint_best.pth`` is written
+    until the retrain beats it -- leaving the run with no best-checkpoint
+    fallback at all.
+    """
+    mgr, _ = _make_manager(tmp_path)
+    mgr.save_checkpoint(epoch=1, val_loss_avg=0.01)
+    assert mgr.best_ckpt_path.exists()
+
+    # A driver reusing the run directory: the new manager seeds from disk.
+    mgr2, _ = _make_manager(tmp_path)
+    assert mgr2.best_val_loss == pytest.approx(0.01)
+    mgr2.save_checkpoint(epoch=2, val_loss_avg=0.9)
+    assert mgr2.last_saved_epoch == 2
+
+    mgr2.cleanup(train_from_scratch=True)
+
+    assert not mgr2.last_ckpt_path.exists()
+    assert not mgr2.best_ckpt_path.exists()
+    assert mgr2.best_val_loss == math.inf
+    assert mgr2.last_saved_epoch is None
+
+    # The fresh run's first epoch is its best, and a best checkpoint exists.
+    assert mgr2.save_checkpoint(epoch=1, val_loss_avg=0.5) is True
+    assert mgr2.best_ckpt_path.exists()
+
+
+# ---------------------------------------------------------------------------
+# R07 -- checkpoint debris (.tmp.<pid>, .corrupt) does not accumulate
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_from_scratch_removes_stale_debris(tmp_path):
+    """A from-scratch cleanup clears checkpoint debris, not just the two files.
+
+    ``_atomic_save``'s temp file survives a process kill (its unlink only runs
+    on a Python-level exception) and ``_quarantine_corrupt``'s ``.corrupt``
+    rename is never undone. Both are full-checkpoint-sized, so a cleanup that
+    claims to have cleared the checkpoints while leaving them behind keeps
+    multi-GB files on the shared filesystem.
+    """
+    mgr, _ = _make_manager(tmp_path)
+    mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)
+
+    stale_tmp = tmp_path / "checkpoint_last.pth.tmp.999999"
+    stale_tmp.write_bytes(b"partial checkpoint")
+    quarantined = tmp_path / "checkpoint_best.pth.corrupt"
+    quarantined.write_bytes(b"truncated checkpoint")
+
+    mgr.cleanup(train_from_scratch=True)
+
+    assert not stale_tmp.exists()
+    assert not quarantined.exists()
+    assert list(tmp_path.iterdir()) == []
+
+
+def test_init_sweeps_orphaned_tmp_files(tmp_path):
+    """Constructing a manager sweeps temp files left by killed writes.
+
+    Repeated walltime kills of a long run take the *resume* path, never the
+    from-scratch cleanup, so without this sweep one stranded temp file per
+    killed pid piles up in the run directory. A temp file belonging to this
+    process is left alone (another manager here may still be writing it), and
+    the quarantined ``.corrupt`` file is kept: at most two can ever exist and
+    they are the only evidence of what a resume discarded.
+    """
+    orphan = tmp_path / "checkpoint_last.pth.tmp.999999"
+    orphan.write_bytes(b"partial checkpoint")
+    own = tmp_path / f"checkpoint_last.pth.tmp.{os.getpid()}"
+    own.write_bytes(b"possibly in flight")
+    quarantined = tmp_path / "checkpoint_last.pth.corrupt"
+    quarantined.write_bytes(b"truncated checkpoint")
+
+    _make_manager(tmp_path)
+
+    assert not orphan.exists()
+    assert own.exists()
+    assert quarantined.exists()
+
+
+# ---------------------------------------------------------------------------
+# R08 -- an improving epoch serializes the state dict once, not twice
+# ---------------------------------------------------------------------------
+
+
+def _count_torch_saves(monkeypatch):
+    """Count ``torch.save`` calls made by the checkpoint writer."""
+    calls = []
+    real_save = torch.save
+
+    def counting_save(obj, f, *args, **kwargs):
+        calls.append(str(getattr(f, "name", f)))
+        return real_save(obj, f, *args, **kwargs)
+
+    monkeypatch.setattr(torch, "save", counting_save)
+    return calls
+
+
+def test_best_checkpoint_copied_not_reserialized(tmp_path, monkeypatch):
+    """The best checkpoint reuses the bytes just written to 'last'.
+
+    An improving epoch used to pickle *and fsync* the identical state dict a
+    second time, doubling both the serialization CPU and the checkpoint bytes
+    pushed at the shared filesystem. Copying the file that was just committed
+    costs one read (usually from page cache) and one write instead.
+    """
+    mgr, model = _make_manager(tmp_path)
+    saves = _count_torch_saves(monkeypatch)
+
+    assert mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) is True
+    assert len(saves) == 1
+    assert saves == [str(mgr.last_ckpt_path) + f".tmp.{os.getpid()}"]
+
+    # A second improving epoch: still exactly one serialization.
+    saves.clear()
+    assert mgr.save_checkpoint(epoch=2, val_loss_avg=0.25) is True
+    assert len(saves) == 1
+
+    # A non-improving epoch writes 'last' only and leaves 'best' alone.
+    best_bytes = mgr.best_ckpt_path.read_bytes()
+    saves.clear()
+    assert mgr.save_checkpoint(epoch=3, val_loss_avg=0.9) is False
+    assert len(saves) == 1
+    assert mgr.best_ckpt_path.read_bytes() == best_bytes
+
+    # And the best checkpoint is still a real, loadable checkpoint holding the
+    # epoch-2 state -- byte-identical to the 'last' file it was copied from.
+    monkeypatch.undo()
+    best = torch.load(mgr.best_ckpt_path, map_location="cpu", weights_only=False)
+    assert best["epoch"] == 2
+    assert best["val_loss_avg"] == pytest.approx(0.25)
+    for name, tensor in model.state_dict().items():
+        assert torch.equal(best["model_state_dict"][name], tensor)
+
+
+def test_best_copy_failure_is_not_silent(tmp_path, monkeypatch):
+    """A failed best-checkpoint copy surfaces and leaves no partial file."""
+    mgr, _ = _make_manager(tmp_path)
+
+    def boom(*args, **kwargs):
+        raise OSError("[Errno 28] No space left on device")
+
+    monkeypatch.setattr("ScaFFold.utils.checkpointing.shutil.copyfileobj", boom)
+
+    with pytest.raises(CheckpointSaveError, match="No space left on device"):
+        mgr.save_checkpoint(epoch=1, val_loss_avg=0.5)
+
+    assert not mgr.best_ckpt_path.exists()
+    assert list(tmp_path.glob("checkpoint_*.tmp.*")) == []
+
+
+# ---------------------------------------------------------------------------
+# VA-1/VA-2/VA-3 -- the remaining rank-0 filesystem windows are fenced
+#
+# Rank 0 is the only rank that touches the run directory, and it does so while
+# its peers are already committed to the next collective. Every such window must
+# therefore report its failure *through* that collective; raising inside it
+# leaves the peers in an unmatched collective, where a plain disk error
+# resurfaces as an opaque gloo transport error or an NCCL watchdog timeout.
+# ---------------------------------------------------------------------------
+
+
+def _record_collectives(monkeypatch):
+    """Stub the process-group collectives, recording what a rank posts."""
+    posted = {"broadcasts": [], "barriers": 0}
+
+    def fake_broadcast(objs, src=0):
+        posted["broadcasts"].append(objs[0])
+
+    def fake_barrier(*args, **kwargs):
+        posted["barriers"] += 1
+
+    monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast)
+    monkeypatch.setattr(dist, "barrier", fake_barrier)
+    return posted
+
+
+def _raise_stale(*args, **kwargs):
+    raise OSError("[Errno 116] Stale file handle")
+
+
+def test_cleanup_rank0_fs_error_travels_through_the_broadcast(tmp_path, monkeypatch):
+    """A failure while clearing checkpoints fails every rank, not just rank 0.
+
+    ``_remove_checkpoint_files`` globs and stats the run directory; on a shared
+    filesystem those can raise (ESTALE, EACCES) even though each individual
+    unlink is already tolerated. That happens between the drain and the
+    broadcast, so an unfenced raise re-creates exactly the hazard R05 closed.
+    """
+    mgr, _ = _make_manager(tmp_path)
+    mgr.dist_enabled = True
+    posted = _record_collectives(monkeypatch)
+    monkeypatch.setattr(CheckpointManager, "_remove_checkpoint_files", _raise_stale)
+
+    with pytest.raises(CheckpointSaveError) as excinfo:
+        mgr.cleanup(train_from_scratch=True)
+
+    assert "Stale file handle" in str(excinfo.value)
+    # The peers' collectives were posted before the raise, so they fail with the
+    # same error instead of hanging.
+    assert len(posted["broadcasts"]) == 1
+    assert "Stale file handle" in str(posted["broadcasts"][0])
+    assert posted["barriers"] == 1
+
+
+def test_load_rank0_selection_error_travels_through_the_broadcast(
+    tmp_path, monkeypatch
+):
+    """A failure inside ``_select_and_load`` is a broadcast decision too.
+
+    The load path folded only the *drain* error into its decision broadcast, so
+    rank 0 stat-ing or renaming a checkpoint candidate could still die before
+    the broadcast the peers were already waiting in.
+    """
+    mgr, _ = _make_manager(tmp_path)
+    mgr.dist_enabled = True
+    posted = _record_collectives(monkeypatch)
+    monkeypatch.setattr(CheckpointManager, "_select_and_load", _raise_stale)
+
+    with pytest.raises(CheckpointSaveError) as excinfo:
+        mgr.load_from_checkpoint()
+
+    assert "Stale file handle" in str(excinfo.value)
+    assert len(posted["broadcasts"]) == 1
+    assert posted["broadcasts"][0][0] == "load_failed"
+
+
+def test_init_survives_an_unlistable_run_dir(tmp_path, monkeypatch, capsys):
+    """A run directory that cannot be listed warns instead of killing __init__.
+
+    The orphan sweep runs on rank 0 only and outside any collective, so a raise
+    here aborts rank 0 alone and strands the peers at the manager's first
+    collective (``cleanup``). ``pathlib`` already swallows PermissionError, but
+    a stale NFS/Lustre handle propagates; the sweep is opportunistic, so it must
+    degrade to a warning.
+    """
+    base = tmp_path / "checkpoints"
+    base.mkdir()
+    (base / "checkpoint_last.pth.tmp.999999").write_bytes(b"partial checkpoint")
+    monkeypatch.setattr(Path, "glob", _raise_stale)
+
+    mgr, _ = _make_manager(base)
+
+    assert mgr.base_dir == base
+    assert "sweep" in capsys.readouterr().out.lower()
+
+
 # ---------------------------------------------------------------------------
 # F71 -- CPU tensors are cloned into the snapshot
 # ---------------------------------------------------------------------------
@@ -264,6 +705,115 @@ def test_cpu_tensors_cloned(tmp_path):
     assert torch.equal(snapshot["nested"][0], torch.ones(3))
 
 
+# ---------------------------------------------------------------------------
+# R09 -- the warmup snapshot is host-resident, not a device-side copy
+# ---------------------------------------------------------------------------
+
+
+def _all_tensors(obj):
+    """Yield every tensor reachable from a snapshot payload."""
+    if torch.is_tensor(obj):
+        yield obj
+    elif isinstance(obj, dict):
+        for value in obj.values():
+            yield from _all_tensors(value)
+    elif isinstance(obj, (list, tuple)):
+        for value in obj:
+            yield from _all_tensors(value)
+
+
+class _StubOptimizer:
+    """Stand-in exposing only the state_dict the snapshot reads."""
+
+    def __init__(self, state):
+        self._state = state
+
+    def state_dict(self):
+        return self._state
+
+
+def test_snapshot_holds_no_device_resident_tensors(tmp_path):
+    """``snapshot_training_state`` copies to the host, not device-to-device.
+
+    The snapshot used to clone model and optimizer state on whatever device
+    they lived on, pinning ~3x parameter bytes of accelerator memory (model
+    clone + Adam's two moment buffers) from before the first warmup batch
+    until the restore -- i.e. straight across warmup's own peak, which is
+    where a memory-marginal configuration OOMs. Device tensors are simulated
+    with FakeTensorMode so this runs without a GPU.
+    """
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    mgr, _ = _make_manager(tmp_path)
+
+    with FakeTensorMode():
+        mgr.model = torch.nn.Linear(8, 4, device="cuda")
+        mgr.optimizer = _StubOptimizer(
+            {
+                "state": {
+                    0: {
+                        "step": torch.zeros(1, device="cuda"),
+                        "exp_avg": torch.zeros(4, 8, device="cuda"),
+                        "exp_avg_sq": torch.zeros(4, 8, device="cuda"),
+                    }
+                },
+                "param_groups": [{"params": [0], "lr": 0.1}],
+            }
+        )
+        snapshot = mgr.snapshot_training_state()
+
+        devices = {t.device.type for t in _all_tensors(snapshot)}
+
+    assert devices == {"cpu"}, f"snapshot kept device-resident tensors: {devices}"
+
+
+def test_snapshot_restore_round_trips_values_and_devices(tmp_path):
+    """A host-resident snapshot still restores exact state on its own device.
+
+    ``load_state_dict`` copies into the live parameters and moves optimizer
+    state back to each parameter's device, so nothing downstream has to know
+    the snapshot took a detour through the host.
+    """
+    mgr, model = _make_manager(tmp_path)
+    optimizer = mgr.optimizer
+
+    # One applied step so the optimizer carries real per-parameter state.
+    model(torch.randn(2, 8)).sum().backward()
+    optimizer.step()
+    optimizer.zero_grad(set_to_none=True)
+
+    reference_params = {k: v.detach().clone() for k, v in model.state_dict().items()}
+    reference_state = {
+        pid: {k: v.detach().clone() for k, v in state.items() if torch.is_tensor(v)}
+        for pid, state in optimizer.state_dict()["state"].items()
+    }
+    param_devices = {k: v.device for k, v in model.state_dict().items()}
+
+    snapshot = mgr.snapshot_training_state()
+
+    # Warmup-shaped mutation: more steps, then roll back.
+    for _ in range(3):
+        model(torch.randn(2, 8)).sum().backward()
+        optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
+    assert not torch.equal(model.state_dict()["weight"], reference_params["weight"])
+
+    mgr.restore_training_state(snapshot)
+
+    for name, tensor in model.state_dict().items():
+        assert torch.equal(tensor, reference_params[name]), name
+        assert tensor.device == param_devices[name], name
+    restored_state = optimizer.state_dict()["state"]
+    for pid, state in reference_state.items():
+        for key, value in state.items():
+            assert torch.equal(restored_state[pid][key], value), (pid, key)
+    for group, param in zip(optimizer.param_groups, model.parameters()):
+        del group
+        for value in optimizer.state[param].values():
+            if torch.is_tensor(value) and value.dim() > 0:
+                assert value.device == param.device
+
+
 # ---------------------------------------------------------------------------
 # F50 -- a final checkpoint is written when the run exits between intervals
 # ---------------------------------------------------------------------------
@@ -307,6 +857,54 @@ def fake_evaluate(*args, **kwargs):
     assert saved["epoch"] == 2
 
 
+# ---------------------------------------------------------------------------
+# R03 -- a diverged epoch aborts instead of checkpointing NaN weights
+# ---------------------------------------------------------------------------
+
+
+def test_divergence_aborts_before_poisoning_checkpoint(tiny_trainer, monkeypatch):
+    """Non-finite epoch losses abort the run before any checkpoint is written.
+
+    The dice score is computed from a hard argmax, so an all-NaN model still
+    produces a *finite* dice (argmax of NaN logits is 0 and the one-hots are
+    finite): the dice check can never fire on divergence. Left unguarded, a
+    diverged ``epochs: -1`` run keeps looping on a finite plateau below target
+    while every checkpoint interval overwrites ``checkpoint_last.pth`` with NaN
+    weights, poisoning the next ``--restart``. The reduced losses are the
+    values that actually go non-finite, and being reductions they are identical
+    on every rank, so the check fires on all ranks together.
+    """
+    trainer = tiny_trainer(
+        config_overrides={
+            "checkpoint_interval": 1,
+            "epochs": 3,
+            "target_dice": 0.95,
+        }
+    )
+
+    # A diverged step: NaN loss, and a dice that stays finite.
+    monkeypatch.setattr(
+        trainer,
+        "_run_training_batch",
+        lambda batch, **kw: (1, torch.tensor(float("nan")), torch.tensor(0.0)),
+    )
+
+    def diverged_evaluate(*args, **kwargs):
+        # Exactly what evaluate() returns for an all-NaN model: a tiny but
+        # finite hard-argmax dice sum alongside a NaN validation loss.
+        return (7.4e-10, float("nan"), float("nan"), 2, 2)
+
+    monkeypatch.setattr(trainer_mod, "evaluate", diverged_evaluate)
+
+    trainer.cleanup_or_resume()
+    with pytest.raises(ValueError, match="[Nn]on-finite"):
+        trainer.train()
+
+    # The run died before the diverged epoch could be checkpointed.
+    assert not trainer.checkpoint_manager.last_ckpt_path.exists()
+    assert not trainer.checkpoint_manager.best_ckpt_path.exists()
+
+
 # ---------------------------------------------------------------------------
 # F49 -- GradScaler-skipped steps do not advance the optimizer-step counter
 # ---------------------------------------------------------------------------
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..f7d83ca
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,658 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Tests for the ``scaffold`` CLI entry point (``ScaFFold.cli.main``).
+
+``cli.main`` makes every job-wide decision on MPI rank 0 and broadcasts it, so
+these tests drive it with a *fake* communicator: the real ``MPI.COMM_WORLD`` in
+this environment is always a one-rank singleton, which cannot express the
+multi-rank shapes the CLI must get right (rank-0-decides-and-broadcasts, and
+the mismatch between the MPI world and the launcher's environment).
+
+``_FakeComm`` records what rank 0 broadcasts and, for a non-zero rank, replays
+a scripted sequence of values as if rank 0 had sent them. That makes it
+possible to assert that a non-zero rank *uses the broadcast decision* instead
+of consulting its own view of the filesystem.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+import yaml
+
+import ScaFFold.cli as cli
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = REPO_ROOT / "ScaFFold" / "configs" / "benchmark_default.yml"
+
+
+# ---------------------------------------------------------------------------
+# harness
+# ---------------------------------------------------------------------------
+
+
+class _FakeComm:
+    """A stand-in for ``MPI.COMM_WORLD`` with a settable rank and size.
+
+    ``bcast`` returns the caller's object on rank 0 (recording it); on a
+    non-zero rank it pops the next scripted value from ``bcast_returns``,
+    falling back to the caller's object when the script is exhausted.
+    """
+
+    def __init__(self, rank: int = 0, size: int = 1, bcast_returns=None):
+        self._rank = rank
+        self._size = size
+        self._scripted = list(bcast_returns or [])
+        self.broadcast = []
+        self.barriers = 0
+
+    def Get_rank(self) -> int:
+        return self._rank
+
+    def Get_size(self) -> int:
+        return self._size
+
+    def Barrier(self) -> None:
+        self.barriers += 1
+
+    def bcast(self, obj, root=0):
+        self.broadcast.append(obj)
+        if self._rank == root:
+            return obj
+        if self._scripted:
+            return self._scripted.pop(0)
+        return obj
+
+
+class _FakeMPI:
+    """Minimal ``mpi4py.MPI`` stand-in exposing only ``COMM_WORLD``."""
+
+    def __init__(self, comm):
+        self.COMM_WORLD = comm
+
+
+def write_config(tmp_path, updates=None, name="bench.yml"):
+    """Write a complete benchmark config into ``tmp_path`` and return its path."""
+    config = yaml.safe_load(DEFAULT_CONFIG.read_text())
+    config["base_run_dir"] = str(tmp_path / "runs")
+    config["dataset_dir"] = str(tmp_path / "datasets")
+    config["fract_base_dir"] = str(tmp_path / "fractals")
+    if updates:
+        config.update(updates)
+    path = tmp_path / name
+    path.write_text(yaml.dump(config))
+    return path
+
+
+def run_cli(monkeypatch, argv, *, comm=None, sync_env=True):
+    """Run ``cli.main`` with a fake communicator and stubbed subcommand drivers.
+
+    ``sync_env`` makes the launcher environment agree with the fake
+    communicator, which is what a correctly launched job looks like; tests of
+    the mismatch check itself pass ``sync_env=False`` and set the environment
+    themselves.
+
+    Returns ``(comm, calls)`` where ``calls`` maps the subcommand name to the
+    list of config dicts its driver was invoked with.
+    """
+    import ScaFFold.benchmark as benchmark_mod
+    import ScaFFold.generate_fractals as generate_fractals_mod
+
+    comm = comm if comm is not None else _FakeComm()
+    calls = {"benchmark": [], "generate_fractals": []}
+
+    if sync_env:
+        monkeypatch.setenv("WORLD_SIZE", str(comm.Get_size()))
+        monkeypatch.setenv("RANK", str(comm.Get_rank()))
+    monkeypatch.setattr(sys, "argv", list(argv))
+    monkeypatch.setattr(cli, "MPI", _FakeMPI(comm))
+    monkeypatch.setattr(
+        benchmark_mod,
+        "main",
+        lambda kwargs_dict={}: calls["benchmark"].append(dict(kwargs_dict)),
+    )
+    monkeypatch.setattr(
+        generate_fractals_mod,
+        "main",
+        lambda kwargs_dict={}: calls["generate_fractals"].append(dict(kwargs_dict)),
+    )
+    cli.main()
+    return comm, calls
+
+
+# ---------------------------------------------------------------------------
+# R13: the MPI world must span the whole job
+# ---------------------------------------------------------------------------
+
+
+def test_mpi_singleton_under_multirank_launcher_aborts(monkeypatch, tmp_path):
+    """A 1-rank MPI world inside a 2-rank launcher job aborts with an explanation.
+
+    This is the plain-``torchrun`` shape: every process is an mpi4py singleton
+    while the launcher says WORLD_SIZE=2. Left unchecked, every process runs
+    cli.py's rank-0 block, claims its own run directory, and the job then hangs
+    in the first real collective.
+    """
+    monkeypatch.setenv("WORLD_SIZE", "2")
+    monkeypatch.setenv("RANK", "0")
+    cfg = write_config(tmp_path)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        run_cli(
+            monkeypatch,
+            ["scaffold", "benchmark", "-c", str(cfg)],
+            comm=_FakeComm(rank=0, size=1),
+            sync_env=False,
+        )
+
+    message = str(excinfo.value)
+    assert "1" in message and "2" in message
+    assert "torchrun" in message.lower()
+    # The abort happens before any run directory is claimed.
+    assert not (tmp_path / "runs").exists()
+
+
+def test_matching_world_sizes_are_accepted(monkeypatch, tmp_path):
+    """An MPI world that matches the launcher environment runs normally."""
+    monkeypatch.setenv("WORLD_SIZE", "4")
+    monkeypatch.setenv("RANK", "0")
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(
+        monkeypatch,
+        ["scaffold", "benchmark", "-c", str(cfg)],
+        comm=_FakeComm(rank=0, size=4),
+    )
+
+    assert len(calls["benchmark"]) == 1
+
+
+def test_help_works_under_a_mismatched_launcher_env(monkeypatch, tmp_path, capsys):
+    """``--help`` is answered even when the launcher environment disagrees.
+
+    The cross-check ran before ``parse_args``, so asking what the flags are
+    raised the launcher-mismatch error -- in exactly the environment (a
+    half-configured shell) where someone is most likely to be asking.
+    """
+    monkeypatch.setenv("WORLD_SIZE", "2")
+    monkeypatch.setenv("RANK", "0")
+
+    with pytest.raises(SystemExit) as excinfo:
+        run_cli(
+            monkeypatch,
+            ["scaffold", "--help"],
+            comm=_FakeComm(rank=0, size=1),
+            sync_env=False,
+        )
+
+    assert excinfo.value.code == 0
+    assert "usage" in capsys.readouterr().out.lower()
+
+
+def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path):
+    """With no launcher variables set, the MPI world alone defines the size."""
+    for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "SLURM_NTASKS", "FLUX_JOB_SIZE"):
+        monkeypatch.delenv(var, raising=False)
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(
+        monkeypatch,
+        ["scaffold", "benchmark", "-c", str(cfg)],
+        comm=_FakeComm(rank=0, size=1),
+        sync_env=False,
+    )
+
+    assert len(calls["benchmark"]) == 1
+
+
+# ---------------------------------------------------------------------------
+# R17: the restart script is generated at the true job scale
+# ---------------------------------------------------------------------------
+
+
+def test_restart_script_gets_the_mpi_world_size(monkeypatch, tmp_path):
+    """The CLI passes its communicator size to the restart-script generator.
+
+    Without it the generator falls back to sniffing the environment, which
+    misses launcher variables the rank side honors (e.g. PALS_NRANKS) and
+    silently emits a single-process restart script for a multi-rank job.
+    """
+    recorded = {}
+
+    def _recorder(run_dir, world_size=None):
+        recorded["run_dir"] = run_dir
+        recorded["world_size"] = world_size
+        return Path(run_dir) / "restart.sh"
+
+    monkeypatch.setattr(cli, "create_restart_script", _recorder)
+    monkeypatch.setenv("WORLD_SIZE", "4")
+    monkeypatch.setenv("RANK", "0")
+    cfg = write_config(tmp_path)
+
+    run_cli(
+        monkeypatch,
+        ["scaffold", "benchmark", "-c", str(cfg)],
+        comm=_FakeComm(rank=0, size=4),
+    )
+
+    assert recorded["world_size"] == 4
+
+
+# ---------------------------------------------------------------------------
+# R18: the restart pre-check is a rank-0 decision, broadcast to everyone
+# ---------------------------------------------------------------------------
+
+
+def _restart_argv(cfg, run_dir):
+    return [
+        "scaffold",
+        "benchmark",
+        "-c",
+        str(cfg),
+        "--restart",
+        "--run-dir",
+        str(run_dir),
+    ]
+
+
+def _make_checkpoint(run_dir):
+    ckpt_dir = run_dir / "checkpoints"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    (ckpt_dir / "checkpoint_last.pth").write_bytes(b"")
+    return ckpt_dir
+
+
+def test_restart_without_checkpoint_is_rejected_on_rank0(monkeypatch, tmp_path):
+    """Rank 0 still rejects a restart with no checkpoint, naming the paths."""
+    cfg = write_config(tmp_path)
+    run_dir = tmp_path / "prior"
+    run_dir.mkdir()
+
+    with pytest.raises(FileNotFoundError) as excinfo:
+        run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=_FakeComm(rank=0))
+
+    assert "checkpoint_last.pth" in str(excinfo.value)
+
+
+def test_restart_precheck_follows_the_broadcast_decision(monkeypatch, tmp_path):
+    """A non-zero rank trusts rank 0's verdict instead of stat-ing the FS itself.
+
+    Simulates a stale attribute cache: rank 0 saw the checkpoint and broadcast
+    "go", while this rank's view of the shared filesystem shows nothing. A rank
+    that re-decides locally raises alone and strands its peers in the next
+    barrier.
+    """
+    cfg = write_config(tmp_path)
+    run_dir = tmp_path / "prior"
+    run_dir.mkdir()  # deliberately empty: this rank sees no checkpoint
+    rank0_config = {
+        "restart": True,
+        "run_dir": str(run_dir),
+        "checkpoint_dir": "checkpoints",
+        "verbose": 0,
+    }
+
+    # Rank 0 broadcasts, in order: no config error, the config, no pre-check
+    # error.
+    comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, None])
+    _, calls = run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm)
+
+    assert len(calls["benchmark"]) == 1
+
+
+def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path):
+    """Rank 0's rejection is broadcast, so non-zero ranks raise too.
+
+    Here the local filesystem view *does* show a checkpoint; the rank must
+    still fail, because rank 0 -- the only rank whose verdict counts -- did not
+    find one. Otherwise the job splits: rank 0 aborts and the rest run on.
+    """
+    cfg = write_config(tmp_path)
+    run_dir = tmp_path / "prior"
+    run_dir.mkdir()
+    _make_checkpoint(run_dir)
+    rank0_config = {
+        "restart": True,
+        "run_dir": str(run_dir),
+        "checkpoint_dir": "checkpoints",
+        "verbose": 0,
+    }
+    rank0_error = "Restart requested but no checkpoint was found. Expected /nope."
+
+    comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, rank0_error])
+    with pytest.raises(FileNotFoundError, match="no checkpoint"):
+        run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm)
+
+
+# ---------------------------------------------------------------------------
+# VC-1: the restart state is resolved once, before anything acts on it.
+#
+# ``restart`` and ``run_dir`` can come from the config file as well as the
+# command line (the generated restart.sh replays a dumped config.yaml, which
+# carries both), and an absent ``--restart`` cannot outrank a file that sets
+# it. Resolving the run directory from the command line alone while the
+# pre-check read the merged config made the two disagree.
+# ---------------------------------------------------------------------------
+
+
+def test_yaml_restart_without_run_dir_fails_before_creating_a_run_dir(
+    monkeypatch, tmp_path
+):
+    """``restart: true`` with no run dir aborts without claiming a directory.
+
+    The run directory was resolved from the command line, which said nothing
+    about a restart, so a fresh timestamped directory was created and populated
+    -- and only then did the merged config's ``restart`` trip the pre-check.
+    """
+    cfg = write_config(tmp_path, {"restart": True})
+
+    with pytest.raises(ValueError, match="run directory"):
+        run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)])
+
+    assert not (tmp_path / "runs").exists(), "a run dir was created before the abort"
+
+
+def test_reused_restart_config_resolves_to_one_run_dir(monkeypatch, tmp_path):
+    """A config.yaml from a restarted run cannot split the run across two dirs.
+
+    Reusing such a file as a base config left ``run_dir`` pointing at the run
+    it was dumped by while a *fresh* directory was created to train in. The
+    pre-check then passed on the old run's checkpoints, and the job died hours
+    later with its dataset already generated. Whatever the file resolves to,
+    the directory the pre-check judges and the directory the run uses must be
+    the same one.
+    """
+    cfg = write_config(tmp_path)
+    _, first = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)])
+    first_run = Path(first["benchmark"][0]["benchmark_run_dir"])
+    _make_checkpoint(first_run)
+
+    # Restart it once, so its config.yaml records the restart state.
+    run_cli(monkeypatch, _restart_argv(first_run / "config.yaml", first_run))
+    runs_before = sorted(p.name for p in (tmp_path / "runs").iterdir())
+
+    # Now reuse that config.yaml as a plain base config, with no flags at all.
+    _, reused = run_cli(
+        monkeypatch, ["scaffold", "benchmark", "-c", str(first_run / "config.yaml")]
+    )
+
+    (config,) = reused["benchmark"]
+    assert config["run_dir"] == config["benchmark_run_dir"], (
+        "the pre-check's run_dir and the run's directory disagree"
+    )
+    assert Path(config["benchmark_run_dir"]) == first_run
+    assert sorted(p.name for p in (tmp_path / "runs").iterdir()) == runs_before
+
+
+def test_fresh_run_records_no_restart_state(monkeypatch, tmp_path):
+    """A fresh run's config.yaml carries no restart state to inherit."""
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)])
+
+    (config,) = calls["benchmark"]
+    assert config["restart"] is False
+    assert config["run_dir"] is None
+    dumped = yaml.safe_load(
+        (Path(config["benchmark_run_dir"]) / "config.yaml").read_text()
+    )
+    assert dumped["restart"] is False and dumped["run_dir"] is None
+
+
+# ---------------------------------------------------------------------------
+# R19: generate_fractals is not a benchmark run
+# ---------------------------------------------------------------------------
+
+
+def test_generate_fractals_creates_no_benchmark_run_dir(monkeypatch, tmp_path):
+    """Fractal generation leaves no benchmark run dir and no restart script.
+
+    The rank-0 block used to run for every subcommand, so a generation job
+    littered base_run_dir with a timestamped benchmark directory holding a
+    restart.sh that replays ``generate_fractals --restart --run-dir ...`` --
+    flags the generate_fractals subparser rejects, so the script exits 2.
+    """
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)])
+
+    assert len(calls["generate_fractals"]) == 1
+    assert not (tmp_path / "runs").exists(), "generation created a benchmark run dir"
+    assert list(tmp_path.rglob("restart.sh")) == []
+    assert list(tmp_path.rglob("overrides.yaml")) == []
+
+
+def test_generate_fractals_config_reaches_the_driver(monkeypatch, tmp_path):
+    """The merged config still reaches the generation driver (control)."""
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(
+        monkeypatch,
+        ["scaffold", "generate_fractals", "-c", str(cfg), "--n-categories", "3"],
+    )
+
+    (config,) = calls["generate_fractals"]
+    assert config["n_categories"] == 3
+    assert config["fract_base_dir"] == str(tmp_path / "fractals")
+
+
+def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path):
+    """The benchmark subcommand keeps its run dir, dumps and restart script."""
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)])
+
+    (config,) = calls["benchmark"]
+    run_dir = Path(config["benchmark_run_dir"])
+    assert run_dir.is_dir()
+    assert (run_dir / "config.yaml").exists()
+    assert (run_dir / "overrides.yaml").exists()
+    assert (run_dir / "restart.sh").exists()
+
+
+# ---------------------------------------------------------------------------
+# R20: auxiliary keys set in YAML must survive; CLI > YAML > argparse default
+# ---------------------------------------------------------------------------
+
+
+def test_yaml_aux_keys_reach_the_driver(monkeypatch, tmp_path):
+    """Auxiliary keys set in the config file are not replaced by defaults.
+
+    ``Config`` accepts ``datagen_batch_size``/``verbose`` but does not store
+    them, so they used to vanish from the merged config and the argparse
+    default was installed instead -- making them settable only on the command
+    line despite being documented, validated config keys.
+    """
+    cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 1})
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)])
+
+    (config,) = calls["generate_fractals"]
+    assert config["datagen_batch_size"] == 500
+    assert config["verbose"] == 1
+
+
+def test_cli_flag_outranks_yaml_aux_key(monkeypatch, tmp_path):
+    """An explicit command-line flag still wins over the config file."""
+    cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 0})
+
+    _, calls = run_cli(
+        monkeypatch,
+        [
+            "scaffold",
+            "-v",
+            "generate_fractals",
+            "-c",
+            str(cfg),
+            "--datagen-batch-size",
+            "250",
+        ],
+    )
+
+    (config,) = calls["generate_fractals"]
+    assert config["datagen_batch_size"] == 250
+    assert config["verbose"] == 1
+
+
+def test_argparse_default_used_when_yaml_is_silent(monkeypatch, tmp_path):
+    """With neither a flag nor a config entry, the argparse default applies."""
+    cfg = write_config(tmp_path)
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)])
+
+    (config,) = calls["generate_fractals"]
+    assert config["datagen_batch_size"] == 10000
+    assert config["verbose"] == 0
+
+
+def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path):
+    """The run dir's config.yaml records what the run actually used."""
+    cfg = write_config(tmp_path, {"verbose": 1})
+
+    _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)])
+
+    (config,) = calls["benchmark"]
+    dumped = yaml.safe_load(
+        (Path(config["benchmark_run_dir"]) / "config.yaml").read_text()
+    )
+    assert dumped["verbose"] == 1
+
+
+# ---------------------------------------------------------------------------
+# R25: an out-of-range bottleneck is rejected before any work starts
+# ---------------------------------------------------------------------------
+
+
+def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path):
+    """A command-line override that empties the U-Net is caught at config time.
+
+    The CLI recomputes unet_layers after applying overrides, so the check has
+    to run there too -- not only inside Config.
+    """
+    cfg = write_config(tmp_path)
+
+    with pytest.raises(ValueError) as excinfo:
+        run_cli(
+            monkeypatch,
+            [
+                "scaffold",
+                "benchmark",
+                "-c",
+                str(cfg),
+                "--problem-scale",
+                "4",
+                "--unet-bottleneck-dim",
+                "4",
+            ],
+        )
+
+    message = str(excinfo.value)
+    assert "unet_bottleneck_dim" in message
+    assert "problem_scale" in message
+
+
+# ---------------------------------------------------------------------------
+# VC-2: a rank-0 config failure is broadcast, not left to the barrier.
+#
+# Rank 0 builds the whole job's config while every peer waits in a barrier, so
+# anything that raises in there -- the config-file validation, the recomputed
+# bottleneck check, the run-dir resolution -- has to travel to the peers as a
+# decision. Otherwise the user gets a hang instead of the error message.
+# ---------------------------------------------------------------------------
+
+
+def test_config_failure_is_broadcast_before_the_barrier(monkeypatch, tmp_path):
+    """Rank 0 posts its collectives and broadcasts the error it hit."""
+    cfg = write_config(tmp_path)
+    comm = _FakeComm(rank=0, size=2)
+
+    with pytest.raises(ValueError):
+        run_cli(
+            monkeypatch,
+            [
+                "scaffold",
+                "benchmark",
+                "-c",
+                str(cfg),
+                "--problem-scale",
+                "4",
+                "--unet-bottleneck-dim",
+                "4",
+            ],
+            comm=comm,
+        )
+
+    # The peers' barrier was matched, and what they receive names the failure.
+    assert comm.barriers == 1
+    errors = [
+        payload
+        for payload in comm.broadcast
+        if isinstance(payload, tuple) and payload[0] == "ValueError"
+    ]
+    assert errors, f"no error sentinel was broadcast: {comm.broadcast}"
+    assert "unet_bottleneck_dim" in errors[0][1]
+
+
+def test_peer_raises_the_broadcast_config_error(monkeypatch, tmp_path):
+    """A peer rebuilds rank 0's error instead of running with no config."""
+    cfg = write_config(tmp_path)
+    rank0_error = ("ValueError", "unet_bottleneck_dim (4) must be < problem_scale (4)")
+
+    comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_error])
+    with pytest.raises(ValueError, match="unet_bottleneck_dim"):
+        run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=comm)
+
+
+def test_unknown_error_types_degrade_to_runtime_error():
+    """A non-builtin exception type is still reported, as a RuntimeError."""
+    assert type(cli.rebuild_error("FileNotFoundError", "gone")) is FileNotFoundError
+    rebuilt = cli.rebuild_error("SomeSiteSpecificError", "boom")
+    assert isinstance(rebuilt, RuntimeError)
+    assert "SomeSiteSpecificError" in str(rebuilt) and "boom" in str(rebuilt)
+
+
+# ---------------------------------------------------------------------------
+# The whole config path survives a restart (R20/R22 together)
+# ---------------------------------------------------------------------------
+
+
+def test_config_round_trips_through_a_restart(monkeypatch, tmp_path):
+    """A restart driven by the run dir's config.yaml reproduces the run.
+
+    This is exactly what the generated restart.sh does: ``-c
+    $RUN_DIR/config.yaml --restart --run-dir $RUN_DIR``. It exercises the merge
+    in both directions -- CLI overrides and auxiliary config keys have to come
+    back out of the dumped config, and the resume flags have to win.
+    """
+    cfg = write_config(tmp_path, {"verbose": 1, "datagen_batch_size": 500})
+
+    _, calls = run_cli(
+        monkeypatch,
+        ["scaffold", "benchmark", "-c", str(cfg), "--local-batch-size", "2"],
+    )
+    run_dir = Path(calls["benchmark"][0]["benchmark_run_dir"])
+    _make_checkpoint(run_dir)
+
+    _, resumed = run_cli(monkeypatch, _restart_argv(run_dir / "config.yaml", run_dir))
+
+    (config,) = resumed["benchmark"]
+    assert config["local_batch_size"] == 2  # CLI override from the first run
+    assert config["verbose"] == 1  # auxiliary key set in YAML
+    assert config["datagen_batch_size"] == 500
+    assert config["restart"] is True
+    assert config["train_from_scratch"] is False
+    assert config["benchmark_run_dir"] == str(run_dir)
diff --git a/tests/test_config.py b/tests/test_config.py
index e16f261..086a4f5 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -12,7 +12,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 
-"""Tests for config loading/validation, override merging, and benchmark sweeps."""
+"""Tests for config loading/validation, override merging, and the benchmark driver."""
 
 from pathlib import Path
 
@@ -87,7 +87,7 @@ def test_unknown_key_warns_when_not_strict(capsys):
 )
 def test_documented_keys_accepted(config_path):
     """Every shipped config file loads cleanly under strict validation."""
-    config_utils.load_config(str(config_path), "sweep")
+    config_utils.load_config(str(config_path), "benchmark")
 
 
 def test_invalid_type_message_names_type(tmp_path):
@@ -98,6 +98,39 @@ def test_invalid_type_message_names_type(tmp_path):
         config_utils.load_config(str(path), "bogus")
 
 
+def test_activation_checkpointing_is_a_real_option():
+    """R44: activation checkpointing is reachable from a config, defaulting off.
+
+    The U-Net has always had ``use_checkpointing``, but with no config key and
+    no caller it could not be turned on: any attempt was rejected as an unknown
+    key.
+    """
+    cfg = config_utils.Config(dict(BASE))
+    assert cfg.activation_checkpointing is False
+    assert (
+        config_utils.Config({**BASE, "activation_checkpointing": 1})
+    ).activation_checkpointing is True
+    assert (
+        config_utils.Config({**BASE, "activation_checkpointing": True})
+    ).activation_checkpointing is True
+    assert (
+        config_utils.Config({**BASE, "activation_checkpointing": 0})
+    ).activation_checkpointing is False
+
+
+@pytest.mark.parametrize("value", [2, -1, "yes", 0.0])
+def test_activation_checkpointing_rejects_non_flag_values(value):
+    """Anything that is not a 0/1 (or bool) toggle is rejected by name."""
+    with pytest.raises(ValueError, match="activation_checkpointing"):
+        config_utils.Config({**BASE, "activation_checkpointing": value})
+
+
+def test_activation_checkpointing_documented_in_the_default_config():
+    """The shipped config is the parameter reference, so the key lives there."""
+    text = (CONFIG_DIR / "benchmark_default.yml").read_text()
+    assert "activation_checkpointing:" in text
+
+
 def test_async_save_is_real_option():
     """async_save is an accepted, defaulted option (consumed by the trainer)."""
     cfg = config_utils.Config(dict(BASE))
@@ -167,11 +200,11 @@ def test_config_yaml_roundtrip(tmp_path):
         }
     )
     path = _write_yaml(tmp_path / "config.yaml", combined)
-    config_utils.load_config(path, "sweep")  # must not raise
+    config_utils.load_config(path, "benchmark")  # must not raise
 
 
 # ---------------------------------------------------------------------------
-# Benchmark parameter sweeps
+# Single-run benchmark driver; list-valued (sweep) params are rejected
 # ---------------------------------------------------------------------------
 
 
@@ -203,31 +236,168 @@ def _run_benchmark(monkeypatch, tmp_path, config_updates):
     return calls
 
 
-def test_sweep_cross_product(monkeypatch, tmp_path):
-    """List-valued sweep params produce one scalarized run per combination."""
-    calls = _run_benchmark(
-        monkeypatch,
-        tmp_path,
-        {"local_batch_size": [1, 2], "starting_learning_rate": [0.1, 0.2]},
-    )
-    assert len(calls) == 4
-    combos = sorted((c["local_batch_size"], c["starting_learning_rate"]) for c in calls)
-    assert combos == [(1, 0.1), (1, 0.2), (2, 0.1), (2, 0.2)]
-    # Every run received scalars and its own run directory.
-    run_dirs = {c["run_dir"] for c in calls}
-    assert len(run_dirs) == 4
-
-
-def test_scalar_config_single_run(monkeypatch, tmp_path):
-    """An all-scalar config runs the worker exactly once."""
+def test_benchmark_runs_worker_once(monkeypatch, tmp_path):
+    """One benchmark invocation runs the worker exactly once, in the run dir."""
     calls = _run_benchmark(monkeypatch, tmp_path, {})
     assert len(calls) == 1
     assert calls[0]["local_batch_size"] == BASE["local_batch_size"]
+    # The run uses the benchmark run dir directly: no per-combination subdirs.
+    run_root = tmp_path / "run"
+    assert calls[0]["run_dir"] == str(run_root)
+    assert list(run_root.glob("param_set_*")) == []
+
+
+def test_no_sweep_expansion_helper():
+    """The sweep-expansion entry point is gone from the benchmark driver."""
+    import ScaFFold.benchmark as benchmark_mod
+
+    assert not hasattr(benchmark_mod, "expand_sweep_combinations")
 
 
-def test_list_in_scalar_field_fails_fast():
-    """A list reaching a scalar Config field raises naming the key."""
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("problem_scale", [6, 7]),
+        ("unet_bottleneck_dim", [3, 4]),
+        ("n_categories", [5, 10]),
+        ("local_batch_size", [1, 2]),
+    ],
+)
+def test_list_valued_scalar_key_rejected(key, value):
+    """A list value for a scalar key is rejected by name, not by a stray TypeError."""
     bad = dict(BASE)
-    bad["local_batch_size"] = [1, 2]
+    bad[key] = value
+    with pytest.raises(ValueError) as excinfo:
+        config_utils.Config(bad)
+    message = str(excinfo.value)
+    assert key in message
+    assert "parameter sweeps are no longer supported" in message.lower()
+    assert str(value) in message
+
+
+def test_list_valued_key_rejected_by_load_config(tmp_path):
+    """The benchmark config load path rejects sweeps with the same clear error."""
+    path = _write_yaml(tmp_path / "sweepy.yml", {**BASE, "problem_scale": [6, 7]})
+    with pytest.raises(ValueError) as excinfo:
+        config_utils.load_config(path, "benchmark")
+    message = str(excinfo.value)
+    assert "problem_scale" in message
+    assert "parameter sweeps are no longer supported" in message.lower()
+
+
+def test_all_list_valued_keys_named():
+    """Every offending key is named in one error, not just the first one hit."""
+    bad = {**BASE, "problem_scale": [6, 7], "local_batch_size": [1, 2]}
+    with pytest.raises(ValueError) as excinfo:
+        config_utils.Config(bad)
+    message = str(excinfo.value)
+    assert "problem_scale" in message
+    assert "local_batch_size" in message
+
+
+def test_list_valued_key_rejected_in_run_config(tmp_path):
+    """A list reaching the worker's RunConfig is rejected before any setup."""
+    bad = {
+        **BASE,
+        "local_batch_size": [1, 2],
+        "run_dir": str(tmp_path),
+        "run_iter": str(tmp_path / "run"),
+    }
     with pytest.raises(ValueError, match="local_batch_size"):
+        config_utils.RunConfig(bad)
+
+
+# ---------------------------------------------------------------------------
+# The base config is preserved under a name of its own (R22)
+# ---------------------------------------------------------------------------
+
+
+def _run_benchmark_with_base(monkeypatch, tmp_path, base_name):
+    """Run the driver with a base config named ``base_name``; return the run dir.
+
+    The run dir starts out holding the merged ``config.yaml`` the CLI writes,
+    so the test can tell whether the copy of the base config overwrote it.
+    """
+    import ScaFFold.benchmark as benchmark_mod
+    import ScaFFold.worker as worker_mod
+
+    monkeypatch.setattr(worker_mod, "main", lambda kwargs_dict={}: None)
+    monkeypatch.setattr(benchmark_mod.worker, "main", worker_mod.main)
+
+    base = _write_yaml(tmp_path / base_name, BASE)
+    run_dir = tmp_path / "run"
+    run_dir.mkdir()
+    merged = {**BASE, "local_batch_size": 2, "machine_name": "some-host"}
+    _write_yaml(run_dir / "config.yaml", merged)
+
+    kwargs = {
+        **merged,
+        "command": "benchmark",
+        "restart": False,
+        "verbose": 0,
+        "config": base,
+        "benchmark_run_dir": str(run_dir),
+    }
+    benchmark_mod.main(kwargs_dict=kwargs)
+    return run_dir
+
+
+@pytest.mark.parametrize("base_name", ["config.yaml", "base.yml"])
+def test_base_config_copy_never_clobbers_merged_config(
+    monkeypatch, tmp_path, base_name
+):
+    """The merged config.yaml survives even when the base file has that name.
+
+    Copying the base config into the run dir under its own name silently
+    overwrote the merged dump -- losing the CLI overrides, machine name and
+    scheduler metadata -- and, since restart.sh points -c at
+    ``$RUN_DIR/config.yaml``, a restart then reloaded the raw base config.
+    """
+    run_dir = _run_benchmark_with_base(monkeypatch, tmp_path, base_name)
+
+    merged = yaml.safe_load((run_dir / "config.yaml").read_text())
+    assert merged["local_batch_size"] == 2
+    assert merged["machine_name"] == "some-host"
+
+    preserved = yaml.safe_load((run_dir / "base_config.yaml").read_text())
+    assert preserved["local_batch_size"] == BASE["local_batch_size"]
+    assert "machine_name" not in preserved
+
+
+# ---------------------------------------------------------------------------
+# unet_bottleneck_dim range (R25)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "problem_scale, bottleneck",
+    [(5, -1), (5, 5), (5, 6)],
+    ids=["negative", "zero-layers", "negative-layers"],
+)
+def test_bottleneck_out_of_range_rejected(problem_scale, bottleneck):
+    """An out-of-range bottleneck fails at config time, naming both keys.
+
+    Left unvalidated it produced a U-Net with more pooling levels than the
+    volume has, and the run died hours later inside max_pool3d with "Given
+    input size: (2048x1x1x1)" -- naming no config key at all.
+    """
+    bad = {**BASE, "problem_scale": problem_scale, "unet_bottleneck_dim": bottleneck}
+
+    with pytest.raises(ValueError) as excinfo:
         config_utils.Config(bad)
+
+    message = str(excinfo.value)
+    assert "unet_bottleneck_dim" in message
+    assert "problem_scale" in message
+    assert str(bottleneck) in message
+    assert str(problem_scale) in message
+
+
+@pytest.mark.parametrize("bottleneck", [0, 3, 4])
+def test_bottleneck_in_range_accepted(bottleneck):
+    """The full valid range (at least one U-Net layer) is accepted."""
+    cfg = config_utils.Config(
+        {**BASE, "problem_scale": 5, "unet_bottleneck_dim": bottleneck}
+    )
+    assert cfg.unet_layers == 5 - bottleneck
+    assert cfg.unet_layers >= 1
diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py
index 5f3a623..1f60055 100644
--- a/tests/test_data_loading.py
+++ b/tests/test_data_loading.py
@@ -594,3 +594,223 @@ def counting_load(path, mmap_mode=None):
     assert loaded["image"] == 0
     assert mask_only.dtype == torch.int16
     assert torch.equal(mask_only, expected)
+
+
+# ---------------------------------------------------------------------------
+# R28: a *present but broken* meta.yaml must not be mistaken for a v1 dataset
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "broken_meta",
+    [
+        pytest.param("", id="zero-byte"),
+        pytest.param("{[not: valid: yaml", id="unparseable"),
+        pytest.param("- just\n- a\n- list\n", id="not-a-mapping"),
+        pytest.param("config_id: abc123\n", id="version-key-missing"),
+        pytest.param("dataset_format_version: two\n", id="version-not-an-int"),
+    ],
+)
+def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta):
+    """A corrupt ``meta.yaml`` is an error, never a silent legacy downgrade.
+
+    Treating a broken meta as "no meta" reclassifies a modern dataset as legacy
+    v1: the loader then transposes channels-first volumes (a (3,N,N,N) sample
+    comes back (N,3,N,N)) and remaps already-dense labels. Training proceeds on
+    silently corrupted data. The dataset directory itself is intact here -- only
+    the metadata is damaged -- so the failure must be loud and actionable.
+    """
+    root = _build_v2_constant_dataset(tmp_path / "ds", n_volumes=2)
+    (root / "meta.yaml").write_text(broken_meta)
+
+    with pytest.raises(ValueError) as excinfo:
+        FractalDataset(
+            root / "volumes" / "training",
+            root / "masks" / "training",
+            data_dir=root / "train_unique_mask_vals",
+        )
+
+    message = str(excinfo.value)
+    assert "meta.yaml" in message
+    # The message must point at the offending file so it can be repaired.
+    assert str(root) in message
+
+
+# ---------------------------------------------------------------------------
+# R32: uneven spatial shards are accepted but computed wrong (known, unfixed)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.xfail(
+    reason="uneven spatial shards mishandled; to be fixed upstream in DistConv",
+    strict=True,
+)
+def test_uneven_spatial_shards_pool_like_the_unsharded_volume():
+    """Per-shard pooling must agree with pooling the whole volume.
+
+    ``SpatialShardSpec`` slices with ``torch.chunk`` semantics and only rejects
+    an *empty* shard, so ``vol_size=16`` over 3 shards is accepted as 6/6/4 --
+    nothing anywhere validates ``vol_size % num_shards`` or the per-U-Net-level
+    evenness that pooling needs. DistConv's ``DCTensor`` intercepts only
+    ``aten.convolution``, so ``max_pool3d`` runs independently on each local
+    shard; with unequal (or odd) shards the local pooling windows stop lining up
+    with the global ones, and by the second level whole planes are dropped and
+    values appear that are the max of no global window at all. The same happens
+    for equal-but-odd shards (``vol_size=10`` over 2 shards -> 5/5), and an odd
+    local shard entering a strided conv trips DistConv's own divisibility check
+    with a cryptic error deep in the first forward.
+
+    A volume ramp is used so a pooled value names the plane it came from.
+
+    XFAIL: the fix belongs in DistConv (its sharded ops must handle uneven
+    spatial decompositions), not in the loader, which is why this documents the
+    hazard instead of asserting a workaround. ``strict=True``: the comparison is
+    plain deterministic CPU pooling, so it cannot pass by chance -- if it ever
+    passes, DistConv/ScaFFold has changed and this test must be revisited.
+    """
+    vol_size, num_shards = 16, 3
+    ramp = torch.zeros(1, 1, vol_size, vol_size, vol_size)
+    for plane in range(vol_size):
+        ramp[0, 0, plane] = plane
+
+    # What the dataset hands each rank: uneven shards, accepted without a word.
+    volume = np.arange(vol_size**3, dtype=np.float32).reshape((vol_size,) * 3)
+    shard_sizes = [
+        dl.SpatialShardSpec(
+            shard_dims=(2,), num_shards=(num_shards,), shard_indices=(index,)
+        )
+        .slice_array(volume, {2: 0, 3: 1, 4: 2}, "mask")
+        .shape[0]
+        for index in range(num_shards)
+    ]
+    assert shard_sizes == [6, 6, 4]
+
+    # Two U-Net levels of pooling, globally versus per local shard.
+    pool = torch.nn.MaxPool3d(2)
+    global_pooled = pool(pool(ramp))
+    shards = list(torch.split(ramp, shard_sizes, dim=2))
+    shards = [pool(pool(shard)) for shard in shards]
+    sharded_pooled = torch.cat(shards, dim=2)
+
+    assert sharded_pooled.shape == global_pooled.shape, (
+        f"sharded pooling produced {list(sharded_pooled.shape[2:])} planes vs "
+        f"{list(global_pooled.shape[2:])} globally"
+    )
+    assert torch.equal(sharded_pooled, global_pooled), (
+        f"per-shard planes {sharded_pooled[0, 0, :, 0, 0].tolist()} vs global "
+        f"{global_pooled[0, 0, :, 0, 0].tolist()}"
+    )
+
+
+def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path:
+    """A v2 dataset whose only foreground label is the (large) ``label``.
+
+    v2 masks store raw ``category + 1`` ids and the per-split pickle lists only
+    the categories actually present in that split, so a sparse split can hold a
+    handful of very large ids.
+    """
+    vol_dir = root / "volumes" / "training"
+    mask_dir = root / "masks" / "training"
+    vol_dir.mkdir(parents=True)
+    mask_dir.mkdir(parents=True)
+    np.save(vol_dir / "0.npy", np.zeros((3, 4, 4, 4), dtype=VOLUME_DTYPE))
+    mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE)
+    mask[0, 0, 0] = label
+    np.save(mask_dir / "0_mask.npy", mask)
+    with open(root / "train_unique_mask_vals", "wb") as handle:
+        pickle.dump({"mask_values": [0, label]}, handle)
+    (root / "meta.yaml").write_text("dataset_format_version: 2\n")
+    return root
+
+
+# ---------------------------------------------------------------------------
+# R34: the int16 carrier guard must bound the largest class *id*, not the count
+# ---------------------------------------------------------------------------
+
+
+def test_int16_guard_checks_the_largest_class_id(tmp_path):
+    """A v2 split holding an id above the int16 range is rejected.
+
+    The guard compared ``len(mask_values) - 1`` -- the class *count* -- against
+    the carrier limit, but v2 masks ship raw ids. A split listing only
+    ``[0, 40000]`` passed a two-class check and then wrapped 40000 to -25536 in
+    the int16 carrier, which survives the downstream ``.long()`` cast as a
+    negative label.
+    """
+    root = _build_v2_sparse_label_dataset(tmp_path / "sparse", label=40000)
+
+    with pytest.raises(ValueError) as excinfo:
+        FractalDataset(
+            root / "volumes" / "training",
+            root / "masks" / "training",
+            data_dir=root / "train_unique_mask_vals",
+        )
+
+    message = str(excinfo.value)
+    # The message names the offending id and how many classes the split has.
+    assert "40000" in message
+    assert "2" in message
+    assert str(np.iinfo(np.int16).max) in message
+
+
+def test_int16_guard_accepts_ids_inside_the_range(tmp_path):
+    """A large-but-representable id still loads, and does not wrap negative."""
+    label = int(np.iinfo(np.int16).max)
+    root = _build_v2_sparse_label_dataset(tmp_path / "edge", label=label)
+
+    ds = FractalDataset(
+        root / "volumes" / "training",
+        root / "masks" / "training",
+        data_dir=root / "train_unique_mask_vals",
+    )
+    carrier = ds[0]["mask"]
+    assert carrier.dtype == torch.int16
+    assert int(carrier.min()) >= 0
+    assert int(carrier.max()) == label
+
+
+def test_int16_guard_uses_remapped_ids_for_legacy_datasets(tmp_path):
+    """v1 raw values are remapped to 0..n-1, so huge raw values are fine.
+
+    Guarding on ``max(mask_values)`` alone would reject a legacy dataset that
+    the loader handles perfectly well: its carrier only ever holds the remapped
+    index, not the raw voxel value.
+    """
+    raw_mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE)
+    raw_mask[0, 0, 0] = 40000
+    volume = np.zeros((4, 4, 4, 3), dtype=VOLUME_DTYPE)
+    root = _build_v1_split_dataset(
+        tmp_path / "legacy",
+        raw_mask,
+        volume,
+        train_vals=[0, 40000],
+        val_vals=[0, 40000],
+    )
+
+    ds = FractalDataset(
+        root / "volumes" / "training",
+        root / "masks" / "training",
+        data_dir=root / "train_unique_mask_vals",
+    )
+    assert ds.dataset_format_version == 1
+    carrier = ds[0]["mask"]
+    # 40000 was remapped to class index 1; nothing wrapped.
+    assert int(carrier.max()) == 1
+    assert int(carrier.min()) == 0
+
+
+def test_absent_meta_is_still_legacy_v1(tiny_v1_dataset):
+    """The genuine legacy case (no ``meta.yaml`` at all) is unchanged.
+
+    Control for the test above: v1 datasets predate the metadata file, so a
+    *missing* meta must keep selecting the legacy loader rather than raising.
+    """
+    root = tiny_v1_dataset(n_categories=2, n_train=2, n_val=1, n=8)
+    assert not (root / "meta.yaml").exists()
+
+    ds = FractalDataset(
+        root / "volumes" / "training",
+        root / "masks" / "training",
+        data_dir=root / "train_unique_mask_vals",
+    )
+    assert ds.dataset_format_version == 1
diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py
new file mode 100644
index 0000000..eff199f
--- /dev/null
+++ b/tests/test_groupnorm.py
@@ -0,0 +1,1605 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Tests for the GroupNorm fast paths (``ScaFFold.unet.group_norm``).
+
+The optimization must be invisible everywhere except in the profile: the same
+state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable
+in both directions), the same numbers within reduction-order noise, and a
+fallback for every input the fast kernels cannot or should not take (CPU,
+unknown tensor subclasses, a broken Triton or Inductor install).  The ladder is
+Triton -> compiled -> eager, and a failure at any rung latches that rung off and
+drops to the next, never to the bottom.
+
+DistConv's ``DCTensor`` is not in the rejection list: ``forward`` unwraps it to
+its local shard around both fast kernels, so the wrapped production path is
+served too.  That unwrap is *not* the bare attribute read DistConv's own
+dispatch does -- dispatch runs below autograd, where a bare read is safe, while
+``forward`` runs above it and must go through DistConv's
+``_ToTensor``/``_FromTensor`` pair to keep the graph connected.
+
+The ReLU that used to follow every GroupNorm now lives inside it
+(``activation="relu"``), fused into the Triton store and applied explicitly on
+the other two paths.  ``DoubleConv`` keeps an ``nn.Identity`` in the vacated
+``nn.Sequential`` slot, so the state dict does not move by one key -- which is
+what the checkpoint tests here pin.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+from ScaFFold.unet import group_norm as gn_mod
+from ScaFFold.unet.group_norm import FastGroupNorm
+from ScaFFold.unet.unet_model import UNet
+from tests.helpers import mpi_runner
+
+_N = 16
+_N_CHANNELS = 3
+_N_CLASSES = 2
+_GROUPS = 8
+
+
+@pytest.fixture(autouse=True)
+def _restore_compile_state():
+    """Keep per-test overrides of the module-level routing state contained."""
+    previous_compile = gn_mod.set_compile_enabled(None)
+    previous_triton = gn_mod.set_triton_enabled(None)
+    compile_failed = gn_mod._compile_failed
+    triton_failed = gn_mod._triton_failed
+    yield
+    gn_mod._compile_override = previous_compile
+    gn_mod._triton_override = previous_triton
+    gn_mod._compile_failed = compile_failed
+    gn_mod._triton_failed = triton_failed
+
+
+def _make_unet(seed: int):
+    """Build the worker.py-shaped UNet."""
+    torch.manual_seed(seed)
+    return UNet(
+        n_channels=_N_CHANNELS,
+        n_classes=_N_CLASSES,
+        trilinear=False,
+        layers=2,
+        group_norm_groups=_GROUPS,
+    )
+
+
+def _make_plain_unet(seed: int):
+    """The pre-fusion build: stock ``nn.GroupNorm`` followed by ``nn.ReLU``.
+
+    Built by *converting* a normal UNet rather than by patching the norm class
+    at construction time, because the fusion moved the ReLU into the norm: a
+    class swap alone would leave ``DoubleConv``'s ``nn.Identity`` placeholders
+    in place and produce a model with no activations at all, which would make
+    every numeric comparison below vacuous.  Converting reproduces exactly the
+    module graph this branch replaced -- ``nn.GroupNorm`` where the fast norm
+    sits, an in-place ``nn.ReLU`` where the placeholder sits -- and consumes no
+    RNG (``nn.GroupNorm`` initializes to ones/zeros), so the parameters are
+    bit-identical to what ``_make_unet(seed)`` draws.
+    """
+    model = _make_unet(seed)
+    for parent in [m for m in model.modules() if isinstance(m, nn.Sequential)]:
+        for index, child in enumerate(list(parent)):
+            if isinstance(child, FastGroupNorm):
+                plain = nn.GroupNorm(
+                    child.num_groups,
+                    child.num_channels,
+                    eps=child.eps,
+                    affine=child.affine,
+                )
+                if child.affine:
+                    with torch.no_grad():
+                        plain.weight.copy_(child.weight)
+                        plain.bias.copy_(child.bias)
+                parent[index] = plain
+            elif isinstance(child, nn.Identity):
+                parent[index] = nn.ReLU(inplace=True)
+    return model
+
+
+def _make_input(seed: int = 0, channels: int = _N_CHANNELS, size: int = _N):
+    generator = torch.Generator().manual_seed(seed)
+    return torch.randn(1, channels, size, size, size, generator=generator)
+
+
+# ---------------------------------------------------------------------------
+# state dict compatibility
+# ---------------------------------------------------------------------------
+
+
+def test_state_dict_matches_plain_groupnorm_model():
+    """Names, shapes and dtypes must be unchanged from the nn.GroupNorm build.
+
+    A checkpoint written before this optimization has to keep loading, so the
+    parameter inventory of the model may not shift by even one key.
+    """
+    new_model = _make_unet(seed=0)
+    old_model = _make_plain_unet(seed=0)
+
+    new_sd = new_model.state_dict()
+    old_sd = old_model.state_dict()
+    assert list(new_sd.keys()) == list(old_sd.keys())
+    for key in old_sd:
+        assert new_sd[key].shape == old_sd[key].shape, key
+        assert new_sd[key].dtype == old_sd[key].dtype, key
+    # The optimization must not have introduced buffers either.
+    assert [name for name, _ in new_model.named_buffers()] == [
+        name for name, _ in old_model.named_buffers()
+    ]
+
+
+def test_checkpoint_round_trip_both_directions(tmp_path):
+    """An old checkpoint loads into the new model and vice versa, strict=True.
+
+    Both directions matter: runs resumed onto the new code must accept old
+    checkpoints, and checkpoints written by the new code must stay readable by
+    anything still building plain ``nn.GroupNorm`` (e.g. an older analysis
+    script).  After each load the two models must agree bit for bit.
+    """
+    new_model = _make_unet(seed=0)
+    old_model = _make_plain_unet(seed=1)
+    x = _make_input(seed=3)
+
+    old_path = tmp_path / "old.pth"
+    torch.save({"model_state_dict": old_model.state_dict()}, old_path)
+    loaded = torch.load(old_path, weights_only=True)
+    missing = new_model.load_state_dict(loaded["model_state_dict"], strict=True)
+    assert not missing.missing_keys and not missing.unexpected_keys
+
+    new_model.eval()
+    old_model.eval()
+    with torch.no_grad():
+        assert torch.equal(new_model(x), old_model(x))
+
+    # ... and the reverse: new checkpoint into the plain-GroupNorm model.
+    fresh_new = _make_unet(seed=2)
+    new_path = tmp_path / "new.pth"
+    torch.save({"model_state_dict": fresh_new.state_dict()}, new_path)
+    reloaded = torch.load(new_path, weights_only=True)
+    result = old_model.load_state_dict(reloaded["model_state_dict"], strict=True)
+    assert not result.missing_keys and not result.unexpected_keys
+
+    fresh_new.eval()
+    with torch.no_grad():
+        assert torch.equal(old_model(x), fresh_new(x))
+
+
+def test_unet_uses_fast_group_norm():
+    """Every norm layer in the model is the fast one -- no half-converted build."""
+    model = _make_unet(seed=0)
+    norms = [m for m in model.modules() if isinstance(m, nn.GroupNorm)]
+    assert norms, "UNet should contain GroupNorm layers"
+    assert all(isinstance(m, FastGroupNorm) for m in norms)
+
+
+def test_state_dict_bytes_identical_to_plain_groupnorm_model():
+    """Not just the same keys: the serialized checkpoint must be byte identical.
+
+    ``test_state_dict_matches_plain_groupnorm_model`` compares names, shapes and
+    dtypes; this compares the actual bytes ``torch.save`` writes, which is the
+    thing that has to stay interchangeable.  It is the direct guard on the
+    ``nn.ReLU`` -> ``nn.Identity`` swap: ``nn.Sequential`` names its children by
+    position, so *removing* the activation slot rather than holding it open
+    would renumber ``3.weight`` and ``4.weight``/``4.bias`` and silently
+    invalidate every checkpoint on disk.
+    """
+    import io
+
+    new_model = _make_unet(seed=0)
+    old_model = _make_plain_unet(seed=0)
+
+    def blob(model):
+        buffer = io.BytesIO()
+        torch.save(model.state_dict(), buffer)
+        return buffer.getvalue()
+
+    assert blob(new_model) == blob(old_model)
+
+
+def test_double_conv_keeps_the_activation_slots():
+    """The fused build keeps six positional slots, with nothing in the spares.
+
+    Pins both halves of the fusion design: the ReLU is *in* the norm
+    (``activation == "relu"`` at positions 1 and 4) and its old slots (2 and 5)
+    are parameterless placeholders rather than deletions.
+    """
+    from ScaFFold.unet.unet_parts import DoubleConv
+
+    block = DoubleConv(3, 16, _GROUPS)
+    children = list(block.double_conv)
+    assert len(children) == 6
+    for norm_index, spare_index in ((1, 2), (4, 5)):
+        norm = children[norm_index]
+        assert isinstance(norm, FastGroupNorm)
+        assert norm.activation == "relu"
+        spare = children[spare_index]
+        assert isinstance(spare, nn.Identity)
+        assert list(spare.parameters()) == []
+        assert list(spare.buffers()) == []
+    # And the positional key numbering is exactly the pre-fusion one.
+    assert list(block.state_dict().keys()) == [
+        "double_conv.0.weight",
+        "double_conv.1.weight",
+        "double_conv.1.bias",
+        "double_conv.3.weight",
+        "double_conv.4.weight",
+        "double_conv.4.bias",
+    ]
+
+
+def test_double_conv_output_matches_the_explicit_relu_build():
+    """Folding the ReLU into the norm may not change a single bit of the output.
+
+    The fused module applies the ReLU itself on every path, so on CPU (eager)
+    the block must reproduce ``conv -> GroupNorm -> ReLU`` exactly, gradients
+    included.
+    """
+    from ScaFFold.unet.unet_parts import DoubleConv
+
+    torch.manual_seed(4)
+    fused = DoubleConv(3, 16, _GROUPS)
+    reference = DoubleConv(3, 16, _GROUPS)
+    reference.load_state_dict(fused.state_dict())
+    for index in (1, 4):
+        reference.double_conv[index].activation = None
+    for index in (2, 5):
+        reference.double_conv[index] = nn.ReLU(inplace=True)
+
+    x_fused = _make_input(seed=8, channels=3, size=8).requires_grad_(True)
+    x_reference = x_fused.detach().clone().requires_grad_(True)
+
+    out_fused = fused(x_fused)
+    out_reference = reference(x_reference)
+    assert torch.equal(out_fused, out_reference)
+    # A block whose activation silently vanished would still pass an
+    # output-equality test against another activation-free block, so assert the
+    # ReLU is really there.
+    assert (out_fused < 0).sum() == 0
+    assert out_fused.max() > 0
+
+    out_fused.pow(2).sum().backward()
+    out_reference.pow(2).sum().backward()
+    assert torch.equal(x_fused.grad, x_reference.grad)
+    for (name, a), (_, b) in zip(
+        fused.named_parameters(), reference.named_parameters()
+    ):
+        assert torch.equal(a.grad, b.grad), name
+
+
+def test_activation_argument_is_validated():
+    """An unknown activation must fail at construction, not at the first step."""
+    with pytest.raises(ValueError, match="activation"):
+        FastGroupNorm(_GROUPS, 16, activation="gelu")
+
+
+def test_supported_activations_match_the_kernels():
+    """The module's activation list may not drift from the kernel's.
+
+    ``group_norm`` spells the tuple out rather than importing it (importing the
+    kernel module has to stay off the CPU path), so nothing but this test stops
+    the two copies from diverging into a runtime ``ValueError`` from inside the
+    custom op.
+    """
+    from ScaFFold.unet import triton_group_norm as triton_mod
+
+    assert set(gn_mod.SUPPORTED_ACTIVATIONS) <= set(triton_mod.SUPPORTED_ACTIVATIONS)
+
+
+# ---------------------------------------------------------------------------
+# CPU behavior: identical numerics, and no compilation at all
+# ---------------------------------------------------------------------------
+
+
+def test_cpu_output_bit_identical_to_eager():
+    """On CPU the fast module is literally the stock kernel, so bits must match."""
+    fast = FastGroupNorm(_GROUPS, 64)
+    plain = nn.GroupNorm(_GROUPS, 64)
+    with torch.no_grad():
+        plain.weight.copy_(fast.weight)
+        plain.bias.copy_(fast.bias)
+    x = _make_input(seed=5, channels=64, size=8)
+    assert torch.equal(fast(x), plain(x))
+
+
+def test_cpu_never_invokes_torch_compile(monkeypatch):
+    """The CPU unit suite must not pay Inductor's compile latency.
+
+    Guards the ``input.is_cuda`` check: if it ever regresses, a CPU-only test
+    run would start building C++ kernels for every GroupNorm shape.
+    """
+    calls = []
+
+    def _boom(*a, **kw):
+        calls.append(a)
+        raise AssertionError("torch.compile must not be called for CPU tensors")
+
+    monkeypatch.setattr(torch, "compile", _boom)
+    monkeypatch.setattr(gn_mod, "_compiled_group_norm", None)
+    gn_mod.set_compile_enabled(True)  # even when explicitly forced on
+
+    model = _make_unet(seed=0)
+    with torch.no_grad():
+        model(_make_input(seed=6))
+    assert not calls
+
+
+def test_tensor_subclass_input_stays_eager():
+    """Unknown tensor subclasses must stay on the eager path.
+
+    Dynamo cannot trace ``__torch_dispatch__`` wrappers, so the predicate must
+    reject anything that is not exactly ``torch.Tensor`` before a compile is
+    attempted.  DistConv's ``DCTensor`` is handled separately -- ``forward``
+    unwraps it before consulting the predicate -- but any other wrapper has
+    unknown semantics and keeps the stock kernel.
+    """
+
+    class _Wrapper(torch.Tensor):
+        pass
+
+    plain = torch.randn(1, 8, 4, 4, 4)
+    assert gn_mod._use_compiled(plain) is False  # CPU
+    assert gn_mod._use_compiled(plain.as_subclass(_Wrapper)) is False
+
+
+def test_compile_failure_falls_back_to_eager(monkeypatch, caplog):
+    """A broken compiler degrades to the stock kernel instead of killing the run.
+
+    Simulated by making the compiled callable raise the real thing Dynamo and
+    Inductor raise (``BackendCompilerFailed``/``Unsupported`` share the
+    ``TorchDynamoException`` root the ladder allowlists); the module must return
+    the eager result, warn once, and stop trying for the rest of the process.
+    """
+    import torch._dynamo.exc
+
+    def _raises(*args, **kwargs):
+        raise torch._dynamo.exc.Unsupported("simulated Inductor failure")
+
+    monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises)
+    gn_mod._compile_failed = False
+
+    fast = FastGroupNorm(_GROUPS, 64)
+    x = _make_input(seed=7, channels=64, size=8)
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        out = fast(x)
+    assert torch.equal(
+        out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias)
+    )
+    assert any("falling back to the eager kernel" in r.message for r in caplog.records)
+    assert gn_mod._compile_failed is True
+    # Latched off: the predicate now refuses even a would-be eligible tensor.
+    monkeypatch.undo()
+    assert gn_mod._use_compiled(torch.randn(1, 8, 4, 4, 4)) is False
+
+
+def test_triton_failure_falls_back_to_the_compiled_kernel(monkeypatch, caplog):
+    """A broken Triton install drops to *compiled*, not all the way to eager.
+
+    The distinction is worth 10x on the shapes that dominate the step, so the
+    ladder must have three rungs and not two.  Simulated by forcing the Triton
+    predicate on and making its kernel raise; the compiled stand-in must then be
+    the one that answers, exactly once, with the Triton path latched off.
+    """
+    from ScaFFold.unet.triton_group_norm import TritonKernelError
+
+    compiled_calls = []
+
+    def _raises(*args, **kwargs):
+        raise TritonKernelError("simulated Triton failure")
+
+    def _recording(input, num_groups, weight, bias, eps):
+        compiled_calls.append(type(input))
+        return nn.functional.group_norm(input, num_groups, weight, bias, eps)
+
+    monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True)
+    monkeypatch.setattr(gn_mod, "_get_triton_module", _raises)
+    monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording)
+    gn_mod._triton_failed = False
+    gn_mod._compile_failed = False
+
+    fast = FastGroupNorm(_GROUPS, 64)
+    x = _make_input(seed=44, channels=64, size=8)
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        out = fast(x)
+
+    assert compiled_calls == [torch.Tensor], "compiled kernel was not the fallback"
+    assert torch.equal(
+        out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias)
+    )
+    assert any(
+        "falling back to the compiled kernel" in r.message for r in caplog.records
+    )
+    assert gn_mod._triton_failed is True
+    assert gn_mod._compile_failed is False
+    # Latched off: the predicate now refuses even a would-be eligible tensor.
+    monkeypatch.undo()
+    assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False
+
+
+@pytest.mark.parametrize("rung", ["triton", "compiled"])
+def test_checkpoint_recompute_stop_is_re_raised(monkeypatch, rung):
+    """``_StopRecomputationError`` is control flow, not a kernel failure.
+
+    ``torch.utils.checkpoint``'s non-reentrant recompute stops itself by raising
+    it from a saved-tensor *pack hook*, i.e. from inside whichever op is saving
+    a tensor at that moment -- which, now that the ReLU is fused and nothing
+    follows GroupNorm in a ``DoubleConv``, is this module.  A blanket
+    ``except Exception`` would swallow it, latch the fast kernel off and drop the
+    whole model to eager mid-run.  (Observed exactly that on
+    ``test_gpu_activation_checkpointing_matches_eager`` before the re-raise.)
+    """
+    import torch.utils.checkpoint as checkpoint_mod
+
+    stop = checkpoint_mod._StopRecomputationError
+
+    def _raises(*args, **kwargs):
+        raise stop()
+
+    if rung == "triton":
+        monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_triton_module", _raises)
+    else:
+        monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises)
+    gn_mod._triton_failed = False
+    gn_mod._compile_failed = False
+
+    fast = FastGroupNorm(_GROUPS, 16)
+    with pytest.raises(stop):
+        fast(_make_input(seed=45, channels=16, size=4))
+    assert gn_mod._triton_failed is False
+    assert gn_mod._compile_failed is False
+
+
+def test_cpu_activation_checkpointing_keeps_the_fast_path(monkeypatch):
+    """End-to-end version of the above, on the real model.
+
+    The fast path is forced on for CPU tensors (with the stock kernel standing
+    in for the compiled one, so only the *routing* is under test) and the model
+    is run with activation checkpointing.  Gradients must match the
+    non-checkpointed run and the fast path must still be live afterwards -- a
+    swallowed recompute-stop shows up as a latched-off kernel here.
+    """
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(
+        gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm
+    )
+    gn_mod._compile_failed = False
+
+    x = _make_input(seed=46).requires_grad_(True)
+
+    def grads(checkpointing):
+        model = _make_unet(seed=0)
+        if checkpointing:
+            model.use_checkpointing()
+        model.zero_grad(set_to_none=True)
+        model(x).pow(2).sum().backward()
+        return {n: p.grad.detach().clone() for n, p in model.named_parameters()}
+
+    direct = grads(False)
+    checkpointed = grads(True)
+    assert gn_mod._compile_failed is False, "recompute-stop was swallowed"
+    for name in direct:
+        assert torch.allclose(direct[name], checkpointed[name]), name
+
+
+def test_triton_rejects_unknown_tensor_subclasses():
+    """``is_supported`` only asks ``isinstance``, so the type check lives here.
+
+    A ``__torch_dispatch__`` wrapper other than DCTensor has unknown semantics
+    and must keep the stock kernel, exactly as it does for the compiled path --
+    but ``triton_group_norm.is_supported`` would happily accept one, so
+    ``_use_triton`` has to reject it itself rather than delegating.
+    """
+
+    class _Wrapper(torch.Tensor):
+        pass
+
+    plain = torch.randn(1, 8, 4, 4, 4)
+    assert gn_mod._use_triton(plain, 8, None, None, None) is False  # CPU
+    assert gn_mod._use_triton(plain.as_subclass(_Wrapper), 8, None, None, None) is False
+
+
+def test_triton_env_opt_out_skips_the_kernel_module_entirely(monkeypatch):
+    """``SCAFFOLD_GROUPNORM_TRITON=0`` is checked before anything is imported."""
+
+    def _boom():
+        raise AssertionError("the kernel module must not be imported when opted out")
+
+    monkeypatch.setattr(gn_mod, "_get_triton_module", _boom)
+    gn_mod.set_triton_enabled(False)
+    assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False
+
+
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        ("0", False),
+        ("false", False),
+        ("OFF", False),
+        ("no", False),
+        ("1", True),
+        ("true", True),
+        ("On", True),
+        ("yes", True),
+        ("maybe", None),
+    ],
+)
+def test_env_var_controls_the_fast_path(monkeypatch, value, expected):
+    """``SCAFFOLD_GROUPNORM_COMPILE`` is the documented run-time opt-out."""
+    monkeypatch.setenv(gn_mod.COMPILE_ENV_VAR, value)
+    gn_mod.set_compile_enabled(None)
+    assert gn_mod._compile_override is expected
+
+
+def test_env_var_unset_means_auto(monkeypatch):
+    monkeypatch.delenv(gn_mod.COMPILE_ENV_VAR, raising=False)
+    gn_mod.set_compile_enabled(None)
+    assert gn_mod._compile_override is None
+
+
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        ("0", False),
+        ("false", False),
+        ("OFF", False),
+        ("no", False),
+        ("1", True),
+        ("true", True),
+        ("On", True),
+        ("yes", True),
+        ("maybe", None),
+    ],
+)
+def test_triton_env_var_controls_the_fast_path(monkeypatch, value, expected):
+    """``SCAFFOLD_GROUPNORM_TRITON`` is parsed exactly like its compile twin."""
+    monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, value)
+    gn_mod.set_triton_enabled(None)
+    assert gn_mod._triton_override is expected
+
+
+def test_triton_env_var_unset_means_auto(monkeypatch):
+    """Unset means "on wherever is_supported accepts" -- the production default."""
+    monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False)
+    gn_mod.set_triton_enabled(None)
+    assert gn_mod._triton_override is None
+
+
+def test_triton_env_var_garbage_warns(monkeypatch, caplog):
+    """An unparsable value is ignored *loudly*, like the compile variable."""
+    monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, "sometimes")
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        gn_mod.set_triton_enabled(None)
+    assert any(gn_mod.TRITON_ENV_VAR in r.message for r in caplog.records)
+    assert gn_mod._triton_override is None
+
+
+def test_set_triton_enabled_returns_the_previous_setting(monkeypatch):
+    """The save/restore contract tests rely on, matching set_compile_enabled."""
+    monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False)
+    gn_mod.set_triton_enabled(None)
+    assert gn_mod.set_triton_enabled(False) is None
+    assert gn_mod.set_triton_enabled(True) is False
+    assert gn_mod.set_triton_enabled(None) is True
+    assert gn_mod._triton_override is None
+
+
+def test_cpu_never_imports_triton(fresh_python):
+    """A CPU-only process must not import triton, nor the kernel module.
+
+    Two separate costs, both of which the CPU unit suite would otherwise pay on
+    every run: ``import triton`` (seconds, and it is not installed everywhere),
+    and importing ``ScaFFold.unet.triton_group_norm``, which registers two
+    dispatcher ops and an autograd formula at import time.  ``_use_triton``
+    rejects non-CUDA tensors *before* it touches the module, which is what this
+    pins -- run in a fresh interpreter because the test session itself has long
+    since imported the kernel module for the kernel's own tests.
+    """
+    out = fresh_python(
+        "import sys\n"
+        "import torch\n"
+        "from ScaFFold.unet.unet_model import UNet\n"
+        "m = UNet(n_channels=3, n_classes=2, trilinear=False, layers=1, "
+        "group_norm_groups=8)\n"
+        "with torch.no_grad():\n"
+        "    m(torch.randn(1, 3, 16, 16, 16))\n"
+        "print('triton', 'triton' in sys.modules)\n"
+        "print('kernel', 'ScaFFold.unet.triton_group_norm' in sys.modules)\n"
+    )
+    assert "triton False" in out, out
+    assert "kernel False" in out, out
+
+
+def test_cpu_activation_is_applied_on_the_eager_path():
+    """``activation="relu"`` is a promise of the module, not of the kernel."""
+    fast = FastGroupNorm(_GROUPS, 16, activation="relu")
+    plain = nn.GroupNorm(_GROUPS, 16)
+    with torch.no_grad():
+        plain.weight.copy_(fast.weight)
+        plain.bias.copy_(fast.bias)
+    x = _make_input(seed=41, channels=16, size=8)
+    assert torch.equal(fast(x), torch.relu(plain(x)))
+
+
+def test_cpu_activation_none_leaves_the_output_alone():
+    """The default stays a bare GroupNorm -- no accidental global activation."""
+    fast = FastGroupNorm(_GROUPS, 16)
+    x = _make_input(seed=42, channels=16, size=8)
+    assert torch.equal(
+        fast(x), nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias)
+    )
+
+
+def test_activation_does_not_allocate_a_second_output():
+    """The absorbed ReLU keeps ``nn.ReLU(inplace=True)``'s memory behaviour.
+
+    The old ``nn.Sequential`` spelling mutated the GroupNorm output in place;
+    an out-of-place ``F.relu`` here would add a full activation-sized allocation
+    at all 22 sites.  Checked by handing the module a stand-in kernel whose
+    output we still hold: the ReLU must have rewritten *that* tensor.
+    """
+    fast = FastGroupNorm(_GROUPS, 16, activation="relu")
+    x = _make_input(seed=43, channels=16, size=4)
+    produced = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps)
+    assert (produced < 0).any(), "test input must have negatives to clamp"
+    out = fast._activate(produced)
+    assert out.data_ptr() == produced.data_ptr()
+    assert (produced < 0).sum() == 0
+
+
+def test_recompile_limit_is_raised_never_lowered():
+    """Dynamo's stock cap of 8 is below what one UNet needs.
+
+    A scale-7 UNet presents 5 distinct GroupNorm shapes, and evaluation runs the
+    same 5 again under ``no_grad`` -- 10 cache entries (measured).  Past the cap
+    Dynamo gives up and every GroupNorm silently reverts to the slow kernel, so
+    the module raises the limit; it must never lower one a caller chose.
+    """
+    import torch._dynamo
+
+    config = torch._dynamo.config
+    name = (
+        "recompile_limit" if hasattr(config, "recompile_limit") else "cache_size_limit"
+    )
+    original = getattr(config, name)
+    try:
+        setattr(config, name, 8)
+        gn_mod._raise_recompile_limit()
+        assert getattr(config, name) >= 10
+        setattr(config, name, 4096)
+        gn_mod._raise_recompile_limit()
+        assert getattr(config, name) == 4096
+    finally:
+        setattr(config, name, original)
+
+
+def test_the_compiled_region_carries_its_own_recompile_limit(monkeypatch):
+    """The global limit is thread-local, so the region must carry one too.
+
+    ``torch._dynamo.config`` keeps user overrides in a ``ContextVar``
+    ("User overrides are thread-local", ``torch/utils/_config_module.py``), so
+    what :func:`_raise_recompile_limit` writes is invisible from every *other*
+    thread -- and one of those threads matters: ``torch.utils.checkpoint``'s
+    non-reentrant recompute runs inside the backward pass, i.e. on the autograd
+    engine's device worker thread.  On a ``DCTensor`` that recompute has to
+    compile (it reaches this module with ``__torch_function__`` subclass
+    handling disabled, which is part of Dynamo's ``GLOBAL_STATE`` guard, so it
+    misses every entry the forward built), and there the limit read the stock 8
+    however large the global had been set -- ``FailOnRecompileLimitHit``, run
+    over.  ``torch.compile``'s ``recompile_limit=`` is applied by Dynamo around
+    the compile itself, on whichever thread that compile happens on, which is
+    the only spelling that reaches the worker; this pins that we ask for it.
+    """
+    seen = {}
+
+    def _fake_compile(fn, **kwargs):
+        seen.update(kwargs)
+        return fn
+
+    monkeypatch.setattr(torch, "compile", _fake_compile)
+    assert gn_mod._compile_group_norm() is gn_mod._group_norm
+    assert seen.get("recompile_limit") == gn_mod._MIN_RECOMPILE_LIMIT
+    assert seen.get("fullgraph") is True and seen.get("dynamic") is False
+
+
+def test_compiling_still_works_without_a_per_region_limit(monkeypatch):
+    """A torch too old for ``recompile_limit=`` must still get a callable.
+
+    The keyword is the fix for the worker thread, not a requirement for
+    compiling at all; dropping the whole rung on a ``TypeError`` would be a far
+    bigger regression than the case it addresses.
+    """
+    calls = []
+
+    def _fake_compile(fn, **kwargs):
+        calls.append(kwargs)
+        if "recompile_limit" in kwargs:
+            raise TypeError("compile() got an unexpected keyword 'recompile_limit'")
+        return fn
+
+    monkeypatch.setattr(torch, "compile", _fake_compile)
+    assert gn_mod._compile_group_norm() is gn_mod._group_norm
+    assert len(calls) == 2 and "recompile_limit" not in calls[1]
+
+
+def test_a_recompile_limit_hit_is_a_kernel_failure_not_a_crash(monkeypatch, caplog):
+    """``FailOnRecompileLimitHit`` has to land in the ladder, not in the run.
+
+    It is what ``fullgraph=True`` raises when a frame needs more cache entries
+    than the recompile limit allows, and -- unlike every other Dynamo failure --
+    it derives straight from ``Exception`` rather than from
+    ``TorchDynamoException``, so an allowlist that names only the latter lets it
+    escape and kill the step (observed at ``5943389``).  It is raised while
+    compiling, before the callable has run or saved anything, so the eager
+    retry underneath it is safe.
+    """
+    import torch._dynamo.exc
+
+    limit_hit = torch._dynamo.exc.FailOnRecompileLimitHit
+    assert not issubclass(limit_hit, torch._dynamo.exc.TorchDynamoException), (
+        "naming it separately is only needed while it sits outside that root"
+    )
+
+    def _kernel(*args, **kwargs):
+        raise limit_hit("simulated recompile limit hit")
+
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel)
+    gn_mod._compile_failed = False
+
+    fast = _seeded_norm()
+    x = _make_input(seed=51, channels=16, size=4)
+    expected = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps)
+
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        out = fast(x)
+
+    assert torch.allclose(out, expected), "the eager fallback did not run"
+    assert gn_mod._compile_failed is True, "the failure did not latch the rung off"
+    assert fast._compiled_ok is False
+    assert any("falling back" in record.message for record in caplog.records)
+
+
+# ---------------------------------------------------------------------------
+# DCTensor routing: unwrap -> compiled kernel -> rewrap
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def dc_cpu(gloo_group_1rank):
+    """DistConv package plus a CPU ParallelStrategy over the 1-rank group.
+
+    ``num_shards=(1, 1, 1)`` on dims (2, 3, 4) is exactly what worker.py builds
+    for a single-device run; a process group must exist even then.
+    """
+    import distconv
+
+    ps = distconv.ParallelStrategy(
+        num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cpu"
+    )
+    return distconv, ps
+
+
+def _seeded_norm(channels=16):
+    """A FastGroupNorm with non-default affine params (the defaults are 1/0,
+    which would let a kernel that drops weight/bias slip through)."""
+    fast = FastGroupNorm(_GROUPS, channels)
+    generator = torch.Generator().manual_seed(97)
+    with torch.no_grad():
+        fast.weight.normal_(1.0, 0.1, generator=generator)
+        fast.bias.normal_(0.0, 0.1, generator=generator)
+    return fast
+
+
+def test_dctensor_routes_through_compiled_kernel(monkeypatch, dc_cpu):
+    """A DCTensor input reaches the compiled callable as its plain local shard.
+
+    Verified with a recording stand-in for the compiled callable: it must see
+    exactly ``torch.Tensor`` (Dynamo cannot trace the wrapper), and the caller
+    must get a DCTensor back with the same values the stock kernel produces.
+    """
+    distconv, ps = dc_cpu
+    seen = []
+
+    def _recording(input, num_groups, weight, bias, eps):
+        seen.append(type(input))
+        return nn.functional.group_norm(input, num_groups, weight, bias, eps)
+
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording)
+
+    fast = _seeded_norm()
+    x = _make_input(seed=21, channels=16, size=4)
+    out = fast(distconv.DCTensor.from_shard(x, ps))
+
+    assert seen == [torch.Tensor]
+    assert isinstance(out, distconv.DCTensor)
+    reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps)
+    assert torch.equal(out._tensor, reference)
+
+
+def test_dctensor_gradients_reach_the_layer_upstream(monkeypatch, dc_cpu):
+    """Gradients must flow past GroupNorm into the layer that produced its input.
+
+    The unwrap has to go through DistConv's autograd pair rather than a bare
+    ``input._tensor`` read.  The distinction is invisible when the DCTensor
+    wraps a leaf -- the leaf *is* ``_tensor``, so even a bare read reaches it --
+    which is why this test puts a producer in front, as production does
+    (``conv -> GroupNorm`` in every block).  With a bare read the graph is
+    severed there: the input and the producer's weight get no gradient at all
+    while GroupNorm's own weight/bias still look healthy.
+    """
+    distconv, ps = dc_cpu
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(
+        gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm
+    )
+
+    fast = _seeded_norm()
+    reference = nn.GroupNorm(_GROUPS, 16)
+    producer = nn.Conv3d(16, 16, 1, bias=False)
+    reference_producer = nn.Conv3d(16, 16, 1, bias=False)
+    with torch.no_grad():
+        reference.weight.copy_(fast.weight)
+        reference.bias.copy_(fast.bias)
+        reference_producer.weight.copy_(producer.weight)
+
+    x_fast = _make_input(seed=22, channels=16, size=4).requires_grad_(True)
+    x_ref = x_fast.detach().clone().requires_grad_(True)
+
+    # A 1x1x1 conv on a DCTensor takes DistConv's convolution path, so the
+    # DCTensor handed to GroupNorm is a genuine intermediate, not a leaf.
+    out = fast(producer(distconv.DCTensor.from_shard(x_fast, ps)))
+    assert isinstance(out, distconv.DCTensor)
+    # Unwrap the way a downstream consumer would (autograd-aware) and drive a
+    # scalar backward through it.
+    distconv.distconv._ToTensor.apply(out).pow(2).sum().backward()
+    reference(reference_producer(x_ref)).pow(2).sum().backward()
+
+    assert x_fast.grad is not None, "gradient never reached the input"
+    assert producer.weight.grad is not None, "gradient never reached the producer"
+    assert torch.equal(x_fast.grad, x_ref.grad)
+    assert torch.equal(producer.weight.grad, reference_producer.weight.grad)
+    assert torch.equal(fast.weight.grad, reference.weight.grad)
+    assert torch.equal(fast.bias.grad, reference.bias.grad)
+
+
+def test_dctensor_on_cpu_never_invokes_torch_compile(monkeypatch, dc_cpu):
+    """The unwrap route obeys the same CPU guard as plain tensors.
+
+    On CPU the local shard fails the ``is_cuda`` check, so a DCTensor must fall
+    through to the stock eager dispatch -- and still come back wrapped.  The
+    stand-in records instead of only raising: ``forward`` catches ``Exception``
+    to fall back, so a raise alone would be swallowed by the very code path
+    under test and the assertion would never fire.
+    """
+    distconv, ps = dc_cpu
+    calls = []
+
+    def _boom(*a, **kw):
+        calls.append(a)
+        raise AssertionError("torch.compile must not be called for CPU tensors")
+
+    monkeypatch.setattr(torch, "compile", _boom)
+    monkeypatch.setattr(gn_mod, "_compiled_group_norm", None)
+    gn_mod.set_compile_enabled(True)  # even when explicitly forced on
+
+    fast = _seeded_norm()
+    x = _make_input(seed=23, channels=16, size=4)
+    out = fast(distconv.DCTensor.from_shard(x, ps))
+
+    assert not calls
+    assert isinstance(out, distconv.DCTensor)
+    reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps)
+    assert torch.equal(out._tensor, reference)
+
+
+def test_dctensor_compile_failure_falls_back_to_eager(monkeypatch, caplog, dc_cpu):
+    """A broken compiler degrades the wrapped path to eager, like the plain one."""
+    distconv, ps = dc_cpu
+    import torch._dynamo.exc
+
+    def _raises(*args, **kwargs):
+        raise torch._dynamo.exc.Unsupported("simulated Inductor failure")
+
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises)
+    gn_mod._compile_failed = False
+
+    fast = _seeded_norm()
+    x = _make_input(seed=24, channels=16, size=4)
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        out = fast(distconv.DCTensor.from_shard(x, ps))
+
+    assert isinstance(out, distconv.DCTensor)
+    reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps)
+    assert torch.equal(out._tensor, reference)
+    assert gn_mod._compile_failed is True
+    assert any("falling back to the eager kernel" in r.message for r in caplog.records)
+
+
+def test_dctensor_two_shards_matches_eager_and_normalizes_per_shard():
+    """Shard count > 1: same values as the eager route, still per-shard stats.
+
+    Every other DCTensor test here runs ``num_shards=(1, 1, 1)``, where the
+    local shard is the whole tensor and the sharding is a no-op -- so none of
+    them can catch a fast path that quietly reduced over the wrong set of
+    elements.  This one splits a spatial dim over two ranks and asserts both
+    halves of the claim the fast path rests on: bit-identical to the eager
+    wrapped route, and normalizing the local shard rather than the global
+    volume (which is DistConv's existing semantics, not something this change
+    introduces -- see the R32/R39 notes for the upstream discussion).
+    """
+    script = (
+        Path(__file__).resolve().parent
+        / "helpers"
+        / "rank_scripts"
+        / "groupnorm_shards_2rank.py"
+    )
+    rc, out, err = mpi_runner.torchrun_gloo(str(script), n=2, timeout=180)
+    assert rc == 0, f"2-rank job failed rc={rc}\nstdout:\n{out}\nstderr:\n{err[-3000:]}"
+
+    results = {
+        rank: match
+        for rank, *match in re.findall(
+            # Bounded alternatives, not \S+: the ranks' lines can arrive
+            # concatenated, so a greedy final field would swallow the next
+            # line's "RESULT".
+            r"RESULT rank=(\d+) shape=(\S+) identical=(True|False) "
+            r"per_shard=(True|False) global=(True|False)",
+            out,
+        )
+    }
+    assert set(results) == {"0", "1"}, f"missing ranks\nstdout:\n{out}"
+    for rank, (shape, identical, per_shard, matches_global) in results.items():
+        assert identical == "True", f"rank {rank}: compiled route differs from eager"
+        assert per_shard == "True", f"rank {rank}: not per-shard statistics"
+        assert matches_global == "False", f"rank {rank}: matched global statistics"
+        # Each rank holds half of the sharded dim.
+        assert shape == "1x16x4x8x8", f"rank {rank}: unexpected shard shape {shape}"
+
+
+@pytest.fixture
+def dc_cuda():
+    """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group."""
+    import distconv
+    import torch.distributed as dist
+
+    created = False
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29517")
+        dist.init_process_group(backend="nccl", rank=0, world_size=1)
+        created = True
+    ps = distconv.ParallelStrategy(
+        num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda"
+    )
+    yield distconv, ps
+    if created and dist.is_initialized():
+        dist.destroy_process_group()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("autocast", [False, True])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last_3d"])
+def test_gpu_dctensor_matches_eager_dctensor(dc_cuda, autocast, layout):
+    """The compiled unwrap path matches today's eager wrapped path on GPU.
+
+    This is the production configuration: worker.py wraps every activation in
+    a DCTensor (even at dc_num_shards=[1,1,1]), which used to force the eager
+    kernel.  Values and gradients must agree within reduction-order noise, the
+    output must still be a DCTensor, and the compile must actually engage.
+
+    Both layouts are covered because production requests ``channels_last_3d``
+    (worker.py) and, with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` set as it is there,
+    the convolutions really do hand GroupNorm channels-last activations.  Only
+    parity is asserted, not the output layout: both routes compared here return
+    contiguous regardless of the input layout (the Triton kernel, which does
+    not, is pinned off below and covered by its own tests).
+    """
+    distconv, ps = dc_cuda
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(31)
+    x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator)
+    if layout == "channels_last_3d":
+        x = x.to(memory_format=torch.channels_last_3d)
+    grad_out = torch.randn(*x.shape, device=device, generator=generator)
+
+    fast = FastGroupNorm(_GROUPS, 64).to(device)
+    with torch.no_grad():
+        fast.weight.normal_(1.0, 0.1, generator=generator)
+        fast.bias.normal_(0.0, 0.1, generator=generator)
+
+    def plain(t):
+        return t._tensor if isinstance(t, distconv.DCTensor) else t
+
+    def run(compiled):
+        # This test is about the compiled rung of the ladder, which the Triton
+        # one would otherwise pre-empt on the channels-last parametrization.
+        gn_mod.set_triton_enabled(False)
+        gn_mod.set_compile_enabled(compiled)
+        inp = x.clone().requires_grad_(True)
+        fast.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
+            out = fast(distconv.DCTensor.from_shard(inp, ps))
+        assert isinstance(out, distconv.DCTensor)
+        local = distconv.distconv._ToTensor.apply(out)
+        local.backward(grad_out.to(local.dtype))
+        return (
+            local.detach(),
+            inp.grad.detach().clone(),
+            plain(fast.weight.grad).detach().clone(),
+            plain(fast.bias.grad).detach().clone(),
+        )
+
+    eager = run(False)
+    compiled = run(True)
+    assert gn_mod._compiled_group_norm is not None, "compiled path was not taken"
+    assert not gn_mod._compile_failed
+
+    _assert_close(compiled[0], eager[0], 1e-5, "output")
+    _assert_close(compiled[1], eager[1], 1e-4, "d_input")
+    _assert_close(compiled[2], eager[2], 1e-4, "d_weight")
+    _assert_close(compiled[3], eager[3], 1e-4, "d_bias")
+    assert compiled[0].dtype == eager[0].dtype
+
+
+# ---------------------------------------------------------------------------
+# GPU behavior: numerics, single compile, checkpointing
+# ---------------------------------------------------------------------------
+
+
+def _assert_close(actual, expected, tol, what):
+    diff = (actual.float() - expected.float()).abs().max().item()
+    scale = expected.float().abs().max().item()
+    assert diff <= tol * max(scale, 1.0), f"{what}: max|diff|={diff:.3e}"
+    return diff
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape", [(1, 64, 32, 32, 32), (2, 128, 16, 16, 16)])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_gpu_compiled_matches_eager(shape, autocast):
+    """Compiled forward and gradients match eager, fp32 and under bf16 autocast.
+
+    ``(1, 64, 32^3)`` is the hot production shape (``[1, 64, 128^3]``) at
+    reduced size -- same channel count and group count, same reduction
+    structure, small enough for a unit test.  Tolerances are loose enough for
+    reduction-order differences and tight enough to catch a real numerics bug;
+    observed maxima on MI300A are ~1e-6 relative.
+    """
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(11)
+    x = torch.randn(*shape, device=device, generator=generator)
+    grad_out = torch.randn(*shape, device=device, generator=generator)
+
+    fast = FastGroupNorm(_GROUPS, shape[1]).to(device)
+    with torch.no_grad():
+        fast.weight.normal_(1.0, 0.1, generator=generator)
+        fast.bias.normal_(0.0, 0.1, generator=generator)
+
+    def run(compiled):
+        # Compiled-rung test: the inputs here are contiguous, which the Triton
+        # kernel declines anyway, but pin it off so the routing cannot drift.
+        gn_mod.set_triton_enabled(False)
+        gn_mod.set_compile_enabled(compiled)
+        inp = x.clone().requires_grad_(True)
+        fast.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
+            out = fast(inp)
+        out.backward(grad_out.to(out.dtype))
+        return (
+            out.detach(),
+            inp.grad.detach(),
+            fast.weight.grad.detach().clone(),
+            fast.bias.grad.detach().clone(),
+        )
+
+    eager = run(False)
+    compiled = run(True)
+    assert gn_mod._compiled_group_norm is not None, "compiled path was not taken"
+    assert not gn_mod._compile_failed
+
+    _assert_close(compiled[0], eager[0], 1e-5, "output")
+    _assert_close(compiled[1], eager[1], 1e-4, "d_input")
+    _assert_close(compiled[2], eager[2], 1e-4, "d_weight")
+    _assert_close(compiled[3], eager[3], 1e-4, "d_bias")
+    # Autocast policy must be preserved: GroupNorm is an fp32 op, so the
+    # compiled path may not quietly hand back bf16 activations.
+    assert compiled[0].dtype == eager[0].dtype
+
+
+@pytest.mark.gpu
+def test_gpu_steady_state_does_not_recompile():
+    """Fixed shapes must compile once and then never again.
+
+    A recompile inside a timed epoch would show up as a multi-second outlier in
+    ``epoch_duration`` (and therefore the FOM), so the guard set has to be
+    stable across steps.
+    """
+    from torch._dynamo.utils import counters
+
+    gn_mod.set_triton_enabled(False)  # this is the compiled rung's guard set
+    gn_mod.set_compile_enabled(True)
+    device = torch.device("cuda")
+    fast = FastGroupNorm(_GROUPS, 64).to(device)
+    x = torch.randn(1, 64, 16, 16, 16, device=device, requires_grad=True)
+
+    def step():
+        fast.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            out = fast(x)
+        out.sum().backward()
+
+    step()  # first call: compiles
+    before = counters["stats"]["unique_graphs"]
+    for _ in range(5):
+        step()
+    assert counters["stats"]["unique_graphs"] == before, "recompiled in steady state"
+
+
+@pytest.mark.gpu
+def test_gpu_activation_checkpointing_matches_eager():
+    """The compiled kernel must survive recompute under use_checkpointing().
+
+    Non-reentrant checkpointing replays the block's forward inside the backward
+    pass; a compiled region has to produce the same activations both times or
+    the gradients silently change.
+
+    Compared as relative L2 error per gradient tensor, because whole-network
+    agreement is not bitwise even without this change: with cudnn.benchmark on
+    and bf16 autocast, two eager runs of this model differ by ~4e-3 relative
+    (measured), and checkpointing on vs. off differs by the same amount.
+    Measured here: compiled vs. eager 6.4e-3, i.e. the same order as that noise
+    floor -- while a genuinely wrong kernel would be O(1).
+    """
+    device = torch.device("cuda")
+    x = _make_input(seed=9).to(device).requires_grad_(True)
+    tolerance = 5e-2
+
+    def grads(compiled, checkpointing):
+        gn_mod.set_triton_enabled(False)
+        gn_mod.set_compile_enabled(compiled)
+        model = _make_unet(seed=0).to(device)
+        if checkpointing:
+            model.use_checkpointing()
+        model.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            out = model(x)
+        out.float().pow(2).mean().backward()
+        return {n: p.grad.detach().clone() for n, p in model.named_parameters()}
+
+    def assert_agrees(actual, expected, label):
+        for name in expected:
+            reference = expected[name].float()
+            error = (actual[name].float() - reference).norm().item()
+            relative = error / max(reference.norm().item(), 1e-12)
+            assert relative < tolerance, f"{label} {name}: rel L2 {relative:.3e}"
+
+    eager = grads(False, True)
+    compiled = grads(True, True)
+    compiled_nockpt = grads(True, False)
+    assert gn_mod._compiled_group_norm is not None, "compiled path was not taken"
+    assert not gn_mod._compile_failed
+    assert_agrees(compiled, eager, "checkpointed grad")
+    assert_agrees(compiled_nockpt, compiled, "grad")
+
+
+@pytest.mark.gpu
+def test_gpu_the_recompile_limit_holds_on_a_worker_thread():
+    """Past Dynamo's stock 8 entries, compiling from a thread that never set it.
+
+    ``torch._dynamo.config``'s user overrides live in a ``ContextVar``, so the
+    limit :func:`_raise_recompile_limit` writes on the main thread is not the
+    limit another thread reads -- and the compiles that matter happen on
+    another thread, because ``torch.utils.checkpoint``'s recompute runs inside
+    the backward pass, on the autograd engine's device worker.  This drives the
+    same shape of traffic directly: entries live on ``_group_norm``'s code
+    object and are shared between threads, so a worker that pushes the count
+    past 8 is exactly the situation the recompute creates.  Before the
+    per-region ``recompile_limit=``, the ninth compile raised
+    ``FailOnRecompileLimitHit`` here.
+
+    ``torch._dynamo.reset()`` first because those entries also accumulate
+    across the whole test session, which would otherwise decide the outcome.
+    """
+    import threading
+
+    torch._dynamo.reset()
+    gn_mod.set_triton_enabled(False)  # this is the compiled rung's limit
+    gn_mod.set_compile_enabled(True)
+    gn_mod._compile_failed = False
+
+    device = torch.device("cuda")
+    # Nine distinct channel counts: nine cache entries, one more than the stock
+    # limit allows, and the one that overflows must land on the worker thread.
+    norms = [FastGroupNorm(_GROUPS, 8 * n).to(device) for n in range(1, 10)]
+    failures = []
+
+    def run(subset):
+        try:
+            for norm in subset:
+                norm(torch.randn(1, norm.num_channels, 2, 2, 2, device=device))
+        except BaseException as error:  # noqa: BLE001 - re-raised below
+            failures.append(error)
+
+    run(norms[:2])
+    worker = threading.Thread(target=run, args=(norms[2:],))
+    worker.start()
+    worker.join()
+
+    if failures:
+        raise AssertionError(f"compiling off the main thread failed: {failures[0]}")
+    assert gn_mod._compile_failed is False, "the rung latched itself off"
+    assert all(norm._compiled_ok for norm in norms), (
+        "some module never had a call served by the compiled rung"
+    )
+
+
+@pytest.mark.gpu
+def test_gpu_checkpointed_dctensor_recompute_keeps_the_compiled_rung(dc_cuda):
+    """The three-way combination that used to die: ckpt + compiled rung + DCTensor.
+
+    ``activation_checkpointing: true`` with ``SCAFFOLD_GROUPNORM_TRITON=0`` on
+    DistConv activations is a supported configuration and it crashed: the
+    recompute reaches this module with ``__torch_function__`` subclass handling
+    *disabled* (DistConv's backward runs below it), which is part of Dynamo's
+    ``GLOBAL_STATE`` guard, so it misses every cache entry the forward built and
+    compiles a second set beside them -- twice the shapes, past 8 -- on the
+    autograd worker thread, where the module's raised limit was invisible.  Each
+    pair of the three is fine on its own; all three together raised
+    ``FailOnRecompileLimitHit`` (at ``5943389``) or, once the ladder caught it
+    and dropped a *proven* module to eager mid-recompute, ``CheckpointError``.
+
+    Five norms is the smallest count that reproduces it: 5 forward entries plus
+    5 recompute entries is 10, and the ninth compile is the one that overflows.
+    The convolutions are what make the block's backward run below torch-function
+    (a bare unwrap does not), and the loss is taken on the ``DCTensor`` for the
+    same reason the trainer's is.
+    """
+    import threading
+
+    import torch.utils.checkpoint
+
+    distconv, ps = dc_cuda
+    device = torch.device("cuda")
+    channels = (8, 16, 24, 32, 40)
+
+    torch._dynamo.reset()
+    gn_mod.set_triton_enabled(False)
+    gn_mod.set_compile_enabled(True)
+    gn_mod._compile_failed = False
+
+    torch.manual_seed(5)
+    norms = [FastGroupNorm(_GROUPS, c).to(device) for c in channels]
+    convs = [
+        nn.Conv3d(previous, c, 1, bias=False).to(device)
+        for previous, c in zip((1,) + channels[:-1], channels)
+    ]
+    tail = nn.Conv3d(channels[-1], 1, 1, bias=False).to(device)
+
+    # Where each GroupNorm call happens, as Dynamo's GLOBAL_STATE guard sees it.
+    states = set()
+
+    def block(t):
+        for conv, norm in zip(convs, norms):
+            states.add(
+                (
+                    threading.current_thread() is threading.main_thread(),
+                    torch._C._is_torch_function_enabled(),
+                )
+            )
+            t = norm(conv(t))
+        return tail(t)
+
+    x = torch.randn(1, 1, 4, 4, 4, device=device)
+
+    def step(checkpointing):
+        for parameter in [x] + [
+            p for m in convs + norms + [tail] for p in m.parameters()
+        ]:
+            parameter.grad = None
+        x.requires_grad_(True)
+        wrapped = distconv.DCTensor.from_shard(x, ps)
+        if checkpointing:
+            out = torch.utils.checkpoint.checkpoint(block, wrapped, use_reentrant=False)
+        else:
+            out = block(wrapped)
+        out.float().square().mean().backward()
+        return [norm.weight.grad.detach().clone() for norm in norms]
+
+    checkpointed = step(True)
+    step(True)  # a second step must not compile anything new either
+    direct = step(False)
+
+    assert (False, False) in states, (
+        "the recompute did not run below torch-function off the main thread; "
+        "this configuration no longer reproduces the guard split it targets"
+    )
+    assert gn_mod._compile_failed is False, "the compiled rung latched itself off"
+    assert all(norm._compiled_ok for norm in norms), "a norm never ran compiled"
+    for index, (recomputed, plain) in enumerate(zip(checkpointed, direct)):
+        assert torch.isfinite(recomputed).all(), index
+        _assert_close(recomputed, plain, 1e-4, f"norm {index} weight grad")
+
+
+# ---------------------------------------------------------------------------
+# GPU behavior: the Triton rung
+# ---------------------------------------------------------------------------
+
+
+def _channels_last(t):
+    return t.is_contiguous(memory_format=torch.channels_last_3d)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", [None, "relu"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_gpu_triton_matches_eager(activation, autocast):
+    """The Triton kernel is the default for channels-last input and matches eager.
+
+    ``(1, 64, 32^3)`` channels-last is the production shape family at unit-test
+    size.  Three claims at once: the routing really picks Triton when nothing is
+    forced; the values and gradients match the eager reference within
+    reduction-order noise; and the fused activation equals an explicit ReLU on
+    the eager result.
+
+    The rung is established by spying on the entry point rather than by
+    inspecting the output's layout: every rung now returns the *input's* memory
+    format, deliberately (a fallback that returned contiguous re-broke the
+    channels-last chain for every convolution after it), so layout no longer
+    distinguishes them.  The layout is asserted separately, of both.
+    """
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(11)
+    x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator).to(
+        memory_format=torch.channels_last_3d
+    )
+    grad_out = torch.randn(*x.shape, device=device, generator=generator)
+
+    fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device)
+    with torch.no_grad():
+        fast.weight.normal_(1.0, 0.1, generator=generator)
+        fast.bias.normal_(0.0, 0.1, generator=generator)
+
+    calls = []
+    original_triton_forward = FastGroupNorm._triton_forward
+
+    def run(triton):
+        gn_mod.set_triton_enabled(triton)
+        gn_mod.set_compile_enabled(False)  # eager reference, not Inductor
+        inp = x.clone().requires_grad_(True)
+        fast.zero_grad(set_to_none=True)
+        before = len(calls)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
+            out = fast(inp)
+        out.backward(grad_out.to(out.dtype))
+        return (
+            out.detach(),
+            inp.grad.detach().clone(),
+            fast.weight.grad.detach().clone(),
+            fast.bias.grad.detach().clone(),
+            len(calls) - before,
+        )
+
+    def spy(self, local):
+        calls.append(tuple(local.shape))
+        return original_triton_forward(self, local)
+
+    FastGroupNorm._triton_forward = spy
+    try:
+        eager = run(False)
+        triton = run(None)  # None = the production default, no override at all
+    finally:
+        FastGroupNorm._triton_forward = original_triton_forward
+    assert not gn_mod._triton_failed
+
+    assert triton[4] == 1, "Triton path was not taken"
+    # ... and the control: the reference really did *not* take it, so the
+    # comparison below is between two kernels and not one kernel with itself.
+    assert eager[4] == 0
+    # Both preserve the input's channels-last layout; that is the contract now,
+    # not a rung signature.
+    assert _channels_last(triton[0]) and _channels_last(eager[0])
+    _assert_close(triton[0], eager[0], 1e-5, "output")
+    _assert_close(triton[1], eager[1], 1e-4, "d_input")
+    _assert_close(triton[2], eager[2], 1e-4, "d_weight")
+    _assert_close(triton[3], eager[3], 1e-4, "d_bias")
+    # Autocast's fp32 policy for GroupNorm must survive the swap.
+    assert triton[0].dtype == eager[0].dtype
+    if activation == "relu":
+        assert (triton[0] < 0).sum() == 0
+        assert triton[0].max() > 0
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_gpu_triton_dctensor_matches_eager_and_stays_wrapped(dc_cuda, activation):
+    """The production configuration: DCTensor in, DCTensor out, NDHWC preserved.
+
+    worker.py wraps every activation in a DCTensor even at
+    ``dc_num_shards=[1,1,1]``, so this -- not the plain-tensor case -- is the
+    path the benchmark actually runs.  A producing convolution sits in front so
+    that the DCTensor handed to GroupNorm is a genuine intermediate: the unwrap
+    has to be the autograd-aware one or the gradient never reaches the conv.
+    """
+    distconv, ps = dc_cuda
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(53)
+    x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to(
+        memory_format=torch.channels_last_3d
+    )
+
+    fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device)
+    producer = nn.Conv3d(64, 64, 1, bias=False).to(
+        device, memory_format=torch.channels_last_3d
+    )
+    with torch.no_grad():
+        fast.weight.normal_(1.0, 0.1, generator=generator)
+        fast.bias.normal_(0.0, 0.1, generator=generator)
+
+    def plain(t):
+        return t._tensor if isinstance(t, distconv.DCTensor) else t
+
+    def run(triton):
+        gn_mod.set_triton_enabled(triton)
+        gn_mod.set_compile_enabled(False)
+        inp = x.clone().requires_grad_(True)
+        fast.zero_grad(set_to_none=True)
+        producer.zero_grad(set_to_none=True)
+        out = fast(producer(distconv.DCTensor.from_shard(inp, ps)))
+        assert isinstance(out, distconv.DCTensor), "DCTensor did not survive"
+        local = distconv.distconv._ToTensor.apply(out)
+        local.float().pow(2).sum().backward()
+        assert inp.grad is not None, "gradient never reached the input"
+        assert producer.weight.grad is not None, "gradient never reached the producer"
+        return (
+            local.detach().clone(),
+            inp.grad.detach().clone(),
+            plain(producer.weight.grad).detach().clone(),
+            plain(fast.weight.grad).detach().clone(),
+            plain(fast.bias.grad).detach().clone(),
+        )
+
+    eager = run(False)
+    triton = run(None)
+    assert not gn_mod._triton_failed
+    assert _channels_last(triton[0]), "Triton path was not taken (output not NDHWC)"
+
+    _assert_close(triton[0], eager[0], 1e-5, "output")
+    for index, what in (
+        (1, "d_input"),
+        (2, "d_producer"),
+        (3, "d_weight"),
+        (4, "d_bias"),
+    ):
+        _assert_close(triton[index], eager[index], 1e-4, what)
+
+
+@pytest.mark.gpu
+def test_gpu_unet_keeps_the_channels_last_chain(monkeypatch):
+    """The whole point: GroupNorm stops breaking the layout chain in the model.
+
+    Before this kernel, every one of the model's GroupNorms consumed
+    ``channels_last_3d`` and emitted contiguous, forcing the next convolution to
+    convert back -- 22 breaks per scale-8 forward.  A hook census asserts that
+    every ``FastGroupNorm`` invocation now takes NDHWC in *and* hands NDHWC out.
+
+    Needs ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` in the environment for the
+    convolutions to emit channels-last at all; without it there is nothing to
+    preserve and the test skips rather than passing vacuously.
+    """
+    device = torch.device("cuda")
+    model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d)
+    x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d)
+
+    census = []
+
+    def hook(module, inputs, output):
+        census.append((_channels_last(inputs[0]), _channels_last(output)))
+
+    for module in model.modules():
+        if isinstance(module, FastGroupNorm):
+            module.register_forward_hook(hook)
+
+    gn_mod.set_triton_enabled(None)
+    with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad():
+        model(x)
+
+    assert census, "no GroupNorm ran"
+    if not any(seen_in for seen_in, _ in census):
+        pytest.skip(
+            "convolutions did not emit channels_last_3d; set "
+            "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)"
+        )
+    breaks = [i for i, (seen_in, seen_out) in enumerate(census) if seen_in != seen_out]
+    assert not breaks, f"GroupNorm broke the layout chain at sites {breaks}"
+    assert all(seen_out for _, seen_out in census)
+
+
+@pytest.mark.gpu
+def test_gpu_unet_triton_matches_the_compiled_build(monkeypatch):
+    """Whole-model gradients with the Triton kernel vs. without it.
+
+    Compared as relative L2 per parameter against the noise floor
+    ``test_gpu_activation_checkpointing_matches_eager`` documents: with
+    ``cudnn.benchmark`` on and bf16 autocast, two *eager* runs of this model
+    differ by ~4e-3 relative.  Anything of that order is the model's own
+    nondeterminism; a genuinely wrong kernel would be O(1).
+
+    Skips (rather than passing vacuously) when the convolutions are not emitting
+    channels-last, since the Triton kernel would then never engage.
+    """
+    device = torch.device("cuda")
+    x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d)
+    tolerance = 5e-2
+
+    engaged = []
+    original = FastGroupNorm._triton_forward
+
+    def spy(self, local):
+        engaged.append(tuple(local.shape))
+        return original(self, local)
+
+    monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy)
+
+    def grads(triton):
+        gn_mod.set_triton_enabled(triton)
+        gn_mod.set_compile_enabled(True)
+        model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d)
+        model.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            out = model(x)
+        out.float().pow(2).mean().backward()
+        return {n: p.grad.detach().clone() for n, p in model.named_parameters()}
+
+    without = grads(False)
+    assert not engaged
+    with_triton = grads(None)
+    if not engaged:
+        pytest.skip(
+            "Triton kernel never engaged; set PYTORCH_MIOPEN_SUGGEST_NHWC=1 "
+            "(the production setting) so the convolutions emit channels-last"
+        )
+    assert not gn_mod._triton_failed
+
+    worst = 0.0
+    for name, reference in without.items():
+        reference = reference.float()
+        error = (with_triton[name].float() - reference).norm().item()
+        relative = error / max(reference.norm().item(), 1e-12)
+        worst = max(worst, relative)
+        assert relative < tolerance, f"grad {name}: rel L2 {relative:.3e}"
+    assert worst < tolerance
diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py
new file mode 100644
index 0000000..66d3c06
--- /dev/null
+++ b/tests/test_groupnorm_wiring_edge.py
@@ -0,0 +1,1173 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Edge cases of the GroupNorm *wiring* (``FastGroupNorm``'s three-rung ladder).
+
+Written as an adversarial review of commit ``ca58664``; ``tests/test_groupnorm.py``
+covers the happy paths and the routing predicates, this file covers the places
+where the ladder, the latches and the absorbed ReLU interact with the rest of
+torch.  See ``review/gn-dctensor/WIRING_REVIEW.md`` for the full write-up.
+
+The review left ten of these as ``xfail(strict=True)``, one per defect, each
+asserting the behaviour the module *should* have.  All ten are fixed and the
+markers are gone; the tests stay, now as regression guards.  The properties
+they pin, in the order the defects were found:
+
+* the fused activation is bit-identical to ``F.relu`` on NaN, the infinities
+  and ``-0.0``, forward and backward, on all three rungs -- and a NaN produced
+  under it still reaches the trainer's non-finite-loss abort;
+* a latch may not change the rung a checkpointed block is *recomputed* on;
+* the ladder catches "the kernel is broken" and nothing else -- not the
+  checkpoint machinery's control flow, not a user's saved-tensor hook, not an
+  exception from after the rung already saved something;
+* ``torch.func`` is a routing question, not a kernel failure;
+* a latch can be cleared;
+* ``activation`` is validated where it is used, and a module pickled before it
+  existed still runs.
+"""
+
+from __future__ import annotations
+
+import io
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ScaFFold.unet import group_norm as gn_mod
+from ScaFFold.unet.group_norm import FastGroupNorm
+from ScaFFold.unet.unet_model import UNet
+
+_GROUPS = 8
+
+
+@pytest.fixture(autouse=True)
+def _restore_routing_state():
+    """Keep per-test overrides of the module-level routing state contained.
+
+    Same contract as ``tests/test_groupnorm.py``'s fixture: several tests here
+    deliberately trip a latch, which is a process global.
+    """
+    previous_compile = gn_mod.set_compile_enabled(None)
+    previous_triton = gn_mod.set_triton_enabled(None)
+    compile_failed = gn_mod._compile_failed
+    triton_failed = gn_mod._triton_failed
+    yield
+    gn_mod._compile_override = previous_compile
+    gn_mod._triton_override = previous_triton
+    gn_mod._compile_failed = compile_failed
+    gn_mod._triton_failed = triton_failed
+
+
+def _cl(t):
+    return t.is_contiguous(memory_format=torch.channels_last_3d)
+
+
+def _cuda_norm(channels=64, activation="relu", size=16, seed=11):
+    """A seeded ``FastGroupNorm`` plus a channels-last CUDA input for it."""
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(seed)
+    x = torch.randn(1, channels, size, size, size, device=device, generator=generator)
+    x = x.to(memory_format=torch.channels_last_3d)
+    module = FastGroupNorm(_GROUPS, channels, activation=activation).to(device)
+    with torch.no_grad():
+        module.weight.normal_(1.0, 0.1, generator=generator)
+        module.bias.normal_(0.0, 0.1, generator=generator)
+    return module, x
+
+
+def _small_unet(device="cpu", channels_last=False):
+    torch.manual_seed(0)
+    model = UNet(
+        n_channels=3, n_classes=2, trilinear=False, layers=1, group_norm_groups=_GROUPS
+    )
+    if channels_last:
+        return model.to(device, memory_format=torch.channels_last_3d)
+    return model.to(device)
+
+
+def _triton_spy(monkeypatch):
+    """Record every call that actually reached the Triton rung."""
+    calls = []
+    original = FastGroupNorm._triton_forward
+
+    def spy(self, local):
+        calls.append(tuple(local.shape))
+        return original(self, local)
+
+    monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy)
+    return calls
+
+
+# ---------------------------------------------------------------------------
+# activation semantics: every rung must apply the same function
+# ---------------------------------------------------------------------------
+
+
+def test_activate_handles_every_supported_activation():
+    """``_activate`` must implement every activation the module advertises.
+
+    ``SUPPORTED_ACTIVATIONS`` is what the *constructor* accepts and what
+    ``is_supported`` is asked about, but the compiled and eager rungs apply it
+    through ``_activate``, which tests one literal string.  Adding a second
+    activation to both tuples (the only thing
+    ``test_supported_activations_match_the_kernels`` checks) would fuse it into
+    the Triton store and silently drop it everywhere else -- i.e. the network's
+    function would depend on the memory format of its input.  This is the guard
+    on that: for every non-``None`` activation, ``_activate`` has to *change*
+    an input that the identity would leave alone.
+    """
+    x = torch.linspace(-2.0, 2.0, 64).reshape(1, 8, 2, 2, 2)
+    for activation in gn_mod.SUPPORTED_ACTIVATIONS:
+        module = FastGroupNorm(_GROUPS, 8, activation=activation)
+        out = module._activate(x.clone())
+        if activation is None:
+            assert torch.equal(out, x)
+        else:
+            assert not torch.equal(out, x), (
+                f"_activate is a no-op for activation={activation!r}: the "
+                "compiled and eager rungs would silently skip it while the "
+                "Triton rung fused it in"
+            )
+
+
+def test_affine_false_still_applies_the_activation():
+    """``affine=False`` leaves ``weight``/``bias`` ``None`` on every rung."""
+    module = FastGroupNorm(_GROUPS, 16, affine=False, activation="relu")
+    x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(5))
+    assert torch.equal(
+        module(x), F.relu(F.group_norm(x, _GROUPS, None, None, module.eps))
+    )
+    assert list(module.state_dict().keys()) == []
+
+
+def test_activation_is_not_part_of_the_state():
+    """``activation`` may not become a parameter, a buffer or a state-dict key."""
+    with_relu = FastGroupNorm(_GROUPS, 16, activation="relu")
+    without = FastGroupNorm(_GROUPS, 16)
+    assert list(with_relu.state_dict().keys()) == list(without.state_dict().keys())
+    assert list(with_relu.buffers()) == []
+    # ... and a checkpoint written by one loads into the other, strict.
+    result = without.load_state_dict(with_relu.state_dict(), strict=True)
+    assert not result.missing_keys and not result.unexpected_keys
+
+
+def test_double_conv_bytes_match_a_hand_built_pre_fusion_block():
+    """Byte-for-byte state-dict identity against an independently built block.
+
+    ``test_state_dict_bytes_identical_to_plain_groupnorm_model`` compares against
+    a model produced by *converting* the fused one, which shares its
+    construction order by definition.  This builds the pre-fusion
+    ``nn.Sequential`` from scratch -- ``Conv3d, GroupNorm, ReLU, Conv3d,
+    GroupNorm, ReLU`` -- and compares the serialized bytes, which is the
+    independent version of the same claim.
+    """
+    from ScaFFold.unet.unet_parts import DoubleConv
+
+    torch.manual_seed(17)
+    fused = DoubleConv(3, 16, _GROUPS)
+
+    torch.manual_seed(17)
+    reference = nn.Sequential(
+        nn.Conv3d(3, 16, kernel_size=3, padding=1, bias=False),
+        nn.GroupNorm(_GROUPS, 16),
+        nn.ReLU(inplace=True),
+        nn.Conv3d(16, 16, kernel_size=3, padding=1, bias=False),
+        nn.GroupNorm(_GROUPS, 16),
+        nn.ReLU(inplace=True),
+    )
+
+    def blob(state_dict):
+        buffer = io.BytesIO()
+        torch.save(state_dict, buffer)
+        return buffer.getvalue()
+
+    assert list(fused.double_conv.state_dict().keys()) == list(
+        reference.state_dict().keys()
+    )
+    assert blob(fused.double_conv.state_dict()) == blob(reference.state_dict())
+
+
+def test_module_pickled_before_the_fusion_still_runs():
+    """A whole-module pickle predates ``self.activation``; forward must cope.
+
+    ``nn.Module.__setstate__`` replaces ``__dict__`` wholesale, so an instance
+    restored from a ``torch.save(model)`` written before the fusion has no
+    ``activation`` at all -- and none of the routing state added since either.
+    Every attribute this module reads outside ``__init__`` therefore needs a
+    class-level default.
+    """
+    module = FastGroupNorm(_GROUPS, 16, activation="relu")
+    state = module.__dict__.copy()
+    for added_since in ("activation", "_triton_ok", "_compiled_ok"):
+        state.pop(added_since, None)  # exactly what a pre-fusion pickle carries
+
+    restored = FastGroupNorm.__new__(FastGroupNorm)
+    nn.Module.__setstate__(restored, state)
+
+    x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(6))
+    out = restored(x)
+    # A pre-fusion pickle had no activation, so it must behave as one.
+    assert torch.equal(
+        out, F.group_norm(x, _GROUPS, restored.weight, restored.bias, restored.eps)
+    )
+
+
+def test_unsupported_activation_assigned_after_construction_is_caught():
+    """``activation`` is validated where it is *used*, not only at construction.
+
+    It is a plain attribute, so it can be assigned afterwards; ``is_supported``
+    would then decline the Triton rung while ``_activate`` silently applied
+    nothing, i.e. the module would quietly become a bare GroupNorm.  The same
+    hole is the forward-looking risk: adding a third entry to both
+    ``SUPPORTED_ACTIVATIONS`` tuples without implementing it in ``_activate``
+    must not produce a network whose activation depends on its input's memory
+    format.  Failing loudly on the rung that cannot apply it closes both.
+    """
+    module = FastGroupNorm(_GROUPS, 16, activation="relu")
+    module.activation = "gelu"
+    x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(7))
+    with pytest.raises(ValueError, match="activation must be one of"):
+        module(x)
+
+
+# ---------------------------------------------------------------------------
+# the ladder must not swallow torch's own control flow
+# ---------------------------------------------------------------------------
+
+
+def test_base_exceptions_are_not_caught():
+    """``KeyboardInterrupt``/``SystemExit`` must escape the ladder untouched."""
+    for exception in (KeyboardInterrupt, SystemExit):
+
+        def _raises(*args, **kwargs):
+            raise exception()
+
+        previous = gn_mod._get_triton_module
+        gn_mod._get_triton_module = _raises
+        original_use = gn_mod._use_triton
+        gn_mod._use_triton = lambda *a, **kw: True
+        gn_mod._triton_failed = False
+        try:
+            module = FastGroupNorm(_GROUPS, 16)
+            with pytest.raises(exception):
+                module(torch.randn(1, 16, 4, 4, 4))
+            assert gn_mod._triton_failed is False
+        finally:
+            gn_mod._get_triton_module = previous
+            gn_mod._use_triton = original_use
+
+
+@pytest.mark.parametrize("rung", ["triton", "compiled"])
+def test_checkpoint_error_is_re_raised(monkeypatch, rung):
+    """``CheckpointError`` is the checkpoint machinery talking, not a kernel.
+
+    It is raised by the recompute pack hook -- i.e. from inside whichever op is
+    saving a tensor -- exactly like ``_StopRecomputationError``, and it is a
+    ``RuntimeError`` subclass, so any handler wide enough to catch "a broken
+    kernel" by type catches it too.  Swallowing it latches the rung off, retries
+    on the next one and leaves the checkpoint frame in a state the machinery
+    never expected.  The allowlist has to be narrow enough that this propagates
+    untouched and nothing latches.
+    """
+    import torch.utils.checkpoint as checkpoint_mod
+
+    def _raises(*args, **kwargs):
+        raise checkpoint_mod.CheckpointError("simulated recompute mismatch")
+
+    if rung == "triton":
+        monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_triton_module", _raises)
+    else:
+        monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises)
+    gn_mod._triton_failed = False
+    gn_mod._compile_failed = False
+
+    module = FastGroupNorm(_GROUPS, 16)
+    with pytest.raises(checkpoint_mod.CheckpointError):
+        module(torch.randn(1, 16, 4, 4, 4))
+    assert gn_mod._triton_failed is False
+    assert gn_mod._compile_failed is False
+
+
+def test_a_rung_failure_does_not_re_fire_saved_tensor_hooks(monkeypatch):
+    """The retry has to be idempotent with respect to saved-tensor hooks.
+
+    Under non-reentrant activation checkpointing the recompute counts pack-hook
+    firings and requires the count *and* the metadata to match the forward's, so
+    a rung that packed some tensors and then failed -- with the fallback packing
+    its own set on top -- corrupts the frame.  A user's offloading hook has the
+    same problem in a less dramatic way (an offload failure retried by
+    offloading a second, larger set).
+
+    The property is structural rather than defensive: the ladder catches only
+    failures that are raised *before* their rung saves anything (the Triton op
+    saves in ``_setup_context``, after the launch region its ``TritonKernelError``
+    comes from; a Dynamo/Inductor error is a compile-time error, before
+    execution).  Both halves are asserted here -- a caught failure packs exactly
+    what a clean fallback packs, and an exception from *after* the packing is
+    not the ladder's to swallow.
+    """
+    import torch._dynamo.exc
+
+    class _Boom(Exception):
+        pass
+
+    def _fails_before_packing(input, num_groups, weight, bias, eps):
+        raise torch._dynamo.exc.Unsupported("failed while compiling")
+
+    def _fails_after_packing(input, num_groups, weight, bias, eps):
+        F.group_norm(input, num_groups, weight, bias, eps)
+        raise _Boom("failed after packing")
+
+    monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+    module = FastGroupNorm(_GROUPS, 16)
+    x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True)
+
+    def packs_during(run):
+        packed = []
+        with torch.autograd.graph.saved_tensors_hooks(
+            lambda t: (packed.append(1), t)[1], lambda t: t
+        ):
+            run()
+        return len(packed)
+
+    baseline = packs_during(lambda: module._eager_forward(x))
+
+    monkeypatch.setattr(
+        gn_mod, "_get_compiled_group_norm", lambda: _fails_before_packing
+    )
+    gn_mod._compile_failed = False
+    retried = packs_during(lambda: module(x))
+    assert retried == baseline, (
+        f"the failed rung packed {retried - baseline} extra tensors before the "
+        "fallback ran"
+    )
+
+    # ... and a failure that *did* have observable effects is not retried at all.
+    monkeypatch.setattr(
+        gn_mod, "_get_compiled_group_norm", lambda: _fails_after_packing
+    )
+    gn_mod._compile_failed = False
+    with pytest.raises(_Boom):
+        module(x)
+
+
+# ---------------------------------------------------------------------------
+# latches
+# ---------------------------------------------------------------------------
+
+
+def test_forcing_a_rung_on_clears_its_failure_latch():
+    """``set_*_enabled(True)`` is the documented way to retry after a failure.
+
+    Without this a one-off failure (a transient OOM, a cache-directory hiccup)
+    costs the rung for the rest of the process with no recovery at all, and the
+    function's own docstring -- "forcing it on is overridden only by the
+    correctness checks" -- is false.  ``None`` deliberately does *not* clear it:
+    that restores a preference, it does not assert that the kernel works again.
+    """
+    for setter, latch in (
+        (gn_mod.set_triton_enabled, "_triton_failed"),
+        (gn_mod.set_compile_enabled, "_compile_failed"),
+    ):
+        setattr(gn_mod, latch, True)
+        setter(True)
+        assert getattr(gn_mod, latch) is False
+
+        setattr(gn_mod, latch, True)
+        setter(None)
+        assert getattr(gn_mod, latch) is True
+        setter(False)
+        assert getattr(gn_mod, latch) is True
+        setattr(gn_mod, latch, False)
+
+
+@pytest.mark.parametrize("rung", ["triton", "compiled"])
+def test_out_of_memory_is_not_recorded_as_a_kernel_failure(monkeypatch, rung):
+    """A transient OOM must propagate, and must not latch a rung off forever.
+
+    ``torch.OutOfMemoryError`` is a resource condition, not a defect: every
+    fallback allocates an output of the same size, so retrying one is a second,
+    differently-shaped OOM at a call site the caller never asked about.
+    Latching on it is worse still -- a per-rank, nondeterministic event that
+    permanently changes which kernel that rank runs, and therefore (measured)
+    the all-reduced gradients of the whole job.
+    """
+    from ScaFFold.unet.triton_group_norm import TritonKernelError
+
+    def _oom(*args, **kwargs):
+        raise torch.OutOfMemoryError("simulated OOM")
+
+    if rung == "triton":
+        monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_triton_module", _oom)
+    else:
+        monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True)
+        monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _oom)
+    gn_mod._triton_failed = False
+    gn_mod._compile_failed = False
+
+    module = FastGroupNorm(_GROUPS, 16)
+    with pytest.raises(torch.OutOfMemoryError):
+        module(torch.randn(1, 16, 4, 4, 4))
+    assert gn_mod._triton_failed is False
+    assert gn_mod._compile_failed is False
+    # It is a RuntimeError, so a handler that caught the kernel's own error by
+    # base class would have swallowed it; the allowlist is the tagged type.
+    assert issubclass(torch.OutOfMemoryError, RuntimeError)
+    assert not issubclass(torch.OutOfMemoryError, TritonKernelError)
+
+
+def test_a_global_latch_does_not_demote_a_module_that_already_used_the_rung(
+    monkeypatch, caplog
+):
+    """The unit of the latch is the module, not the process.
+
+    A rung that has already served a module keeps serving it; only modules that
+    have never used it are steered away.  That is what makes a checkpointed
+    block's forward and its recompute agree (they save different tensors on
+    different rungs, so a mid-graph change is fatal), and it is why a broken
+    install still costs exactly one attempt per module rather than one per call.
+
+    The corollary tested here too: because a proven module keeps retrying, the
+    warning has to be emitted on the latch's edge rather than per call, or a
+    persistently broken kernel floods the log for the rest of the run.
+    """
+    import logging
+
+    import torch._dynamo.exc
+
+    calls = []
+
+    def _kernel(input, num_groups, weight, bias, eps):
+        calls.append(1)
+        if len(calls) > 1:
+            raise torch._dynamo.exc.Unsupported("simulated Inductor failure")
+        return F.group_norm(input, num_groups, weight, bias, eps)
+
+    monkeypatch.setattr(gn_mod, "_use_compiled", gn_mod._use_compiled)
+    monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel)
+    monkeypatch.setattr(
+        gn_mod,
+        "_use_compiled",
+        lambda t, proven=False, **kw: (not gn_mod._compile_failed) or proven,
+    )
+    gn_mod._compile_failed = False
+
+    proven = FastGroupNorm(_GROUPS, 16)
+    fresh = FastGroupNorm(_GROUPS, 16)
+    x = torch.randn(1, 16, 4, 4, 4)
+
+    proven(x)  # succeeds: this module is now proven on the compiled rung
+    assert proven._compiled_ok is True
+
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        proven(x)  # fails, latches, falls back to eager
+        assert gn_mod._compile_failed is True
+        warned = sum("falling back" in r.message for r in caplog.records)
+        proven(x)  # ... and still *tries* the rung, because it is proven
+        proven(x)
+        assert sum("falling back" in r.message for r in caplog.records) == warned, (
+            "a persistently failing rung warned once per call"
+        )
+    assert len(calls) == 4, "the proven module stopped trying its rung"
+
+    fresh(x)  # never used it, so the global latch keeps it away entirely
+    assert len(calls) == 4
+    assert fresh._compiled_ok is False
+
+
+@pytest.mark.parametrize("rung", ["triton", "compiled"])
+def test_a_proven_rung_does_not_degrade_while_a_backward_replays_it(monkeypatch, rung):
+    """A fallback *during a recompute* corrupts rather than degrades.
+
+    The latch already refuses to demote a module that has used a rung, but the
+    fallback itself sidestepped that: the failing call still got answered from
+    the next rung down, and if that call is ``torch.utils.checkpoint``'s
+    recompute of a forward that ran on the failing rung, the recomputed forward
+    saves a different set of tensors than the original did and torch rejects
+    the whole step (``CheckpointError``; on one measured shape a GPU memory
+    fault instead).  Neither is a degradation, so this one case re-raises.
+
+    It is narrow on purpose -- ``_replaying_a_forward()`` is false in an
+    ordinary forward, where ``test_a_global_latch_does_not_demote_a_module_...``
+    still requires the fallback -- and the second half here pins the other side
+    of the narrowness: a module that has *not* used the rung degrades even
+    inside the backward, because the forward it is replaying went down the
+    ladder too and the two agree.
+    """
+    import torch._dynamo.exc
+    import torch.utils.checkpoint as checkpoint_mod
+
+    from ScaFFold.unet.triton_group_norm import TritonKernelError
+
+    failure = TritonKernelError if rung == "triton" else torch._dynamo.exc.Unsupported
+    latch = "_triton_failed" if rung == "triton" else "_compile_failed"
+    proven_flag = "_triton_ok" if rung == "triton" else "_compiled_ok"
+
+    def make_kernel(fail_always):
+        def _kernel(input, num_groups, weight, bias, eps, *activation):
+            if fail_always or gn_mod._replaying_a_forward():
+                raise failure("simulated kernel failure")
+            return F.group_norm(input, num_groups, weight, bias, eps)
+
+        return _kernel
+
+    def run(fail_always):
+        gn_mod._triton_failed = False
+        gn_mod._compile_failed = False
+        kernel = make_kernel(fail_always)
+        if rung == "triton":
+            monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True)
+            module_stub = type("_Stub", (), {"triton_group_norm": staticmethod(kernel)})
+            monkeypatch.setattr(gn_mod, "_get_triton_module", lambda: module_stub)
+        else:
+            monkeypatch.setattr(gn_mod, "_use_compiled", lambda t, **kw: True)
+            monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: kernel)
+        module = FastGroupNorm(_GROUPS, 16)
+        x = torch.randn(1, 16, 4, 4, 4, requires_grad=True)
+        out = checkpoint_mod.checkpoint(module, x, use_reentrant=False)
+        out.pow(2).sum().backward()
+        return module, x
+
+    # Proven in the forward, failing in the recompute: answering from another
+    # rung would be the metadata mismatch, so the failure has to come back out.
+    with pytest.raises(failure):
+        run(fail_always=False)
+
+    # Never served by the rung: the forward already went down the ladder, so
+    # the recompute doing the same agrees with it and the step survives.
+    module, x = run(fail_always=True)
+    assert getattr(module, proven_flag) is False
+    assert getattr(gn_mod, latch) is True
+    assert torch.isfinite(x.grad).all()
+
+
+def test_a_latch_flip_mid_forward_is_not_a_numerics_error(monkeypatch):
+    """Two rungs inside one forward still compose (values, not bits, agree)."""
+    monkeypatch.setattr(
+        gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor
+    )
+    monkeypatch.setattr(
+        gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm
+    )
+    gn_mod._compile_failed = False
+
+    module = FastGroupNorm(_GROUPS, 16, activation="relu")
+    x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True)
+    first = module(x)
+    gn_mod._compile_failed = True  # latch flips between the two calls
+    second = module(first)
+    second.pow(2).sum().backward()
+    assert torch.isfinite(x.grad).all()
+
+
+# ---------------------------------------------------------------------------
+# predicates
+# ---------------------------------------------------------------------------
+
+
+def test_predicates_reject_a_parameter_input():
+    """``nn.Parameter`` is a subclass, so both fast rungs decline it.
+
+    Not a bug -- the model never feeds a Parameter to a norm -- but it is the
+    documented consequence of the ``type(input) is torch.Tensor`` policy, and a
+    regression that loosened it to ``isinstance`` would route real
+    ``__torch_dispatch__`` wrappers into the kernel.
+    """
+    parameter = nn.Parameter(torch.randn(1, 8, 4, 4, 4))
+    assert gn_mod._use_triton(parameter, _GROUPS, None, None, None) is False
+    assert gn_mod._use_compiled(parameter) is False
+
+
+def test_use_triton_is_side_effect_free_for_rejected_inputs(monkeypatch):
+    """The predicate may look, but it may not allocate, launch or mutate."""
+    module = FastGroupNorm(_GROUPS, 16)
+    x = torch.randn(1, 16, 4, 4, 4)
+    before = x.clone()
+    assert gn_mod._use_triton(x, _GROUPS, module.weight, module.bias, None) is False
+    assert torch.equal(x, before)
+
+
+# ---------------------------------------------------------------------------
+# GPU: the Triton rung in situ
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+def test_gpu_double_backward_fails_loudly():
+    """The kernel is first-order only; a second derivative must *raise*.
+
+    ``triton_group_norm``'s backward is itself a custom op with no autograd
+    formula, so a gradient penalty or an HVP through the wired model has to die
+    with a clear message rather than silently return a wrong number -- and the
+    failure must not be mistaken for a broken kernel and latch the rung off.
+    """
+    module, x = _cuda_norm(activation=None)
+    x = x.clone().requires_grad_(True)
+    out = module(x)
+    assert _cl(out), "Triton rung was not taken; the test would be vacuous"
+
+    (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True)
+    with pytest.raises(RuntimeError, match="no autograd formula"):
+        torch.autograd.grad(first.pow(2).sum(), x)
+    assert gn_mod._triton_failed is False, "a caller error latched the kernel off"
+
+
+@pytest.mark.gpu
+def test_gpu_model_double_backward_fails_loudly():
+    """Same, through the wired model, which is where a user would hit it."""
+    model = _small_unet("cuda", channels_last=True)
+    x = (
+        torch.randn(1, 3, 16, 16, 16, device="cuda")
+        .contiguous(memory_format=torch.channels_last_3d)
+        .requires_grad_(True)
+    )
+    gn_mod.set_triton_enabled(None)
+    out = model(x)
+    (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True)
+    with pytest.raises(RuntimeError, match="no autograd formula"):
+        first.pow(2).sum().backward()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("fullgraph", [True, False])
+def test_gpu_triton_rung_inside_a_compiled_region(monkeypatch, fullgraph):
+    """The Triton rung is *allowed* inside ``torch.compile``; prove it works.
+
+    ``_use_compiled`` bails out when ``torch.compiler.is_compiling()`` so the
+    functional GroupNorm inlines, but ``_use_triton`` has no such guard: an
+    enclosing compiled region traces straight into the custom op.  Nothing in
+    ScaFFold compiles ``FastGroupNorm.forward`` today, so this was untested in
+    situ.  Values, gradients *and* the channels-last output must survive
+    Dynamo/AOTAutograd unchanged.
+    """
+    import torch._dynamo
+
+    module, x = _cuda_norm(activation="relu")
+    reference_input = x.clone().requires_grad_(True)
+    reference = module(reference_input)
+    reference.pow(2).sum().backward()
+    reference_grad = reference_input.grad.detach().clone()
+    module.zero_grad(set_to_none=True)
+    assert _cl(reference), "Triton rung was not taken; the test would be vacuous"
+
+    torch._dynamo.reset()
+    calls = _triton_spy(monkeypatch)
+    compiled = torch.compile(lambda t: module(t), fullgraph=fullgraph, dynamic=False)
+
+    compiled_input = x.clone().requires_grad_(True)
+    out = compiled(compiled_input)
+    out.pow(2).sum().backward()
+
+    assert calls, "the Triton rung was not traced inside the compiled region"
+    assert _cl(out), "the compiled region lost the channels-last output"
+    assert torch.equal(out, reference)
+    assert torch.equal(compiled_input.grad, reference_grad)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("proven", [False, True])
+def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch, proven):
+    """A rung failure *while Dynamo is tracing* must still fall back, not die.
+
+    The handler used to call ``logger.warning``, which Dynamo cannot trace
+    ("Unsupported: logging.Logger method not supported for non-export cases"),
+    so a caller compiling this forward with ``fullgraph=True`` got a hard error
+    instead of the fallback -- the one caller for whom the fallback matters
+    most, since the thing it is reacting to is usually a compile-time failure.
+    Nothing in ScaFFold compiles ``FastGroupNorm.forward`` today; this pins the
+    claim that it can.
+
+    Both halves of the handler's guard have to trace, which is why ``proven``
+    is parametrized: with ``_triton_ok`` false Dynamo folds the ``and`` away
+    without ever looking at ``_replaying_a_forward()``, so only the ``True``
+    arm reaches it -- and a probe Dynamo cannot trace there would be the same
+    defect as the logging call, reintroduced.
+    """
+    import torch._dynamo
+
+    from ScaFFold.unet.triton_group_norm import TritonKernelError
+
+    module, x = _cuda_norm(activation="relu")
+    reference = module(x).detach().clone()
+    assert _cl(reference), "Triton rung was not taken; the test would be vacuous"
+
+    real = gn_mod._get_triton_module()
+
+    class _BrokenKernelModule:
+        def __getattr__(self, name):
+            if name == "triton_group_norm":
+
+                def _raises(*args, **kwargs):
+                    raise TritonKernelError("simulated Triton failure")
+
+                return _raises
+            return getattr(real, name)
+
+    monkeypatch.setattr(gn_mod, "_get_triton_module", _BrokenKernelModule)
+    gn_mod._triton_failed = False
+    module._triton_ok = proven
+    torch._dynamo.reset()
+
+    compiled = torch.compile(lambda t: module(t), fullgraph=True, dynamic=False)
+    out = compiled(x)
+
+    assert gn_mod._triton_failed is True, "the latch was not recorded"
+    assert _cl(out), "the fallback rung dropped the channels-last chain"
+    assert (out - reference).abs().max().item() < 1e-5
+
+
+@pytest.mark.gpu
+def test_gpu_inference_mode_takes_the_triton_rung(monkeypatch):
+    """``evaluate()`` runs the whole model under ``torch.inference_mode``."""
+    module, x = _cuda_norm(activation="relu")
+    module.eval()
+    reference = F.relu(F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps))
+    calls = _triton_spy(monkeypatch)
+    with torch.inference_mode():
+        out = module(x)
+    assert calls, "inference_mode fell off the Triton rung"
+    assert _cl(out)
+    assert out.is_inference()
+    assert (out.float() - reference.float()).abs().max().item() < 1e-5
+
+
+@pytest.mark.gpu
+def test_gpu_evaluation_shaped_forward_matches_training_shaped_one(monkeypatch):
+    """``eval()`` + ``inference_mode`` + autocast is the evaluate() combination."""
+    model = _small_unet("cuda", channels_last=True)
+    x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous(
+        memory_format=torch.channels_last_3d
+    )
+    gn_mod.set_triton_enabled(None)
+    model.eval()
+    calls = _triton_spy(monkeypatch)
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        out = model(x)
+    if not calls:
+        pytest.skip(
+            "convolutions did not emit channels_last_3d; set "
+            "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)"
+        )
+    assert torch.isfinite(out.float()).all()
+    assert gn_mod._triton_failed is False
+
+
+@pytest.mark.gpu
+def test_gpu_vmap_over_the_module_still_works():
+    """``torch.func`` is a routing question, not a kernel failure.
+
+    ``is_supported``'s ``is_contiguous(memory_format=...)`` raises outright
+    under a ``vmap`` layer ("NYI"), so a predicate that reaches it -- or a
+    relayout helper that does -- turns a plain ``nn.GroupNorm`` drop-in into a
+    hard error for any caller using ``torch.func``.  Both fast rungs must
+    decline while a transform is active and let the stock kernel answer.
+    """
+    module, x = _cuda_norm(activation=None)
+    batched = torch.stack([x[0], x[0]])
+    out = torch.func.vmap(lambda t: module(t.unsqueeze(0)).squeeze(0))(batched)
+    assert out.shape == batched.shape
+    assert torch.allclose(out[0], module(x[None, 0]).squeeze(0), atol=1e-5)
+
+
+@pytest.mark.gpu
+def test_gpu_a_predicate_that_cannot_answer_falls_back_without_latching(
+    monkeypatch, caplog
+):
+    """``is_supported`` raising is a routing miss, and "no" is always a valid answer.
+
+    The predicate runs *outside* the ladder's try, so anything it raises escapes
+    ``forward()`` -- which is how a ``torch.func`` transform used to turn a
+    drop-in ``nn.GroupNorm`` into a hard error.  The functorch check upstream
+    covers the one caller known to trip it; this covers the shape of the
+    problem, because ``is_supported`` inspects an *arbitrary* tensor and the set
+    of wrappers that can make an attribute read raise is not closed.  A broad
+    catch is right here and nowhere else in this module: the predicate has done
+    no work anyone can observe and a correct answer ("use the stock kernel") is
+    always available -- so it must fall back, and must not latch, because
+    nothing about the kernel has been learned.
+    """
+    import logging
+
+    class _Unanswerable:
+        def __getattr__(self, name):
+            if name == "is_supported":
+
+                def _raises(*args, **kwargs):
+                    raise RuntimeError("NYI: querying is_contiguous inside of vmap")
+
+                return _raises
+            return getattr(gn_mod._get_triton_module(), name)
+
+    module, x = _cuda_norm(activation="relu")
+    reference = module(x).detach().clone()
+    monkeypatch.setattr(gn_mod, "_get_triton_module", _Unanswerable)
+    monkeypatch.setattr(gn_mod, "_predicate_warned", False)
+    gn_mod._triton_failed = False
+    module._triton_ok = False
+
+    with caplog.at_level(logging.WARNING, logger=gn_mod.__name__):
+        out = module(x)
+    assert (out - reference).abs().max().item() < 1e-5
+    assert gn_mod._triton_failed is False, "a routing miss latched the kernel off"
+    assert gn_mod._compile_failed is False
+    assert any("routing check failed" in r.message for r in caplog.records)
+
+
+@pytest.mark.gpu
+def test_gpu_torch_func_grad_does_not_latch_the_rungs_off():
+    """A ``torch.func`` call anywhere must not demote the whole process.
+
+    ``torch.func.grad`` used to reach the kernel, fail, and latch *both* fast
+    rungs off permanently -- i.e. one transform anywhere in a process silently
+    dropped every GroupNorm in the model to the stock kernel for the rest of
+    the run.  Recording a routing miss as a kernel failure is the general shape
+    of the bug; this pins the specific instance.
+    """
+    module, x = _cuda_norm(activation=None)
+    gn_mod._triton_failed = False
+    gn_mod._compile_failed = False
+
+    grad = torch.func.grad(lambda t: module(t).pow(2).sum())(x)
+
+    assert torch.isfinite(grad).all()
+    assert gn_mod._triton_failed is False
+    assert gn_mod._compile_failed is False
+    # ... and the module is still on the Triton rung afterwards.
+    assert _cl(module(x))
+
+
+@pytest.mark.gpu
+def test_gpu_a_triton_failure_during_a_checkpointed_step_degrades_not_dies():
+    """A latch may not change the rung a checkpointed block is recomputed on.
+
+    Non-reentrant checkpointing compares the metadata of every tensor the
+    recomputed forward saves against the forward's, and the three rungs save
+    *different* tensors -- Triton ``(input, weight, bias, mean, rstd)``, the
+    other two ``(input, weight, mean, rstd, relu_output)``.  So a rung change
+    between a block's forward and its recompute kills the step with a
+    ``CheckpointError``, which is the exact opposite of the ladder's contract
+    ("a broken Triton install must degrade a multi-node run, not kill it") and
+    is reachable whenever the ``activation_checkpointing`` option is on
+    (``worker.py:230``).  Matching the output memory format is *not* enough on
+    its own -- measured; the saved sets still differ -- so the fix is that a
+    global latch does not demote a module that has already used the rung.
+
+    The same hazard predates the Triton rung: flipping ``_compile_failed``
+    between forward and recompute dies too, which is why both latches are
+    checked here.
+    """
+    for latch in ("_triton_failed", "_compile_failed"):
+        model = _small_unet("cuda", channels_last=True)
+        model.use_checkpointing()
+        x = (
+            torch.randn(1, 3, 16, 16, 16, device="cuda")
+            .contiguous(memory_format=torch.channels_last_3d)
+            .requires_grad_(True)
+        )
+        gn_mod.set_triton_enabled(None)
+        gn_mod._triton_failed = False
+        gn_mod._compile_failed = False
+
+        out = model(x)
+        # A one-off failure at any *later* GroupNorm site latches the rung off
+        # while the blocks already run are waiting to be recomputed.
+        setattr(gn_mod, latch, True)
+        out.pow(2).sum().backward()
+
+        assert torch.isfinite(x.grad).all(), latch
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("channels_last", [True, False])
+def test_gpu_every_rung_returns_the_inputs_memory_format(channels_last):
+    """All three rungs must agree on the output layout, not just the values.
+
+    ``F.group_norm`` -- eager or Inductor-compiled -- returns a *contiguous*
+    tensor whatever it was given, so a single fallback used to re-break the
+    channels-last chain for every convolution after it, which is the exact
+    thing this module exists to prevent.  It also made the rungs distinguishable
+    to anything that inspects metadata (``torch.utils.checkpoint``, a compiled
+    caller's guards), which is a correctness problem rather than a speed one.
+    """
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(9)
+    x = torch.randn(1, 64, 8, 8, 8, device=device, generator=generator)
+    x = x.to(memory_format=torch.channels_last_3d) if channels_last else x.contiguous()
+    module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device)
+
+    outputs = _run_on_every_rung(module, x)
+    for label, out in outputs.items():
+        assert _cl(out) is channels_last, (
+            f"the {label} rung returned "
+            f"{'channels_last_3d' if _cl(out) else 'contiguous'} for a "
+            f"{'channels_last_3d' if channels_last else 'contiguous'} input"
+        )
+    reference = outputs["eager"].detach().float()
+    for label, out in outputs.items():
+        assert (out.detach().float() - reference).abs().max().item() < 1e-5, label
+
+
+@pytest.mark.gpu
+def test_gpu_triton_rejects_a_cuda_tensor_subclass():
+    """The subclass check has to be tested on a tensor that would otherwise pass.
+
+    ``test_triton_rejects_unknown_tensor_subclasses`` hands ``_use_triton`` a
+    *CPU* subclass, which the ``is_cuda`` check rejects one line later -- so it
+    cannot tell whether the ``type(input) is torch.Tensor`` test exists at all
+    (a mutation deleting that line survives the whole suite).  The check is
+    load-bearing: ``is_supported`` only asks ``isinstance``, so without it every
+    unknown ``__torch_dispatch__`` wrapper would be routed into the kernel.
+    """
+
+    class _Wrapper(torch.Tensor):
+        pass
+
+    x = torch.randn(1, 64, 8, 8, 8, device="cuda").to(
+        memory_format=torch.channels_last_3d
+    )
+    wrapped = x.as_subclass(_Wrapper)
+    # The control: everything *except* the subclass test accepts this tensor.
+    from ScaFFold.unet import triton_group_norm as triton_mod
+
+    assert triton_mod.is_supported(wrapped, _GROUPS, None, None, None) is True
+    assert gn_mod._use_triton(wrapped, _GROUPS, None, None, None) is False
+    assert gn_mod._use_compiled(wrapped) is False
+
+
+#: NaN, +Inf, -Inf, -0.0 and four ordinary values -- everything the fused
+#: activation has to agree with ``F.relu`` on.  ``tl.maximum(y, 0)`` and
+#: ``tl.where(y > 0, y, 0)`` both map NaN to 0.0 (the first returns the non-NaN
+#: operand, the second because ``NaN > 0`` is False); ``F.relu`` propagates it.
+_SPECIAL_VALUES = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0]
+
+
+def _run_on_every_rung(module, x):
+    """``{rung: output}`` for the same module and input on all three rungs."""
+    results = {}
+    for label, triton, compiled in (
+        ("triton", True, False),
+        ("compiled", False, True),
+        ("eager", False, False),
+    ):
+        gn_mod.set_triton_enabled(triton)
+        gn_mod.set_compile_enabled(compiled)
+        results[label] = module(x)
+    return results
+
+
+def _bits(t):
+    return t.detach().float().cpu().contiguous().view(torch.int32)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", ["relu", None])
+@pytest.mark.parametrize("poison", [float("nan"), float("inf"), float("-inf")])
+def test_gpu_all_rungs_agree_on_nan_and_inf(poison, activation):
+    """One non-finite input value must poison the same elements on every rung.
+
+    The Triton store used ``tl.maximum(y, 0.0)``, which returns the *non*-NaN
+    operand, so a diverging activation came back from the Triton rung as a
+    finite 0.0 while ``F.relu`` on the other two kept it NaN.  That is worse
+    than a numerics discrepancy: the forward looks finite while the backward is
+    still NaN, so the run sails past ScaFFold's non-finite-loss abort and
+    checkpoints a broken model -- and the model's output becomes a function of
+    its input's memory format.  ``activation=None`` is the control: all three
+    agreed there even before the fix, which is what localizes the divergence to
+    the fused activation.
+    """
+    device = torch.device("cuda")
+    x = torch.randn(1, 64, 4, 4, 4, device=device)
+    x.view(-1)[0] = poison
+    x = x.to(memory_format=torch.channels_last_3d)
+    module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device)
+
+    results = {
+        label: out.detach().float().cpu().contiguous().isnan()
+        for label, out in _run_on_every_rung(module, x).items()
+    }
+    assert int(results["eager"].sum()) > 0, "the poison did not reach the output"
+    assert torch.equal(results["compiled"], results["eager"])
+    assert torch.equal(results["triton"], results["eager"]), (
+        f"the fused activation turned {int(results['eager'].sum())} NaNs into "
+        f"{int(results['triton'].sum())}"
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", ["relu", None])
+def test_gpu_fused_activation_is_bit_identical_to_relu(activation):
+    """Every special value of the *pre-activation*, on every rung, bit for bit.
+
+    Poisoning the input can only produce NaN pre-activations (one NaN or Inf
+    makes the whole group's statistics NaN), so the four values that actually
+    distinguish the spellings of ReLU are reached the other way round: a zero
+    ``weight`` makes the pre-activation exactly ``bias``, elementwise, so the
+    bias vector chooses what the activation sees.  Expected, per
+    ``F.relu``: NaN stays NaN, ``+Inf`` stays ``+Inf``, ``-Inf`` and both zeros
+    become ``+0.0`` (never ``-0.0``).
+
+    With ``activation=None`` the zeros are normalized (``+ 0.0`` maps ``-0.0``
+    to ``+0.0`` and leaves NaN, the infinities and every normal value alone)
+    before the same bitwise comparison: a ``-0.0`` bias survives to the output
+    there, and whether ``xhat * 0 + (-0.0)`` keeps the sign depends on whether
+    the kernel contracted the multiply-add into an FMA -- true of the Triton
+    *and* the Inductor rung, false of eager, and nothing to do with the
+    activation.  ``torch.equal`` is no use for either case: it reports NaN as
+    unequal to itself.
+    """
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(3)
+    x = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to(
+        memory_format=torch.channels_last_3d
+    )
+    module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device)
+    with torch.no_grad():
+        module.weight.zero_()
+        module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device))
+
+    reference = F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps)
+    if activation == "relu":
+        reference = F.relu(reference)
+
+    for label, out in _run_on_every_rung(module, x).items():
+        if activation == "relu":
+            assert torch.equal(_bits(out), _bits(reference)), (
+                f"{label} rung differs from F.relu(F.group_norm(...)) in the "
+                "bit pattern of at least one special value"
+            )
+        else:
+            assert torch.equal(_bits(out + 0.0), _bits(reference + 0.0)), label
+
+
+@pytest.mark.gpu
+def test_gpu_fused_relu_backward_gates_like_threshold_backward():
+    """ReLU's backward passes the gradient where the output is NaN, too.
+
+    ``threshold_backward(grad, result, 0)`` zeroes where ``result <= 0``, and
+    ``NaN <= 0`` is False -- so a NaN pre-activation passes its gradient.  The
+    kernel recomputes the pre-activation and must gate with the same
+    complement; ``pre > 0 ? dy : 0`` would silently zero it.
+    """
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(4)
+    base = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to(
+        memory_format=torch.channels_last_3d
+    )
+    module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device)
+    with torch.no_grad():
+        module.weight.zero_()
+        module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device))
+
+    grads = {}
+    for label, triton in (("triton", True), ("eager", False)):
+        gn_mod.set_triton_enabled(triton)
+        gn_mod.set_compile_enabled(False)
+        x = base.clone().requires_grad_(True)
+        module.zero_grad(set_to_none=True)
+        module(x).sum().backward()
+        grads[label] = (x.grad.detach().clone(), module.bias.grad.detach().clone())
+
+    # d_bias is exactly the gate: one per element that passed.
+    assert torch.equal(_bits(grads["triton"][1]), _bits(grads["eager"][1]))
+    assert grads["eager"][1][0].item() > 0, "the NaN lane's gradient was gated off"
+    assert torch.equal(
+        _bits(grads["triton"][0].float()), _bits(grads["eager"][0].float())
+    )
+
+
+@pytest.mark.gpu
+def test_gpu_fused_relu_nan_still_trips_the_trainers_non_finite_guard(
+    monkeypatch, tiny_trainer
+):
+    """End to end: a NaN produced under the fused path reaches the abort.
+
+    ScaFFold aborts a run whose reduced epoch losses are non-finite, precisely
+    so a diverged run stops instead of overwriting ``checkpoint_last.pth`` with
+    NaN weights.  A fused activation that ate the NaN would hand that guard a
+    finite loss and let the run continue on a model whose *gradients* are still
+    NaN.  The loss below is the real one: a real UNet, on the GPU, with the
+    Triton rung verified to have served every GroupNorm in it.
+    """
+    from ScaFFold.utils import trainer as trainer_mod
+
+    model = _small_unet("cuda", channels_last=True)
+    poisoned = torch.randn(1, 3, 16, 16, 16, device="cuda")
+    poisoned.view(-1)[0] = float("nan")
+    x = poisoned.contiguous(memory_format=torch.channels_last_3d).requires_grad_(True)
+    gn_mod.set_triton_enabled(None)
+    calls = _triton_spy(monkeypatch)
+    loss = model(x).float().mean()
+    if not calls:
+        pytest.skip(
+            "convolutions did not emit channels_last_3d; set "
+            "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)"
+        )
+    assert not torch.isfinite(loss).item(), (
+        "the fused activation swallowed the NaN: the forward is finite while "
+        "the backward is not, which is exactly what hides divergence"
+    )
+
+    trainer = tiny_trainer(config_overrides={"checkpoint_interval": 1, "epochs": 3})
+    monkeypatch.setattr(
+        trainer,
+        "_run_training_batch",
+        lambda batch, **kw: (1, loss.detach().cpu(), torch.tensor(0.0)),
+    )
+    # A *finite* validation loss, so the abort can only come from the model's.
+    monkeypatch.setattr(
+        trainer_mod, "evaluate", lambda *a, **kw: (7.4e-10, 0.5, 0.5, 2, 2)
+    )
+    trainer.cleanup_or_resume()
+    with pytest.raises(ValueError, match="[Nn]on-finite"):
+        trainer.train()
+    assert not trainer.checkpoint_manager.last_ckpt_path.exists()
+
+
+@pytest.mark.gpu
+def test_gpu_ddp_wrapped_model_takes_the_triton_rung(monkeypatch):
+    """DDP's module-tree walk must not be confused by the nn.Identity slots."""
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel
+
+    created = False
+    if not dist.is_initialized():
+        import os
+
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29623")
+        dist.init_process_group(backend="nccl", rank=0, world_size=1)
+        created = True
+    try:
+        model = _small_unet("cuda", channels_last=True)
+        wrapped = DistributedDataParallel(model, device_ids=[0])
+        x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous(
+            memory_format=torch.channels_last_3d
+        )
+        gn_mod.set_triton_enabled(None)
+        calls = _triton_spy(monkeypatch)
+        wrapped(x).pow(2).sum().backward()
+        if not calls:
+            pytest.skip(
+                "convolutions did not emit channels_last_3d; set "
+                "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)"
+            )
+        assert all(torch.isfinite(p.grad).all() for p in model.parameters())
+    finally:
+        if created and dist.is_initialized():
+            dist.destroy_process_group()
diff --git a/tests/test_infra.py b/tests/test_infra.py
index bf0fbeb..aaf3a73 100644
--- a/tests/test_infra.py
+++ b/tests/test_infra.py
@@ -22,13 +22,16 @@
 
 from __future__ import annotations
 
+import logging
 import os
 
 import numpy as np
+import pytest
 import torch
 
 from ScaFFold.utils.data_loading import FractalDataset
 from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE
+from ScaFFold.utils.utils import gather_and_print_mem, mem_stats
 from tests.helpers import mpi_runner
 
 # ---------------------------------------------------------------------------
@@ -230,3 +233,54 @@ def test_torchrun_gloo_two_ranks(tmp_path):
     # Both ranks reported; all_reduce of ranks {0,1} sums to 1.
     assert "RANK 0/2 sum=1.0" in out
     assert "RANK 1/2 sum=1.0" in out
+
+
+# ---------------------------------------------------------------------------
+# R21: memory diagnostics on a CPU-only run
+# ---------------------------------------------------------------------------
+
+
+def _debug_logger(name):
+    log = logging.getLogger(name)
+    log.setLevel(logging.DEBUG)
+    return log
+
+
+# The bug was a CUDA-free host taking the CUDA path, so these only mean
+# something where CUDA is genuinely unavailable.
+_requires_no_cuda = pytest.mark.skipif(
+    torch.cuda.is_available(), reason="covers the CPU-only path"
+)
+
+
+@_requires_no_cuda
+def test_mem_stats_without_cuda(caplog):
+    """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host."""
+    stats = mem_stats()
+
+    assert stats["cuda_available"] is False
+    assert "rank" in stats
+
+
+@_requires_no_cuda
+def test_gather_and_print_mem_without_cuda(caplog):
+    """A DEBUG-level CPU run logs a fallback instead of crashing.
+
+    ``BaseTrainer.__init__`` calls this unconditionally, so a CPU/gloo run with
+    ``-v`` used to die in trainer construction with "No CUDA GPUs are
+    available".
+    """
+    log = _debug_logger("test_gather_and_print_mem_without_cuda")
+    with caplog.at_level(logging.DEBUG, logger=log.name):
+        gather_and_print_mem(log, "after_trainer_setup")
+
+    messages = " ".join(record.getMessage() for record in caplog.records)
+    assert "after_trainer_setup" in messages
+    assert "cuda" in messages.lower() or "gpu" in messages.lower()
+
+
+def test_trainer_constructs_with_debug_logging(tiny_trainer):
+    """The real call site survives: a trainer builds with a DEBUG logger."""
+    trainer = tiny_trainer(log_level=logging.DEBUG)
+
+    assert trainer.log.getEffectiveLevel() == logging.DEBUG
diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py
index 840bc3d..a683c59 100644
--- a/tests/test_perf_hotpath.py
+++ b/tests/test_perf_hotpath.py
@@ -188,3 +188,160 @@ def test_training_batch_gathers_mem_only_first_batch():
     # first-batch predicate.
     assert "gather_mem_stats=True" not in src
     assert "first_batch" in src
+
+
+# ---------------------------------------------------------------------------
+# R41: warmup covers the ragged final batch
+#
+# Warmup only ever runs the *leading* batches of the train loader, which are
+# all local_batch_size wide, and neither loader drops its last batch. When a
+# rank's sample count is not a multiple of the batch size, the narrower final
+# batch is therefore a set of convolution shapes nothing has warmed, and with
+# cudnn.benchmark on the algorithm search for it runs inside the first *timed*
+# epoch (measured: 95 s) -- straight into epoch_duration, the FOM denominator.
+# ---------------------------------------------------------------------------
+
+
+def _stub_warmup_steps(trainer, monkeypatch, *, mutate=False):
+    """Record the batch sizes warmup runs; optionally mutate model state.
+
+    The real training step is hardwired through DistConv and cannot run on the
+    CPU ``ps=None`` fixture, so the step itself is stubbed; what is under test
+    is which batches warmup feeds it.
+    """
+    import ScaFFold.utils.trainer as tr
+
+    sizes = []
+
+    def fake_training_batch(batch, **kwargs):
+        sizes.append(int(batch["image"].shape[0]))
+        if mutate:
+            with torch.no_grad():
+                for param in trainer.model.parameters():
+                    param.add_(1.0)
+        return int(batch["image"].shape[0]), torch.tensor(0.0), torch.tensor(0.0)
+
+    monkeypatch.setattr(trainer, "_run_training_batch", fake_training_batch)
+    monkeypatch.setattr(tr, "evaluate", lambda *args, **kwargs: (0.0, 0.0))
+    return sizes
+
+
+def test_warmup_covers_the_ragged_train_batch(tiny_trainer, monkeypatch):
+    # 5 local training samples at local_batch_size 2: every epoch ends with a
+    # 1-sample batch that the leading warmup batches never present.
+    trainer = tiny_trainer(
+        n_train=5,
+        n_val=4,
+        config_overrides={"local_batch_size": 2, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch)
+
+    trainer.warmup()
+
+    assert sizes == [2, 2, 1]
+
+
+def test_warmup_adds_no_extra_batch_when_shards_divide_evenly(
+    tiny_trainer, monkeypatch
+):
+    # local_batch_size 1 can never produce a ragged batch: no extra work.
+    trainer = tiny_trainer(
+        n_train=4,
+        n_val=2,
+        config_overrides={"local_batch_size": 1, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch)
+
+    trainer.warmup()
+
+    assert sizes == [1, 1]
+
+
+def test_warmup_covers_a_ragged_validation_batch(tiny_trainer, monkeypatch):
+    # Training divides evenly (4 / 2) but validation does not (3 / 2), so the
+    # 1-sample shape still has to be warmed.
+    trainer = tiny_trainer(
+        n_train=4,
+        n_val=3,
+        config_overrides={"local_batch_size": 2, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch)
+
+    trainer.warmup()
+
+    assert sizes == [2, 2, 1]
+
+
+def test_warmup_ragged_sizes_are_agreed_across_ranks(tiny_trainer, monkeypatch):
+    # Validation is sharded unpadded, so peers can end their epoch with a
+    # different partial size. Every rank must run the same extra steps or the
+    # collectives inside them diverge; the peer's remainder (2) is gathered and
+    # warmed here even though this rank's own is 1.
+    import torch.distributed as dist
+
+    trainer = tiny_trainer(
+        n_train=6,
+        n_val=4,
+        config_overrides={"local_batch_size": 3, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch)
+    trainer.world_size = 2
+
+    def fake_all_gather(tensor_list, tensor, *args, **kwargs):
+        tensor_list[0].copy_(tensor)
+        tensor_list[1].fill_(2)
+
+    monkeypatch.setattr(dist, "all_gather", fake_all_gather)
+
+    trainer.warmup()
+
+    assert sizes == [3, 3, 1, 2]
+
+
+def test_warmup_ragged_all_gather_is_posted_even_without_a_batch(
+    tiny_trainer, monkeypatch
+):
+    # The ragged-size agreement is a collective: a rank whose warmup fetched
+    # no batch at all (empty train loader) must still post the all_gather, or
+    # its peers block in it. The no-batch early return has to come after the
+    # gather, even though such a rank then runs no extra step itself.
+    import torch.distributed as dist
+
+    trainer = tiny_trainer(
+        n_train=4,
+        n_val=3,
+        config_overrides={"local_batch_size": 2, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch)
+    gathers = []
+
+    def fake_all_gather(tensor_list, tensor, *args, **kwargs):
+        gathers.append(int(tensor.item()))
+        for out in tensor_list:
+            out.copy_(tensor)
+
+    monkeypatch.setattr(dist, "all_gather", fake_all_gather)
+
+    trainer._warmup_ragged_batches(None)
+
+    assert gathers, "rank skipped the ragged-size all_gather when batch was None"
+    assert sizes == []
+
+
+def test_warmup_rolls_back_state_including_the_ragged_batch(tiny_trainer, monkeypatch):
+    # The extra ragged iteration stays inside warmup's snapshot/restore
+    # envelope, so nothing it touches survives into training.
+    trainer = tiny_trainer(
+        n_train=5,
+        n_val=4,
+        config_overrides={"local_batch_size": 2, "warmup_batches": 2},
+    )
+    sizes = _stub_warmup_steps(trainer, monkeypatch, mutate=True)
+    before = {k: v.detach().clone() for k, v in trainer.model.state_dict().items()}
+
+    trainer.warmup()
+
+    assert sizes == [2, 2, 1]
+    after = trainer.model.state_dict()
+    for name, tensor in before.items():
+        assert torch.equal(after[name], tensor), name
diff --git a/tests/test_reporting.py b/tests/test_reporting.py
index a6f0106..af50a04 100644
--- a/tests/test_reporting.py
+++ b/tests/test_reporting.py
@@ -13,11 +13,13 @@
 # SPDX-License-Identifier: (Apache-2.0)
 
 import csv
+import logging
 from pathlib import Path
 from types import SimpleNamespace
 
 import matplotlib
 import numpy as np
+import pytest
 
 matplotlib.use("Agg")
 
@@ -55,6 +57,50 @@ def test_figures_dir_idempotent(self, tmp_path):
         assert (run_dir / "figures" / "train_loss.png").exists()
 
 
+class TestFigureLifetime:
+    """R43: standard_viz closes every figure it opens.
+
+    ``worker.main`` calls ``standard_viz.main`` in-process once per sweep
+    combination, and pyplot keeps a strong reference to every unclosed figure,
+    so the canvases (and their Figure/Axes/Line objects) accumulate for the
+    whole sweep -- 36 live figures / 42 MiB after 12 combinations, plus
+    matplotlib's max_open_warning from the seventh on.
+    """
+
+    def _config(self, tmp_path, name):
+        run_dir = tmp_path / name
+        run_dir.mkdir()
+        (run_dir / "train_stats.csv").write_text(
+            "epoch,overall_loss,val_dice,val_loss_avg\n1,0.9,0.40,0.8\n2,0.5,0.70,0.4\n"
+        )
+        return SimpleNamespace(
+            run_dir=str(run_dir), vol_size=32, n_categories=5, unet_layers=2
+        )
+
+    def test_no_figures_left_open(self, tmp_path):
+        """Repeated calls (a sweep) leave no figure behind."""
+        plt.close("all")
+        for i in range(3):
+            standard_viz.main(self._config(tmp_path, f"run{i}"))
+            assert plt.get_fignums() == [], f"figures leaked after call {i}"
+
+    def test_no_figures_left_open_when_plotting_fails(self, tmp_path, monkeypatch):
+        """A failure between figure() and savefig() does not strand a figure.
+
+        ``main`` logs and swallows plotting errors, so without an unconditional
+        close the leak survives exactly the case it is hardest to notice.
+        """
+        plt.close("all")
+
+        def boom(*args, **kwargs):
+            raise OSError("[Errno 28] No space left on device")
+
+        monkeypatch.setattr(plt, "savefig", boom)
+        standard_viz.main(self._config(tmp_path, "failing_run"))
+
+        assert plt.get_fignums() == []
+
+
 class TestDiceFigure:
     """F70: Validation Dice figure saved as val_dice.png, not val_loss.png."""
 
@@ -270,3 +316,274 @@ def test_torch_profiler_independent_of_caliper(self, monkeypatch):
                 )
         finally:
             importlib.reload(perf_measure)
+
+
+class TestProfilerTraceExport:
+    """R15: a failed trace export must not strand the other ranks."""
+
+    @staticmethod
+    def _unstepped_profiler():
+        """A profiler whose window never opened (a run with zero batches)."""
+        from torch.profiler import ProfilerActivity, profile, schedule
+
+        prof = profile(
+            activities=[ProfilerActivity.CPU],
+            schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
+        )
+        with prof:
+            pass  # no prof.step(): the schedule never leaves its wait phase
+        return prof
+
+    @staticmethod
+    def _stepped_profiler():
+        """A profiler with a completed capture window."""
+        from torch.profiler import ProfilerActivity, profile, schedule
+
+        prof = profile(
+            activities=[ProfilerActivity.CPU],
+            schedule=schedule(wait=1, warmup=1, active=1, repeat=1),
+        )
+        with prof:
+            for _ in range(4):
+                prof.step()
+        return prof
+
+    @staticmethod
+    def _config(run_dir):
+        return SimpleNamespace(
+            problem_scale=4,
+            epochs=1,
+            n_instances_used_per_fractal=2,
+            run_dir=str(run_dir),
+        )
+
+    def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog):
+        """Exporting an unstepped profiler logs an error instead of raising.
+
+        The export runs before the ``dist.barrier()`` that precedes rank-0
+        post-processing, so a raise here kills the profiling rank and leaves
+        every other rank blocked in that barrier until the collective timeout.
+        """
+        import ScaFFold.worker as worker
+
+        log = logging.getLogger("test_zero_step_export")
+        with caplog.at_level(logging.DEBUG, logger=log.name):
+            result = worker.export_profiler_trace(
+                self._unstepped_profiler(),
+                self._config(tmp_path),
+                log,
+                rank=0,
+                world_size=1,
+                ranks_per_node=1,
+            )
+
+        assert result is None
+        messages = " ".join(record.getMessage() for record in caplog.records)
+        assert "trace" in messages.lower()
+
+    def test_successful_export_writes_a_trace(self, tmp_path, caplog):
+        """A profiler with a completed window still writes its trace (control)."""
+        from torch.profiler import ProfilerActivity, profile, schedule
+
+        import ScaFFold.worker as worker
+
+        prof = profile(
+            activities=[ProfilerActivity.CPU],
+            schedule=schedule(wait=1, warmup=1, active=1, repeat=1),
+        )
+        with prof:
+            for _ in range(4):
+                prof.step()
+
+        log = logging.getLogger("test_successful_export")
+        path = worker.export_profiler_trace(
+            prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1
+        )
+
+        assert path is not None
+        assert Path(path).exists()
+
+    def test_trace_lands_in_the_run_dir(self, tmp_path, caplog):
+        """R23: the trace goes to the run dir, not whatever CWD happens to be."""
+        import ScaFFold.worker as worker
+
+        prof = self._stepped_profiler()
+        log = logging.getLogger("test_trace_lands_in_the_run_dir")
+
+        path = worker.export_profiler_trace(
+            prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1
+        )
+
+        assert Path(path).parent == tmp_path
+        assert list(tmp_path.glob("torch-*.json")) == [Path(path)]
+
+    @pytest.mark.parametrize(
+        "world_size, ranks_per_node, expected",
+        [(8, 4, "-N2-n8-"), (6, 4, "-N2-n6-"), (1, 1, "-N1-n1-")],
+        ids=["even", "ragged-last-node", "singleton"],
+    )
+    def test_trace_name_counts_nodes_not_ranks(
+        self, tmp_path, world_size, ranks_per_node, expected
+    ):
+        """R23: the N field is a node count, and never rounds a node away."""
+        import ScaFFold.worker as worker
+
+        prof = self._stepped_profiler()
+        log = logging.getLogger("test_trace_name_counts_nodes")
+
+        path = worker.export_profiler_trace(
+            prof,
+            self._config(tmp_path),
+            log,
+            rank=0,
+            world_size=world_size,
+            ranks_per_node=ranks_per_node,
+        )
+
+        assert expected in Path(path).name
+
+
+class TestProfileTorchGate:
+    """R24: PROFILE_TORCH is parsed like every other profiler flag."""
+
+    @staticmethod
+    def _reload_with(monkeypatch_context, value):
+        import importlib
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        if value is None:
+            monkeypatch_context.delenv("PROFILE_TORCH", raising=False)
+        else:
+            monkeypatch_context.setenv("PROFILE_TORCH", value)
+        monkeypatch_context.delenv("CALI_CONFIG", raising=False)
+        importlib.reload(perf_measure)
+        return perf_measure
+
+    @pytest.mark.parametrize("value", [None, "", "0", "false", "no", "off", "OFF"])
+    def test_disabled_values(self, monkeypatch, value):
+        """Anything that is not an affirmative value leaves profiling off.
+
+        ``PROFILE_TORCH=0`` used to *enable* the profiler: the gate only
+        rejected the literal "off", so every conventional way of saying "no"
+        silently turned profiling on.
+        """
+        import importlib
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            with monkeypatch.context() as m:
+                assert not self._reload_with(m, value).TORCH_PERF_ENABLED
+        finally:
+            importlib.reload(perf_measure)
+
+    @pytest.mark.parametrize("value", ["1", "true", "on", "ON", "yes", "TRUE"])
+    def test_enabled_values(self, monkeypatch, value):
+        """The affirmative spellings still enable the profiler."""
+        import importlib
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            with monkeypatch.context() as m:
+                assert self._reload_with(m, value).TORCH_PERF_ENABLED
+        finally:
+            importlib.reload(perf_measure)
+
+    def test_gate_matches_the_sub_option_parser(self, monkeypatch):
+        """The master switch and the sub-option flags agree on every spelling."""
+        import importlib
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            for value in ("1", "true", "on", "yes", "0", "false", "no", "off", ""):
+                with monkeypatch.context() as m:
+                    module = self._reload_with(m, value)
+                    assert module.TORCH_PERF_ENABLED == module._profiler_env_flag(
+                        "PROFILE_TORCH"
+                    ), value
+        finally:
+            importlib.reload(perf_measure)
+
+
+class TestProfilerSchedule:
+    """R26: the schedule must not record everything before the first step."""
+
+    @staticmethod
+    def _context_with(monkeypatch_context, env):
+        """Reload perf_measure with ``env`` applied and build a profiler context."""
+        import importlib
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        monkeypatch_context.setenv("PROFILE_TORCH", "1")
+        monkeypatch_context.delenv("CALI_CONFIG", raising=False)
+        for name in ("PROFILE_TORCH_WAIT", "PROFILE_TORCH_WARMUP"):
+            monkeypatch_context.delenv(name, raising=False)
+        for key, value in env.items():
+            monkeypatch_context.setenv(key, value)
+        importlib.reload(perf_measure)
+        assert perf_measure.TORCH_PERF_ENABLED
+        ctx, is_local = perf_measure.get_torch_context(1, 0)
+        assert is_local
+        return ctx
+
+    def test_wait_zero_does_not_record_step_zero(self, monkeypatch, caplog):
+        """PROFILE_TORCH_WAIT=0 is clamped so step 0 records nothing.
+
+        worker.main enters the profiler context around checkpoint cleanup and
+        every warmup batch, and ``prof.step()`` only advances once per training
+        batch -- so a schedule that is already active at step 0 buffers all of
+        that as a single unbounded step, which is exactly what the bounded
+        window exists to prevent.
+        """
+        import importlib
+
+        from torch.profiler import ProfilerAction
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            with caplog.at_level(logging.WARNING, logger=perf_measure.logger.name):
+                with monkeypatch.context() as m:
+                    ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"})
+                    assert ctx.schedule(0) == ProfilerAction.NONE
+            messages = " ".join(record.getMessage() for record in caplog.records)
+            assert "PROFILE_TORCH_WAIT" in messages
+        finally:
+            importlib.reload(perf_measure)
+
+    def test_default_schedule_skips_step_zero(self, monkeypatch):
+        """The default window already skips step 0 (control)."""
+        import importlib
+
+        from torch.profiler import ProfilerAction
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            with monkeypatch.context() as m:
+                ctx = self._context_with(m, {})
+                assert ctx.schedule(0) == ProfilerAction.NONE
+        finally:
+            importlib.reload(perf_measure)
+
+    def test_larger_wait_is_preserved(self, monkeypatch):
+        """A wait longer than the minimum is left alone."""
+        import importlib
+
+        from torch.profiler import ProfilerAction
+
+        import ScaFFold.utils.perf_measure as perf_measure
+
+        try:
+            with monkeypatch.context() as m:
+                ctx = self._context_with(
+                    m, {"PROFILE_TORCH_WAIT": "3", "PROFILE_TORCH_WARMUP": "1"}
+                )
+                assert ctx.schedule(2) == ProfilerAction.NONE
+                assert ctx.schedule(3) == ProfilerAction.WARMUP
+        finally:
+            importlib.reload(perf_measure)
diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py
index 1059c7e..291c57a 100644
--- a/tests/test_restart_script.py
+++ b/tests/test_restart_script.py
@@ -25,6 +25,7 @@
 
 from __future__ import annotations
 
+import os
 import shutil
 import subprocess
 import sys
@@ -45,8 +46,10 @@
     "SLURM_JOB_NUM_NODES",
     "SLURM_NNODES",
     "WORLD_SIZE",
+    "MV2_COMM_WORLD_SIZE",
     "OMPI_COMM_WORLD_SIZE",
     "PMI_SIZE",
+    "PALS_NRANKS",
 )
 
 # Profiling variables re-exported only when set in the generating run.
@@ -182,3 +185,171 @@ def test_generated_script_is_valid_bash(monkeypatch, tmp_path):
         assert result.returncode == 0, (
             f"bash -n failed for variant {i}:\n{result.stderr}"
         )
+
+
+# ---------------------------------------------------------------------------
+# R14: combined ``--flag=value`` tokens
+# ---------------------------------------------------------------------------
+
+
+def test_config_equals_form_is_substituted(monkeypatch, tmp_path):
+    """``--config=PATH`` is repointed at the run dir, not left as a placeholder.
+
+    The rewriter emits the placeholder as part of a combined token, so a
+    substitution that only matches whole tokens leaves ``--config=__CFG__`` in
+    the script and the restart dies with "Config file '__CFG__' not found".
+    """
+    _isolate_env(monkeypatch)
+    argv = [
+        "/usr/bin/scaffold",
+        "benchmark",
+        "--config=/some/where/config.yml",
+        "--epochs",
+        "10",
+    ]
+
+    script = _generate(monkeypatch, tmp_path / "run", argv=argv)
+
+    assert "__CFG__" not in script
+    assert '--config="$RUN_DIR/config.yaml"' in script
+    assert "/some/where/config.yml" not in script
+
+
+@pytest.mark.parametrize(
+    "config_argv",
+    [
+        ["-c", "/some/where/config.yml"],
+        ["--config", "/some/where/config.yml"],
+        ["--config=/some/where/config.yml"],
+    ],
+    ids=["short", "long-space", "long-equals"],
+)
+@pytest.mark.skipif(shutil.which("bash") is None, reason="bash not available")
+def test_every_config_spelling_expands_to_the_run_dir_config(
+    monkeypatch, tmp_path, config_argv
+):
+    """Bash expands every ``--config`` spelling to ``$RUN_DIR/config.yaml``.
+
+    The generated PY array is sourced and expanded by a real shell so the test
+    asserts on the arguments the restarted CLI actually receives.
+    """
+    _isolate_env(monkeypatch)
+    argv = ["/usr/bin/scaffold", "benchmark"] + config_argv
+
+    script = _generate(monkeypatch, tmp_path / "run", argv=argv)
+
+    py_decl = next(line for line in script.splitlines() if line.startswith("PY=("))
+    probe = tmp_path / f"probe_{config_argv[0][-1]}.sh"
+    probe.write_text(f'RUN_DIR=/run/dir\n{py_decl}\nprintf "%s\\n" "${{PY[@]}}"\n')
+    result = subprocess.run(
+        ["bash", str(probe)], capture_output=True, text=True, check=True
+    )
+
+    tokens = result.stdout.split("\n")
+    assert "/run/dir/config.yaml" in tokens or (
+        "--config=/run/dir/config.yaml" in tokens
+    )
+    assert not any("__CFG__" in tok for tok in tokens)
+
+
+def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path):
+    """The appended ``--run-dir`` placeholder still resolves (control)."""
+    _isolate_env(monkeypatch)
+
+    script = _generate(monkeypatch, tmp_path / "run")
+
+    assert "__RUN_DIR__" not in script
+    assert '--run-dir "$RUN_DIR"' in script
+
+
+# ---------------------------------------------------------------------------
+# R17: launch-shape sniffing must match the rank side
+# ---------------------------------------------------------------------------
+
+# Every variable ``ScaFFold.utils.distributed.get_world_size`` honors. The
+# restart generator must derive the same world size from each of them, or a
+# restart script silently relaunches the job at the wrong scale.
+_WORLD_SIZE_ENV_VARS = (
+    "WORLD_SIZE",
+    "MV2_COMM_WORLD_SIZE",
+    "OMPI_COMM_WORLD_SIZE",
+    "PMI_SIZE",
+    "PALS_NRANKS",
+    "SLURM_NTASKS",
+    "FLUX_JOB_SIZE",
+)
+
+
+@pytest.mark.parametrize("var", _WORLD_SIZE_ENV_VARS)
+def test_sniffed_world_size_matches_rank_side(monkeypatch, var):
+    """The generator and ``get_world_size`` agree on every launcher variable."""
+    from ScaFFold.utils.distributed import get_world_size
+
+    _isolate_env(monkeypatch)
+    for other in _WORLD_SIZE_ENV_VARS:
+        monkeypatch.delenv(other, raising=False)
+    monkeypatch.setenv(var, "8")
+
+    _, _, sniffed = crs._sniff_launch_shape(os.environ)
+
+    assert sniffed == 8, f"{var} not recognized by the restart generator"
+    assert sniffed == get_world_size()
+
+
+def test_pals_job_gets_a_multirank_restart_script(monkeypatch, tmp_path):
+    """A Cray PALS launch (PALS_NRANKS) emits the multi-rank template."""
+    _isolate_env(monkeypatch)
+    monkeypatch.setenv("PALS_NRANKS", "8")
+
+    script = _generate(monkeypatch, tmp_path / "run")
+
+    assert "torchrun-hpc" in script
+    assert 'exec "${PY[@]}"' not in script
+
+
+# ---------------------------------------------------------------------------
+# VC-5: the node shape comes from the local rank count when no scheduler
+# reports one. Flux and Slurm state their node count; PALS and plain torchrun
+# do not, and assuming one node relaunched an 8-rank/2-node job as
+# NODES=1 TASKS_PER_NODE=8 -- an oversubscribed node, or a rejected job.
+# ---------------------------------------------------------------------------
+
+
+def test_pals_multinode_shape_uses_the_local_rank_count(monkeypatch, tmp_path):
+    """8 PALS ranks, 4 per node -> NODES=2, TASKS_PER_NODE=4."""
+    _isolate_env(monkeypatch)
+    monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False)
+    monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False)
+    monkeypatch.setenv("PALS_NRANKS", "8")
+    monkeypatch.setenv("PALS_LOCAL_SIZE", "4")
+
+    script = _generate(monkeypatch, tmp_path / "run")
+
+    assert 'NODES="2"' in script
+    assert 'TASKS_PER_NODE="4"' in script
+
+
+def test_single_node_torchrun_shape_is_unchanged(monkeypatch, tmp_path):
+    """8 torchrun ranks all on one node stay NODES=1, TASKS_PER_NODE=8."""
+    _isolate_env(monkeypatch)
+    monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False)
+    monkeypatch.setenv("WORLD_SIZE", "8")
+    monkeypatch.setenv("LOCAL_WORLD_SIZE", "8")
+
+    script = _generate(monkeypatch, tmp_path / "run")
+
+    assert 'NODES="1"' in script
+    assert 'TASKS_PER_NODE="8"' in script
+
+
+def test_unknown_local_size_keeps_the_single_node_assumption(monkeypatch, tmp_path):
+    """With nothing reporting a per-node count, the old assumption stands."""
+    _isolate_env(monkeypatch)
+    for var in ("PALS_LOCAL_SIZE", "LOCAL_WORLD_SIZE", "PMI_LOCAL_SIZE"):
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("PALS_NRANKS", "8")
+
+    script = _generate(monkeypatch, tmp_path / "run")
+
+    assert 'NODES="1"' in script
+    assert 'TASKS_PER_NODE="8"' in script
diff --git a/tests/test_resume.py b/tests/test_resume.py
index 209a6eb..dfe7b2c 100644
--- a/tests/test_resume.py
+++ b/tests/test_resume.py
@@ -28,6 +28,7 @@
 
 from __future__ import annotations
 
+import logging
 from pathlib import Path
 
 import numpy as np
@@ -132,6 +133,7 @@ def cfg():
 
 from types import SimpleNamespace  # noqa: E402
 
+import ScaFFold.utils.trainer as trainer_mod  # noqa: E402
 from ScaFFold.utils.checkpointing import CheckpointManager  # noqa: E402
 from ScaFFold.utils.trainer import PyTorchTrainer  # noqa: E402
 
@@ -142,19 +144,27 @@ def cfg():
 )
 
 
-def _stub_trainer(run_dir, *, train_from_scratch, log):
+def _stub_trainer(run_dir, *, train_from_scratch, log, **config_overrides):
     """A PyTorchTrainer carrying only what cleanup_or_resume/train touch.
 
     Built via ``object.__new__`` (as the checkpointing tests do) so no dataset,
     model, or process group is needed. A real CheckpointManager over a tiny
     linear model backs the resume path so load/save round-trip faithfully.
+
+    ``config_overrides`` set additional config fields (``epochs``,
+    ``target_dice``, ``checkpoint_interval``, ...) for tests that drive
+    ``train()``'s loop-entry and final-save control flow.
     """
     t = object.__new__(PyTorchTrainer)
-    t.config = SimpleNamespace(
-        train_from_scratch=train_from_scratch,
-        run_dir=str(run_dir),
-        checkpoint_interval=-1,
-    )
+    config_fields = {
+        "train_from_scratch": train_from_scratch,
+        "run_dir": str(run_dir),
+        "checkpoint_interval": -1,
+        "epochs": -1,
+        "target_dice": 0.95,
+    }
+    config_fields.update(config_overrides)
+    t.config = SimpleNamespace(**config_fields)
     t.world_rank = 0
     t.global_step = 0
     t.total_optimizer_steps = 0
@@ -346,6 +356,160 @@ def test_step_counters_roundtrip(tmp_path):
     assert t2.total_optimizer_steps == 37
 
 
+def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog):
+    """Restarting a run whose checkpoint already covers ``epochs`` exits cleanly.
+
+    The epoch loop breaks at the max-epoch check before any new epoch runs, so
+    no epoch metric exists to checkpoint: the final-save block must be skipped
+    (the last checkpoint already covers the completed epochs) rather than
+    saving with an unbound ``val_loss_avg``. ``train()`` has to return normally
+    so the worker's post-processing still runs off the CSV already on disk, and
+    the run must say plainly that it had no epoch left to run.
+    """
+    log = logging.getLogger("resume.r01")
+    run = tmp_path / "run"
+    run.mkdir()
+
+    # Artifacts of a completed epochs=2 run: a checkpoint recording epoch 2
+    # (written by the in-loop save) plus its CSV rows.
+    t1 = _stub_trainer(
+        run,
+        train_from_scratch=False,
+        log=log,
+        epochs=2,
+        checkpoint_interval=1,
+    )
+    t1.checkpoint_manager.save_checkpoint(
+        epoch=2,
+        val_loss_avg=0.5,
+        extras={
+            "train_mask_values": None,
+            "global_step": 8,
+            "total_optimizer_steps": 8,
+        },
+    )
+    csv = run / "train_stats.csv"
+    csv.write_text(_HEADER + "\n")
+    _write_rows(csv, [1, 2])
+
+    ckpt = t1.checkpoint_manager.last_ckpt_path
+    before = ckpt.read_bytes()
+
+    # The user reruns the generated restart command against the same dir.
+    t2 = _stub_trainer(
+        run,
+        train_from_scratch=False,
+        log=log,
+        epochs=2,
+        checkpoint_interval=1,
+    )
+    t2.config.restart = True
+    t2.cleanup_or_resume()
+    assert t2.start_epoch == 3  # past the last epoch: nothing left to train
+
+    with caplog.at_level(logging.WARNING):
+        t2.train()  # must not raise
+
+    # Nothing new was written: the existing checkpoint is byte-identical (a
+    # fresh save would serialize this trainer's own randomly-initialized model)
+    # and the CSV still holds exactly the original run's epochs.
+    assert ckpt.read_bytes() == before
+    epochs = [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]]
+    assert epochs == ["1", "2"]
+    assert "no epoch left to run" in caplog.text.lower()
+
+
+def test_fresh_run_with_no_epochs_does_not_claim_a_resume(tmp_path, caplog):
+    """A fresh ``epochs: 0`` run reports the truth: nothing was resumed.
+
+    The same message covers both ways of entering ``train()`` with no epoch to
+    run, so it must not describe one of them as the other. This run has no
+    checkpoint and never asked for one -- telling its user "there was nothing
+    to resume" sends them looking for a checkpoint that was never part of the
+    story.
+    """
+    log = logging.getLogger("resume.va4")
+    run = tmp_path / "run"
+    run.mkdir()
+
+    trainer = _stub_trainer(
+        run,
+        train_from_scratch=True,
+        log=log,
+        epochs=0,
+        checkpoint_interval=1,
+    )
+    trainer.cleanup_or_resume()
+    assert trainer.start_epoch == 1  # a fresh run: nothing was resumed
+
+    with caplog.at_level(logging.WARNING):
+        trainer.train()
+
+    assert "no epoch left to run" in caplog.text.lower()
+    assert "nothing to resume" not in caplog.text.lower()
+    assert not trainer.checkpoint_manager.last_ckpt_path.exists()
+
+
+def test_converged_resume_does_not_retrain(tiny_trainer, monkeypatch):
+    """A restart of an already-converged run must not train another epoch.
+
+    The validation dice the checkpointed epoch achieved is persisted with the
+    checkpoint and restored on resume, so the ``while dice_score_train <
+    target_dice`` loop is never entered. Without it a converged ``epochs: -1``
+    run re-trains, re-logs and re-checkpoints one full epoch on every restart,
+    inflating ``sum(epoch_duration)`` -- the FOM denominator -- and the
+    reported epoch count relative to the same run left un-restarted.
+    """
+    overrides = {
+        "checkpoint_interval": 1,
+        "epochs": -1,
+        "target_dice": 0.95,
+        "train_from_scratch": 0,
+    }
+
+    def converged_evaluate(*args, **kwargs):
+        # Two validation samples at hard dice 0.96, i.e. above target.
+        return (0.96 * 2, 0.1 * 2, 0.1, 2, 2)
+
+    monkeypatch.setattr(trainer_mod, "evaluate", converged_evaluate)
+
+    def stub_batches(trainer):
+        """Replace the DistConv forward path and count the batches it runs."""
+        calls = {"n": 0}
+
+        def _step(batch, **kwargs):
+            calls["n"] += 1
+            return 1, torch.tensor(0.1), torch.tensor(0.96)
+
+        monkeypatch.setattr(trainer, "_run_training_batch", _step)
+        return calls
+
+    # The original run: converges in epoch 1 and checkpoints it.
+    first = tiny_trainer(config_overrides=overrides)
+    first_calls = stub_batches(first)
+    first.cleanup_or_resume()
+    first.train()
+    assert first_calls["n"] > 0  # it really did train
+
+    ckpt_path = first.checkpoint_manager.last_ckpt_path
+    saved = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    assert saved["epoch"] == 1
+    csv = Path(first.outfile_path)
+    assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"]
+
+    # The user reruns the restart command: nothing is left to train.
+    second = tiny_trainer(config_overrides=overrides)
+    second_calls = stub_batches(second)
+    second.cleanup_or_resume()
+    assert second.start_epoch == 2
+    second.train()
+
+    assert second_calls["n"] == 0, "a converged run re-trained an extra epoch"
+    assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"]
+    reloaded = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    assert reloaded["epoch"] == 1
+
+
 def test_total_steps_resume_predates_dedicated_key(tmp_path):
     """A checkpoint recording only global_step still resumes the step total.
 
diff --git a/tests/test_triton_group_norm.py b/tests/test_triton_group_norm.py
new file mode 100644
index 0000000..ff4d0ee
--- /dev/null
+++ b/tests/test_triton_group_norm.py
@@ -0,0 +1,868 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Tests for the channels-last Triton GroupNorm (``ScaFFold.unet.triton_group_norm``).
+
+The kernel replaces a stock op, so almost every test here is a *parity* test:
+values and gradients against ``F.group_norm``, but with the reference computed
+in **float64** rather than against another fp32 result -- an fp32-vs-fp32
+comparison cannot tell a correct kernel from one that has merely made the same
+mistake, and it cannot see the variance-formula failure the Welford
+implementation exists to fix (``test_welford_survives_large_mean``).  Each
+parity test reports the measured relative error so a regression shows up as a
+number, not just a boolean.
+
+The other three things being pinned down:
+
+* the **contract** -- output dtype exactly matches ``F.group_norm``'s
+  (including its fp32 autocast policy), output *memory format* matches the
+  input's (which is where the kernel deliberately differs from stock, and the
+  entire reason it exists), and ``is_supported`` accepts exactly the inputs the
+  native kernel serves;
+* **determinism** -- the same call twice is bitwise identical, forward and
+  backward, because the split count and tiling are pure functions of the shape;
+* **composition** -- the op is a real dispatcher op, so it must survive
+  ``torch.compile(fullgraph=True)`` without a graph break and a ``DCTensor``
+  round trip through ``__torch_dispatch__`` with the autograd graph intact.
+
+CPU runs never touch Triton: the module defers ``import triton`` to the first
+call that reaches a kernel, which ``test_import_does_not_pull_in_triton``
+checks in a fresh interpreter.
+"""
+
+import os
+import subprocess
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from ScaFFold.unet import triton_group_norm as tgn
+from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm
+
+CL = torch.channels_last_3d
+GROUPS = 8
+EPS = 1e-5
+
+#: Relative-error ceilings against a float64 reference, by input dtype.  The
+#: fp32 numbers observed on MI300A at these (small) test shapes are ~2e-07 for
+#: y/dx/dweight/dbias; the ceiling leaves room for the 3e-05 that a 134M-element
+#: fp32 reduction shows at the largest production shape.  The low-precision
+#: ceilings are set just above the output's own rounding: 2^-8 for bf16 and
+#: 2^-11 for fp16, measured 3.3e-03 and 3.8e-04.
+_TOL = {
+    torch.float32: 1e-4,
+    torch.bfloat16: 2e-2,
+    torch.float16: 3e-3,
+}
+
+
+# ---------------------------------------------------------------------------
+# helpers
+# ---------------------------------------------------------------------------
+
+
+def _rel(actual, expected):
+    """max|actual - expected| / max|expected|, computed in float64."""
+    a = actual.detach().double()
+    e = expected.detach().double()
+    scale = e.abs().max().clamp_min(1e-30)
+    return ((a - e).abs().max() / scale).item()
+
+
+def _tensors(shape, dtype, device, affine=True, seed=0, mean=0.0, std=1.0):
+    """Channels-last input plus (optionally) affine parameters and a cotangent."""
+    gen = torch.Generator(device=device).manual_seed(seed)
+    x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL)
+    x.normal_(mean, std, generator=gen)
+    channels = shape[1]
+    if affine:
+        weight = torch.empty(channels, device=device, dtype=dtype)
+        weight.normal_(1.0, 0.25, generator=gen)
+        bias = torch.empty(channels, device=device, dtype=dtype)
+        bias.normal_(0.0, 0.25, generator=gen)
+    else:
+        weight = bias = None
+    grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL)
+    grad_out.normal_(generator=gen)
+    return x, weight, bias, grad_out
+
+
+def _run(fn, x, weight, bias, grad_out, activation=None, eps=EPS):
+    """Forward + backward through ``fn``, returning detached results."""
+    x = x.detach().clone().requires_grad_(True)
+    weight = None if weight is None else weight.detach().clone().requires_grad_(True)
+    bias = None if bias is None else bias.detach().clone().requires_grad_(True)
+    out = fn(x, GROUPS, weight, bias, eps, activation)
+    out.backward(grad_out.to(out.dtype))
+    return (
+        out.detach(),
+        x.grad,
+        None if weight is None else weight.grad,
+        None if bias is None else bias.grad,
+    )
+
+
+def _reference(x, weight, bias, grad_out, activation=None, eps=EPS):
+    """``F.group_norm`` (+ optional ReLU) evaluated entirely in float64."""
+
+    def fn(x, groups, weight, bias, eps, activation):
+        out = F.group_norm(x, groups, weight, bias, eps)
+        return F.relu(out) if activation == "relu" else out
+
+    return _run(
+        fn,
+        x.double(),
+        None if weight is None else weight.double(),
+        None if bias is None else bias.double(),
+        grad_out.double(),
+        activation,
+        eps,
+    )
+
+
+def _assert_parity(got, ref, dtype, label, tol=None):
+    """Compare (y, dx, dweight, dbias) against the float64 reference."""
+    tol = _TOL[dtype] if tol is None else tol
+    errors = {}
+    for name, a, e in zip(("y", "dx", "dweight", "dbias"), got, ref):
+        if a is None:
+            assert e is None or True  # no parameter -> no gradient to compare
+            continue
+        errors[name] = _rel(a, e)
+    print(
+        f"[{label}] "
+        + " ".join(f"{k}={v:.3e}" for k, v in errors.items())
+        + f"  (tol {tol:.1e})"
+    )
+    for name, err in errors.items():
+        assert err <= tol, f"{label}: {name} relative error {err:.3e} > {tol:.1e}"
+    return errors
+
+
+def _cuda_shapes():
+    """Shapes covering N>1, non-power-of-two extents and a wide channel count."""
+    return [
+        (1, 64, 8, 8, 8),  # the canonical UNet shape, shrunk
+        (2, 64, 9, 7, 5),  # N>1, all three extents non-power-of-two
+        (1, 128, 5, 6, 7),
+        (3, 256, 4, 4, 4),
+        (1, 2048, 6, 6, 6),  # widest UNet channel count
+    ]
+
+
+# ---------------------------------------------------------------------------
+# CPU-only behaviour (no Triton, no GPU)
+# ---------------------------------------------------------------------------
+
+
+def test_import_does_not_pull_in_triton():
+    """Importing the module must not import Triton.
+
+    Run in a fresh interpreter because any earlier GPU test in this session
+    would already have built the kernels.  The guarantee matters twice over: a
+    CPU-only unit run must not pay Triton's import, and the module must stay
+    importable on a build that has no Triton at all.
+    """
+    script = (
+        "import sys; import ScaFFold.unet.triton_group_norm as m; "
+        "assert m.tl is None, 'kernels built at import time'; "
+        "print('triton' in sys.modules)"
+    )
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        capture_output=True,
+        text=True,
+        cwd=repo_root,
+        timeout=300,
+    )
+    assert result.returncode == 0, result.stderr
+    assert result.stdout.strip() == "False", (
+        f"triton was imported at module import time: {result.stdout!r}"
+    )
+
+
+def test_cpu_input_falls_back_bitwise():
+    """A CPU tensor is not supported, and the fallback is the stock kernel."""
+    gen = torch.Generator().manual_seed(0)
+    x = torch.randn(1, 64, 4, 4, 4, generator=gen).requires_grad_(True)
+    weight = torch.randn(64, generator=gen).requires_grad_(True)
+    bias = torch.randn(64, generator=gen).requires_grad_(True)
+    assert is_supported(x, GROUPS, weight, bias) is False
+
+    out = triton_group_norm(x, GROUPS, weight, bias, EPS)
+    assert torch.equal(out, F.group_norm(x, GROUPS, weight, bias, EPS))
+    out.pow(2).sum().backward()
+    assert x.grad is not None and weight.grad is not None and bias.grad is not None
+
+
+def test_cpu_fused_relu_falls_back_bitwise():
+    gen = torch.Generator().manual_seed(1)
+    x = torch.randn(2, 32, 3, 4, 5, generator=gen)
+    weight = torch.randn(32, generator=gen)
+    bias = torch.randn(32, generator=gen)
+    got = triton_group_norm(x, 8, weight, bias, EPS, "relu")
+    assert torch.equal(got, F.relu(F.group_norm(x, 8, weight, bias, EPS)))
+
+
+def test_unknown_activation_raises():
+    x = torch.randn(1, 8, 2, 2, 2)
+    with pytest.raises(ValueError, match="activation"):
+        triton_group_norm(x, 2, activation="gelu")
+    assert is_supported(x, 2, activation="gelu") is False
+
+
+def test_select_strategy_is_a_pure_function_of_shape():
+    """The dispatch hook must be deterministic -- the reduction order, and so
+    the bits of the result, depend on it."""
+    for args in ((1, 64, 8**3, 8), (2, 2048, 16**3, 8), (1, 128, 7 * 5 * 3, 4)):
+        first = tgn.select_strategy(*args)
+        assert first in tgn.STRATEGIES
+        assert all(tgn.select_strategy(*args) == first for _ in range(3))
+
+
+def test_tuning_table_covers_the_scale8_shapes():
+    """The frozen table is what makes the kernel reproducible; keep it honest."""
+    for channels, edge in ((64, 256), (128, 128), (256, 64), (512, 32), (1024, 16)):
+        assert tgn.default_config(channels, edge**3) is tgn._TUNED[(channels, edge)]
+    # An unlisted shape falls back to the generic config rather than failing.
+    assert tgn.default_config(96, 11**3) == tgn.GNConfig()
+
+
+def test_plan_depends_only_on_shape():
+    """Two plans for the same shape must be identical objects of identical
+    content, or the split count could drift between calls and break bitwise
+    reproducibility."""
+    a = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3)
+    b = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3)
+    assert (a.nsplit, a.chunk, a.block_s_stats, a.block_s_elem, a.int64) == (
+        b.nsplit,
+        b.chunk,
+        b.block_s_stats,
+        b.block_s_elem,
+        b.int64,
+    )
+    # int64 addressing turns on exactly when a linear index can overflow int32.
+    small = tgn._plan(1, 64, 128**3, 8, 64 * 128**3)
+    big = tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3)
+    assert small.int64 is False
+    assert big.int64 is True
+
+
+# ---------------------------------------------------------------------------
+# is_supported
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+def test_is_supported_accepts_the_fast_path():
+    device = torch.device("cuda")
+    x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_()
+    weight = torch.randn(64, device=device)
+    assert is_supported(x, GROUPS, weight, weight) is True
+    assert is_supported(x, GROUPS) is True
+    assert is_supported(x, GROUPS, activation="relu") is True
+
+
+@pytest.mark.gpu
+def test_is_supported_rejections():
+    """Everything ``is_supported`` rejects must be something a caller can hand
+    to ``F.group_norm`` instead -- so the rejections are the contract's edge."""
+    device = torch.device("cuda")
+    cl = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_()
+    cases = {
+        "cpu tensor": (torch.randn(1, 64, 6, 6, 6), GROUPS, None, None, None),
+        "contiguous (NCDHW)": (
+            torch.randn(1, 64, 6, 6, 6, device=device),
+            GROUPS,
+            None,
+            None,
+            None,
+        ),
+        "float64": (
+            torch.empty(
+                1, 64, 6, 6, 6, device=device, dtype=torch.float64, memory_format=CL
+            ),
+            GROUPS,
+            None,
+            None,
+            None,
+        ),
+        "4-D": (torch.randn(1, 64, 6, 6, device=device), GROUPS, None, None, None),
+        "channels not divisible": (cl, 7, None, None, None),
+        "num_groups=0": (cl, 0, None, None, None),
+        "bad activation": (cl, GROUPS, None, None, "gelu"),
+        "weight wrong size": (cl, GROUPS, torch.randn(32, device=device), None, None),
+        "weight on cpu": (cl, GROUPS, torch.randn(64), None, None),
+        "weight dtype mismatch": (
+            cl,
+            GROUPS,
+            torch.randn(64, device=device, dtype=torch.bfloat16),
+            None,
+            None,
+        ),
+        "sliced (non-contiguous)": (
+            torch.empty(1, 64, 6, 6, 12, device=device, memory_format=CL)[..., ::2],
+            GROUPS,
+            None,
+            None,
+            None,
+        ),
+        "not a tensor": (None, GROUPS, None, None, None),
+    }
+    for label, args in cases.items():
+        assert is_supported(*args) is False, f"{label} should be rejected"
+
+
+@pytest.mark.gpu
+def test_rejected_inputs_still_produce_stock_results():
+    """``triton_group_norm`` stays total: rejects go to ``F.group_norm``."""
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(4)
+    x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen).requires_grad_(True)
+    weight = torch.randn(64, device=device, generator=gen).requires_grad_(True)
+    bias = torch.randn(64, device=device, generator=gen).requires_grad_(True)
+    assert is_supported(x, GROUPS, weight, bias) is False
+    got = triton_group_norm(x, GROUPS, weight, bias, EPS)
+    assert torch.equal(got, F.group_norm(x, GROUPS, weight, bias, EPS))
+
+
+# ---------------------------------------------------------------------------
+# value / gradient parity against float64
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape", _cuda_shapes())
+def test_parity_fp32(shape):
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors(
+        shape, torch.float32, device, seed=hash(shape) % 1000
+    )
+    assert is_supported(x, GROUPS, weight, bias)
+    got = _run(triton_group_norm, x, weight, bias, grad_out)
+    ref = _reference(x, weight, bias, grad_out)
+    _assert_parity(got, ref, torch.float32, f"fp32 {shape}")
+    assert got[0].is_contiguous(memory_format=CL)
+    assert got[1].is_contiguous(memory_format=CL)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_parity_every_dtype_and_activation(dtype, activation):
+    device = torch.device("cuda")
+    shape = (2, 128, 7, 6, 5)
+    x, weight, bias, grad_out = _tensors(shape, dtype, device, seed=11)
+    assert is_supported(x, GROUPS, weight, bias, activation)
+    got = _run(triton_group_norm, x, weight, bias, grad_out, activation)
+    ref = _reference(x, weight, bias, grad_out, activation)
+    _assert_parity(got, ref, dtype, f"{dtype} act={activation}")
+    assert got[0].dtype == dtype
+    assert got[1].dtype == dtype
+    assert got[2].dtype == dtype and got[3].dtype == dtype
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("affine", ["both", "weight_only", "bias_only", "neither"])
+def test_parity_without_affine_parameters(affine):
+    """``weight=None`` / ``bias=None`` are separate kernel constexpr paths."""
+    device = torch.device("cuda")
+    shape = (2, 64, 5, 5, 5)
+    x, weight, bias, grad_out = _tensors(shape, torch.float32, device, seed=19)
+    if affine in ("bias_only", "neither"):
+        weight = None
+    if affine in ("weight_only", "neither"):
+        bias = None
+    assert is_supported(x, GROUPS, weight, bias)
+    got = _run(triton_group_norm, x, weight, bias, grad_out)
+
+    reference_weight = weight
+    if weight is None and bias is not None:
+        # Upstream limitation, not a difference in this kernel:
+        # ``F.group_norm(x, g, None, bias).backward()`` raises "tensor does not
+        # have a device" on both CPU and CUDA (torch 2.13.0+rocm7.2), so the
+        # float64 reference has to spell the same computation with weight=1.
+        # This kernel handles the combination directly.
+        reference_weight = torch.ones_like(bias)
+    ref = _reference(x, reference_weight, bias, grad_out)
+    # ``_assert_parity`` skips outputs this configuration does not produce.
+    _assert_parity(got, ref, torch.float32, f"affine={affine}")
+
+
+@pytest.mark.gpu
+def test_partial_gradient_requirements():
+    """Only some inputs requiring grad must not change the ones that do."""
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors(
+        (1, 64, 5, 5, 5), torch.float32, device, seed=23
+    )
+    full = _run(triton_group_norm, x, weight, bias, grad_out)
+
+    frozen_w = weight.detach().clone()
+    frozen_b = bias.detach().clone()
+    xi = x.detach().clone().requires_grad_(True)
+    out = triton_group_norm(xi, GROUPS, frozen_w, frozen_b, EPS)
+    out.backward(grad_out)
+    assert torch.equal(xi.grad, full[1])
+    assert frozen_w.grad is None and frozen_b.grad is None
+
+    # ... and the mirror image: parameters only.
+    xn = x.detach().clone()
+    wn = weight.detach().clone().requires_grad_(True)
+    bn = bias.detach().clone().requires_grad_(True)
+    triton_group_norm(xn, GROUPS, wn, bn, EPS).backward(grad_out)
+    assert torch.equal(wn.grad, full[2])
+    assert torch.equal(bn.grad, full[3])
+
+
+# ---------------------------------------------------------------------------
+# fused activation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_fused_relu_matches_unfused(dtype):
+    """The fused store and ``F.relu`` on the unfused output must agree bitwise,
+    forward *and* backward -- the backward recomputes the pre-activation rather
+    than reading it back, so this is the test that recomputation is exact."""
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors((2, 64, 6, 7, 8), dtype, device, seed=29)
+    fused = _run(triton_group_norm, x, weight, bias, grad_out, "relu")
+
+    def unfused(x, groups, weight, bias, eps, _activation):
+        return F.relu(triton_group_norm(x, groups, weight, bias, eps))
+
+    separate = _run(unfused, x, weight, bias, grad_out)
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), fused, separate):
+        assert torch.equal(a, b), f"fused vs unfused+relu differ in {name}"
+    # A ReLU that never fires would make this test vacuous.
+    assert (fused[0] == 0).any() and (fused[0] > 0).any()
+
+
+# ---------------------------------------------------------------------------
+# determinism
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_bitwise_determinism(activation):
+    """Same input twice => bitwise-equal output and gradients.
+
+    Guaranteed by construction (no float atomics; grid, split count and tile
+    sizes are pure functions of the shape) and asserted here because a future
+    run-time autotuner would silently break it.
+    """
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors(
+        (2, 128, 9, 11, 13), torch.float32, device, seed=31
+    )
+    first = _run(triton_group_norm, x, weight, bias, grad_out, activation)
+    second = _run(triton_group_norm, x, weight, bias, grad_out, activation)
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), first, second):
+        assert torch.equal(a, b), f"{name} is not bitwise reproducible"
+
+
+# ---------------------------------------------------------------------------
+# numerics: Welford vs E[x^2] - E[x]^2
+# ---------------------------------------------------------------------------
+
+
+def _naive_group_norm(x, num_groups, weight, bias, eps):
+    """The prototype's variance formula, reproduced in fp32 torch ops.
+
+    ``var = E[x^2] - E[x]^2`` is split-friendly and cheap, and it is what the
+    kernel used before the Welford rewrite; this is the thing the test below
+    must show is broken so that "the new one passes" means something.
+    """
+    n, channels = x.shape[0], x.shape[1]
+    flat = x.reshape(n, num_groups, -1)
+    mean = flat.mean(-1)
+    mean_sq = (flat * flat).mean(-1)
+    var = mean_sq - mean * mean
+    rstd = 1.0 / torch.sqrt(var + eps)
+    out = (flat - mean[..., None]) * rstd[..., None]
+    out = out.reshape(x.shape)
+    shape = (1, channels) + (1,) * (x.dim() - 2)
+    return out * weight.reshape(shape) + bias.reshape(shape)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "mean,std,naive_floor",
+    [
+        (0.0, 1.0, None),  # both formulations are fine here
+        (100.0, 1.0, 1e-4),  # E[x^2]-E[x]^2 already an order of magnitude off
+        (1e3, 1e-2, 1e-1),  # ... and here it has lost the variance entirely
+    ],
+)
+def test_welford_survives_large_mean(mean, std, naive_floor):
+    """Large-mean / small-variance input: the regression case for the rewrite.
+
+    Measured here on MI300A at ``[1, 256, 24^3]`` (relative error of the output
+    against a float64 reference computed from the same fp32 samples), with the
+    production shape ``[1, 256, 64^3]`` in parentheses:
+
+        mean=0,   std=1     this 1.3e-07 (1.6e-07)  E[x^2]-E[x]^2 1.4e-07 (1.8e-07)
+        mean=1e2, std=1     this 1.1e-06 (8.0e-07)  E[x^2]-E[x]^2 9.5e-04 (5.6e-04)
+        mean=1e3, std=1e-2  this 4.4e-04 (1.1e-04)  E[x^2]-E[x]^2 2.3e+00 (2.3e+00)
+
+    i.e. at ``mean/std = 1e5`` the old formulation loses the variance outright
+    (the difference of the two ~1e6-sized fp32 terms is below one ulp, so
+    ``rstd`` saturates on ``eps`` and the output is meaningless) while Welford
+    is still correct to 4.4e-04 -- itself *5x better* than ATen's own fp32
+    GroupNorm on the same input (2.2e-03), and dominated by the fp32
+    representation of a mean of 1e3 rather than by anything the kernel does.
+    """
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors(
+        (1, 256, 24, 24, 24), torch.float32, device, seed=37, mean=mean, std=std
+    )
+    got = triton_group_norm(x, GROUPS, weight, bias, EPS)
+    ref = F.group_norm(x.double(), GROUPS, weight.double(), bias.double(), EPS)
+    stock = F.group_norm(x, GROUPS, weight, bias, EPS)
+    naive = _naive_group_norm(x, GROUPS, weight, bias, EPS)
+
+    err = _rel(got, ref)
+    err_stock = _rel(stock, ref)
+    err_naive = _rel(naive, ref)
+    print(
+        f"[welford mean={mean:g} std={std:g}] triton={err:.3e} "
+        f"aten_fp32={err_stock:.3e} naive_Ex2={err_naive:.3e}"
+    )
+    # Never worse than ATen's own fp32 kernel by more than a small factor.
+    assert err <= max(4.0 * err_stock, 1e-5), (
+        f"triton {err:.3e} vs aten fp32 {err_stock:.3e}"
+    )
+    if naive_floor is not None:
+        assert err_naive > naive_floor, (
+            "the naive formulation was expected to fail here "
+            f"but only reached {err_naive:.3e}"
+        )
+        assert err < err_naive / 10.0, (
+            f"Welford ({err:.3e}) is not clearly better than "
+            f"E[x^2]-E[x]^2 ({err_naive:.3e})"
+        )
+
+
+# ---------------------------------------------------------------------------
+# dtypes, layouts, autocast
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("autocast_dtype", [None, torch.bfloat16, torch.float16])
+def test_output_dtype_matches_stock(dtype, autocast_dtype):
+    """The dtype contract, including autocast's fp32 policy for GroupNorm.
+
+    Stock behaviour on this build (measured, not assumed): without autocast the
+    output dtype is the input dtype; inside *any* enabled CUDA autocast region
+    it is fp32, because ``at::group_norm`` carries the fp32 cast policy.
+    """
+    device = torch.device("cuda")
+    x, weight, bias, _ = _tensors((1, 64, 5, 5, 5), dtype, device, seed=41)
+    if autocast_dtype is not None:
+        # Autocast casts the parameters itself, and production keeps them fp32.
+        weight = weight.float()
+        bias = bias.float()
+    ctx = (
+        torch.autocast("cuda", dtype=autocast_dtype)
+        if autocast_dtype is not None
+        else torch.autocast("cuda", enabled=False)
+    )
+    with ctx:
+        assert is_supported(x, GROUPS, weight, bias)
+        stock = F.group_norm(x, GROUPS, weight, bias, EPS)
+        got = triton_group_norm(x, GROUPS, weight, bias, EPS)
+    expected = torch.float32 if autocast_dtype is not None else dtype
+    assert stock.dtype == expected, "assumption about stock GroupNorm broke"
+    assert got.dtype == stock.dtype
+    # ... and the one deliberate difference: stock always returns contiguous.
+    assert stock.is_contiguous() and not stock.is_contiguous(memory_format=CL)
+    assert got.is_contiguous(memory_format=CL)
+    print(
+        f"[dtype in={dtype} autocast={autocast_dtype}] "
+        f"stock={stock.dtype}/CONT mine={got.dtype}/CL rel={_rel(got, stock):.3e}"
+    )
+    assert _rel(got.float(), stock.float()) <= _TOL[stock.dtype]
+
+
+@pytest.mark.gpu
+def test_autocast_gradient_dtypes_match_stock():
+    """Under autocast, ``d_input`` keeps the input's dtype and the parameter
+    gradients stay fp32 -- exactly what the cast nodes around stock GroupNorm
+    produce."""
+    device = torch.device("cuda")
+    x, _, _, grad_out = _tensors((1, 64, 5, 5, 5), torch.bfloat16, device, seed=43)
+    weight = torch.randn(64, device=device, requires_grad=True)
+    bias = torch.randn(64, device=device, requires_grad=True)
+
+    def run(fn):
+        xi = x.detach().clone().requires_grad_(True)
+        w = weight.detach().clone().requires_grad_(True)
+        b = bias.detach().clone().requires_grad_(True)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            out = fn(xi, GROUPS, w, b, EPS)
+        out.backward(grad_out.to(out.dtype))
+        return out, xi.grad, w.grad, b.grad
+
+    stock = run(F.group_norm)
+    got = run(triton_group_norm)
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, stock):
+        assert a.dtype == b.dtype, f"{name}: {a.dtype} != {b.dtype}"
+        assert _rel(a.float(), b.float()) <= 5e-2, name
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("layout", ["channels_last_3d", "contiguous"])
+def test_memory_format_is_preserved(layout):
+    """Both layouts round-trip their own format; contiguous input takes the
+    documented ``F.group_norm`` fallback rather than silently changing layout."""
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(47)
+    x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen)
+    if layout == "channels_last_3d":
+        x = x.contiguous(memory_format=CL)
+    weight = torch.randn(64, device=device, generator=gen)
+    bias = torch.randn(64, device=device, generator=gen)
+    grad_out = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen)
+    if layout == "channels_last_3d":
+        grad_out = grad_out.contiguous(memory_format=CL)
+
+    assert is_supported(x, GROUPS, weight, bias) is (layout == "channels_last_3d")
+    got = _run(triton_group_norm, x, weight, bias, grad_out)
+    ref = _reference(x, weight, bias, grad_out)
+    _assert_parity(got, ref, torch.float32, f"layout={layout}")
+    if layout == "channels_last_3d":
+        assert got[0].is_contiguous(memory_format=CL)
+        assert got[1].is_contiguous(memory_format=CL)
+    else:
+        assert got[0].is_contiguous()
+        assert got[1].is_contiguous()
+
+
+# ---------------------------------------------------------------------------
+# int64 offsets
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+def test_int64_switch_flips_at_int32_max():
+    """The switch is a pure function of the element count, so pin the boundary.
+
+    ``[2, 64, 256^3]`` is *exactly* 2^31 elements: the shape that made an
+    int64 path mandatory before batch>1 or scale 16.
+    """
+    assert tgn._plan(1, 64, 255**3, 8, 64 * 255**3).int64 is False
+    assert tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3).int64 is True
+
+
+@pytest.mark.gpu
+@pytest.mark.slow
+def test_correct_above_int32_max_elements():
+    """Correctness at a shape whose linear element count exceeds INT32_MAX.
+
+    ``[2, 64, 256, 256, 257]`` is 2_155_872_256 elements -- 8.4M past 2^31, and
+    non-power-of-two in the fastest spatial dimension so a truncated offset
+    cannot accidentally land on the right address.  fp32 (8.03 GiB per tensor)
+    keeps the comparison sharp; the reference needs an NCDHW copy, so the peak
+    is ~48 GiB and the test skips, loudly, if the device cannot hold that.
+    """
+    device = torch.device("cuda")
+    shape = (2, 64, 256, 256, 257)
+    numel = 1
+    for dim in shape:
+        numel *= dim
+    assert numel > 2**31 - 1
+    needed = 6 * numel * 4  # x, y, x_contig, reference, and slack for the diff
+    free, total = torch.cuda.mem_get_info()
+    if free < needed:
+        pytest.skip(
+            f"needs ~{needed / 2**30:.0f} GiB free, device has "
+            f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB"
+        )
+
+    gen = torch.Generator(device=device).manual_seed(53)
+    x = torch.empty(shape, device=device, memory_format=CL)
+    x.normal_(generator=gen)
+    weight = torch.randn(64, device=device, generator=gen)
+    bias = torch.randn(64, device=device, generator=gen)
+
+    assert tgn._plan(
+        shape[0], shape[1], shape[2] * shape[3] * shape[4], GROUPS, numel
+    ).int64
+    got = triton_group_norm(x, GROUPS, weight, bias, EPS)
+    assert got.is_contiguous(memory_format=CL)
+
+    contiguous = x.contiguous()
+    del x
+    torch.cuda.empty_cache()
+    reference = F.group_norm(contiguous, GROUPS, weight, bias, EPS)
+    del contiguous
+    torch.cuda.empty_cache()
+
+    # Compare both batch items separately: a truncated 32-bit offset wraps
+    # partway through, so the second half would be wrong while the first is not.
+    errors = [_rel(got[i], reference[i]) for i in range(shape[0])]
+    print(f"[int64 {shape}] per-sample relative error {errors}")
+    for i, err in enumerate(errors):
+        assert err < 1e-4, f"sample {i}: relative error {err:.3e}"
+    del got, reference
+    torch.cuda.empty_cache()
+
+
+# ---------------------------------------------------------------------------
+# composition: torch.compile and DCTensor
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+def test_custom_op_is_registered_with_a_fake_kernel():
+    """A meta/fake kernel is what lets Dynamo trace the op without running it."""
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    assert hasattr(torch.ops.scaffold_gn, "group_norm")
+    assert hasattr(torch.ops.scaffold_gn, "group_norm_backward")
+    with FakeTensorMode():
+        x = torch.empty(2, 64, 5, 6, 7, device="cuda", memory_format=CL)
+        weight = torch.empty(64, device="cuda")
+        out, mean, rstd = torch.ops.scaffold_gn.group_norm(
+            x, GROUPS, weight, weight, EPS, "relu", None
+        )
+        assert out.shape == x.shape and out.dtype == x.dtype
+        assert out.is_contiguous(memory_format=CL)
+        assert mean.shape == (2, GROUPS) and rstd.dtype == torch.float32
+        # ... and the dtype override autocast uses.
+        bf16 = torch.empty(
+            2, 64, 5, 6, 7, device="cuda", dtype=torch.bfloat16, memory_format=CL
+        )
+        out32, _, _ = torch.ops.scaffold_gn.group_norm(
+            bf16, GROUPS, None, None, EPS, None, torch.float32
+        )
+        assert out32.dtype == torch.float32
+        assert out32.is_contiguous(memory_format=CL)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_torch_compile_fullgraph(activation):
+    """``fullgraph=True`` raises on a graph break, so this *is* the no-break
+    test; the compiled result must additionally be bitwise equal to eager,
+    because the op is opaque to Inductor and so cannot be re-associated."""
+    device = torch.device("cuda")
+    x, weight, bias, grad_out = _tensors(
+        (2, 64, 6, 6, 6), torch.float32, device, seed=59
+    )
+
+    def fn(x, weight, bias):
+        return triton_group_norm(x, GROUPS, weight, bias, EPS, activation) * 2.0
+
+    def wrapped(x, groups, weight, bias, eps, _activation):
+        return fn(x, weight, bias)
+
+    eager = _run(wrapped, x, weight, bias, grad_out)
+
+    compiled_fn = torch.compile(fn, fullgraph=True, dynamic=False)
+
+    def wrapped_compiled(x, groups, weight, bias, eps, _activation):
+        return compiled_fn(x, weight, bias)
+
+    compiled = _run(wrapped_compiled, x, weight, bias, grad_out)
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager):
+        assert torch.equal(a, b), f"compiled and eager differ in {name}"
+    assert compiled[0].is_contiguous(memory_format=CL)
+
+
+@pytest.fixture
+def dc_cuda():
+    """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group.
+
+    Mirrors ``tests/test_groupnorm.py``'s fixture (``num_shards=(1, 1, 1)`` on
+    dims (2, 3, 4) is what worker.py builds for a single-device run) on its own
+    rendezvous port so the two suites can run in one session.
+    """
+    import torch.distributed as dist
+
+    distconv = pytest.importorskip("distconv")
+
+    created = False
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29519")
+        dist.init_process_group(backend="nccl", rank=0, world_size=1)
+        created = True
+    strategy = distconv.ParallelStrategy(
+        num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda"
+    )
+    yield distconv, strategy
+    if created and dist.is_initialized():
+        dist.destroy_process_group()
+
+
+@pytest.mark.gpu
+def test_dctensor_round_trip(dc_cuda):
+    """A DCTensor must go in and come out, with the graph back to its producer
+    intact.
+
+    The op is a real dispatcher op, so DistConv's generic
+    ``__torch_dispatch__`` unwraps to the local shard, runs it, and rewraps --
+    no GroupNorm-specific handling needed on either side.  The producer in
+    front matters: with a bare ``input._tensor`` read the graph would be severed
+    there and only GroupNorm's own parameters would see gradients.
+    """
+    distconv, strategy = dc_cuda
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(61)
+    x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL)
+    x.normal_(generator=gen)
+    grad_out = torch.empty_like(x)
+    grad_out.normal_(generator=gen)
+    weight = torch.randn(64, device=device, generator=gen)
+    bias = torch.randn(64, device=device, generator=gen)
+    producer = torch.nn.Conv3d(64, 64, 1, bias=False).to(device)
+
+    def run(fn, wrap):
+        xi = x.detach().clone().requires_grad_(True)
+        conv = torch.nn.Conv3d(64, 64, 1, bias=False).to(device)
+        with torch.no_grad():
+            conv.weight.copy_(producer.weight)
+        w = weight.detach().clone().requires_grad_(True)
+        b = bias.detach().clone().requires_grad_(True)
+        inp = distconv.DCTensor.from_shard(xi, strategy) if wrap else xi
+        # The conv is the producer; the explicit channels-last conversion is
+        # what PYTORCH_MIOPEN_SUGGEST_NHWC=1 gives production for free.
+        hidden = conv(inp).contiguous(memory_format=CL)
+        out = fn(hidden, GROUPS, w, b, EPS, "relu")
+        if wrap:
+            assert isinstance(out, distconv.DCTensor)
+            assert out.is_contiguous(memory_format=CL)
+            out = distconv.distconv._ToTensor.apply(out)
+        out.backward(grad_out)
+        return out.detach(), xi.grad, conv.weight.grad, w.grad, b.grad
+
+    def stock(x, groups, weight, bias, eps, _activation):
+        return F.relu(F.group_norm(x, groups, weight, bias, eps))
+
+    got = run(triton_group_norm, wrap=True)
+    ref = run(stock, wrap=False)
+    for name, a, b in zip(("y", "dx", "dconv", "dweight", "dbias"), got, ref):
+        assert a is not None, f"{name} never received a gradient"
+        err = _rel(a, b)
+        print(f"[dctensor] {name}={err:.3e}")
+        assert err <= 1e-4, f"{name}: relative error {err:.3e}"
diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py
new file mode 100644
index 0000000..5998c5f
--- /dev/null
+++ b/tests/test_triton_group_norm_edge.py
@@ -0,0 +1,1701 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+"""Adversarial edge-case tests for the channels-last Triton GroupNorm.
+
+Companion to ``tests/test_triton_group_norm.py``, written independently during
+an audit of the kernel.  It covers the ground the author's suite does not, and
+pins the divergences that audit found.  See
+``review/gn-dctensor/KERNEL_REVIEW.md`` for the full write-up.
+
+The two structural gaps this file closes:
+
+* **The masked channel axis is never exercised upstream.**  Every GPU test in
+  ``test_triton_group_norm.py`` uses ``num_groups=8`` with a channel count of
+  64/128/256/2048, so ``G`` and ``C/G`` are *always* powers of two and
+  ``_Plan.masked_c`` is always ``False``.  The entire ``MASKED_C=True`` code
+  path -- the ``cmask``/``wbm`` predicates in all four kernels, and the
+  ``inner`` offsets that deliberately run past the end of a voxel -- ships
+  untested.  :func:`test_masked_channel_axis_parity` and friends run it.
+
+* **Uninitialised split-K scratch is never checked.**  ``_forward`` and
+  ``_backward`` allocate their partial buffers with ``torch.empty``, so a slot
+  that is read before it is written would surface as *plausible* numbers, not
+  as a crash.  :func:`test_scratch_slots_are_all_written` poisons every
+  ``torch.empty`` with NaN for the duration of the call, which turns that class
+  of bug into a hard failure.
+
+The audit's six findings -- no device guard, the backward fake kernel's stride
+promise, silently-differentiable ``mean``/``rstd``, accepting a shape
+``F.group_norm`` rejects, a non-zero ``d_input`` for single-element groups, and
+undocumented double backward -- were tested here as ``xfail(strict=True)``
+first and fixed afterwards; the tests remain, without the markers, as the
+regression pins.  The last section adds the coverage a mutation sweep of the
+kernels found thinnest: the ``INT64=True`` branch (which a default run never
+compiled), the split-K Welford merge on *unequal* split counts, ``eps``
+placement, and the tile-mean correction term.
+"""
+
+import contextlib
+import os
+import subprocess
+import sys
+import textwrap
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from ScaFFold.unet import triton_group_norm as tgn
+from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm
+
+CL = torch.channels_last_3d
+EPS = 1e-5
+
+#: Relative-error ceiling against the float64 reference below.  fp32 parity at
+#: these (small) shapes measures ~1e-07; the ceiling leaves room for the
+#: reduction noise a production-sized split-K reduction shows.
+FP32_TOL = 1e-4
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+# ---------------------------------------------------------------------------
+# independent float64 reference (deliberately *not* F.group_norm, and
+# deliberately not the helper the author's suite uses)
+# ---------------------------------------------------------------------------
+
+
+def _ref64(x, groups, weight, bias, eps, activation=None):
+    """GroupNorm written from scratch in float64, in the (N, G, ...) view."""
+    xd = x.double()
+    n, channels = xd.shape[0], xd.shape[1]
+    flat = xd.reshape(n, groups, -1)
+    mu = flat.mean(-1, keepdim=True)
+    var = ((flat - mu) ** 2).mean(-1, keepdim=True)
+    y = ((flat - mu) / torch.sqrt(var + eps)).reshape(xd.shape)
+    shape = (1, channels) + (1,) * (xd.dim() - 2)
+    if weight is not None:
+        y = y * weight.double().reshape(shape)
+    if bias is not None:
+        y = y + bias.double().reshape(shape)
+    return torch.relu(y) if activation == "relu" else y
+
+
+def _rel(actual, expected):
+    a = actual.detach().double()
+    e = expected.detach().double()
+    return ((a - e).abs().max() / e.abs().max().clamp_min(1e-300)).item()
+
+
+def _make(shape, groups, dtype=torch.float32, affine=True, seed=0, mean=0.0, std=1.0):
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(seed)
+    x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL)
+    x.normal_(mean, std, generator=gen)
+    channels = shape[1]
+    if affine:
+        weight = torch.empty(channels, device=device, dtype=dtype)
+        weight.normal_(1.0, 0.25, generator=gen)
+        bias = torch.empty(channels, device=device, dtype=dtype)
+        bias.normal_(0.0, 0.25, generator=gen)
+    else:
+        weight = bias = None
+    grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL)
+    grad_out.normal_(generator=gen)
+    return x, weight, bias, grad_out
+
+
+def _parity(shape, groups, activation=None, affine=True, seed=0, eps=EPS, label=""):
+    """Forward + backward against the float64 reference.  Returns the errors."""
+    x, weight, bias, grad_out = _make(shape, groups, affine=affine, seed=seed)
+    assert is_supported(x, groups, weight, bias, activation), (
+        f"is_supported rejected {shape} groups={groups}"
+    )
+
+    xi = x.detach().clone().requires_grad_(True)
+    wi = None if weight is None else weight.detach().clone().requires_grad_(True)
+    bi = None if bias is None else bias.detach().clone().requires_grad_(True)
+    triton_group_norm(xi, groups, wi, bi, eps, activation).backward(grad_out)
+
+    xd = x.detach().clone().double().requires_grad_(True)
+    wd = (
+        None
+        if weight is None
+        else weight.detach().clone().double().requires_grad_(True)
+    )
+    bd = None if bias is None else bias.detach().clone().double().requires_grad_(True)
+    _ref64(xd, groups, wd, bd, eps, activation).backward(grad_out.double())
+
+    errors = {"dx": _rel(xi.grad, xd.grad)}
+    if wi is not None:
+        errors["dw"] = _rel(wi.grad, wd.grad)
+    if bi is not None:
+        errors["db"] = _rel(bi.grad, bd.grad)
+    xj = x.detach().clone()
+    errors["y"] = _rel(
+        triton_group_norm(xj, groups, weight, bias, eps, activation),
+        _ref64(x, groups, weight, bias, eps, activation),
+    )
+    print(f"[{label or shape}] " + " ".join(f"{k}={v:.2e}" for k, v in errors.items()))
+    for name, err in errors.items():
+        assert err <= FP32_TOL, f"{label or shape}: {name} rel err {err:.3e}"
+    return errors
+
+
+# ---------------------------------------------------------------------------
+# 1. the masked channel axis (MASKED_C=True) -- never reached upstream
+# ---------------------------------------------------------------------------
+
+#: ``(shape, num_groups)`` pairs for which ``_Plan.masked_c`` is True, i.e.
+#: ``num_groups`` and/or ``num_channels // num_groups`` is not a power of two,
+#: so ``GP``/``CGP`` over-cover the channel axis and every load, store and
+#: reduction in all four kernels has to be predicated.
+_MASKED_CASES = [
+    ((1, 6, 4, 4, 4), 3),  # G=3 -> GP=4, CG=2
+    ((1, 15, 5, 5, 5), 3),  # G=3, CG=5 -> both padded
+    ((2, 12, 7, 5, 3), 3),  # N>1 with a padded group axis
+    ((1, 24, 9, 9, 9), 6),  # G=6 -> GP=8, CG=4
+    ((1, 20, 4, 4, 4), 5),  # G=5 -> GP=8, CG=4
+    ((1, 20, 4, 4, 4), 4),  # G=4, CG=5 -> only the inner axis padded
+    ((3, 20, 3, 5, 7), 5),
+    ((1, 63, 5, 5, 5), 7),  # G=7, CG=9
+    ((2, 63, 5, 5, 5), 7),
+    ((1, 96, 5, 5, 5), 6),  # G=6, CG=16
+    ((1, 10, 3, 3, 3), 10),  # G == C, neither a power of two
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape,groups", _MASKED_CASES)
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_masked_channel_axis_parity(shape, groups, activation):
+    """``MASKED_C=True``: the padded (G, C/G) tile must be fully predicated.
+
+    ``inner = g * CG + j`` deliberately runs past the end of a voxel for the
+    padding lanes, so a wrong ``cmask``/``wbm`` predicate reads (or writes) the
+    *next* voxel's channels, and a wrong ``other=`` poisons the Welford sums.
+    Neither shows up anywhere in the author's suite, which only ever runs
+    ``num_groups=8`` over 64/128/256/2048 channels.
+    """
+    plan = tgn._plan(shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, 0)
+    assert plan.masked_c, "case is supposed to exercise the padded channel axis"
+    _parity(shape, groups, activation, seed=abs(hash((shape, groups))) % 997)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [
+        ((1, 64, 4, 4, 4), 64),  # instance norm, C/G == 1
+        ((1, 7, 3, 3, 3), 7),  # instance norm, prime channel count
+        ((1, 64, 4, 4, 4), 1),  # layer norm, G == 1
+        ((1, 7, 3, 3, 3), 1),  # layer norm, prime channel count
+        ((1, 1, 4, 4, 4), 1),  # single channel
+        ((2, 1, 4, 4, 4), 1),
+        ((1, 3, 5, 5, 5), 3),
+    ],
+)
+def test_extreme_group_counts(shape, groups):
+    """``num_groups == num_channels`` (instance norm) and ``== 1`` (layer norm).
+
+    Both collapse one axis of the ``(BLOCK_S, GP, CGP)`` tile to length 1 and
+    are accepted by ``is_supported``; neither appears upstream.
+    """
+    _parity(shape, groups, seed=abs(hash((shape, groups))) % 997)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 64, 1, 1, 1),  # S == 1: a single voxel, far below one tile
+        (2, 64, 1, 1, 1),
+        (4, 64, 2, 1, 1),
+        (1, 64, 1, 1, 127),  # just under a 128-voxel stats tile
+        (1, 64, 1, 1, 128),  # exactly one tile
+        (1, 64, 1, 1, 129),  # just over
+        (1, 64, 1, 1, 257),
+        (1, 64, 13, 17, 19),  # three primes
+        (5, 64, 3, 3, 3),  # N not a power of two
+        (1, 2048, 1, 1, 2),  # widest channel count, two voxels
+    ],
+)
+def test_ragged_spatial_tails(shape):
+    """Spatial extents that are prime, or sit just either side of a tile edge.
+
+    The ragged tail is where ``offs_s < S - s0`` in ``_normalize_kernel`` /
+    ``_dx_kernel`` and ``nvalid = min(BLOCK_S, s_end - s0)`` in the two partial
+    kernels have to agree; ``cnt_t = nvalid * CG`` also has to be the *valid*
+    lane count or the Welford mean is scaled wrong.
+    """
+    _parity(shape, 8, seed=abs(hash(shape)) % 997)
+
+
+# ---------------------------------------------------------------------------
+# 2. split-K scratch
+# ---------------------------------------------------------------------------
+
+
+def _empty_split_count(n, channels, spatial, groups):
+    plan = tgn._plan(n, channels, spatial, groups, n * channels * spatial)
+    return sum(1 for sp in range(plan.nsplit) if sp * plan.chunk >= spatial), plan
+
+
+#: Shapes whose ``ceil(S / nsplit)`` chunking leaves at least one split with
+#: ``s_begin >= S``, i.e. a program that writes an all-zero ``(cnt, mean, M2)``
+#: partial that the finalize tree then has to absorb.  Found by search over the
+#: plan; ``_welford_combine``'s ``cnt == 0`` guard is what makes them harmless.
+_EMPTY_SPLIT_CASES = [
+    ((1, 64, 1, 1, 32775), 8),
+    ((2, 64, 1, 1, 32775), 8),
+    ((1, 128, 1, 1, 8198), 8),
+    ((1, 256, 1, 1, 2049), 8),
+    ((1, 2048, 1, 1, 33), 8),
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape,groups", _EMPTY_SPLIT_CASES)
+def test_empty_split_slots(shape, groups):
+    """A split whose whole chunk lies past ``S`` still has to combine cleanly.
+
+    ``chunk = ceil(S / nsplit)`` can leave trailing splits entirely empty; that
+    program's loop never runs, so it stores ``(0, 0, 0)``.  Chan's combine is
+    only exact for those because of its ``cnt == 0`` guard, and no upstream
+    shape produces one.
+    """
+    empties, plan = _empty_split_count(
+        shape[0], shape[1], shape[2] * shape[3] * shape[4], groups
+    )
+    assert empties > 0, (
+        f"expected an empty split for {shape}; plan has nsplit={plan.nsplit} "
+        f"chunk={plan.chunk}"
+    )
+    _parity(shape, groups, seed=abs(hash(shape)) % 997)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [
+        ((1, 64, 8, 8, 8), 8),
+        ((2, 64, 8, 8, 8), 8),
+        ((3, 15, 5, 5, 5), 3),
+        ((1, 64, 1, 1, 32775), 8),  # has an empty split
+        ((2, 2048, 1, 1, 1), 8),
+    ],
+)
+def test_scratch_slots_are_all_written(shape, groups):
+    """Every split-K partial slot must be written before it is read.
+
+    ``_forward``/``_backward`` allocate ``pcnt/pmean/pm2`` and
+    ``ps1/ps2/pdw/pdb`` with ``torch.empty``.  A slot that is read but never
+    written would inherit whatever the caching allocator last left there --
+    usually finite, plausible numbers, which no parity test can be relied on to
+    catch.  Poisoning every ``torch.empty``/``empty_like`` with NaN for the
+    duration of the call turns that into a hard failure, and also proves the
+    output buffer itself is fully covered by the store masks.
+    """
+    x, weight, bias, grad_out = _make(shape, groups, seed=5)
+    real_empty, real_empty_like = torch.empty, torch.empty_like
+
+    def poisoned_empty(*args, **kwargs):
+        t = real_empty(*args, **kwargs)
+        return t.fill_(float("nan")) if t.is_floating_point() else t
+
+    def poisoned_empty_like(*args, **kwargs):
+        t = real_empty_like(*args, **kwargs)
+        return t.fill_(float("nan")) if t.is_floating_point() else t
+
+    xi = x.detach().clone().requires_grad_(True)
+    wi = weight.detach().clone().requires_grad_(True)
+    bi = bias.detach().clone().requires_grad_(True)
+    torch.empty, torch.empty_like = poisoned_empty, poisoned_empty_like
+    try:
+        out = triton_group_norm(xi, groups, wi, bi, EPS)
+        out.backward(grad_out)
+    finally:
+        torch.empty, torch.empty_like = real_empty, real_empty_like
+
+    for name, t in (("y", out), ("dx", xi.grad), ("dw", wi.grad), ("db", bi.grad)):
+        assert torch.isfinite(t).all(), (
+            f"{name} contains NaN with poisoned scratch: a split-K slot (or an "
+            f"output element) is read/returned without ever being written"
+        )
+    ref = _ref64(x, groups, weight, bias, EPS)
+    assert _rel(out, ref) <= FP32_TOL
+
+
+# ---------------------------------------------------------------------------
+# 3. numerics
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("value", [0.0, 3.0, 1e3])
+@pytest.mark.parametrize("eps", [1e-5, 1e-12])
+def test_all_equal_input_has_exactly_zero_variance(value, eps):
+    """Variance exactly 0 => ``rstd = 1/sqrt(eps)`` and ``xhat`` exactly 0.
+
+    This is the sharpest possible statement of the Welford claim: with
+    ``weight=1, bias=0`` the output must be *identically* zero, with no
+    tolerance at all.  ATen's fp32 GroupNorm does not manage it (it forms the
+    variance by cancellation and leaves ~1e-05 of noise at ``value=3`` and
+    ~3e-03 at ``value=1e3``), which is asserted here so the comparison stays
+    honest if ATen ever changes.
+    """
+    device = torch.device("cuda")
+    shape = (2, 64, 8, 8, 8)
+    x = torch.full(shape, value, device=device).contiguous(memory_format=CL)
+    weight = torch.ones(64, device=device)
+    bias = torch.zeros(64, device=device)
+
+    got = triton_group_norm(x, 8, weight, bias, eps)
+    assert torch.equal(got, torch.zeros_like(got)), (
+        f"all-equal input must normalise to exactly 0, got max "
+        f"{got.abs().max().item():.3e}"
+    )
+    # ... and rstd really is 1/sqrt(eps), which only the output scale can show.
+    _out, _mean, rstd = torch.ops.scaffold_gn.group_norm(
+        x, 8, None, None, eps, None, None
+    )
+    assert torch.allclose(rstd, torch.full_like(rstd, 1.0 / eps**0.5), rtol=1e-6), (
+        f"rstd={rstd.flatten()[0].item():.6e} != 1/sqrt(eps)={1.0 / eps**0.5:.6e}"
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "mean,std,naive_floor",
+    [(1e4, 1e-2, 1e-1), (1e6, 1.0, 1e1)],
+)
+def test_welford_at_extreme_mean_to_std_ratio(mean, std, naive_floor):
+    """Past the ratios the author's suite tests (mu/sigma up to 1e5).
+
+    At ``mu/sigma = 1e6`` the ``E[x^2]-E[x]^2`` formulation is off by ~2e0 to
+    ~3e2 relative while this kernel holds ~5e-04, and ATen's own fp32 kernel is
+    50-70x worse than this one.  Measured on MI300A at ``[1, 256, 24^3]``.
+    """
+    x, weight, bias, _ = _make((1, 256, 24, 24, 24), 8, seed=37, mean=mean, std=std)
+    got = triton_group_norm(x, 8, weight, bias, EPS)
+    ref = _ref64(x, 8, weight, bias, EPS)
+    stock = F.group_norm(x, 8, weight, bias, EPS)
+
+    flat = x.reshape(1, 8, -1)
+    mu = flat.mean(-1)
+    var = (flat * flat).mean(-1) - mu * mu
+    naive = (flat - mu[..., None]) / torch.sqrt(var + EPS)[..., None]
+    naive = naive.reshape(x.shape) * weight.reshape(1, 256, 1, 1, 1) + bias.reshape(
+        1, 256, 1, 1, 1
+    )
+
+    err, err_stock, err_naive = _rel(got, ref), _rel(stock, ref), _rel(naive, ref)
+    print(
+        f"[mu={mean:g} sd={std:g}] triton={err:.3e} aten={err_stock:.3e} "
+        f"naive={err_naive:.3e}"
+    )
+    assert err_naive > naive_floor, "the naive formulation was supposed to fail here"
+    assert err < err_naive / 100.0
+    assert err <= err_stock, (
+        f"triton {err:.3e} is worse than ATen fp32 {err_stock:.3e} at mu/sigma="
+        f"{mean / std:g}"
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("eps", [1e-5, 1e-8, 1e-12])
+def test_tiny_eps_with_tiny_variance(eps):
+    """``eps`` far below the default with data whose std is ~1e-4.
+
+    ``rstd = 1/sqrt(var + eps)`` reaches ~1e4 here, so any error in the
+    variance is amplified by that factor before it reaches the output.
+    """
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(7)
+    x = torch.empty(1, 64, 16, 16, 16, device=device, memory_format=CL)
+    x.normal_(0.0, 1e-4, generator=gen)
+    weight = torch.ones(64, device=device)
+    bias = torch.zeros(64, device=device)
+    got = triton_group_norm(x, 8, weight, bias, eps)
+    assert _rel(got, _ref64(x, 8, weight, bias, eps)) <= FP32_TOL
+
+
+# ---------------------------------------------------------------------------
+# 4. autograd plumbing
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "kind", ["contiguous", "sliced", "expanded", "transposed", "channels_last"]
+)
+def test_grad_out_layout_variants(kind):
+    """A cotangent that is not channels-last-contiguous.
+
+    ``_group_norm_backward_op`` relayouts it; the kernels index it with the
+    *input's* channels-last stride pattern, so a missed relayout silently
+    permutes the gradient rather than raising.  The author's suite only ever
+    feeds a channels-last-contiguous cotangent to the fast path.
+    """
+    shape = (2, 64, 5, 6, 7)
+    x, weight, bias, _ = _make(shape, 8, seed=71)
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(77)
+    base = torch.empty(shape, device=device)
+    base.normal_(generator=gen)
+    if kind == "contiguous":
+        grad_out = base.contiguous()
+    elif kind == "channels_last":
+        grad_out = base.contiguous(memory_format=CL)
+    elif kind == "sliced":
+        wide = torch.empty(
+            (shape[0], shape[1], shape[2], shape[3], shape[4] * 2), device=device
+        )
+        wide.normal_(generator=gen)
+        grad_out = wide.contiguous(memory_format=CL)[..., ::2]
+    elif kind == "expanded":
+        col = torch.empty((shape[0], shape[1], shape[2], shape[3], 1), device=device)
+        col.normal_(generator=gen)
+        grad_out = col.expand(shape)
+    else:  # transposed
+        swapped = torch.empty(
+            (shape[0], shape[1], shape[2], shape[4], shape[3]), device=device
+        )
+        swapped.normal_(generator=gen)
+        grad_out = swapped.contiguous(memory_format=CL).transpose(3, 4)
+
+    xi = x.detach().clone().requires_grad_(True)
+    wi = weight.detach().clone().requires_grad_(True)
+    bi = bias.detach().clone().requires_grad_(True)
+    triton_group_norm(xi, 8, wi, bi, EPS).backward(grad_out)
+
+    xd = x.detach().clone().double().requires_grad_(True)
+    wd = weight.detach().clone().double().requires_grad_(True)
+    bd = bias.detach().clone().double().requires_grad_(True)
+    _ref64(xd, 8, wd, bd, EPS).backward(grad_out.double())
+
+    assert _rel(xi.grad, xd.grad) <= FP32_TOL
+    assert _rel(wi.grad, wd.grad) <= FP32_TOL
+    assert _rel(bi.grad, bd.grad) <= FP32_TOL
+    # d_input keeps the *input's* format regardless of the cotangent's.
+    assert xi.grad.is_contiguous(memory_format=CL)
+
+
+@pytest.mark.gpu
+def test_affine_parameters_may_be_non_contiguous_views():
+    """``weight``/``bias`` sliced out of a bigger parameter tensor.
+
+    ``is_supported`` only checks rank, numel, device and dtype, so a strided or
+    offset 1-D parameter reaches the op, which is why it calls ``.contiguous()``
+    on both.  Nothing upstream tests that.
+    """
+    x, _weight, _bias, _ = _make((1, 64, 4, 4, 4), 8, seed=83)
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(83)
+    strided = torch.empty(128, device=device)
+    strided.normal_(1.0, 0.25, generator=gen)
+    weight = strided[::2]  # stride 2
+    pack = torch.empty(4, 64, device=device)
+    pack.normal_(0.0, 0.25, generator=gen)
+    bias = pack[2]  # storage offset
+    assert not weight.is_contiguous()
+    assert is_supported(x, 8, weight, bias)
+    got = triton_group_norm(x, 8, weight, bias, EPS)
+    assert _rel(got, _ref64(x, 8, weight, bias, EPS)) <= FP32_TOL
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("lo,hi", [(0, 2), (1, 3), (3, 4)])
+def test_channels_last_views_with_a_storage_offset(lo, hi):
+    """A batch slice of a bigger channels-last tensor stays channels-last
+    contiguous but has a non-zero storage offset -- the kernels must address
+    from ``data_ptr()``, not from the storage base."""
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(89)
+    big = torch.empty((4, 64, 5, 6, 7), device=device, memory_format=CL)
+    big.normal_(generator=gen)
+    weight = torch.empty(64, device=device)
+    weight.normal_(1.0, 0.25, generator=gen)
+    bias = torch.empty(64, device=device)
+    bias.normal_(0.0, 0.25, generator=gen)
+    view = big[lo:hi]
+    assert view.is_contiguous(memory_format=CL) and is_supported(view, 8, weight, bias)
+    got = triton_group_norm(view, 8, weight, bias, EPS)
+    assert _rel(got, _ref64(view, 8, weight, bias, EPS)) <= FP32_TOL
+
+
+@pytest.mark.gpu
+def test_double_backward_raises_instead_of_returning_garbage():
+    """Higher-order gradients are *not* supported, and must say so.
+
+    ``scaffold_gn::group_norm_backward`` has no autograd formula of its own, so
+    a second ``torch.autograd.grad`` through the kernel raises.  Stock
+    ``F.group_norm`` supports double backward, so this is a real (if narrow)
+    behavioural difference from the op it replaces -- anything that needs a
+    gradient penalty or a Hessian-vector product cannot use this kernel.  The
+    test pins "raises loudly", which is the safe half of the story.
+    """
+    x, weight, bias, grad_out = _make((1, 64, 4, 4, 4), 8, seed=97)
+    xi = x.detach().clone().requires_grad_(True)
+    wi = weight.detach().clone().requires_grad_(True)
+    bi = bias.detach().clone().requires_grad_(True)
+    y = triton_group_norm(xi, 8, wi, bi, EPS)
+    (gx,) = torch.autograd.grad(y, xi, grad_out, create_graph=True)
+    with pytest.raises(RuntimeError, match="no autograd formula was registered"):
+        torch.autograd.grad(gx.sum(), xi)
+
+    # ... and stock really does support it, so this is a divergence not a law.
+    xr = x.detach().clone().requires_grad_(True)
+    yr = F.group_norm(xr, 8, weight, bias, EPS)
+    (gxr,) = torch.autograd.grad(yr, xr, grad_out, create_graph=True)
+    (ggr,) = torch.autograd.grad(gxr.sum(), xr)
+    assert torch.isfinite(ggr).all()
+
+
+# ---------------------------------------------------------------------------
+# 5. fake / meta kernel
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("out_dtype", [None, torch.float32])
+@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)])
+def test_fake_forward_matches_real_in_every_branch(dtype, out_dtype, has_w, has_b):
+    """The fake kernel must promise the real shape, dtype, stride *and* device.
+
+    A meta mismatch is invisible in eager and silently corrupts
+    ``torch.compile``; the author's suite spot-checks two combinations, this
+    walks the whole cross product of dtype x out_dtype override x affine.
+    """
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    device = torch.device("cuda")
+    shape = (2, 64, 5, 6, 7)
+    x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL).normal_()
+    weight = torch.randn(64, device=device, dtype=dtype) if has_w else None
+    bias = torch.randn(64, device=device, dtype=dtype) if has_b else None
+
+    real = torch.ops.scaffold_gn.group_norm(x, 8, weight, bias, EPS, "relu", out_dtype)
+    with FakeTensorMode() as mode:
+        args = [None if t is None else mode.from_tensor(t) for t in (x, weight, bias)]
+        fake = torch.ops.scaffold_gn.group_norm(
+            args[0], 8, args[1], args[2], EPS, "relu", out_dtype
+        )
+    for i, (r, f) in enumerate(zip(real, fake)):
+        assert r.shape == f.shape, f"out[{i}] shape"
+        assert r.dtype == f.dtype, f"out[{i}] dtype {r.dtype} != {f.dtype}"
+        assert r.stride() == f.stride(), f"out[{i}] stride {r.stride()} != {f.stride()}"
+        assert r.device.type == f.device.type, f"out[{i}] device"
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "layout", ["contiguous", "channels_last", "sliced", "degenerate"]
+)
+@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)])
+def test_fake_backward_matches_real_in_every_branch(layout, has_w, has_b):
+    """The backward's fake kernel must promise what the real op returns.
+
+    The real op relayouts a non-channels-last ``input`` and *always* returns a
+    channels-last ``d_input``; ``torch.empty_like(input)`` would instead
+    preserve the input's own format, so for a plain contiguous NCDHW input the
+    two disagree ((13440, 1, 2688, 448, 64) against (13440, 210, 42, 7, 1)).  A
+    meta mismatch is invisible in eager and silently corrupts
+    ``torch.compile``, so every branch of the promise -- both layouts, a
+    non-contiguous view, the shape where the two formats coincide, and each
+    affine combination -- is checked here rather than only the CL case.
+    """
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    device = torch.device("cuda")
+    if layout == "degenerate":
+        shape = (2, 64, 1, 1, 1)  # contiguous *is* channels_last_3d here
+        x = torch.randn(shape, device=device)
+    else:
+        shape = (2, 64, 5, 6, 7)
+        if layout == "contiguous":
+            x = torch.randn(shape, device=device)
+        elif layout == "channels_last":
+            x = torch.randn(shape, device=device).contiguous(memory_format=CL)
+        else:  # sliced: neither contiguous nor channels-last contiguous
+            x = torch.randn((2, 64, 5, 6, 14), device=device)[..., ::2]
+    grad_out = torch.randn(shape, device=device).contiguous(memory_format=CL)
+    weight = torch.randn(64, device=device) if has_w else None
+    bias = torch.randn(64, device=device) if has_b else None
+    mean = torch.zeros(2, 8, device=device)
+    rstd = torch.ones(2, 8, device=device)
+
+    real = torch.ops.scaffold_gn.group_norm_backward(
+        grad_out, x, weight, bias, mean, rstd, 8, None
+    )
+    with FakeTensorMode() as mode:
+        a = [
+            None if t is None else mode.from_tensor(t)
+            for t in (grad_out, x, weight, bias, mean, rstd)
+        ]
+        fake = torch.ops.scaffold_gn.group_norm_backward(
+            a[0], a[1], a[2], a[3], a[4], a[5], 8, None
+        )
+    names = ("d_input", "d_weight", "d_bias")
+    for name, r, f in zip(names, real, fake):
+        assert r.shape == f.shape, f"{name} shape {r.shape} != {f.shape}"
+        assert r.dtype == f.dtype, f"{name} dtype {r.dtype} != {f.dtype}"
+        assert r.stride() == f.stride(), f"{name} stride {r.stride()} != {f.stride()}"
+        assert r.device.type == f.device.type, f"{name} device"
+    assert real[0].is_contiguous(memory_format=CL)
+
+
+@pytest.mark.gpu
+def test_mean_and_rstd_are_not_silently_differentiable():
+    """``mean``/``rstd`` are backward state, so they must refuse, not lie.
+
+    They are documented as "not differentiable".  Before they were marked as
+    such, they came back with ``requires_grad=True`` and differentiating
+    through them *succeeded*: autograd materialised an all-zero cotangent for
+    the unused ``out``, ran the entire backward (a full-size zeros allocation
+    plus four kernels) and returned zeros -- a plausible wrong answer where the
+    true value is ~6e-04.  ``ctx.mark_non_differentiable`` turns that into an
+    error, which is the only safe outcome short of a real formula.
+    """
+    device = torch.device("cuda")
+    shape = (2, 64, 5, 6, 7)
+    x = torch.empty(shape, device=device, memory_format=CL).normal_()
+    xi = x.clone().requires_grad_(True)
+    out, mean, rstd = torch.ops.scaffold_gn.group_norm(
+        xi, 8, None, None, EPS, None, None
+    )
+    assert out.requires_grad, "the forward output must still be differentiable"
+    assert not mean.requires_grad, "mean must be marked non-differentiable"
+    assert not rstd.requires_grad, "rstd must be marked non-differentiable"
+    for name, t in (("mean", mean), ("rstd", rstd)):
+        with pytest.raises(RuntimeError, match="does not require grad"):
+            torch.autograd.grad(t.sum(), xi)
+        assert xi.grad is None, f"differentiating {name} left a gradient behind"
+    # The value that used to come back silently wrong is genuinely non-zero,
+    # so "returns zeros" was never defensible as an answer.
+    xd = x.clone().double().requires_grad_(True)
+    (want,) = torch.autograd.grad(xd.reshape(2, 8, -1).mean(-1).sum(), xd)
+    assert want.abs().max() > 0
+
+
+# ---------------------------------------------------------------------------
+# 6. contract / drop-in divergences
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [((1, 64, 1, 1, 1), 8), ((2, 64, 1, 1, 1), 8), ((1, 1, 4, 5, 6), 1)],
+)
+def test_is_supported_accepts_layout_ambiguous_contiguous_input(shape, groups):
+    """``is_supported`` is *not* simply "False for contiguous input".
+
+    For shapes whose spatial or channel extents are all 1 the contiguous and
+    channels-last-3d stride patterns coincide, so a plain ``torch.randn``
+    tensor is accepted by the fast path.  That is benign -- the two layouts are
+    the same bytes -- but it means callers cannot use ``is_supported`` as a
+    layout *classifier*.  Pinned here so the behaviour is deliberate.
+    """
+    device = torch.device("cuda")
+    x = torch.randn(shape, device=device)  # never asked for channels_last
+    assert x.is_contiguous()
+    assert x.is_contiguous(memory_format=CL)
+    assert is_supported(x, groups) is True
+    got = triton_group_norm(x, groups, None, None, EPS)
+    assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape,groups", [((1, 8, 1, 1, 1), 8), ((1, 1, 1, 1, 1), 1)])
+def test_one_value_per_channel_matches_stock_rejection(shape, groups):
+    """``N*(C/G)*D*H*W == 1`` is a shape ``F.group_norm`` refuses to run.
+
+    The kernel *can* compute it (every group has zero variance, so the answer
+    is ``bias``), and it used to: ``is_supported`` returned True and
+    ``triton_group_norm`` returned a value where the op it is a drop-in for
+    raises ``ValueError``.  A caller branching on ``is_supported`` would then
+    get a different answer from the reference path, which is worse than being
+    slower, so all three of ``is_supported``, the public wrapper and the raw op
+    now reject it the same way stock does.
+    """
+    device = torch.device("cuda")
+    x = torch.empty(shape, device=device, memory_format=CL).normal_()
+    with pytest.raises(ValueError, match="more than 1 value per channel"):
+        F.group_norm(x, groups, None, None, EPS)
+    assert is_supported(x, groups) is False, (
+        "is_supported accepts a shape F.group_norm rejects"
+    )
+    # The public wrapper reaches the same rejection through its fallback...
+    with pytest.raises(ValueError, match="more than 1 value per channel"):
+        triton_group_norm(x, groups, None, None, EPS)
+    # ... and the op itself refuses too, for anyone calling it directly.
+    with pytest.raises(ValueError, match="more than 1 value per channel"):
+        torch.ops.scaffold_gn.group_norm(x, groups, None, None, EPS, None, None)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [((2, 8, 1, 1, 1), 8), ((1, 8, 2, 1, 1), 8), ((1, 16, 1, 1, 1), 8)],
+)
+def test_neighbours_of_the_one_value_per_channel_shape_are_still_served(shape, groups):
+    """The rejection must be exactly stock's, not a shape family around it.
+
+    ``_verify_batch_size`` rejects ``N*(C/G)*spatial == 1`` and nothing else, so
+    bumping *any one* of N, C/G or the spatial extent to 2 has to come back to
+    the fast path -- including ``(2, 8, 1, 1, 1)``, which still has a single
+    element per group.
+    """
+    device = torch.device("cuda")
+    x = torch.empty(shape, device=device, memory_format=CL).normal_()
+    F.group_norm(x, groups, None, None, EPS)  # stock accepts it
+    assert is_supported(x, groups) is True
+    got = triton_group_norm(x, groups, None, None, EPS)
+    assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups,activation",
+    [
+        ((2, 8, 1, 1, 1), 8, None),
+        ((2, 8, 1, 1, 1), 8, "relu"),
+        ((1, 8, 2, 1, 1), 8, None),  # 2 elements per group: NOT the degenerate case
+        ((3, 16, 1, 1, 1), 16, None),
+    ],
+)
+def test_single_element_group_gradient_is_exactly_zero(shape, groups, activation):
+    """One element per group => y is constant in x => dx must be identically 0.
+
+    ``mean == x`` and ``var == 0`` identically, so ``xhat`` is the constant 0
+    and nothing downstream depends on ``x``.  ``_dx_kernel`` used to answer
+    2.2e-05 instead: the compiler contracts ``dy*w - c1`` to
+    ``fma(dy, w, -c1)`` while ``c1`` was accumulated from the *rounded*
+    product, so what survives is the product's rounding error (7.0e-08, well
+    under one ulp of ``dyw``), amplified by ``rstd = 1/sqrt(eps) = 316``.
+    ``_backward`` now recognises the degenerate case and returns the exact
+    zero; ATen, on the shapes where it will run at all, leaves ~3e-05 there.
+
+    The ``(1, 8, 2, 1, 1)`` case is the control: two elements per group, so the
+    gradient is *not* identically zero and the kernel must not zero it.
+    """
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(3)
+    channels = shape[1]
+    x = torch.empty(shape, device=device, memory_format=CL)
+    x.normal_(generator=gen)
+    weight = torch.empty(channels, device=device)
+    weight.normal_(1.0, 0.25, generator=gen)
+    bias = torch.empty(channels, device=device)
+    bias.normal_(0.0, 0.25, generator=gen)
+    grad_out = torch.empty(shape, device=device, memory_format=CL)
+    grad_out.normal_(generator=gen)
+
+    assert is_supported(x, groups, weight, bias, activation)
+    xi = x.clone().requires_grad_(True)
+    wi = weight.clone().requires_grad_(True)
+    bi = bias.clone().requires_grad_(True)
+    triton_group_norm(xi, groups, wi, bi, EPS, activation).backward(grad_out)
+
+    xd = x.clone().double().requires_grad_(True)
+    wd = weight.clone().double().requires_grad_(True)
+    bd = bias.clone().double().requires_grad_(True)
+    _ref64(xd, groups, wd, bd, EPS, activation).backward(grad_out.double())
+
+    if channels // groups * shape[2] * shape[3] * shape[4] == 1:
+        assert torch.equal(xi.grad, torch.zeros_like(xi.grad)), (
+            f"dx should be exactly 0, got {xi.grad.abs().max().item():.3e}"
+        )
+        assert xd.grad.abs().max() == 0, "the float64 reference disagrees"
+        # d_weight is exactly 0 too (xhat is exactly 0); d_bias is not.
+        assert torch.equal(wi.grad, torch.zeros_like(wi.grad))
+        assert _rel(bi.grad, bd.grad) <= FP32_TOL
+    else:
+        assert xd.grad.abs().max() > 0, "control case is supposed to be non-trivial"
+        assert xi.grad.abs().max() > 0, "the kernel zeroed a non-degenerate gradient"
+        # A *two*-element group is merely ill-conditioned, not degenerate:
+        # xhat is +-1/sqrt(1+eps/var) and dx is a difference of near-equal
+        # terms, so every fp32 implementation loses digits here -- 4.9e-04
+        # relative for this kernel and 1.4e-04 for ATen on this input.  The
+        # bound is therefore loose against float64, and tight against ATen,
+        # which suffers the same cancellation.
+        assert _rel(xi.grad, xd.grad) <= 1e-3
+        xa = x.clone().requires_grad_(True)
+        F.group_norm(xa, groups, weight, bias, EPS).backward(grad_out)
+        assert _rel(xi.grad, xa.grad) <= 1e-3
+
+
+# ---------------------------------------------------------------------------
+# 7. composition
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_torch_compile_with_dynamic_shapes(activation):
+    """``dynamic=True`` as well as the author's ``dynamic=False``.
+
+    With dynamic shapes the fake kernel is invoked on *symbolic* sizes, so a
+    shape/stride promise that only happens to hold for a concrete size shows up
+    here and nowhere else.  ``fullgraph=True`` is the no-graph-break assertion.
+    """
+    x, weight, bias, grad_out = _make((2, 64, 6, 6, 6), 8, seed=59)
+
+    def fn(x, weight, bias):
+        return triton_group_norm(x, 8, weight, bias, EPS, activation) * 2.0
+
+    def run(f):
+        xi = x.detach().clone().requires_grad_(True)
+        wi = weight.detach().clone().requires_grad_(True)
+        bi = bias.detach().clone().requires_grad_(True)
+        out = f(xi, wi, bi)
+        out.backward(grad_out)
+        return out.detach(), xi.grad, wi.grad, bi.grad
+
+    torch._dynamo.reset()
+    eager = run(fn)
+    compiled = run(torch.compile(fn, fullgraph=True, dynamic=True))
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager):
+        assert torch.equal(a, b), f"dynamic-shape compile differs from eager in {name}"
+    assert compiled[0].is_contiguous(memory_format=CL)
+
+
+# ---------------------------------------------------------------------------
+# 8. determinism, across processes
+# ---------------------------------------------------------------------------
+
+_DETERMINISM_SCRIPT = textwrap.dedent(
+    """
+    import hashlib, sys, torch
+    from ScaFFold.unet.triton_group_norm import triton_group_norm
+    CL = torch.channels_last_3d
+
+    def h(t):
+        b = t.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
+        return hashlib.sha256(b).hexdigest()
+
+    if sys.argv[1] == "warm":
+        # Different JIT order, different lru_cache occupancy, different
+        # allocator state and different free memory before the real work.
+        junk = []
+        for shape, g in (((3, 128, 7, 7, 7), 8), ((1, 15, 5, 5, 5), 3)):
+            a = torch.empty(shape, device="cuda", memory_format=CL).normal_()
+            triton_group_norm(a, g, None, None, 1e-5, "relu")
+            junk.append(torch.empty(1 << 25, device="cuda"))
+        del junk
+        torch.cuda.empty_cache()
+
+    for shape, groups, act, dtype in (
+        ((2, 128, 9, 11, 13), 8, None, torch.float32),
+        ((2, 64, 6, 7, 8), 8, "relu", torch.bfloat16),
+        ((3, 15, 5, 5, 5), 3, None, torch.float32),
+    ):
+        gen = torch.Generator(device="cuda").manual_seed(31)
+        x = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL)
+        x.normal_(generator=gen)
+        w = torch.empty(shape[1], device="cuda", dtype=dtype)
+        w.normal_(1.0, 0.25, generator=gen)
+        b = torch.empty(shape[1], device="cuda", dtype=dtype)
+        b.normal_(0.0, 0.25, generator=gen)
+        go = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL)
+        go.normal_(generator=gen)
+        xi = x.clone().requires_grad_(True)
+        wi = w.clone().requires_grad_(True)
+        bi = b.clone().requires_grad_(True)
+        y = triton_group_norm(xi, groups, wi, bi, 1e-5, act)
+        y.backward(go)
+        print(shape, groups, act, dtype,
+              h(y), h(xi.grad), h(wi.grad), h(bi.grad), flush=True)
+    """
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.slow
+@pytest.mark.timeout(900)
+def test_bitwise_determinism_across_processes():
+    """Process-to-process bitwise reproducibility, which is half the claim.
+
+    ``test_bitwise_determinism`` upstream only calls the kernel twice in *one*
+    process, where the plan is already memoised and the JIT cache already warm.
+    This runs three fresh interpreters -- one of which first JITs other shapes,
+    churns the caching allocator and changes how much memory is free -- and
+    compares SHA-256 of the raw output bytes.  Anything that made the split
+    count, tile size or launch geometry depend on device state rather than on
+    the shape would show up only here.
+    """
+    outputs = []
+    for mode in ("plain", "plain", "warm"):
+        result = subprocess.run(
+            [sys.executable, "-c", _DETERMINISM_SCRIPT, mode],
+            capture_output=True,
+            text=True,
+            cwd=REPO_ROOT,
+            timeout=600,
+        )
+        assert result.returncode == 0, result.stderr[-3000:]
+        outputs.append(result.stdout)
+    assert outputs[0] == outputs[1], "two identical processes disagree"
+    assert outputs[0] == outputs[2], (
+        "a process that JITted other shapes first disagrees:\n"
+        f"{outputs[0]}\n--- vs ---\n{outputs[2]}"
+    )
+    assert outputs[0].count("\n") >= 3
+
+
+# ---------------------------------------------------------------------------
+# 9. multi-device
+# ---------------------------------------------------------------------------
+
+_DEVICE_GUARD_SCRIPT = textwrap.dedent(
+    """
+    import sys, torch
+    import torch.nn.functional as F
+    from ScaFFold.unet.triton_group_norm import triton_group_norm
+    CL = torch.channels_last_3d
+    torch.cuda.set_device(0)                      # current device = 0
+    other = "cuda:1"
+    g = torch.Generator(device=other).manual_seed(1)
+    x = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL)
+    x.normal_(generator=g)
+    w = torch.empty(64, device=other); w.normal_(1.0, 0.25, generator=g)
+    b = torch.empty(64, device=other); b.normal_(0.0, 0.25, generator=g)
+    go = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL)
+    go.normal_(generator=g)
+
+    def rel(a, e):
+        return ((a.double() - e.double()).abs().max()
+                / e.double().abs().max().clamp_min(1e-30)).item()
+
+    # ATen carries a DeviceGuard, so this is the behaviour to match.
+    xr = x.clone().requires_grad_(True)
+    wr = w.clone().requires_grad_(True)
+    br = b.clone().requires_grad_(True)
+    F.group_norm(xr, 8, wr, br, 1e-5).backward(go)
+
+    xi = x.clone().requires_grad_(True)
+    wi = w.clone().requires_grad_(True)
+    bi = b.clone().requires_grad_(True)
+    y = triton_group_norm(xi, 8, wi, bi, 1e-5)    # tensors on 1, current is 0
+    y.backward(go)                                # ... and so is the backward
+    torch.cuda.synchronize()
+    assert torch.cuda.current_device() == 0, "the guard leaked the device"
+    for name, got, want in (("y", y, F.group_norm(xr.detach(), 8, w, b, 1e-5)),
+                            ("dx", xi.grad, xr.grad),
+                            ("dw", wi.grad, wr.grad),
+                            ("db", bi.grad, br.grad)):
+        assert got.device == torch.device(other), f"{name} on {got.device}"
+        e = rel(got, want)
+        assert e < 1e-4, f"{name}: rel err {e}"
+    print("OK")
+    """
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.slow
+@pytest.mark.timeout(600)
+def test_kernel_runs_on_the_inputs_device_not_the_current_one():
+    """Tensors on cuda:1 while cuda:0 is current, forward *and* backward.
+
+    A Triton launch goes to whatever device is *current*, so without a device
+    guard the kernel dereferences another device's pointers and the process
+    dies with ``Memory access fault by GPU node-N``.  ``F.group_norm`` carries
+    ATen's ``DeviceGuard`` and handles the identical call, so this is a
+    divergence from the op being replaced, not a PyTorch limitation.
+
+    Run in a subprocess because the failure mode is an unrecoverable GPU memory
+    fault, which would take the whole pytest session with it.
+    """
+    if torch.cuda.device_count() < 2:
+        pytest.skip("needs 2 visible CUDA devices")
+    result = subprocess.run(
+        [sys.executable, "-c", _DEVICE_GUARD_SCRIPT],
+        capture_output=True,
+        text=True,
+        cwd=REPO_ROOT,
+        timeout=480,
+    )
+    assert result.returncode == 0 and "OK" in result.stdout, (
+        f"returncode={result.returncode}\nstdout={result.stdout}\n"
+        f"stderr={result.stderr[-2000:]}"
+    )
+
+
+@pytest.mark.gpu
+def test_device_guard_helper_is_a_no_op_on_the_current_device():
+    """The guard must be free on the hot path and real off it.
+
+    ``_device_guard`` skips ``torch.cuda.device`` when the tensor already lives
+    on the current device (1.55 us against 0.51 us of host time per call, which
+    is 0.5% of the two smallest scale-8 shapes' 0.65 ms fwd+bwd because they
+    are host-dispatch bound).  This pins both halves of that shortcut so a
+    future edit cannot quietly turn it into "no guard at all"; the multi-device
+    behaviour itself is covered by the subprocess test above.
+    """
+    device = torch.device("cuda", torch.cuda.current_device())
+    guard = tgn._device_guard(device)
+    assert guard is tgn._NO_GUARD, "should not build a guard for the current device"
+    # Constructing a guard for another index does not touch that device.
+    other = torch.device("cuda", device.index + 1)
+    assert tgn._device_guard(other) is not tgn._NO_GUARD, (
+        "a foreign device must get a real guard"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 10. addressing at the int32 boundary
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.slow
+@pytest.mark.timeout(1800)
+def test_int32_addressing_at_its_documented_maximum():
+    """``numel = INT32_MAX - 127`` with N=2, i.e. ``plan.int64 is False``.
+
+    The author's suite tests the shape *above* the switch
+    (``test_correct_above_int32_max_elements``) but never the largest shape the
+    **int32** path itself has to serve, which is where a missing term in the
+    ``numel + channels > INT32_MAX`` guard would bite.  ``65 * 63 * 4097`` is
+    ``2^24 - 1`` voxels, so nothing about the extents is a power of two.
+
+    Verified without materialising an NCDHW reference: the statistics are
+    checked against a chunked float64 reduction over the physical (N, S, C)
+    view, and the output against an elementwise recomputation done per batch
+    item (a truncated offset wraps partway through, so sample 1 would break
+    while sample 0 did not).
+    """
+    device = torch.device("cuda")
+    shape = (2, 64, 65, 63, 4097)
+    n, channels = shape[0], shape[1]
+    spatial = shape[2] * shape[3] * shape[4]
+    numel = n * channels * spatial
+    assert numel == 2**31 - 128, numel
+
+    plan = tgn._plan(n, channels, spatial, 8, numel)
+    assert plan.int64 is False, "this shape is supposed to use the int32 path"
+
+    free, total = torch.cuda.mem_get_info()
+    needed = 4 * numel * 4
+    if free < needed:
+        pytest.skip(
+            f"needs ~{needed / 2**30:.0f} GiB free, device has "
+            f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB"
+        )
+
+    gen = torch.Generator(device=device).manual_seed(53)
+    x = torch.empty(shape, device=device, memory_format=CL)
+    x.normal_(generator=gen)
+    weight = torch.randn(channels, device=device, generator=gen)
+    bias = torch.randn(channels, device=device, generator=gen)
+    out, mean, rstd = torch.ops.scaffold_gn.group_norm(
+        x, 8, weight, bias, EPS, None, None
+    )
+
+    group_channels = channels // 8
+    flat = x.permute(0, 2, 3, 4, 1).reshape(n, spatial, channels)  # no copy
+    chunk = 1 << 20
+    for i in range(n):
+        acc = torch.zeros(8, dtype=torch.float64, device=device)
+        for s in range(0, spatial, chunk):
+            acc += (
+                flat[i, s : s + chunk]
+                .double()
+                .reshape(-1, 8, group_channels)
+                .sum(dim=(0, 2))
+            )
+        mu = acc / (spatial * group_channels)
+        acc2 = torch.zeros(8, dtype=torch.float64, device=device)
+        for s in range(0, spatial, chunk):
+            d = (
+                flat[i, s : s + chunk].double().reshape(-1, 8, group_channels)
+                - mu[None, :, None]
+            )
+            acc2 += (d * d).sum(dim=(0, 2))
+        var = acc2 / (spatial * group_channels)
+        assert _rel(mean[i], mu) <= 1e-5, f"sample {i} mean"
+        assert _rel(rstd[i], 1.0 / torch.sqrt(var + EPS)) <= 1e-5, f"sample {i} rstd"
+    del flat
+
+    mv = (
+        mean.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1)
+    )
+    rv = (
+        rstd.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1)
+    )
+    for i in range(n):
+        recomputed = (x[i : i + 1] - mv[i : i + 1]) * rv[i : i + 1] * weight.reshape(
+            1, channels, 1, 1, 1
+        ) + bias.reshape(1, channels, 1, 1, 1)
+        # fp32 subtraction of two near-equal fp32 values is exact.
+        err = (out[i : i + 1] - recomputed).abs().max().item()
+        scale = recomputed.abs().max().item()
+        print(f"[int32-max sample {i}] elementwise rel err {err / scale:.3e}")
+        assert err / scale < 1e-5, f"sample {i}"
+        del recomputed
+    del x, out, mean, rstd, mv, rv
+    torch.cuda.empty_cache()
+
+
+# ---------------------------------------------------------------------------
+# 11. coverage the mutation sweep found thin
+# ---------------------------------------------------------------------------
+
+
+@contextlib.contextmanager
+def _force_int64_addressing():
+    """Make every plan take the int64 tile-base path, whatever the shape.
+
+    ``_Plan`` sets ``int64 = numel + channels > _INT32_MAX``, so dropping the
+    threshold turns the wide path on for a shape that fits in a few MiB.  The
+    plan cache is keyed on the shape, not on the threshold, so it has to be
+    cleared on the way in *and* on the way out.
+    """
+    real = tgn._INT32_MAX
+    tgn._plan.cache_clear()
+    tgn._INT32_MAX = -1
+    try:
+        yield
+    finally:
+        tgn._INT32_MAX = real
+        tgn._plan.cache_clear()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [
+        ((2, 64, 5, 6, 7), 8),  # ragged tail, several splits
+        ((1, 2048, 6, 6, 6), 8),  # widest channel count
+        ((3, 15, 5, 5, 5), 3),  # masked channel axis as well
+        ((1, 64, 1, 1, 32775), 8),  # has an empty split
+    ],
+)
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_int64_addressing_path_is_behaviourally_correct(shape, groups, activation):
+    """Run the ``INT64=True`` branch of all seven kernels on a small shape.
+
+    ``INT64`` is a ``tl.constexpr``, so the wide and narrow paths are *different
+    compiled kernels*; only shapes above 2^31 elements reach the wide one
+    naturally, and the one test that does is ``@pytest.mark.slow`` and needs
+    8 GiB.  In a default ``-m "not slow"`` run the int64 branch therefore has no
+    behavioural coverage at all -- forcing ``self.int64`` gives it some for the
+    price of a few MiB.
+
+    The two paths differ only in the *type* of the scalar tile base, so the
+    results must be **bitwise** identical, which is a far sharper assertion than
+    a tolerance and would catch a widened offset that lost or duplicated a tile.
+    """
+    x, weight, bias, grad_out = _make(shape, groups, seed=abs(hash(shape)) % 997)
+
+    def run():
+        xi = x.detach().clone().requires_grad_(True)
+        wi = weight.detach().clone().requires_grad_(True)
+        bi = bias.detach().clone().requires_grad_(True)
+        out = triton_group_norm(xi, groups, wi, bi, EPS, activation)
+        out.backward(grad_out)
+        return out.detach(), xi.grad, wi.grad, bi.grad
+
+    plan32 = tgn._plan(
+        shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel()
+    )
+    assert plan32.int64 is False, "shape is supposed to fit the int32 path"
+    narrow = run()
+
+    with _force_int64_addressing():
+        plan64 = tgn._plan(
+            shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel()
+        )
+        assert plan64.int64 is True, "the int64 path was not forced on"
+        wide = run()
+
+    for name, a, b in zip(("y", "dx", "dweight", "dbias"), wide, narrow):
+        assert torch.equal(a, b), f"int64 path differs from int32 in {name}"
+    # ... and both are actually right, not identically wrong.
+    ref = _ref64(x, groups, weight, bias, EPS, activation)
+    assert _rel(wide[0], ref) <= FP32_TOL
+
+
+#: ``(shape, groups)`` whose split-K partials have *unequal* counts, because
+#: ``chunk = ceil(S / nsplit)`` does not divide ``S``.  Chan's combine weights
+#: the delta by ``cnt_b / (cnt_a + cnt_b)``; with equal counts every level of
+#: the reduction tree has ``cnt_a == cnt_b``, so weighting by the wrong one is
+#: invisible.  Only a ragged (or empty) trailing split exposes it -- which is
+#: why the mutation sweep killed that bug with exactly two parametrizations of
+#: one test upstream.
+_UNEVEN_SPLIT_CASES = [
+    ((2, 64, 9, 7, 5), 8),
+    ((1, 2048, 6, 6, 6), 8),
+    ((1, 64, 1, 1, 32775), 8),
+    ((1, 256, 1, 1, 2049), 8),
+    ((2, 128, 11, 13, 17), 8),
+    ((2, 15, 9, 9, 9), 3),  # masked channel axis as well
+    ((1, 20, 17, 17, 17), 5),  # 16 splits, trailing split 15 voxels short
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape,groups", _UNEVEN_SPLIT_CASES)
+@pytest.mark.parametrize("eps", [1e-5, 0.5])
+def test_group_statistics_match_float64_with_uneven_splits(shape, groups, eps):
+    """Assert ``mean``/``rstd`` themselves, not just the output they feed.
+
+    Two things hide inside the output's 1e-4 tolerance and show up here:
+
+    * **the Welford merge.**  The shapes above all have at least one split with
+      a different element count from its neighbours, which is the only
+      configuration in which mis-weighting Chan's delta changes the answer.
+    * **where ``eps`` goes.**  Every parity test in both files uses
+      ``eps=1e-5`` against a variance of ~1, where ``1/sqrt(var+eps)`` and
+      ``1/(sqrt(var)+eps)`` agree to ~1e-5 -- inside that tolerance.  At
+      ``eps=0.5`` they are 0.816 and 0.667, a 22% difference that no tolerance
+      can absorb.
+    """
+    spatial = shape[2] * shape[3] * shape[4]
+    plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial)
+    counts = {
+        max(0, min(sp * plan.chunk + plan.chunk, spatial) - sp * plan.chunk)
+        for sp in range(plan.nsplit)
+    }
+    assert plan.nsplit > 1 and len(counts) > 1, (
+        f"{shape} was supposed to give unequal split counts; nsplit="
+        f"{plan.nsplit} chunk={plan.chunk} counts={sorted(counts)}"
+    )
+
+    x, _weight, _bias, _ = _make(shape, groups, seed=abs(hash(shape)) % 997)
+    _out, mean, rstd = torch.ops.scaffold_gn.group_norm(
+        x, groups, None, None, eps, None, None
+    )
+    flat = x.double().reshape(shape[0], groups, -1)
+    mean64 = flat.mean(-1)
+    var64 = ((flat - mean64[..., None]) ** 2).mean(-1)
+    rstd64 = 1.0 / torch.sqrt(var64 + eps)
+    assert _rel(mean, mean64) <= 1e-5, "group mean"
+    assert _rel(rstd, rstd64) <= 1e-5, "group rstd (eps placement / Welford merge)"
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("groups", [1, 2, 4])
+@pytest.mark.parametrize("seed", [0, 1, 2, 3])
+def test_welford_correction_recovers_rstd_in_a_single_tile_reduction(groups, seed):
+    """The third reduction pass (``corr``) is load-bearing, and here is where.
+
+    ``mean0 = sum(x)/n`` loses digits in proportion to the tile's element count
+    times ``mu/sigma``; ``corr = sum(x-mean0)/n`` recovers them, and ``M2`` is
+    then formed around the corrected mean.  The effect is largest when one tile
+    carries a whole group's reduction, which is this shape: ``block_s_stats``
+    covers all 128 voxels and ``nsplit == 1``, so 8192/``G`` elements per group
+    go through a single ``mean0``.
+
+    At ``mu/sigma = 1e6`` the correction is worth **216x** (G=1), **580x**
+    (G=2) and **1472x** (G=4) on the relative error of ``rstd`` -- measured by
+    running a copy of this module with the term deleted.  Corrected lands at
+    ~1e-07 for every seed and group count; without it, at 2.1e-05 to 1.6e-04.
+    The 1e-06 ceiling below sits an order of magnitude above the first and an
+    order of magnitude below the second.
+
+    The *output* is not a witness for this: ``y`` moves by at most ~1.4x with
+    or without the term, because it is dominated by the fp32 representation of
+    the mean.  That is why this asserts ``rstd`` directly.
+    """
+    device = torch.device("cuda")
+    shape = (2, 64, 8, 4, 4)
+    spatial = shape[2] * shape[3] * shape[4]
+    plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial)
+    assert plan.nsplit == 1 and plan.block_s_stats >= spatial, (
+        f"case is supposed to be a single-tile reduction; nsplit={plan.nsplit} "
+        f"block_s_stats={plan.block_s_stats} spatial={spatial}"
+    )
+
+    gen = torch.Generator(device=device).manual_seed(seed)
+    x = torch.empty(shape, device=device, memory_format=CL)
+    x.normal_(1e4, 1e-2, generator=gen)  # mu/sigma = 1e6
+    _out, mean, rstd = torch.ops.scaffold_gn.group_norm(
+        x, groups, None, None, EPS, None, None
+    )
+    flat = x.double().reshape(shape[0], groups, -1)
+    mean64 = flat.mean(-1)
+    var64 = ((flat - mean64[..., None]) ** 2).mean(-1)
+    err = _rel(rstd, 1.0 / torch.sqrt(var64 + EPS))
+    print(f"[corr G={groups} seed={seed}] rstd rel err {err:.3e}")
+    assert err <= 1e-6, (
+        f"rstd rel err {err:.3e} at mu/sigma=1e6 with G={groups}: the tile mean "
+        f"correction is not doing its job"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 12. the fused finalize and the capped elementwise grid
+# ---------------------------------------------------------------------------
+#
+# ``_stats_finalize``/``_bwd_finalize``/``_dwdb_reduce`` are no longer their own
+# launches: each is recomputed inside the elementwise kernel that consumes it.
+# Two consequences need pinning.
+#
+# * The elementwise grid is capped at ``GNConfig.elem_progs`` and each program
+#   *strides* over its share of the tiles, so that the fused finalize costs
+#   ``nprog_elem * nsplit`` and not ``nblk_elem * nsplit`` reads.  No scale-8
+#   shape and no shape in either suite reaches that path with the shipped
+#   table -- ``nprog_elem == nblk_elem`` at every small shape -- so it has to be
+#   reached deliberately.
+# * The cap is a *performance* knob.  If it could change a single bit of the
+#   output it would break the module's reproducibility contract, since it is
+#   the one plan field that does not follow from the shape alone.
+
+
+@contextlib.contextmanager
+def _forced_config(channels, spatial, **overrides):
+    """Temporarily install a tiling config for one ``(channels, spatial)`` key.
+
+    ``default_config`` keys the frozen table by ``(num_channels, cube-root
+    spatial extent)``, so the spatial extent has to be a perfect cube here.
+    Restores the previous entry (or its absence) and clears the plan cache on
+    the way out, so no other test can see it.
+    """
+    edge = round(spatial ** (1.0 / 3.0))
+    assert edge**3 == spatial, "forced configs need a cube spatial extent"
+    key = (channels, edge)
+    cfg = tgn.GNConfig(*tgn.default_config(channels, spatial).key())
+    for name, value in overrides.items():
+        assert hasattr(cfg, name), name
+        setattr(cfg, name, value)
+    sentinel = object()
+    saved = tgn._TUNED.get(key, sentinel)
+    tgn._TUNED[key] = cfg
+    tgn._plan.cache_clear()
+    try:
+        yield cfg
+    finally:
+        if saved is sentinel:
+            del tgn._TUNED[key]
+        else:
+            tgn._TUNED[key] = saved
+        tgn._plan.cache_clear()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups",
+    [
+        ((1, 64, 16, 16, 16), 8),
+        ((2, 64, 16, 16, 16), 8),  # N > 1: the stride is per (blk, n) program
+        ((1, 20, 8, 8, 8), 5),  # capped grid *and* a padded channel axis
+    ],
+)
+@pytest.mark.parametrize("elem_progs", [1, 3, 8])
+@pytest.mark.parametrize("activation", [None, "relu"])
+def test_capped_elementwise_grid_strides_over_its_tiles(
+    shape, groups, elem_progs, activation
+):
+    """Fewer elementwise programs than tiles: each must cover several tiles.
+
+    A grid-stride loop that got its start, stride or trip count wrong leaves
+    part of the output (and of ``d_input``) unwritten -- which, since both are
+    ``torch.empty``, surfaces as plausible stale numbers rather than as a
+    crash.  ``elem_progs=1`` is the extreme: one program per sample walks every
+    tile, so it also pins that the fused statistics are hoisted out of the loop
+    correctly rather than being recomputed per iteration from stale state.
+    """
+    spatial = shape[2] * shape[3] * shape[4]
+    with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs):
+        plan = tgn._plan(shape[0], shape[1], spatial, groups, 0)
+        assert plan.nprog_elem == min(plan.nblk_elem, elem_progs)
+        assert plan.nprog_elem < plan.nblk_elem, (
+            f"the cap has to actually bite: nprog={plan.nprog_elem} "
+            f"nblk={plan.nblk_elem}"
+        )
+        _parity(
+            shape,
+            groups,
+            activation,
+            seed=abs(hash((shape, elem_progs))) % 997,
+            label=f"{shape} elem_progs={elem_progs}",
+        )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups", [((2, 64, 16, 16, 16), 8), ((1, 20, 8, 8, 8), 5)]
+)
+def test_elementwise_grid_cap_is_bitwise_neutral(shape, groups):
+    """``elem_progs`` may not change a single bit of any output.
+
+    It is the only field of ``_Plan`` that is a free parameter rather than a
+    consequence of the shape, and the module promises bitwise reproducibility.
+    That promise holds only because the elementwise kernels carry nothing
+    across loop iterations: the fused finalize is computed once per program
+    from the *same* partials with the *same* tile shape, and the tile bodies
+    are pure elementwise.  If tuning this knob ever moved a result, the frozen
+    table would have become part of the numerical contract.
+    """
+    spatial = shape[2] * shape[3] * shape[4]
+    x, weight, bias, grad_out = _make(shape, groups, seed=11)
+    results = []
+    for elem_progs in (0, 1, 5, 64, 4096):
+        with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs):
+            xi = x.detach().clone().requires_grad_(True)
+            wi = weight.detach().clone().requires_grad_(True)
+            bi = bias.detach().clone().requires_grad_(True)
+            y = triton_group_norm(xi, groups, wi, bi, EPS)
+            y.backward(grad_out)
+            results.append(
+                (y.detach().clone(), xi.grad.clone(), wi.grad.clone(), bi.grad.clone())
+            )
+    for elem_progs, got in zip((1, 5, 64, 4096), results[1:]):
+        for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, results[0]):
+            assert torch.equal(a, b), (
+                f"elem_progs={elem_progs} changed {name} bitwise; the grid cap "
+                f"is supposed to be a pure performance knob"
+            )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups,elem_progs",
+    [
+        ((1, 64, 16, 16, 16), 8, 0),
+        ((1, 64, 16, 16, 16), 8, 3),
+        ((2, 96, 8, 8, 8), 6, 2),  # padded channel axis, N > 1, capped grid
+    ],
+)
+def test_fused_finalize_publishes_the_statistics(shape, groups, elem_progs):
+    """``mean``/``rstd`` are published by program 0 of the *normalize* kernel.
+
+    There is no separate finalize launch any more: every elementwise program
+    re-derives the statistics from the split-K Welford partials, and program 0
+    is the one that stores them for the backward pass.  A wrong publishing
+    program, a wrong partials index, or a group-mask slip in that fused
+    reduction would hand the backward garbage while leaving the forward -- which
+    uses its own locally computed copy -- perfectly correct.  So check the
+    published tensors directly against float64.
+    """
+    spatial = shape[2] * shape[3] * shape[4]
+    with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs):
+        x, weight, bias, _grad = _make(shape, groups, seed=3)
+        _y, mean, rstd = torch.ops.scaffold_gn.group_norm(
+            x, groups, weight, bias, EPS, None, None
+        )
+        flat = x.double().reshape(shape[0], groups, -1)
+        mean64 = flat.mean(-1)
+        var64 = ((flat - mean64[..., None]) ** 2).mean(-1)
+        assert _rel(mean, mean64) <= FP32_TOL
+        assert _rel(rstd, 1.0 / torch.sqrt(var64 + EPS)) <= FP32_TOL
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "shape,groups", [((2, 2048, 1, 1, 1), 8), ((1, 1024, 2, 1, 1), 8)]
+)
+def test_dweight_blocks_are_covered_when_there_are_more_of_them_than_tiles(
+    shape, groups
+):
+    """The dweight/dbias reduction rides in ``_dx_kernel``'s first NDW programs.
+
+    Those blocks are per-*channel*, the elementwise tiles are per-*voxel*, and
+    nothing makes the first outnumber the second: at ``(2, 2048, 1, 1, 1)``
+    there is one elementwise tile and eight dweight blocks.  The grid is
+    ``max(nprog_elem, dwdb_progs)`` for exactly that reason, and a grid of
+    ``nprog_elem`` alone would silently leave 7/8 of ``d_weight`` unwritten.
+    """
+    spatial = shape[2] * shape[3] * shape[4]
+    plan = tgn._plan(shape[0], shape[1], spatial, groups, 0)
+    assert plan.dwdb_progs > plan.nprog_elem, (
+        f"case is supposed to have more dweight blocks ({plan.dwdb_progs}) than "
+        f"elementwise programs ({plan.nprog_elem})"
+    )
+    assert plan.grid_dx == plan.dwdb_progs
+    _parity(shape, groups, seed=13)
+
+
+# ---------------------------------------------------------------------------
+# the kernel-failure boundary
+# ---------------------------------------------------------------------------
+
+
+def test_kernel_failures_are_tagged_and_carry_their_cause():
+    """Everything the launch region raises comes out as ``TritonKernelError``.
+
+    The tag is what lets a caller with a fallback (``FastGroupNorm``'s ladder)
+    catch *exactly* "the kernel is broken" instead of catching ``Exception`` and
+    then trying to enumerate every framework mechanism -- saved-tensor pack
+    hooks, ``torch.utils.checkpoint``'s recompute control flow, functorch --
+    that legitimately raises through a forward.  The region it wraps is closed
+    (allocations and launches, no autograd-observable op), so a blanket catch
+    inside it is sound where one at the call site is not.
+
+    The tag must survive the *type* of the original error, whatever it was: a
+    mismatched Triton release raises ``TypeError``/``AttributeError`` from a
+    changed signature, an unwritable JIT cache ``OSError``, a bad launch
+    ``RuntimeError``.
+    """
+    for original in (
+        RuntimeError("launch failed"),
+        TypeError("triton API changed"),
+        AttributeError("no such attribute"),
+        OSError("unwritable cache dir"),
+        ImportError("no module named triton"),
+    ):
+
+        @tgn._tag_kernel_failures
+        def _boom():
+            raise original
+
+        with pytest.raises(tgn.TritonKernelError) as caught:
+            _boom()
+        assert caught.value.__cause__ is original
+        assert type(original).__name__ in str(caught.value)
+
+
+def test_out_of_memory_is_not_tagged_as_a_kernel_failure():
+    """An OOM is a resource condition, and every fallback allocates as much.
+
+    Tagging it would make the ladder retry on a rung that is about to OOM in
+    the same place, and would latch a rung off for the rest of the process on a
+    transient, per-rank event.  It has to come out unchanged.
+    """
+
+    @tgn._tag_kernel_failures
+    def _oom():
+        raise torch.OutOfMemoryError("simulated OOM")
+
+    with pytest.raises(torch.OutOfMemoryError):
+        _oom()
+    assert not issubclass(torch.OutOfMemoryError, tgn.TritonKernelError)
+
+
+def test_contract_violations_are_not_tagged():
+    """``_validate``'s ``ValueError``s are caller errors, and stay loud.
+
+    ``is_supported`` accepts exactly what ``_validate`` accepts, so a caller
+    that branches on the predicate can never see one; if the two ever disagree,
+    the failure must not be laundered into "the kernel is broken" and silently
+    fall back.
+    """
+    with pytest.raises(ValueError, match="activation must be one of"):
+        tgn._validate(torch.zeros(1, 8, 2, 2, 2), 8, None, None, "gelu")
+    with pytest.raises(ValueError, match="expected a 5-D"):
+        tgn._validate(torch.zeros(1, 8, 2, 2), 8, None, None, None)
+
+
+@pytest.mark.gpu
+def test_a_real_launch_failure_is_tagged(monkeypatch):
+    """End to end: break the launch and the public op raises the tagged type."""
+    x = torch.randn(1, 64, 4, 4, 4, device="cuda").to(memory_format=CL)
+
+    def _broken(*args, **kwargs):
+        raise RuntimeError("simulated HIP launch failure")
+
+    tgn._ensure_kernels()
+    monkeypatch.setattr(tgn, "_stats_partial_kernel", _Unlaunchable(_broken))
+    with pytest.raises(tgn.TritonKernelError):
+        torch.ops.scaffold_gn.group_norm(x, 8, None, None, EPS, None, None)
+
+
+class _Unlaunchable:
+    """A stand-in for a ``triton.jit`` kernel whose launch raises."""
+
+    def __init__(self, fn):
+        self._fn = fn
+
+    def __getitem__(self, grid):
+        return self._fn
+
+
+# ---------------------------------------------------------------------------
+# the fused activation on non-finite values
+# ---------------------------------------------------------------------------
+
+#: NaN, +Inf, -Inf, -0.0 and four ordinary values.  ``tl.maximum(y, 0)`` returns
+#: the non-NaN operand and ``tl.where(y > 0, y, 0)`` fails ``NaN > 0``, so both
+#: of the obvious spellings map NaN to 0.0 where ``F.relu`` propagates it.
+_SPECIALS = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0]
+
+
+def _zero_weight_case(activation, seed=3):
+    """A case whose pre-activation is exactly ``bias``, elementwise.
+
+    Poisoning the *input* can only produce NaN pre-activations -- one non-finite
+    value makes the whole group's statistics NaN -- so the values that actually
+    distinguish the spellings of ReLU have to be placed directly.  A zero
+    ``weight`` does that: ``xhat * 0 + bias == bias``.
+    """
+    x = torch.randn(
+        1,
+        64,
+        4,
+        4,
+        4,
+        device="cuda",
+        generator=torch.Generator("cuda").manual_seed(seed),
+    ).to(memory_format=CL)
+    weight = torch.zeros(64, device="cuda")
+    bias = torch.tensor(_SPECIALS * 8, device="cuda")
+    reference = F.group_norm(x, 8, weight, bias, EPS)
+    if activation == "relu":
+        reference = F.relu(reference)
+    return x, weight, bias, reference
+
+
+@pytest.mark.gpu
+def test_fused_relu_matches_f_relu_on_nan_inf_and_negative_zero():
+    """The fused store must be ``F.relu``, bit for bit, on every special value.
+
+    NaN in, NaN out -- and that matters beyond numerics: ScaFFold aborts a run
+    whose loss goes non-finite, so an activation that turns a diverging NaN into
+    a finite 0.0 makes the forward look healthy while the backward is still NaN,
+    and the run checkpoints a broken model.  ``-Inf`` and both signed zeros must
+    come out as ``+0.0``, never ``-0.0``.
+    """
+    x, weight, bias, reference = _zero_weight_case("relu")
+    out, _mean, _rstd = torch.ops.scaffold_gn.group_norm(
+        x, 8, weight, bias, EPS, "relu", None
+    )
+    assert torch.equal(out.cpu().view(torch.int32), reference.cpu().view(torch.int32))
+    # ... and the control: without the fusion the same values pass through.
+    plain, _m, _r = torch.ops.scaffold_gn.group_norm(
+        x, 8, weight, bias, EPS, None, None
+    )
+    assert plain[0, 0, 0, 0, 0].isnan() and plain[0, 1, 0, 0, 0].isinf()
+
+
+@pytest.mark.gpu
+def test_fused_relu_backward_gates_like_threshold_backward():
+    """ReLU's backward is ``result <= 0 ? 0 : grad``, so a NaN passes.
+
+    The kernel recomputes the pre-activation and must gate with the same
+    complement: ``pre > 0 ? dy : 0`` reads identically on every finite value and
+    silently zeroes the NaN lane, which is the backward half of the same defect.
+    """
+    x, weight, bias, _reference = _zero_weight_case("relu")
+    grad_out = torch.ones_like(x)
+
+    weight = weight.requires_grad_(True)
+    bias = bias.requires_grad_(True)
+    xg = x.clone().requires_grad_(True)
+    reference = F.relu(F.group_norm(xg, 8, weight, bias, EPS))
+    reference.backward(grad_out)
+    ref_dbias = bias.grad.detach().clone()
+    ref_dx = xg.grad.detach().clone()
+
+    weight.grad = bias.grad = xg.grad = None
+    triton_group_norm(xg, 8, weight, bias, EPS, "relu").backward(grad_out)
+
+    # d_bias counts exactly the elements whose gradient the gate let through.
+    assert ref_dbias[0].item() == 64, "the reference gated the NaN lane off"
+    assert torch.equal(bias.grad.cpu(), ref_dbias.cpu())
+    assert torch.equal(xg.grad.cpu(), ref_dx.cpu())
diff --git a/tests/test_unet.py b/tests/test_unet.py
index b61ae21..fa5f47b 100644
--- a/tests/test_unet.py
+++ b/tests/test_unet.py
@@ -198,3 +198,144 @@ def counting_pad(tensor, pad, *args, **kwargs):
             f"Guard not yet in place: {exact_match_pad_calls} pad calls with diffs=0 "
             f"(expected 0 when fixed). This is the RED baseline."
         )
+
+
+# --------------------------------------------------------------------------- #
+# The decoder skip concatenation.
+#
+# ``Up.forward`` no longer calls ``torch.cat`` directly; it goes through
+# ``unet_parts._skip_concat``, which may legitimately emit a narrower dtype
+# than ``torch.cat`` would when autocast is on (see the ``Up`` docstring).
+# Everything below pins the part that must NOT change: outside autocast the
+# block is bitwise what it was, the ``F.pad`` path still works, and both the
+# ``trilinear`` and ``ConvTranspose3d`` branches agree with an explicit
+# ``torch.cat`` reference.
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("trilinear", [False, True])
+def test_up_matches_an_explicit_torch_cat_reference(trilinear):
+    """``Up.forward`` must equal ``conv(cat([x2, up(x1)]))``, bitwise, on CPU."""
+    from ScaFFold.unet.unet_parts import Up
+
+    up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=trilinear)
+    up.eval()
+    generator = torch.Generator().manual_seed(11)
+    # Either branch must hand ``self.conv`` ``in_channels`` channels: the
+    # transposed convolution halves 32 -> 16, while ``nn.Upsample`` changes no
+    # channels, so its input already carries 16.
+    x1 = torch.randn(1, 16 if trilinear else 32, 8, 8, 8, generator=generator)
+    x2 = torch.randn(1, 16, 16, 16, 16, generator=generator)
+
+    with torch.no_grad():
+        got = up(x1, x2)
+        reference = up.conv(torch.cat([x2, up.up(x1)], dim=1))
+
+    assert got.shape == reference.shape
+    assert torch.equal(got, reference), (
+        "the skip concatenation must be bitwise torch.cat outside autocast"
+    )
+
+
+def test_up_still_pads_and_concatenates_when_shapes_disagree():
+    """The non-power-of-two path: ``F.pad`` fires and the result still matches."""
+    import torch.nn.functional as F
+
+    from ScaFFold.unet.unet_parts import Up
+
+    up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False)
+    up.eval()
+    generator = torch.Generator().manual_seed(12)
+    x1 = torch.randn(1, 32, 7, 7, 7, generator=generator)  # -> 14^3 after up
+    x2 = torch.randn(1, 16, 16, 16, 16, generator=generator)  # 16^3: diff = 2
+
+    with torch.no_grad():
+        got = up(x1, x2)
+        padded = F.pad(up.up(x1), [1, 1, 1, 1, 1, 1])
+        reference = up.conv(torch.cat([x2, padded], dim=1))
+
+    assert got.shape == (1, 16, 16, 16, 16)
+    assert torch.equal(got, reference)
+
+
+def test_up_gradients_match_an_explicit_torch_cat_reference():
+    """Backward through the skip concatenation, bitwise, on CPU."""
+    from ScaFFold.unet.unet_parts import Up
+
+    up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False)
+    generator = torch.Generator().manual_seed(13)
+    x1 = torch.randn(1, 32, 8, 8, 8, generator=generator)
+    x2 = torch.randn(1, 16, 16, 16, 16, generator=generator)
+
+    a, b = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True)
+    up.zero_grad(set_to_none=True)
+    up(a, b).pow(2).sum().backward()
+    got = (a.grad.clone(), b.grad.clone())
+    got_params = {n: p.grad.clone() for n, p in up.named_parameters()}
+
+    c, d = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True)
+    up.zero_grad(set_to_none=True)
+    up.conv(torch.cat([d, up.up(c)], dim=1)).pow(2).sum().backward()
+
+    assert torch.equal(got[0], c.grad)
+    assert torch.equal(got[1], d.grad)
+    for name, param in up.named_parameters():
+        assert torch.equal(got_params[name], param.grad), name
+
+
+def test_up_concatenation_keeps_channels_last():
+    """The concatenation must not break the layout chain it exists to preserve.
+
+    Asserted on ``_skip_concat`` with two channels-last halves rather than on a
+    whole ``Up`` block: on CPU ``nn.ConvTranspose3d`` returns a *contiguous*
+    tensor whatever it is handed, so the block's own inputs to the
+    concatenation are not both channels-last there and the block-level
+    assertion would be measuring the convolution's layout policy, not this
+    one's.  On GPU with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- the production
+    configuration -- both halves are channels-last and this is the property
+    ``Up`` relies on.
+    """
+    from ScaFFold.unet.unet_parts import _skip_concat as skip_concat
+
+    generator = torch.Generator().manual_seed(14)
+    x1 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+    x2 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+    out = skip_concat(x2, x1)
+    assert out.shape == (1, 32, 16, 16, 16)
+    assert out.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(out, torch.cat([x2, x1], dim=1))
+
+
+def test_whole_model_forward_and_backward_still_agree_with_a_cat_based_up():
+    """End to end: swapping the concatenation back must change nothing on CPU."""
+    import torch as _torch
+
+    from ScaFFold.unet import unet_parts
+
+    def cat_forward(self, x1, x2):
+        x1 = self.up(x1)
+        return self.conv(_torch.cat([x2, x1], dim=1))
+
+    model = UNet(
+        n_channels=_N_CHANNELS, n_classes=_N_CLASSES, trilinear=False, layers=2
+    )
+    x = _make_input(seed=15).requires_grad_(True)
+
+    model.zero_grad(set_to_none=True)
+    model(x).pow(2).sum().backward()
+    grads = {n: p.grad.clone() for n, p in model.named_parameters()}
+    x_grad = x.grad.clone()
+
+    original = unet_parts.Up.forward
+    try:
+        unet_parts.Up.forward = cat_forward
+        x2 = _make_input(seed=15).requires_grad_(True)
+        model.zero_grad(set_to_none=True)
+        model(x2).pow(2).sum().backward()
+        for name, param in model.named_parameters():
+            assert torch.equal(grads[name], param.grad), name
+        assert torch.equal(x_grad, x2.grad)
+    finally:
+        unet_parts.Up.forward = original
diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py
index 40be05b..b5c6b32 100644
--- a/tests/test_worker_dist.py
+++ b/tests/test_worker_dist.py
@@ -21,6 +21,7 @@
 
 import logging
 
+import pytest
 import torch
 
 import ScaFFold.utils.distributed as distributed_mod
@@ -171,21 +172,21 @@ def fake_ddp(model, parallel_strategy=None, **kwargs):
 # ---------------------------------------------------------------------------
 
 
-def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset):
-    """worker.main completes end to end as a one-rank gloo job on CPU.
+def _run_singleton_worker(
+    monkeypatch, tiny_config, tiny_dataset, *, port, config_overrides=None
+):
+    """Run ``worker.main`` as a one-rank gloo job on CPU; return (rc, trainer).
 
-    ScaFFold always runs distributed; the supported singleton case is a
-    one-rank launch. The worker initializes the (gloo) process group itself,
-    builds a real unsharded ParallelStrategy, and tears the group down before
-    rank-0 post-processing.
+    Training itself is stubbed (one synthetic epoch row so post-processing has
+    data), so what this exercises is the worker's own setup path.
     """
     monkeypatch.setenv("MASTER_ADDR", "127.0.0.1")
-    monkeypatch.setenv("MASTER_PORT", "29513")
+    monkeypatch.setenv("MASTER_PORT", str(port))
     # Force the CPU path so initialize_dist selects gloo: this test must not
     # depend on a working GPU/NCCL stack.
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
 
-    cfg = tiny_config()
+    cfg = tiny_config(**(config_overrides or {}))
     kwargs = dict(vars(cfg))
     kwargs.update(
         {
@@ -212,7 +213,20 @@ def fake_train(self, profiler=None):
 
     monkeypatch.setattr(worker_mod.PyTorchTrainer, "train", fake_train)
     result = worker_mod.main(kwargs_dict=kwargs)
-    trainer = seen.get("trainer")
+    return result, seen.get("trainer")
+
+
+def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset):
+    """worker.main completes end to end as a one-rank gloo job on CPU.
+
+    ScaFFold always runs distributed; the supported singleton case is a
+    one-rank launch. The worker initializes the (gloo) process group itself,
+    builds a real unsharded ParallelStrategy, and tears the group down before
+    rank-0 post-processing.
+    """
+    result, trainer = _run_singleton_worker(
+        monkeypatch, tiny_config, tiny_dataset, port=29513
+    )
 
     assert result == 0
     assert trainer is not None
@@ -226,3 +240,133 @@ def fake_train(self, profiler=None):
     assert trainer.config.global_batch_size == trainer.config.local_batch_size
     # The worker destroyed the process group before rank-0 post-processing.
     assert not torch.distributed.is_initialized()
+
+
+# ---------------------------------------------------------------------------
+# R44: the activation-checkpointing config flag reaches the model
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("flag, expected", [(0, False), (1, True)])
+def test_activation_checkpointing_flag_reaches_the_model(
+    monkeypatch, tiny_config, tiny_dataset, flag, expected
+):
+    """``activation_checkpointing: 1`` turns the U-Net's flag on.
+
+    ``use_checkpointing`` had no caller at all, and it has to be invoked
+    *before* the DDP wrap: afterwards the model is only reachable through
+    ``.module``, which is exactly why this asserts on the wrapped model's
+    inner module.
+    """
+    _result, trainer = _run_singleton_worker(
+        monkeypatch,
+        tiny_config,
+        tiny_dataset,
+        port=29520 + flag,
+        config_overrides={"activation_checkpointing": flag},
+    )
+
+    model = getattr(trainer.model, "module", trainer.model)
+    assert model.checkpointing is expected
+
+
+# ---------------------------------------------------------------------------
+# Local size detection (R23)
+# ---------------------------------------------------------------------------
+
+_LOCAL_SIZE_CASES = [
+    # torchrun exports LOCAL_WORLD_SIZE alongside LOCAL_RANK.
+    ({"LOCAL_WORLD_SIZE": "4"}, 4),
+    ({"MV2_COMM_WORLD_LOCAL_SIZE": "4"}, 4),
+    ({"OMPI_COMM_WORLD_LOCAL_SIZE": "4"}, 4),
+    ({"PMI_LOCAL_SIZE": "4"}, 4),
+    ({"PALS_LOCAL_SIZE": "4"}, 4),
+    ({"SLURM_NTASKS": "8", "SLURM_NNODES": "2"}, 4),
+    ({"FLUX_JOB_SIZE": "8", "FLUX_JOB_NNODES": "2"}, 4),
+]
+
+_LOCAL_SIZE_VARS = [
+    "LOCAL_WORLD_SIZE",
+    "MV2_COMM_WORLD_LOCAL_SIZE",
+    "OMPI_COMM_WORLD_LOCAL_SIZE",
+    "PMI_LOCAL_SIZE",
+    "PALS_LOCAL_SIZE",
+    "SLURM_NTASKS",
+    "SLURM_NNODES",
+    "FLUX_JOB_SIZE",
+    "FLUX_JOB_NNODES",
+]
+
+
+def test_local_size_detection_matrix(monkeypatch):
+    """Every launcher that reports a local rank has its local size honored too.
+
+    An unrecognized variable silently yields 1, which makes the per-node
+    profiler gate ``rank % ranks_per_node == 0`` select *every* rank and
+    mislabels the trace's node count.
+    """
+    for env, want_local_size in _LOCAL_SIZE_CASES:
+        for var in _LOCAL_SIZE_VARS:
+            monkeypatch.delenv(var, raising=False)
+        for key, value in env.items():
+            monkeypatch.setenv(key, value)
+        assert distributed_mod.get_local_size() == want_local_size, env
+
+
+def test_local_size_defaults_to_one(monkeypatch):
+    """With nothing to go on, one rank per node is still the assumption."""
+    for var in _LOCAL_SIZE_VARS:
+        monkeypatch.delenv(var, raising=False)
+    assert distributed_mod.get_local_size() == 1
+
+
+# ---------------------------------------------------------------------------
+# VC-3: an unusable launcher variable is ignored, not fatal
+#
+# These helpers run at the top of every entry point, so a bare int() on a
+# variable a site wrapper exported empty ("WORLD_SIZE=") or as a placeholder
+# ("auto") killed the invocation -- ``scaffold --help`` included -- with a
+# ValueError naming neither the variable nor a remedy.
+# ---------------------------------------------------------------------------
+
+
+def _clear_launcher_env(monkeypatch):
+    for var in set(_LAUNCHER_VARS) | set(_LOCAL_SIZE_VARS):
+        monkeypatch.delenv(var, raising=False)
+
+
+@pytest.mark.parametrize("value", ["", "   ", "auto"])
+def test_unusable_launcher_values_fall_through(monkeypatch, value):
+    """Empty and non-numeric values are treated as absent, never raise."""
+    _clear_launcher_env(monkeypatch)
+    for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "LOCAL_WORLD_SIZE"):
+        monkeypatch.setenv(var, value)
+
+    # Falls through to the next source -- here the (singleton) communicator and
+    # the documented defaults.
+    assert distributed_mod.get_world_size() == 1
+    assert distributed_mod.get_world_rank() == 0
+    assert distributed_mod.get_local_rank() == 0
+    assert distributed_mod.get_local_size() == 1
+
+
+def test_unusable_value_defers_to_the_next_launcher_variable(monkeypatch, caplog):
+    """A garbage value does not mask a usable variable further down the order."""
+    _clear_launcher_env(monkeypatch)
+    monkeypatch.setenv("WORLD_SIZE", "auto")
+    monkeypatch.setenv("PALS_NRANKS", "8")
+
+    with caplog.at_level(logging.WARNING, logger=distributed_mod.logger.name):
+        assert distributed_mod.get_world_size() == 8
+
+    messages = " ".join(record.getMessage() for record in caplog.records)
+    assert "WORLD_SIZE" in messages, "the ignored value was not reported"
+
+
+def test_zero_node_count_does_not_divide_by_zero(monkeypatch):
+    """A nonsense node count falls through instead of raising."""
+    _clear_launcher_env(monkeypatch)
+    monkeypatch.setenv("SLURM_NTASKS", "8")
+    monkeypatch.setenv("SLURM_NNODES", "0")
+
+    assert distributed_mod.get_local_size() == 1