diff --git a/AGENTS.md b/AGENTS.md
index 2698812..f21c0ef 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -106,9 +106,17 @@ below the ensemble member-to-member spread (docs/validation.md).
(zero precipitation, land-filled SST), so rounding-level forward differences flip signs. Between
two runs of the same comparison (official vs XLA-autodiff path, H200, `highest`) the gradient
difference went from 1.1e-3 to 4.3e-3. Judge implementations with the smooth `--loss linear`.
-- **XLA autotuning is timing-based**: kernel picks (and last bits) differ between runs; once, an
- A100 fp32 run deviated at TF32 level (1 of 6, not reproduced). Use `--autotune-cache` for
- reproducible runs and `--no-autotune` for strict fp32.
+- **XLA autotuning is timing-based**: kernel picks (and last bits) differ between runs. Use
+ `--autotune-cache` for reproducible runs.
+- **The grid↔mesh graphs depend on the host CPU** (official code): they are built with NumPy from
+ float32 lat/lon at trace time, and NumPy's math code path, picked from the host CPU (AVX-512,
+ AVX2, baseline), changes last bits and with them which grid points near the ball-query radius
+ get an edge. An AVX2 host
+ (Modal's A100-40GB) and an AVX-512 host differ in 13 of 1.6M grid→mesh edges and 517 mesh→grid
+ triangles: 1.3e-2 max rel RMS at `highest`, which looks like TF32 and is not. Before reading a
+ cross-machine fp32 difference as GPU numerics, compare the graph hashes in `cloud/canary_run.py`'s
+ JSON (or `scripts/graph_edges.py`); `--edges` pins a graph. Modal may serve "A100-40GB" requests
+ on 80GB cards: check `device`.
- **In-place block updates** need the new block wrapped in `optimization_barrier` before the
`dynamic_update_slice`, or XLA fuses the read into the update and copies the whole buffer.
- **Transformer layers as a loop** (`hk.scan` + `hk.switch`) avoid XLA heap fragmentation
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4ad627b..445a838 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,18 @@
# Changelog
+## Unreleased
+
+- The A100 fp32 deviation in the canaries (1.3e-2 max rel RMS vs the RTX 5090) is explained, and
+ it is not GPU precision. The official code builds its grid↔mesh graphs with NumPy on the host
+ CPU, and Modal's A100-40GB hosts (AVX2 only) build a slightly different graph than AVX-512
+ hosts. `cloud/canary_run.py` records the graph hashes and NumPy's SIMD targets, runs on
+ another machine's graph with `--edges`, and skips the rollout with `--steps 0`. New
+ `scripts/graph_edges.py` counts and saves the graphs on the CPU; new `cloud/modal_a100_tf32.py`
+ and `cloud/modal_graph_hosts.py` are the probes. Validation: docs/validation.md §4, A100 note.
+- README FAQ and `--no-autotune` help: no longer recommend `--no-autotune` for strict fp32
+ (`--precision highest` gives fp32-level results; autotuning was not the cause). The FAQ now
+ says why results can differ between machines.
+
## 0.2.0 — 2026-11-01
Fine-tuning (experimental), IFS open-data mirror fallback, Intel Arc. Inference is unchanged:
diff --git a/README.md b/README.md
index b85c439..da23078 100644
--- a/README.md
+++ b/README.md
@@ -168,7 +168,9 @@ No. Weights are the official fp32 checkpoints, attention is exact, and every lay
Why aren't results bit-identical between runs?
-XLA picks GPU kernels by timing them, so kernel choices — and the last bits of the result — vary between runs, with or without this project. Use `fwn forecast --autotune-cache kernels.txt` to save the choices on the first run and reuse them. For strict fp32, use `--precision highest --no-autotune`.
+XLA picks GPU kernels by timing them, so kernel choices — and the last bits of the result — vary between runs, with or without this project. Use `fwn forecast --autotune-cache kernels.txt` to save the choices on the first run and reuse them (`--no-autotune` also fixes them; neither is needed for fp32-level results, which `--precision highest` gives).
+
+Between machines, results can differ by more than that: the official code builds its grid↔mesh graph on the CPU with NumPy, and CPUs with and without AVX-512 get a few different edges. In one fp32 step that is up to 1.3 × 10⁻² relative RMS in the most sensitive field (median 6 × 10⁻⁴), around those edges ([details](docs/validation.md#4-other-gpus-modal-vastai)).
diff --git a/cloud/canary_run.py b/cloud/canary_run.py
index 2786d91..3b689cf 100644
--- a/cloud/canary_run.py
+++ b/cloud/canary_run.py
@@ -6,12 +6,20 @@
1. one step at precision "highest", saved on every 4th grid point (--npz) for cross-GPU
comparison (scripts/compare_npz.py);
- 2. an N-step rollout at default precision: per-step wall time, finite check, pool peak.
+ 2. an N-step rollout at default precision: per-step wall time, finite check, pool peak
+ (--steps 0 skips it, e.g. for precision probes).
+
+The JSON also records the grid<->mesh graphs (edge count and hash) and NumPy's SIMD targets:
+the official code builds the graphs on the host CPU, and the grid->mesh edge set depends on
+NumPy's code path there (docs/validation.md §4, A100 note). Compare fp32 outputs across machines
+only when the graph hashes match, or pass --edges (indices saved by scripts/graph_edges.py
+--save-edges on another machine) to run on that machine's graphs.
python cloud/canary_run.py --data SAMPLE.nc --json out.json [--npz out.npz] [--steps 4]
"""
import argparse
import dataclasses
+import hashlib
import json
import os
import time
@@ -24,6 +32,7 @@
p.add_argument("--npz", default="")
p.add_argument("--steps", type=int, default=4)
p.add_argument("--label", default="")
+p.add_argument("--edges", default="", help="grid<->mesh edge indices from scripts/graph_edges.py --save-edges")
args = p.parse_args()
res = dict(label=args.label, env={k: v for k, v in os.environ.items() if k.startswith(("XLA_", "JAX_"))}, ok=False)
@@ -34,8 +43,25 @@
import xarray
import faster_weathernext
from faster_weathernext import model, patch
+ from numpy._core._multiarray_umath import __cpu_dispatch__, __cpu_features__
from weathernext.utils import data_utils
+ from weathernext.utils import points_mesh_gnn
+
+ res.update(cpu=next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), ""),
+ numpy=np.__version__, numpy_simd=[f for f in __cpu_dispatch__ if __cpu_features__.get(f)], graphs={})
+ get_edge_indices = points_mesh_gnn.PointsMeshTypedGraphGNN._get_edge_indices
+
+ saved = dict(np.load(args.edges)) if args.edges else None
+
+ def record_edges(self, *a):
+ name = self._connectivity_type.name
+ edges = ((saved[f"{name}_points"].astype(int), saved[f"{name}_mesh"].astype(int)) if saved
+ else get_edge_indices(self, *a))
+ res["graphs"][name] = dict(n_edges=int(len(edges[0])), sha1=hashlib.sha1(
+ np.stack(edges).astype(np.int64).tobytes()).hexdigest()[:16], loaded=bool(saved))
+ return edges
+ points_mesh_gnn.PointsMeshTypedGraphGNN._get_edge_indices = record_edges
faster_weathernext.enable()
G = 2**30
dev = jax.devices()[0]
@@ -43,7 +69,7 @@
compute_capability=getattr(dev, "compute_capability", ""), attention=patch.attention_impl(),
pool_gib=(dev.memory_stats() or {}).get("bytes_limit", 0) / G)
print(f"{dev.device_kind} cc={res['compute_capability']} attention={res['attention']} "
- f"pool {res['pool_gib']:.2f} GiB", flush=True)
+ f"pool {res['pool_gib']:.2f} GiB; host {res['cpu']}, NumPy SIMD {res['numpy_simd']}", flush=True)
task, forward = model.build("WeatherNextCyclones")
params = model.load_params("WeatherNextCyclones", 1, args.weights_dir)
ds = xarray.load_dataset(args.data)
@@ -67,27 +93,29 @@
print(f"highest-precision step saved ({len(fields)} fields) in {res['highest_first_call_s']:.0f}s", flush=True)
del pred
- lead = pd.to_timedelta(np.arange(1, args.steps + 1) * 6, unit="h")
- tvars = {}
- for v, da in targets1.data_vars.items():
- shape = list(da.shape)
- shape[da.dims.index("time")] = args.steps
- tvars[v] = (da.dims, np.broadcast_to(np.float32(np.nan), shape))
- targets = xarray.Dataset(tvars, coords={**{k: c for k, c in targets1.coords.items() if k != "time"}, "time": lead})
- init = pd.Timestamp(ds["datetime"].values[0, 1]).to_pydatetime()
- fds = xarray.Dataset(coords=dict(time=lead, lon=ds.lon, datetime=(("batch", "time"), np.array(
- [[init + td for td in lead.to_pytimedelta()]], dtype="datetime64[ns]"))))
- data_utils.add_derived_vars(fds)
- forcings = fds[list(task.forcing_variables)].drop_vars("datetime")
- forcings = forcings.astype({v: forcings1[v].dtype for v in forcings1.data_vars})
res["step_s"], res["finite"] = [], []
- t0 = time.time()
- for pred in model.rollout(forward, params, jax.random.PRNGKey(0), inputs, targets, forcings):
- msl = np.asarray(pred["mean_sea_level_pressure"])
- res["step_s"].append(time.time() - t0)
- res["finite"].append(bool(np.isfinite(msl).all()))
- print(f" step {len(res['step_s'])}: {res['step_s'][-1]:.1f}s msl min {np.nanmin(msl) / 100:.1f} hPa", flush=True)
+ if args.steps > 0:
+ lead = pd.to_timedelta(np.arange(1, args.steps + 1) * 6, unit="h")
+ tvars = {}
+ for v, da in targets1.data_vars.items():
+ shape = list(da.shape)
+ shape[da.dims.index("time")] = args.steps
+ tvars[v] = (da.dims, np.broadcast_to(np.float32(np.nan), shape))
+ targets = xarray.Dataset(tvars, coords={**{k: c for k, c in targets1.coords.items() if k != "time"}, "time": lead})
+ init = pd.Timestamp(ds["datetime"].values[0, 1]).to_pydatetime()
+ fds = xarray.Dataset(coords=dict(time=lead, lon=ds.lon, datetime=(("batch", "time"), np.array(
+ [[init + td for td in lead.to_pytimedelta()]], dtype="datetime64[ns]"))))
+ data_utils.add_derived_vars(fds)
+ forcings = fds[list(task.forcing_variables)].drop_vars("datetime")
+ forcings = forcings.astype({v: forcings1[v].dtype for v in forcings1.data_vars})
t0 = time.time()
+ for pred in model.rollout(forward, params, jax.random.PRNGKey(0), inputs, targets, forcings):
+ msl = np.asarray(pred["mean_sea_level_pressure"])
+ res["step_s"].append(time.time() - t0)
+ res["finite"].append(bool(np.isfinite(msl).all()))
+ print(f" step {len(res['step_s'])}: {res['step_s'][-1]:.1f}s msl min {np.nanmin(msl) / 100:.1f} hPa", flush=True)
+ t0 = time.time()
+ print(f"graphs: {res['graphs']}", flush=True)
res["peak_gib"] = (dev.memory_stats() or {}).get("peak_bytes_in_use", 0) / G
res["ok"] = all(res["finite"])
print(f"peak {res['peak_gib']:.2f} GiB of pool {res['pool_gib']:.2f} GiB; ok={res['ok']}", flush=True)
diff --git a/cloud/modal_a100_tf32.py b/cloud/modal_a100_tf32.py
new file mode 100644
index 0000000..afd8ab4
--- /dev/null
+++ b/cloud/modal_a100_tf32.py
@@ -0,0 +1,123 @@
+"""A100 fp32 probe: why does the canary's "highest" step land at TF32 level on A100-SXM4-40GB?
+
+Runs the canary's fp32 step (cloud/canary_run.py --steps 0) in fresh processes with different
+XLA settings on an A100-40GB and an A100-80GB. The base run on each card also dumps the autotune
+picks, the HLO, the thunk sequence and XLA's device description. Compare each .npz with an
+RTX 5090 run of the same code (scripts/compare_npz.py). The "graph" step saves the host's
+grid<->mesh graphs and "edges5090" runs the step on the RTX 5090's graphs instead (the cause:
+see docs/validation.md §4, A100 note).
+
+ modal run --detach cloud/modal_a100_tf32.py --tag tf32probe01
+ modal run --detach cloud/modal_a100_tf32.py --tag tf32probe02 --v40 graph,edges5090 --v80 ""
+ modal run --detach cloud/modal_a100_tf32.py --tag tf32probe03 --v40 "" --v80 numpy_avx2
+ modal volume get fwn-canary results/tf32probe01 ./tf32probe01
+ modal app stop -y # detached apps can linger after finishing
+
+Cost (2026 prices): about $1 (A100-40GB ~18 min at ~$2.10/h, A100-80GB ~6 min at ~$2.50/h).
+"""
+import os
+import subprocess
+import time
+
+import modal
+
+REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA = "source-hres_forecast_init-2024-10-07 00:00:00_res-0.25_levels-13_steps-01.nc"
+
+app = modal.App("fwn-a100-tf32")
+vol = modal.Volume.from_name("fwn-canary", create_if_missing=True)
+image = (
+ modal.Image.debian_slim(python_version="3.12")
+ .apt_install("git")
+ .add_local_dir(REPO, "/opt/faster-weathernext", copy=True,
+ ignore=[".git", "**/__pycache__", "docs/results", "*.nc", "*.npz"])
+ .run_commands("pip install '/opt/faster-weathernext[cuda12,weathernext]'")
+)
+
+# name -> (extra environment, dump autotune picks + HLO, extra canary_run.py arguments)
+VARIANTS = {
+ "base": ({}, True, []),
+ "at0": ({"XLA_FLAGS": "--xla_gpu_autotune_level=0"}, False, []),
+ "tf32off": ({"NVIDIA_TF32_OVERRIDE": "0"}, False, []),
+ "notritongemm": ({"XLA_FLAGS": "--xla_gpu_enable_triton_gemm=false"}, False, []),
+ # The RTX 5090's grid<->mesh graphs (scripts/graph_edges.py --save-edges, uploaded with
+ # `modal volume put fwn-canary edges_rtx5090.npz edges/edges_rtx5090.npz`).
+ "edges5090": ({}, False, ["--edges", "/data/edges/edges_rtx5090.npz"]),
+ # NumPy's AVX2 code path on an AVX-512 host: same host-side graph and features as the
+ # A100-40GB hosts (AVX2 only), so the GPUs are compared on identical inputs.
+ "numpy_avx2": ({"NPY_DISABLE_CPU_FEATURES": "X86_V4 AVX512_ICL AVX512_SPR"}, False, []),
+}
+
+
+def _run(card: str, tag: str, names: list[str]) -> dict:
+ import json
+ out = f"/data/results/{tag}"
+ os.makedirs(out, exist_ok=True)
+ smi = subprocess.run(["nvidia-smi", "--query-gpu=name,driver_version,memory.total,clocks.max.sm,ecc.mode.current",
+ "--format=csv"], capture_output=True, text=True).stdout
+ summary = {"card": card, "nvidia_smi": smi,
+ "env": {k: v for k, v in os.environ.items() if any(s in k for s in ("TF32", "XLA", "JAX", "CUBLAS", "CUDNN"))}}
+ print(f"[{card}] {smi}", flush=True)
+ for name in names:
+ if name == "graph": # this host's graphs, saved for reuse elsewhere (CPU only)
+ subprocess.run(["python", "/opt/faster-weathernext/scripts/graph_edges.py", "--weights-dir", "/data",
+ "--data", f"/data/{DATA}", "--json", f"{out}/{card}_graph.json",
+ "--save-edges", f"{out}/{card}_edges.npz"], env={**os.environ, "JAX_PLATFORMS": "cpu"})
+ summary[name] = json.load(open(f"{out}/{card}_graph.json"))
+ print(f"[{card}/{name}] {summary[name]['cpu']} {summary[name]['dispatch']} "
+ + "; ".join(f"{g['connectivity']} {g['n_edges']} {g['sha1']}" for g in summary[name]["graphs"]), flush=True)
+ vol.commit()
+ continue
+ extra, dump, extra_args = VARIANTS[name]
+ env = {**os.environ, "XLA_PYTHON_CLIENT_PREALLOCATE": "false", **extra}
+ if dump:
+ d = f"/tmp/dump_{card}_{name}"
+ env["XLA_FLAGS"] = (env.get("XLA_FLAGS", "") + f" --xla_dump_to={d} --xla_dump_hlo_as_text"
+ " --xla_dump_hlo_module_re=.*lambda.*"
+ f" --xla_gpu_dump_autotune_results_to={out}/{card}_{name}_autotune.txt").strip()
+ cmd = ["python", "/opt/faster-weathernext/cloud/canary_run.py", "--weights-dir", "/data",
+ "--data", f"/data/{DATA}", "--json", f"{out}/{card}_{name}.json", "--npz", f"{out}/{card}_{name}.npz",
+ "--steps", "0", "--label", f"{card}_{name}", *extra_args]
+ t0 = time.time()
+ proc = subprocess.run(cmd, env=env, capture_output=True, text=True)
+ with open(f"{out}/{card}_{name}.log", "w") as f:
+ f.write(proc.stdout + proc.stderr)
+ if dump:
+ # HLO before/after optimization, thunk sequence (which op runs as which kernel or library
+ # call), device description and autotune picks; not the large buffer-assignment files.
+ keep = r"optimizations.txt$|thunk_sequence.txt$|gpu_target_config.pbtxt$|autotune_results.pbtxt$|memory-usage-report.txt$"
+ subprocess.run(f"cd {d} && tar czf {out}/{card}_{name}_dump.tgz $(ls | grep -E '{keep}')", shell=True, check=False)
+ path = f"{out}/{card}_{name}.json"
+ summary[name] = json.load(open(path)) if os.path.exists(path) else {"error": f"exit {proc.returncode}"}
+ summary[name].pop("traceback", None)
+ summary[name]["wall_s"] = time.time() - t0
+ print(f"[{card}/{name}] exit {proc.returncode} device {summary[name].get('device')} "
+ f"{summary[name]['wall_s']:.0f}s {summary[name].get('error', '')}", flush=True)
+ vol.commit()
+ with open(f"{out}/{card}_summary.json", "w") as f:
+ json.dump(summary, f, indent=1)
+ vol.commit()
+ return summary
+
+
+@app.function(image=image, gpu="A100-40GB", volumes={"/data": vol}, timeout=3600, cpu=8, memory=32768)
+def a100_40(tag: str, names: list[str]):
+ return _run("a100_40", tag, names)
+
+
+@app.function(image=image, gpu="A100-80GB", volumes={"/data": vol}, timeout=3600, cpu=8, memory=32768)
+def a100_80(tag: str, names: list[str]):
+ return _run("a100_80", tag, names)
+
+
+@app.local_entrypoint()
+def main(tag: str = "tf32probe01", v40: str = "base,at0,tf32off,notritongemm", v80: str = "base"):
+ calls = []
+ if v40:
+ calls.append(a100_40.spawn(tag, v40.split(",")))
+ if v80:
+ calls.append(a100_80.spawn(tag, v80.split(",")))
+ for call in calls:
+ r = call.get()
+ print(f"== {r['card']}: " + ", ".join(f"{k}: {v.get('device')} ok={v.get('ok')} {v.get('error', '')}"
+ for k, v in r.items() if isinstance(v, dict) and k != "env"), flush=True)
diff --git a/cloud/modal_graph_hosts.py b/cloud/modal_graph_hosts.py
new file mode 100644
index 0000000..11234f4
--- /dev/null
+++ b/cloud/modal_graph_hosts.py
@@ -0,0 +1,86 @@
+"""Which grid->mesh graph does each Modal GPU host build? (host CPU, NumPy SIMD dispatch, edges)
+
+Runs scripts/graph_edges.py (CPU only: traces one step, compiles nothing) on the host of each
+GPU type and records the CPU model, the CPU flags the container sees and NumPy's dispatch
+targets next to the edge count and hash. A host whose graph differs gives different fp32
+outputs on any GPU (docs/validation.md §4, A100 note).
+
+ modal run --detach cloud/modal_graph_hosts.py --cards a100,a100_80,a10g,l4,t4 --tag graphhosts01
+ modal volume get fwn-canary results/graphhosts01 ./graphhosts01
+ modal app stop -y
+
+Cost (2026 prices): about $0.05-0.15 per card (~3 min each).
+"""
+import os
+import subprocess
+
+import modal
+
+REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA = "source-hres_forecast_init-2024-10-07 00:00:00_res-0.25_levels-13_steps-01.nc"
+
+app = modal.App("fwn-graph-hosts")
+vol = modal.Volume.from_name("fwn-canary", create_if_missing=True)
+image = (
+ modal.Image.debian_slim(python_version="3.12")
+ .apt_install("git")
+ .add_local_dir(REPO, "/opt/faster-weathernext", copy=True,
+ ignore=[".git", "**/__pycache__", "docs/results", "*.nc", "*.npz"])
+ .run_commands("pip install '/opt/faster-weathernext[cuda12,weathernext]'")
+)
+
+
+def _graph(card: str, tag: str) -> dict:
+ import json
+ out = f"/data/results/{tag}"
+ os.makedirs(out, exist_ok=True)
+ gpu = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True).stdout.strip()
+ flags = next((l.split(":", 1)[1].split() for l in open("/proc/cpuinfo") if l.startswith("flags")), [])
+ path = f"{out}/{card}_graph.json"
+ proc = subprocess.run(["python", "/opt/faster-weathernext/scripts/graph_edges.py", "--weights-dir", "/data",
+ "--data", f"/data/{DATA}", "--json", path],
+ env={**os.environ, "JAX_PLATFORMS": "cpu"}, capture_output=True, text=True)
+ res = json.load(open(path)) if os.path.exists(path) else {"error": (proc.stdout + proc.stderr)[-3000:]}
+ res.update(card=card, gpu=gpu, nproc=os.cpu_count(),
+ cpu_flags={f: f in flags for f in ("sse4_2", "avx", "avx2", "fma", "avx512f", "avx512bw", "avx512vl")})
+ with open(path, "w") as f:
+ json.dump(res, f, indent=1)
+ vol.commit()
+ print(f"[{card}] {gpu} | {res.get('cpu')} | flags {res['cpu_flags']} | numpy dispatch {res.get('dispatch')} | "
+ + "; ".join(f"{g['connectivity']} {g['n_edges']} {g['sha1']}" for g in res.get("graphs", [])), flush=True)
+ return res
+
+
+@app.function(image=image, gpu="A100-40GB", volumes={"/data": vol}, timeout=1200, cpu=8, memory=32768)
+def g_a100(tag: str):
+ return _graph("a100", tag)
+
+
+@app.function(image=image, gpu="A100-80GB", volumes={"/data": vol}, timeout=1200, cpu=8, memory=32768)
+def g_a100_80(tag: str):
+ return _graph("a100_80", tag)
+
+
+@app.function(image=image, gpu="A10G", volumes={"/data": vol}, timeout=1200, cpu=8, memory=32768)
+def g_a10g(tag: str):
+ return _graph("a10g", tag)
+
+
+@app.function(image=image, gpu="L4", volumes={"/data": vol}, timeout=1200, cpu=8, memory=32768)
+def g_l4(tag: str):
+ return _graph("l4", tag)
+
+
+@app.function(image=image, gpu="T4", volumes={"/data": vol}, timeout=1200, cpu=8, memory=32768)
+def g_t4(tag: str):
+ return _graph("t4", tag)
+
+
+CARDS = {"a100": g_a100, "a100_80": g_a100_80, "a10g": g_a10g, "l4": g_l4, "t4": g_t4}
+
+
+@app.local_entrypoint()
+def main(cards: str = "a100,a100_80,a10g,l4,t4", tag: str = "graphhosts01"):
+ calls = {c: CARDS[c].spawn(tag) for c in cards.split(",")}
+ for c, call in calls.items():
+ call.get()
diff --git a/docs/results/a100_host_graph/cross_card_fp32.txt b/docs/results/a100_host_graph/cross_card_fp32.txt
new file mode 100644
index 0000000..bc04464
--- /dev/null
+++ b/docs/results/a100_host_graph/cross_card_fp32.txt
@@ -0,0 +1,102 @@
+# fp32 ("highest") one-step outputs, every 4th grid point, 101 fields; scripts/compare_npz.py A B.
+# rtx5090 = docs/results/canary_modal/v0.2/rtx5090.npz (AVX-512 host graph, sha1 189d06f0830e1c31).
+# A100-40GB host graph: sha1 1d90c1ed7b086815 (1,618,832 grid->mesh edges vs 1,618,841).
+
+## 1. Reproduction and XLA/TF32 knobs on the A100-SXM4-40GB (tf32probe01, cloud/modal_a100_tf32.py)
+== A100-40GB base vs rtx5090
+101 fields min corr 0.999913072909 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130729 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802987 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813396 rel RMS 6.11e-03
+== A100-40GB --xla_gpu_autotune_level=0 vs rtx5090
+101 fields min corr 0.999913072036 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130720 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802981 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813398 rel RMS 6.11e-03
+== A100-40GB NVIDIA_TF32_OVERRIDE=0 vs rtx5090
+101 fields min corr 0.999913071873 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130719 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802993 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813400 rel RMS 6.11e-03
+== A100-40GB --xla_gpu_enable_triton_gemm=false vs rtx5090
+101 fields min corr 0.999913075693 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130757 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802899 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813399 rel RMS 6.11e-03
+== A100 80GB PCIe base vs rtx5090
+101 fields min corr 0.999999999834 max rel RMS 1.83e-05 median rel RMS 1.96e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999998 rel RMS 1.83e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.45e-05
+ vertical_velocity@150 corr 0.9999999999 rel RMS 1.39e-05
+== A100-40GB canary02 (v0.1.0) vs A100-40GB canary03 (v0.2): same host graph, reproducible
+101 fields min corr 0.999999999894 max rel RMS 1.46e-05 median rel RMS 1.47e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.46e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.10e-05
+ vertical_velocity@150 corr 0.9999999999 rel RMS 1.06e-05
+
+## 2. Graph swap (scripts/graph_edges.py --save-edges, cloud/canary_run.py --edges)
+== RTX 5090 on the A100-40GB host graph vs RTX 5090 on its own graph
+101 fields min corr 0.999913080664 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130807 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999803010 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813436 rel RMS 6.11e-03
+== A100-40GB on the RTX 5090 graph vs A100-40GB on its own graph (tf32probe02)
+101 fields min corr 0.999913084050 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130841 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802932 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813383 rel RMS 6.11e-03
+== A100-40GB on the RTX 5090 graph vs rtx5090
+101 fields min corr 0.999999992431 max rel RMS 1.31e-04 median rel RMS 9.69e-06 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999999924 rel RMS 1.31e-04
+ vertical_velocity@500 corr 0.9999999972 rel RMS 7.47e-05
+ vertical_velocity@600 corr 0.9999999972 rel RMS 7.44e-05
+== RTX 5090 on the A100-40GB host graph vs A100-40GB
+101 fields min corr 0.999999991545 max rel RMS 1.37e-04 median rel RMS 1.11e-05 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999999915 rel RMS 1.37e-04
+ cyclone_exists_gaussian_unit_mode corr 0.9999999963 rel RMS 9.04e-05
+ vertical_velocity@600 corr 0.9999999962 rel RMS 8.75e-05
+
+## 3. Host-side features (same graph indices, NumPy code path switched with NPY_DISABLE_CPU_FEATURES on the RTX 5090 desktop)
+== RTX 5090, own graph indices, NumPy AVX2 path vs NumPy AVX-512 path
+101 fields min corr 0.999999992186 max rel RMS 1.33e-04 median rel RMS 1.01e-05 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999999922 rel RMS 1.33e-04
+ cyclone_exists_gaussian_unit_mode corr 0.9999999969 rel RMS 8.36e-05
+ vertical_velocity@150 corr 0.9999999972 rel RMS 7.53e-05
+== RTX 5090, A100-40GB host graph indices, NumPy AVX2 path vs A100-40GB
+101 fields min corr 0.999999997506 max rel RMS 7.06e-05 median rel RMS 6.81e-06 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999999975 rel RMS 7.06e-05
+ vertical_velocity@500 corr 0.9999999982 rel RMS 6.00e-05
+ vertical_velocity@600 corr 0.9999999982 rel RMS 5.95e-05
+== RTX 5090, NumPy baseline path (own graph build: 1,618,832 edges, sha1 97d163811acc313e) vs RTX 5090
+101 fields min corr 0.999943218592 max rel RMS 1.07e-02 median rel RMS 7.12e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999432186 rel RMS 1.07e-02
+ vertical_velocity@50 corr 0.9999815077 rel RMS 6.08e-03
+ cyclone_all_wind_disc corr 0.9999886737 rel RMS 4.76e-03
+== RTX 5090 own graph loaded with --edges vs RTX 5090 (run-to-run level)
+101 fields min corr 0.999999999899 max rel RMS 1.42e-05 median rel RMS 1.55e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.42e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.08e-05
+ vertical_velocity@150 corr 0.9999999999 rel RMS 1.06e-05
+== RTX 5090 today vs rtx5090 (run-to-run level)
+101 fields min corr 0.999999999896 max rel RMS 1.44e-05 median rel RMS 1.54e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.44e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.11e-05
+ vertical_velocity@200 corr 0.9999999999 rel RMS 1.08e-05
+
+## 4. Same host inputs on different GPUs (tf32probe03): A100 80GB PCIe with NPY_DISABLE_CPU_FEATURES="X86_V4 AVX512_ICL AVX512_SPR"
+## (NumPy AVX2 path on an AVX-512 host) builds the A100-40GB host graph (sha1 1d90c1ed7b086815)
+== A100-80GB, NumPy AVX2 path vs A100-40GB (tf32probe01)
+101 fields min corr 0.999999999879 max rel RMS 1.69e-05 median rel RMS 1.47e-06 NaN masks equal: True
+ cyclone_exists_gaussian_unit_mode corr 0.9999999999 rel RMS 1.69e-05
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.50e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.11e-05
+== A100-80GB, NumPy AVX2 path vs A100-40GB canary03
+101 fields min corr 0.999999999892 max rel RMS 1.47e-05 median rel RMS 1.48e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.47e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.10e-05
+ vertical_velocity@150 corr 0.9999999999 rel RMS 1.06e-05
+== A100-80GB, NumPy AVX2 path vs the same card with NumPy's default (AVX-512) path
+101 fields min corr 0.999913071891 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130719 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802987 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813399 rel RMS 6.11e-03
diff --git a/docs/results/a100_host_graph/differing_edges.txt b/docs/results/a100_host_graph/differing_edges.txt
new file mode 100644
index 0000000..1f5ac41
--- /dev/null
+++ b/docs/results/a100_host_graph/differing_edges.txt
@@ -0,0 +1,18 @@
+# grid<->mesh edges that differ between the AVX-512 host graph (RTX 5090 desktop, Modal T4 / A100-80GB;
+# sha1 189d06f0830e1c31) and the A100-SXM4-40GB host graph (NumPy X86_V3; sha1 1d90c1ed7b086815)
+BALL_QUERY: 1618841 vs 1618832 edges; 11 only in the AVX-512 graph, 2 only in the 40GB-host graph; 13 grid points affected
+ grid ( -87.75, 30.50) mesh node 17483 ( -87.538, 13.924) AVX-512 graph only
+ grid ( -87.75, 89.50) mesh node 17486 ( -87.538, 106.076) AVX-512 graph only
+ grid ( -87.75, 150.50) mesh node 17166 ( -87.538, 133.924) AVX-512 graph only
+ grid ( -87.75, 209.50) mesh node 16203 ( -87.538, 226.076) AVX-512 graph only
+ grid ( -87.75, 270.50) mesh node 16202 ( -87.538, 253.924) AVX-512 graph only
+ grid ( -87.75, 329.50) mesh node 17162 ( -87.538, 346.076) AVX-512 graph only
+ grid ( -54.75, 33.50) mesh node 27691 ( -54.092, 33.962) AVX-512 graph only
+ grid ( -54.75, 206.50) mesh node 7200 ( -54.092, 206.038) AVX-512 graph only
+ grid ( -54.75, 326.50) mesh node 16171 ( -54.092, 326.038) AVX-512 graph only
+ grid ( -18.75, 28.75) mesh node 35351 ( -18.056, 28.590) AVX-512 graph only
+ grid ( 9.50, 111.00) mesh node 33731 ( 9.720, 110.315) AVX-512 graph only
+ grid ( 22.75, 121.75) mesh node 30499 ( 22.485, 121.036) 40GB-host graph only
+ grid ( 23.00, 285.75) mesh node 9005 ( 22.890, 284.988) 40GB-host graph only
+IN_TRIANGLE: 3114720 vs 3114720 edges; 517 only in the AVX-512 graph, 517 only in the 40GB-host graph; 517 grid points affected
+ most common longitudes of the affected grid points: 240° (100), 0° (98), 120° (98), 180° (74), 60° (68), 300° (49), 26.25° (1), 18° (1)
diff --git a/docs/results/a100_host_graph/errmap_a100_40gb_vs_rtx5090.png b/docs/results/a100_host_graph/errmap_a100_40gb_vs_rtx5090.png
new file mode 100644
index 0000000..242a0e3
Binary files /dev/null and b/docs/results/a100_host_graph/errmap_a100_40gb_vs_rtx5090.png differ
diff --git a/docs/results/a100_host_graph/graphhosts01/a100_80_graph.json b/docs/results/a100_host_graph/graphhosts01/a100_80_graph.json
new file mode 100644
index 0000000..bd6313a
--- /dev/null
+++ b/docs/results/a100_host_graph/graphhosts01/a100_80_graph.json
@@ -0,0 +1,48 @@
+{
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": true,
+ "dispatch": [
+ "X86_V3",
+ "X86_V4"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404078617691994,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 3,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d"
+ }
+ ],
+ "card": "a100_80",
+ "gpu": "NVIDIA A100-SXM4-80GB",
+ "nproc": 24,
+ "cpu_flags": {
+ "sse4_2": true,
+ "avx": true,
+ "avx2": true,
+ "fma": true,
+ "avx512f": true,
+ "avx512bw": true,
+ "avx512vl": true
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/graphhosts01/a100_graph.json b/docs/results/a100_host_graph/graphhosts01/a100_graph.json
new file mode 100644
index 0000000..6e98036
--- /dev/null
+++ b/docs/results/a100_host_graph/graphhosts01/a100_graph.json
@@ -0,0 +1,48 @@
+{
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": true,
+ "dispatch": [
+ "X86_V3",
+ "X86_V4"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404078617691994,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 3,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d"
+ }
+ ],
+ "card": "a100",
+ "gpu": "NVIDIA A100-SXM4-80GB",
+ "nproc": 24,
+ "cpu_flags": {
+ "sse4_2": true,
+ "avx": true,
+ "avx2": true,
+ "fma": true,
+ "avx512f": true,
+ "avx512bw": true,
+ "avx512vl": true
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/graphhosts01/a10g_graph.json b/docs/results/a100_host_graph/graphhosts01/a10g_graph.json
new file mode 100644
index 0000000..8b01969
--- /dev/null
+++ b/docs/results/a100_host_graph/graphhosts01/a10g_graph.json
@@ -0,0 +1,47 @@
+{
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": false,
+ "dispatch": [
+ "X86_V3"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404080480337143,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 1,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521"
+ }
+ ],
+ "card": "a10g",
+ "gpu": "NVIDIA A10G",
+ "nproc": 24,
+ "cpu_flags": {
+ "sse4_2": true,
+ "avx": true,
+ "avx2": true,
+ "fma": true,
+ "avx512f": false,
+ "avx512bw": false,
+ "avx512vl": false
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/graphhosts01/t4_graph.json b/docs/results/a100_host_graph/graphhosts01/t4_graph.json
new file mode 100644
index 0000000..0399ef2
--- /dev/null
+++ b/docs/results/a100_host_graph/graphhosts01/t4_graph.json
@@ -0,0 +1,48 @@
+{
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": true,
+ "dispatch": [
+ "X86_V3",
+ "X86_V4"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404078617691994,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 3,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d"
+ }
+ ],
+ "card": "t4",
+ "gpu": "Tesla T4",
+ "nproc": 21,
+ "cpu_flags": {
+ "sse4_2": true,
+ "avx": true,
+ "avx2": true,
+ "fma": true,
+ "avx512f": true,
+ "avx512bw": true,
+ "avx512vl": true
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/graph_numpy_avx512.json b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_avx512.json
new file mode 100644
index 0000000..cc9c1dc
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_avx512.json
@@ -0,0 +1,37 @@
+{
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "avx512_skx": true,
+ "dispatch": [
+ "X86_V3",
+ "X86_V4",
+ "AVX512_ICL"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404078617691994,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 3,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/graph_numpy_baseline.json b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_baseline.json
new file mode 100644
index 0000000..4ba7921
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_baseline.json
@@ -0,0 +1,33 @@
+{
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "avx512_skx": false,
+ "dispatch": [],
+ "npy_disable_cpu_features": "X86_V3 X86_V4 AVX512_ICL",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618832,
+ "sha1": "97d163811acc313e",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404092587530613,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 7,
+ "1e-05": 59
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "a1e24853a843cd17"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/graph_numpy_x86_v3.json b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_x86_v3.json
new file mode 100644
index 0000000..0b8739e
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/graph_numpy_x86_v3.json
@@ -0,0 +1,35 @@
+{
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "avx512_skx": false,
+ "dispatch": [
+ "X86_V3"
+ ],
+ "npy_disable_cpu_features": "X86_V4 AVX512_ICL",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618831,
+ "sha1": "2faf55842ee44405",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404080480337143,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 1,
+ "1e-05": 65
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "c86e548279d2dda1"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/numpy_code_paths.txt b/docs/results/a100_host_graph/local_rtx5090/numpy_code_paths.txt
new file mode 100644
index 0000000..f758204
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/numpy_code_paths.txt
@@ -0,0 +1,14 @@
+# NumPy 2.5.3, glibc 2.43, AMD Ryzen 9 9950X3D; code path chosen with NPY_DISABLE_CPU_FEATURES.
+# 1. Grid/mesh lat/lon as passed to PointsMeshTypedGraphGNN._get_edge_indices (float32 degrees): AVX-512 path vs AVX2 path
+grid lat: 0 of 1038240 differ, max |diff| 0 deg
+grid lon: 0 of 1038240 differ, max |diff| 0 deg
+mesh lat: 12454 of 40962 differ, max |diff| 1.53e-05 deg
+mesh lon: 9964 of 40962 differ, max |diff| 3.05e-05 deg
+# 2. sha1 of np.sin / np.cos (float32, of deg2rad of those lat/lon, AVX-512-path inputs) per code path
+['X86_V3', 'X86_V4', 'AVX512_ICL'] plat.sin=d5d8f045 plat.cos=c1f67be5 plon.sin=7d39474b plon.cos=4a982ecb mlat.sin=a182b930 mlat.cos=5a82959f mlon.sin=a041ab3a mlon.cos=8cf23768
+['X86_V3'] plat.sin=d5d8f045 plat.cos=c1f67be5 plon.sin=7d39474b plon.cos=4a982ecb mlat.sin=a182b930 mlat.cos=5a82959f mlon.sin=a041ab3a mlon.cos=8cf23768
+[] plat.sin=7858a558 plat.cos=1a851bd7 plon.sin=045c86ec plon.cos=0bdf7d9f mlat.sin=ff320064 mlat.cos=809d1d43 mlon.sin=c1255eb9 mlon.cos=7ee4dbc1
+# 3. float32 sin/cos of the grid's lat/lon in radians vs glibc sinf/cosf and vs float64-then-rounded, per code path
+'' simd=['X86_V3', 'X86_V4', 'AVX512_ICL'] sin 77c0589a3f cos b298336f92 | != glibc sinf 242 cosf 389 | != f64-rounded sin 223 cos 404
+'X86_V4 AVX512_ICL' simd=['X86_V3'] sin 77c0589a3f cos b298336f92 | != glibc sinf 242 cosf 389 | != f64-rounded sin 223 cos 404
+'X86_V3 X86_V4 AVX512_ICL' simd=[] sin 24493c9f5f cos da72d7cf59 | != glibc sinf 0 cosf 0 | != f64-rounded sin 35 cos 31
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.json b/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.json
new file mode 100644
index 0000000..3e230dd
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.json
@@ -0,0 +1,36 @@
+{
+ "label": "rtx5090_edges_a100_40",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "numpy_simd": [
+ "X86_V3",
+ "X86_V4",
+ "AVX512_ICL"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 72.90617728233337,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.810781002044678
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.log b/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.log
new file mode 100644
index 0000000..fdc8a60
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_a100_40gb_graph.log
@@ -0,0 +1,13 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB; host AMD Ryzen 9 9950X3D 16-Core Processor, NumPy SIMD ['X86_V3', 'X86_V4', 'AVX512_ICL']
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:56:24.012498 463665 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 73s
+graphs: {'BALL_QUERY': {'n_edges': 1618832, 'sha1': '1d90c1ed7b086815', 'loaded': True}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '31e1d7a376c54521', 'loaded': True}}
+peak 5.81 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.json b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.json
new file mode 100644
index 0000000..4d76299
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.json
@@ -0,0 +1,34 @@
+{
+ "label": "rtx5090_numpy_avx2_a100_40_graph",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 70.61653590202332,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937248945236206
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.log b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.log
new file mode 100644
index 0000000..05a99ff
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_a100_40gb_graph.log
@@ -0,0 +1,13 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB; host AMD Ryzen 9 9950X3D 16-Core Processor, NumPy SIMD ['X86_V3']
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:59:43.440247 469836 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 71s
+graphs: {'BALL_QUERY': {'n_edges': 1618832, 'sha1': '1d90c1ed7b086815', 'loaded': True}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '31e1d7a376c54521', 'loaded': True}}
+peak 5.94 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.json b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.json
new file mode 100644
index 0000000..e017b23
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.json
@@ -0,0 +1,34 @@
+{
+ "label": "rtx5090_numpy_avx2_own_graph",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 70.45442032814026,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937260389328003
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.log b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.log
new file mode 100644
index 0000000..3392139
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_avx2_own_graph.log
@@ -0,0 +1,13 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB; host AMD Ryzen 9 9950X3D 16-Core Processor, NumPy SIMD ['X86_V3']
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:58:29.252897 466929 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 70s
+graphs: {'BALL_QUERY': {'n_edges': 1618841, 'sha1': '189d06f0830e1c31', 'loaded': True}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '0bd261acea40807d', 'loaded': True}}
+peak 5.94 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.json b/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.json
new file mode 100644
index 0000000..11f74f2
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.json
@@ -0,0 +1,17 @@
+{
+ "label": "rtx5090_npy_x86v2",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 86.5319275856018,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937266826629639
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.log b/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.log
new file mode 100644
index 0000000..a2229a7
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_numpy_baseline.log
@@ -0,0 +1,12 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:47:27.177475 454271 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 87s
+peak 5.94 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_own_graph.json b/docs/results/a100_host_graph/local_rtx5090/step_own_graph.json
new file mode 100644
index 0000000..23a269c
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_own_graph.json
@@ -0,0 +1,18 @@
+{
+ "label": "rtx5090_base",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_dump_to=/tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/103238b9-ccc5-49c4-8b99-159fd6100067/scratchpad/local5090/dump --xla_dump_hlo_as_text --xla_dump_hlo_module_re=.*lambda.* --xla_gpu_dump_autotune_results_to=/tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/103238b9-ccc5-49c4-8b99-159fd6100067/scratchpad/local5090/autotune.txt"
+ },
+ "ok": true,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 87.90489864349365,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.585663318634033
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_own_graph.log b/docs/results/a100_host_graph/local_rtx5090/step_own_graph.log
new file mode 100644
index 0000000..58a2205
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_own_graph.log
@@ -0,0 +1,12 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:37:26.006912 445268 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 88s
+peak 5.59 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.json b/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.json
new file mode 100644
index 0000000..060b959
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.json
@@ -0,0 +1,36 @@
+{
+ "label": "rtx5090_ownedges",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "AMD Ryzen 9 9950X3D 16-Core Processor",
+ "numpy": "2.5.3",
+ "numpy_simd": [
+ "X86_V3",
+ "X86_V4",
+ "AVX512_ICL"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 70.63282299041748,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.9372663497924805
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.log b/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.log
new file mode 100644
index 0000000..103ed4f
--- /dev/null
+++ b/docs/results/a100_host_graph/local_rtx5090/step_own_graph_loaded.log
@@ -0,0 +1,13 @@
+NVIDIA GeForce RTX 5090 cc=12.0 attention=pallas pool 23.49 GiB; host AMD Ryzen 9 9950X3D 16-Core Processor, NumPy SIMD ['X86_V3', 'X86_V4', 'AVX512_ICL']
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 13:52:39.465119 460127 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest-precision step saved (101 fields) in 71s
+graphs: {'BALL_QUERY': {'n_edges': 1618841, 'sha1': '189d06f0830e1c31', 'loaded': True}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '0bd261acea40807d', 'loaded': True}}
+peak 5.94 GiB of pool 23.49 GiB; ok=True
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_at0.json b/docs/results/a100_host_graph/tf32probe01/a100_40_at0.json
new file mode 100644
index 0000000..49f5637
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_at0.json
@@ -0,0 +1,18 @@
+{
+ "label": "a100_40_at0",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_gpu_autotune_level=0"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 126.01233983039856,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.581749677658081
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_at0.log b/docs/results/a100_host_graph/tf32probe01/a100_40_at0.log
new file mode 100644
index 0000000..cf96a8b
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_at0.log
@@ -0,0 +1,12 @@
+NVIDIA A100-SXM4-40GB cc=8.0 attention=pallas pool 29.64 GiB
+highest-precision step saved (101 fields) in 126s
+peak 5.58 GiB of pool 29.64 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 18:39:15.994254 701 hlo_rematerialization.cc:3363] Can't reduce memory use below 26.91GiB (28895515853 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_base.json b/docs/results/a100_host_graph/tf32probe01/a100_40_base.json
new file mode 100644
index 0000000..7e56668
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_base.json
@@ -0,0 +1,18 @@
+{
+ "label": "a100_40_base",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_dump_to=/tmp/dump_a100_40_base --xla_dump_hlo_as_text --xla_dump_hlo_module_re=.*lambda.* --xla_gpu_dump_autotune_results_to=/data/results/tf32probe01/a100_40_base_autotune.txt"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 132.97908067703247,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268257141113
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_base.log b/docs/results/a100_host_graph/tf32probe01/a100_40_base.log
new file mode 100644
index 0000000..5ec0a9e
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_base.log
@@ -0,0 +1,12 @@
+NVIDIA A100-SXM4-40GB cc=8.0 attention=pallas pool 29.64 GiB
+highest-precision step saved (101 fields) in 133s
+peak 5.94 GiB of pool 29.64 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 18:37:05.174324 6 hlo_rematerialization.cc:3363] Can't reduce memory use below 26.91GiB (28895515853 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.json b/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.json
new file mode 100644
index 0000000..4905541
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.json
@@ -0,0 +1,18 @@
+{
+ "label": "a100_40_notritongemm",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_gpu_enable_triton_gemm=false"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 120.14120078086853,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.7115514278411865
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.log b/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.log
new file mode 100644
index 0000000..25dcbd5
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_notritongemm.log
@@ -0,0 +1,11 @@
+NVIDIA A100-SXM4-40GB cc=8.0 attention=pallas pool 29.64 GiB
+highest-precision step saved (101 fields) in 120s
+peak 5.71 GiB of pool 29.64 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_summary.json b/docs/results/a100_host_graph/tf32probe01/a100_40_summary.json
new file mode 100644
index 0000000..0cfaad9
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_summary.json
@@ -0,0 +1,80 @@
+{
+ "card": "a100_40",
+ "nvidia_smi": "name, driver_version, memory.total [MiB], clocks.max.sm [MHz], ecc.mode.current\nNVIDIA A100-SXM4-40GB, 610.57.04, 40960 MiB, 1410 MHz, Enabled\n",
+ "env": {},
+ "base": {
+ "label": "a100_40_base",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_dump_to=/tmp/dump_a100_40_base --xla_dump_hlo_as_text --xla_dump_hlo_module_re=.*lambda.* --xla_gpu_dump_autotune_results_to=/data/results/tf32probe01/a100_40_base_autotune.txt"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 132.97908067703247,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268257141113,
+ "wall_s": 143.64632391929626
+ },
+ "at0": {
+ "label": "a100_40_at0",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_gpu_autotune_level=0"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 126.01233983039856,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.581749677658081,
+ "wall_s": 133.0238482952118
+ },
+ "tf32off": {
+ "label": "a100_40_tf32off",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 133.15456080436707,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268018722534,
+ "wall_s": 140.29184436798096
+ },
+ "notritongemm": {
+ "label": "a100_40_notritongemm",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_gpu_enable_triton_gemm=false"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 120.14120078086853,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.7115514278411865,
+ "wall_s": 127.63998007774353
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.json b/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.json
new file mode 100644
index 0000000..eb66105
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.json
@@ -0,0 +1,17 @@
+{
+ "label": "a100_40_tf32off",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 133.15456080436707,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268018722534
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.log b/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.log
new file mode 100644
index 0000000..e4bf7ab
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_40_tf32off.log
@@ -0,0 +1,12 @@
+NVIDIA A100-SXM4-40GB cc=8.0 attention=pallas pool 29.64 GiB
+highest-precision step saved (101 fields) in 133s
+peak 5.94 GiB of pool 29.64 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 18:41:44.970353 1100 hlo_rematerialization.cc:3363] Can't reduce memory use below 26.91GiB (28895515853 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_80_base.json b/docs/results/a100_host_graph/tf32probe01/a100_80_base.json
new file mode 100644
index 0000000..b7d297a
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_80_base.json
@@ -0,0 +1,18 @@
+{
+ "label": "a100_80_base",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_dump_to=/tmp/dump_a100_80_base --xla_dump_hlo_as_text --xla_dump_hlo_module_re=.*lambda.* --xla_gpu_dump_autotune_results_to=/data/results/tf32probe01/a100_80_base_autotune.txt"
+ },
+ "ok": true,
+ "device": "NVIDIA A100 80GB PCIe",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 59.455078125,
+ "highest_first_call_s": 127.11047983169556,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937267303466797
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_80_base.log b/docs/results/a100_host_graph/tf32probe01/a100_80_base.log
new file mode 100644
index 0000000..8ea02ed
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_80_base.log
@@ -0,0 +1,11 @@
+NVIDIA A100 80GB PCIe cc=8.0 attention=pallas pool 59.46 GiB
+highest-precision step saved (101 fields) in 127s
+peak 5.94 GiB of pool 59.46 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
diff --git a/docs/results/a100_host_graph/tf32probe01/a100_80_summary.json b/docs/results/a100_host_graph/tf32probe01/a100_80_summary.json
new file mode 100644
index 0000000..b68b462
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/a100_80_summary.json
@@ -0,0 +1,24 @@
+{
+ "card": "a100_80",
+ "nvidia_smi": "name, driver_version, memory.total [MiB], clocks.max.sm [MHz], ecc.mode.current\nNVIDIA A100 80GB PCIe, 610.57.04, 81920 MiB, 1410 MHz, Enabled\n",
+ "env": {},
+ "base": {
+ "label": "a100_80_base",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
+ "XLA_FLAGS": "--xla_dump_to=/tmp/dump_a100_80_base --xla_dump_hlo_as_text --xla_dump_hlo_module_re=.*lambda.* --xla_gpu_dump_autotune_results_to=/data/results/tf32probe01/a100_80_base_autotune.txt"
+ },
+ "ok": true,
+ "device": "NVIDIA A100 80GB PCIe",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 59.455078125,
+ "highest_first_call_s": 127.11047983169556,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937267303466797,
+ "wall_s": 139.12689113616943
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe01/hlo_before_optimizations_40gb_vs_80gb.diff b/docs/results/a100_host_graph/tf32probe01/hlo_before_optimizations_40gb_vs_80gb.diff
new file mode 100644
index 0000000..b8269f4
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe01/hlo_before_optimizations_40gb_vs_80gb.diff
@@ -0,0 +1,14 @@
+831,832c831,832
+< %_pad.26 (Arg_0.35: f32[1618832,1,4], Arg_1.25: s32[]) -> f32[1638400,1,4] {
+< %Arg_0.35 = f32[1618832,1,4]{2,1,0} parameter(0)
+---
+> %_pad.26 (Arg_0.35: f32[1618841,1,4], Arg_1.25: s32[]) -> f32[1638400,1,4] {
+> %Arg_0.35 = f32[1618841,1,4]{2,1,0} parameter(0)
+835c835
+< ROOT %pad.1 = f32[1638400,1,4]{2,1,0} pad(%Arg_0.35, %convert_element_type.5), padding=0_19568x0_0x0_0, metadata={op_name="pad"}
+---
+> ROOT %pad.1 = f32[1638400,1,4]{2,1,0} pad(%Arg_0.35, %convert_element_type.5), padding=0_19559x0_0x0_0, metadata={op_name="pad"}
+5235c5235
+< %constant.976 = f32[1618832,1,4]{2,1,0} constant({...})
+---
+> %constant.976 = f32[1618841,1,4]{2,1,0} constant({...})
diff --git a/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.json b/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.json
new file mode 100644
index 0000000..ce73e98
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.json
@@ -0,0 +1,34 @@
+{
+ "label": "a100_40_edges5090",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 112.11894607543945,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268018722534
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.log b/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.log
new file mode 100644
index 0000000..6ca53c8
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe02/a100_40_edges5090.log
@@ -0,0 +1,13 @@
+NVIDIA A100-SXM4-40GB cc=8.0 attention=pallas pool 29.64 GiB; host unknown, NumPy SIMD ['X86_V3']
+highest-precision step saved (101 fields) in 112s
+graphs: {'BALL_QUERY': {'n_edges': 1618841, 'sha1': '189d06f0830e1c31', 'loaded': True}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '0bd261acea40807d', 'loaded': True}}
+peak 5.94 GiB of pool 29.64 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 18:56:45.863282 95 hlo_rematerialization.cc:3363] Can't reduce memory use below 26.91GiB (28895515853 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
diff --git a/docs/results/a100_host_graph/tf32probe02/a100_40_graph.json b/docs/results/a100_host_graph/tf32probe02/a100_40_graph.json
new file mode 100644
index 0000000..4a887eb
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe02/a100_40_graph.json
@@ -0,0 +1,35 @@
+{
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": false,
+ "dispatch": [
+ "X86_V3"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404080480337143,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 1,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe02/a100_40_summary.json b/docs/results/a100_host_graph/tf32probe02/a100_40_summary.json
new file mode 100644
index 0000000..69cb099
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe02/a100_40_summary.json
@@ -0,0 +1,75 @@
+{
+ "card": "a100_40",
+ "nvidia_smi": "name, driver_version, memory.total [MiB], clocks.max.sm [MHz], ecc.mode.current\nNVIDIA A100-SXM4-40GB, 610.57.04, 40960 MiB, 1410 MHz, Enabled\n",
+ "env": {},
+ "graph": {
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "avx512_skx": false,
+ "dispatch": [
+ "X86_V3"
+ ],
+ "npy_disable_cpu_features": "",
+ "graphs": [
+ {
+ "connectivity": "BALL_QUERY",
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "dtypes": [
+ "float32",
+ "float32",
+ "float32",
+ "float32"
+ ],
+ "radius": 0.012404080480337143,
+ "pairs_near_radius": {
+ "1e-12": 0,
+ "1e-09": 0,
+ "1e-07": 0,
+ "1e-06": 1,
+ "1e-05": 70
+ }
+ },
+ {
+ "connectivity": "IN_TRIANGLE",
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521"
+ }
+ ]
+ },
+ "edges5090": {
+ "label": "a100_40_edges5090",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618841,
+ "sha1": "189d06f0830e1c31",
+ "loaded": true
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "0bd261acea40807d",
+ "loaded": true
+ }
+ },
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 112.11894607543945,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268018722534,
+ "wall_s": 120.11274552345276
+ }
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.json b/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.json
new file mode 100644
index 0000000..6b1ece9
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.json
@@ -0,0 +1,34 @@
+{
+ "label": "a100_80_numpy_avx2",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "loaded": false
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521",
+ "loaded": false
+ }
+ },
+ "device": "NVIDIA A100 80GB PCIe",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 59.455078125,
+ "highest_first_call_s": 137.35874152183533,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268495559692
+}
\ No newline at end of file
diff --git a/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.log b/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.log
new file mode 100644
index 0000000..3601e4a
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe03/a100_80_numpy_avx2.log
@@ -0,0 +1,12 @@
+NVIDIA A100 80GB PCIe cc=8.0 attention=pallas pool 59.46 GiB; host unknown, NumPy SIMD ['X86_V3']
+highest-precision step saved (101 fields) in 137s
+graphs: {'BALL_QUERY': {'n_edges': 1618832, 'sha1': '1d90c1ed7b086815', 'loaded': False}, 'IN_TRIANGLE': {'n_edges': 3114720, 'sha1': '31e1d7a376c54521', 'loaded': False}}
+peak 5.94 GiB of pool 59.46 GiB; ok=True
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
diff --git a/docs/results/a100_host_graph/tf32probe03/a100_80_summary.json b/docs/results/a100_host_graph/tf32probe03/a100_80_summary.json
new file mode 100644
index 0000000..5208312
--- /dev/null
+++ b/docs/results/a100_host_graph/tf32probe03/a100_80_summary.json
@@ -0,0 +1,40 @@
+{
+ "card": "a100_80",
+ "nvidia_smi": "name, driver_version, memory.total [MiB], clocks.max.sm [MHz], ecc.mode.current\nNVIDIA A100 80GB PCIe, 610.57.04, 81920 MiB, 1410 MHz, Enabled\n",
+ "env": {},
+ "numpy_avx2": {
+ "label": "a100_80_numpy_avx2",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "cpu": "unknown",
+ "numpy": "2.5.4",
+ "numpy_simd": [
+ "X86_V3"
+ ],
+ "graphs": {
+ "BALL_QUERY": {
+ "n_edges": 1618832,
+ "sha1": "1d90c1ed7b086815",
+ "loaded": false
+ },
+ "IN_TRIANGLE": {
+ "n_edges": 3114720,
+ "sha1": "31e1d7a376c54521",
+ "loaded": false
+ }
+ },
+ "device": "NVIDIA A100 80GB PCIe",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 59.455078125,
+ "highest_first_call_s": 137.35874152183533,
+ "step_s": [],
+ "finite": [],
+ "peak_gib": 5.937268495559692,
+ "wall_s": 149.64930152893066
+ }
+}
\ No newline at end of file
diff --git a/docs/validation.md b/docs/validation.md
index e3be86c..7a936f9 100644
--- a/docs/validation.md
+++ b/docs/validation.md
@@ -99,7 +99,7 @@ rollout step including host overhead (steps 2–4 of a 4-step rollout, default p
|---|---|---|---|---|---|---|
| RTX 5090 (local) | sm_120 | pallas | 0.7 | 6.40 GiB | yes | — |
| RTX 4060 (vast.ai) | sm_89 | pallas | 5.2 | 6.21 GiB | yes (real 8 GB card) | not measured (sm_89 covered by L4) |
-| A100-SXM4-40GB | sm_80 | pallas | 1.1 | 6.18 GiB | yes (6.21) | 1.6–1.9e-5 in 5 of 6 runs; see note |
+| A100-SXM4-40GB | sm_80 | pallas | 1.1 | 6.18 GiB | yes (6.21) | 1.32e-2: its host builds another graph; 1.7e-5 vs an A100-80GB given the same host inputs (see note) |
| A10G | sm_86 | pallas | 2.5 | 6.53 GiB | yes (6.34) | 1.6e-5 |
| L4 | sm_89 | pallas | 3.5 | 6.53 GiB | yes (6.34) | 1.75e-5 |
| T4 | sm_75 | xla (automatic; no TF32) | 44 | 6.49 GiB | yes (6.30) | 9.7e-5 |
@@ -114,23 +114,74 @@ out of memory. Installed from the repo on a fresh machine. Files:
An earlier build without the fused attention kernel ran at 18.5 s/step on the same card
([`results/canary_vast_rtx4060/`](results/canary_vast_rtx4060/)).
-**A100 note.** One of six A100 fp32 runs deviated from the RTX 5090 at TF32 level (max rel RMS
-1.3e-2, median 6.2e-4) and could not be reproduced (autotune on / off and three fresh repeats:
-1.6–1.9e-5). XLA's autotuner picks kernels by timing and its picks vary between runs (21 of
-149 autotuned ops differed across three repeats); the most likely explanation is one pick that
-did not honour fp32 in that run. For strict fp32 use `--no-autotune` (5/5 matching A100 runs)
-and for reproducible runs `--autotune-cache`. Files: [`results/a100_repeats/`](results/a100_repeats/).
-Update (2026-11-01): the canary below reproduced the deviation almost exactly (max rel RMS 1.32e-2,
-median 6.21e-4, same worst fields, min corr equal to 8 digits), so it is not a one-off pick: on
-A100, `cloud/canary_run.py` lands at TF32 level where the repeat script did not. Not yet explained;
-it predates the attention backward (same numbers on v0.1.0).
+**A100 note: the host CPU, not the GPU (cause found 2026-11-01).** On Modal's A100-SXM4-40GB the
+fp32 step differs from the RTX 5090 by max rel RMS 1.32e-2 (median 6.21e-4; worst fields
+`cyclone_usa_rmw_disc`, `vertical_velocity@50`) in every run (canary02 on v0.1.0, canary03 and
+tf32probe01 on v0.2; the three agree with each other at 1.5–2.1e-5). It looks like TF32 but is not.
+Earlier it was read as "1 of 6 runs, an autotune pick that did not honour fp32"; that was wrong:
+Modal served the five "A100-40GB" probe and repeat requests on A100-SXM4-80GB machines (`device`
+in [`results/a100_repeats/`](results/a100_repeats/)), so every 40GB run deviated and every 80GB run
+matched (1.6–1.9e-5).
+
+The official code builds the grid→mesh and mesh→grid graphs on the host CPU, with NumPy, from
+float32 latitudes and longitudes: the mesh nodes' lat/lon come from Cartesian positions through
+`np.arctan2`/`np.arccos`, then grid and mesh go back to unit-sphere coordinates through
+`np.sin`/`np.cos`, and a `cKDTree` ball query (radius 0.6 × the longest mesh edge, 0.0124) picks the
+grid→mesh edges; mesh→grid uses the mesh triangle that contains each grid point. One float32 ULP
+of a coordinate moves a 0.0124 distance by ~1e-5 (relative), and about 70 (grid point, mesh node)
+pairs lie that close to the radius. NumPy picks its math code at run time from the host CPU, and
+the results differ in the last bit: between its AVX-512 and AVX2 code, 12,454 of the 40,962 mesh
+latitudes (and 9,964 longitudes) come out different; its baseline code (no AVX2) also hands
+`sin`/`cos` to glibc ([`numpy_code_paths.txt`](results/a100_host_graph/local_rtx5090/numpy_code_paths.txt)).
+Modal's A100-40GB hosts offer AVX2 but not AVX-512; the RTX 5090 desktop (Zen 5) and Modal's
+A100-80GB and T4 hosts offer AVX-512. The two build different graphs: 1,618,832 vs 1,618,841
+grid→mesh edges (13 differ), and 517 of 1,038,240 grid points in a different mesh triangle (points
+on mesh edges along the meridians 0°, 60°, …, 300°); the edge-feature and node-position constants
+also differ in the last bits.
+faster-weathernext does not touch this (the patched `_build_typed_graph` calls the stock one,
+which calls the official `PointsMeshTypedGraphGNN._get_edge_indices`), so the official code depends
+on the host in the same way. Each graph is what the official code builds on that CPU.
+
+Evidence (one 0.25° step at `highest`; files in [`results/a100_host_graph/`](results/a100_host_graph/),
+comparisons in its `cross_card_fp32.txt`):
+- The HLO that JAX hands to XLA already differs between a 40GB and an 80GB run, in one place: the
+ grid→mesh edge-feature constant is `f32[1618832,1,4]` vs `f32[1618841,1,4]`
+ (`tf32probe01/hlo_before_optimizations_40gb_vs_80gb.diff`).
+- XLA and TF32 settings change nothing on the 40GB card: autotuning off, `NVIDIA_TF32_OVERRIDE=0`
+ and Triton GEMMs off all give 1.32e-2 / 6.21e-4 (`cloud/modal_a100_tf32.py`).
+- The difference is local. It sits on the 13 changed edges and the 517 re-assigned grid points
+ ([map](results/a100_host_graph/errmap_a100_40gb_vs_rtx5090.png),
+ [`differing_edges.txt`](results/a100_host_graph/differing_edges.txt)); elsewhere the two cards
+ agree at ~1e-6.
+- Graph swap (`scripts/graph_edges.py --save-edges`, `cloud/canary_run.py --edges`): the RTX 5090
+ on the 40GB host's graph differs from the RTX 5090 on its own graph by 1.32e-2 / 6.21e-4 (same
+ worst fields), and so does the A100-40GB on the RTX 5090's graph from the A100-40GB on its own.
+ With the graph swapped, the two cards are 1.3e-4 / 1e-5 apart. The remainder is the host-side
+ features: on the RTX 5090, with the graph held fixed, switching only NumPy's code path from
+ AVX-512 to AVX2 moves the output by the same 1.33e-4 / 1.0e-5.
+- One card, one environment variable: on an A100 80GB PCIe (AVX-512 host), forcing NumPy's AVX2
+ code (`NPY_DISABLE_CPU_FEATURES="X86_V4 AVX512_ICL AVX512_SPR"`) builds exactly the 40GB host's
+ graph (same hash) and moves the output by 1.32e-2 / 6.21e-4 from the same card's default run.
+ That output then matches the A100-40GB at 1.69e-5 / 1.47e-6, the usual cross-GPU level
+ (tf32probe03).
+- Host survey (`cloud/modal_graph_hosts.py`): the T4 and A100-80GB hosts build the RTX 5090's graph
+ (same hash); an A10G host with AVX2 only builds the 40GB host's graph. The canary's A10G
+ requests ran on `NVIDIA A10` cards (another host type; its CPU was not recorded) and matched.
+
+So the fp32 agreement across GPUs holds for the same host-built graph. On machines whose CPUs take
+different NumPy code paths the official model runs on slightly different graphs, and one fp32 step
+then differs by up to ~1e-2 max rel RMS (median 6e-4), around the changed edges. `cloud/canary_run.py`
+now records the graph hashes and NumPy's SIMD targets: compare fp32 outputs across machines only
+when the hashes match, or pass `--edges`. `--no-autotune` is not needed for fp32 (it changed nothing
+here). XLA's autotune picks still vary between runs (21 of 149 ops across three A100 repeats), so
+use `--autotune-cache` for bitwise-reproducible runs.
**Canary after the attention backward (2026-11-01, branch `head-only-finetune`).** Same runner and
cards as the table above (`modal run cloud/modal_canary.py --cards a100,a10g,l4,t4`, ~$1–2): every
card ran and fits the 8 GB pool; steady-state s/step and peaks as in the table (A100 1.10 s / 6.18
GiB, A10G 2.43 / 6.53, L4 3.57 / 6.53, T4 43 / 6.49 on the XLA path). fp32 outputs vs a fresh
-RTX 5090 run of the same code: A10G 1.67e-5, L4 1.85e-5, T4 9.67e-5 max rel RMS; A100 1.32e-2 (the
-note above). Files: [`results/canary_modal/v0.2/`](results/canary_modal/v0.2/).
+RTX 5090 run of the same code: A10G 1.67e-5, L4 1.85e-5, T4 9.67e-5 max rel RMS; A100 1.32e-2 (a
+different host-built graph, not the GPU: note above). Files: [`results/canary_modal/v0.2/`](results/canary_modal/v0.2/).
## 5. Mini model and batch size > 1 (v0.1.1)
@@ -423,6 +474,10 @@ Files: [`results/finetune/finetune_full_smoke_rtx5090.json`](results/finetune/fi
rollouts in §3 and the batch checks in §5.
- Batch sizes > 1 were checked at 2 (§5, §6); larger batches only change memory.
- Without a fixed autotune cache, results are not bitwise reproducible between runs (§4 note).
+- fp32 agreement across machines assumes the same host-built graphs. The official code builds the
+ grid↔mesh graphs with NumPy on the host CPU, and they follow NumPy's SIMD code path: between an
+ AVX2 and an AVX-512 host one fp32 step differs by up to 1.3e-2 max rel RMS (median 6e-4), with or
+ without faster-weathernext (§4, A100 note). Only x86 hosts were checked.
- TPU (`splash_mha`) numerics were not compared; TPU matmuls use bf16 passes by default and
differ from any of the GPU paths above.
- Whole-model gradients were compared with the official path on the Mini model at 1° (§9, on an
diff --git a/scripts/graph_edges.py b/scripts/graph_edges.py
new file mode 100644
index 0000000..4d4cfc6
--- /dev/null
+++ b/scripts/graph_edges.py
@@ -0,0 +1,86 @@
+"""Count the edges of the grid<->mesh graphs that the official code builds on the host CPU.
+
+WN2 builds its grid->mesh edges with a ball query (scipy cKDTree, radius = a fraction of the
+longest mesh edge) on unit-sphere coordinates computed from float32 lat/lon with NumPy (mesh
+lat/lon via np.arctan2 / np.arccos, then np.sin / np.cos). Grid points close to the query radius
+are kept or dropped depending on the last bits of those coordinates, and NumPy's results depend on
+the code path it picks for the host CPU (AVX-512, AVX2, baseline). This script traces one forward
+step on the CPU (no compilation) and prints, per graph, the edge count, a hash of the edge indices
+and how many (grid point, mesh node) pairs lie within a relative 1e-12 ... 1e-5 of the radius.
+
+ JAX_PLATFORMS=cpu python scripts/graph_edges.py --data SAMPLE.nc [--weights-dir DIR] [--json out.json]
+ [--save-edges edges.npz] # reuse elsewhere with cloud/canary_run.py --edges edges.npz
+ # NumPy's code path for a CPU without AVX-512 (NumPy 2.5 names; "X86_V3 ..." too: no AVX2):
+ NPY_DISABLE_CPU_FEATURES="X86_V4 AVX512_ICL AVX512_SPR" JAX_PLATFORMS=cpu python scripts/graph_edges.py ...
+"""
+import argparse
+import dataclasses
+import hashlib
+import json
+import os
+import platform
+
+import numpy as np
+
+p = argparse.ArgumentParser()
+p.add_argument("--data", required=True, help="Google's 0.25° sample (see scripts/verify_equivalence.py)")
+p.add_argument("--model", default="WeatherNextCyclones")
+p.add_argument("--weights-dir")
+p.add_argument("--json", default="")
+p.add_argument("--save-edges", default="", help=".npz with the edge indices (for cloud/canary_run.py --edges)")
+args = p.parse_args()
+
+import jax # noqa: E402
+import xarray # noqa: E402
+from scipy import spatial # noqa: E402
+from weathernext.utils import data_utils # noqa: E402
+from weathernext.utils import model_utils # noqa: E402
+from weathernext.utils import points_mesh_gnn as pmg # noqa: E402
+
+import faster_weathernext # noqa: E402
+from faster_weathernext import model # noqa: E402
+
+from numpy._core._multiarray_umath import __cpu_dispatch__, __cpu_features__ # noqa: E402
+
+cpu = next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), platform.processor())
+res = dict(cpu=cpu, numpy=np.__version__, avx512_skx=bool(__cpu_features__.get("AVX512_SKX")),
+ dispatch=[f for f in __cpu_dispatch__ if __cpu_features__.get(f)],
+ npy_disable_cpu_features=os.environ.get("NPY_DISABLE_CPU_FEATURES", ""), graphs=[])
+
+orig = pmg.PointsMeshTypedGraphGNN._get_edge_indices
+edges = {}
+
+
+def hooked(self, plat, plon, mlat, mlon, faces):
+ pts, mesh = orig(self, plat, plon, mlat, mlon, faces)
+ g = dict(connectivity=self._connectivity_type.name, n_edges=int(len(pts)),
+ sha1=hashlib.sha1(np.stack([pts, mesh]).astype(np.int64).tobytes()).hexdigest()[:16])
+ if self._connectivity_type == pmg.ConnectivityType.BALL_QUERY:
+ pc = np.stack(model_utils.lat_lon_to_cartesian(plat, plon), -1)
+ mc = np.stack(model_utils.lat_lon_to_cartesian(mlat, mlon), -1)
+ r = self._ball_query_radius_fraction * pmg._max_edge_distance(mc, faces)
+ near = spatial.cKDTree(mc).query_ball_point(pc, r=r * (1 + 1e-5))
+ d = np.concatenate([np.linalg.norm(mc[nb] - pc[i], axis=-1) for i, nb in enumerate(near) if nb])
+ g.update(dtypes=[str(np.asarray(a).dtype) for a in (plat, plon, mlat, mlon)], radius=float(r),
+ pairs_near_radius={f"{t:.0e}": int((np.abs(d - r) <= t * r).sum()) for t in (1e-12, 1e-9, 1e-7, 1e-6, 1e-5)})
+ res["graphs"].append(g)
+ edges[f"{g['connectivity']}_points"], edges[f"{g['connectivity']}_mesh"] = pts.astype(np.int32), mesh.astype(np.int32)
+ print(g, flush=True)
+ return pts, mesh
+
+
+pmg.PointsMeshTypedGraphGNN._get_edge_indices = hooked
+faster_weathernext.enable()
+task, forward = model.build(args.model)
+params = model.load_params(args.model, 1, args.weights_dir)
+ds = xarray.load_dataset(args.data)
+inputs, targets, forcings = data_utils.extract_inputs_targets_forcings(
+ ds, target_lead_times=slice("6h", "6h"), **dataclasses.asdict(task))
+# Tracing builds the graphs (host numpy/scipy); nothing is compiled or run on a device.
+jax.eval_shape(forward, params, jax.random.PRNGKey(0), inputs, targets * np.nan, forcings)
+print(json.dumps(res, indent=1))
+if args.json:
+ with open(args.json, "w") as f:
+ json.dump(res, f, indent=1)
+if args.save_edges:
+ np.savez_compressed(args.save_edges, **edges)
diff --git a/src/faster_weathernext/cli.py b/src/faster_weathernext/cli.py
index 2d0d324..2c28aa8 100644
--- a/src/faster_weathernext/cli.py
+++ b/src/faster_weathernext/cli.py
@@ -122,7 +122,7 @@ def main(argv=None):
sub = p.add_subparsers(dest="cmd", required=True)
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--no-autotune", action="store_true",
- help="XLA autotune off (deterministic kernel picks; recommended for strict fp32 on A100)")
+ help="XLA autotune off (deterministic kernel picks)")
common.add_argument("--autotune-cache", metavar="PATH",
help="save XLA's kernel picks here on the first run and reuse them afterwards")
sub.add_parser("info", parents=[common], help="device, memory and compatibility check").set_defaults(fn=cmd_info)