diff --git a/LICENSE/License.txt b/LICENSE
similarity index 100%
rename from LICENSE/License.txt
rename to LICENSE
diff --git a/LICENSE/third_party.txt b/LICENSE/third_party.txt
deleted file mode 100644
index 2ac7c383..00000000
--- a/LICENSE/third_party.txt
+++ /dev/null
@@ -1,38 +0,0 @@
-# Third Party Software Attributions
-
-This file documents the third-party open-source code incorporated into
-nvSubquadratic and its associated license terms. Each entry lists the
-upstream source, the license, the in-repo paths that contain or derive
-from the third-party code, and the verbatim license text.
-
-================================================================================
-
-JiT (Just-in-time Tokenization for Diffusion Models)
-- Upstream: https://github.com/LTH14/JiT
-- License:  MIT
-- Used in:  nvsubquadratic/networks/jit.py (port)
-- License text below:
-
-    MIT License
-
-    Copyright (c) 2023 The JiT authors and contributors.
-
-    Permission is hereby granted, free of charge, to any person obtaining a copy
-    of this software and associated documentation files (the "Software"), to deal
-    in the Software without restriction, including without limitation the rights
-    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-    copies of the Software, and to permit persons to whom the Software is
-    furnished to do so, subject to the following conditions:
-
-    The above copyright notice and this permission notice shall be included in all
-    copies or substantial portions of the Software.
-
-    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-    SOFTWARE.
-
-================================================================================
diff --git a/THIRD_PARTY_NOTICES.txt b/THIRD_PARTY_NOTICES.txt
new file mode 100644
index 00000000..c4080e75
--- /dev/null
+++ b/THIRD_PARTY_NOTICES.txt
@@ -0,0 +1,66 @@
+# Third Party Software Attributions
+
+This file documents the third-party open-source code incorporated into
+nvSubquadratic and its associated license terms. Each entry lists the
+upstream source, the license, the in-repo paths that contain or derive
+from the third-party code, and the verbatim license text.
+
+================================================================================
+
+The Well (PolymathicAI/the_well)
+- Upstream: https://github.com/PolymathicAI/the_well
+- License:  BSD-3-Clause
+- Used in:  nvsubquadratic/networks/baselines/unet_convnext.py and
+            nvsubquadratic/networks/baselines/unet_convnext_v2.py — ported from
+            the_well.benchmark.models.unet_convnext (the V2 file corrects the
+            upstream finest-resolution skip-connection bug).
+- License text below:
+
+    BSD 3-Clause License
+
+    Copyright (c) 2024 Polymathic AI.
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+    1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+    2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+    3. Neither the name of Polymathic AI nor the names of the Well contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+================================================================================
+
+ConvNeXt (facebookresearch/ConvNeXt)
+- Upstream: https://github.com/facebookresearch/ConvNeXt
+- License:  MIT
+- Used in:  nvsubquadratic/networks/baselines/unet_convnext.py — the ConvNeXt
+            block, LayerNorm, and stage implementations that The Well adapted
+            and that this port inherits transitively.
+- License text below:
+
+    MIT License
+
+    Copyright (c) Meta Platforms, Inc. and affiliates.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE.
+
+================================================================================
diff --git a/benchmarks/benchmark_imagenet_diffusion_gpu.py b/benchmarks/benchmark_imagenet_diffusion_gpu.py
deleted file mode 100644
index 6e9ce752..00000000
--- a/benchmarks/benchmark_imagenet_diffusion_gpu.py
+++ /dev/null
@@ -1,410 +0,0 @@
-#!/usr/bin/env python
-
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Benchmark GPU memory/time for the ImageNet diffusion model (batch size = 1).
-
-Targets: H100 SXM 80GB or any Ampere+ GPU.  Runs a single
-forward+backward step at ``batch_size=1`` and reports peak memory and
-wall-clock time, used to compare attention vs Hyena vs JiT diffusion
-backbones at parity.
-
-Usage:
-    PYTHONPATH=. conda run -n nv-subq python \\
-        benchmarks/benchmark_imagenet_diffusion_gpu.py --config <config-path>
-
-Output: stdout summary table; no files written.
-"""
-
-from __future__ import annotations
-
-import argparse
-import copy
-import json
-import time
-from contextlib import nullcontext
-from dataclasses import dataclass, replace
-from functools import partial
-from typing import Callable, Iterable, List, Optional
-
-import torch
-
-from experiments.utils.cli import apply_config_overrides, load_config_from_file
-from nvsubquadratic.lazy_config import instantiate
-
-
-RESOLUTIONS = [64, 128, 256, 1024]
-BATCH_SIZE = 1
-
-
-@dataclass(frozen=True)
-class ModelSpec:
-    """Simple descriptor for each benchmarked model size."""
-
-    name: str
-    hidden_dim: int
-    num_layers: int
-    num_params: int | None = None
-
-
-MODEL_SPECS: tuple[ModelSpec, ModelSpec, ModelSpec] = (
-    ModelSpec(name="tiny", hidden_dim=512, num_layers=8),
-    ModelSpec(name="base", hidden_dim=768, num_layers=12),
-    ModelSpec(name="large", hidden_dim=1024, num_layers=16),
-)
-
-
-@dataclass(frozen=True)
-class CompileOptions:
-    """Holds torch.compile arguments for optional graph compilation."""
-
-    backend: str = "inductor"
-    mode: str = "default"
-
-
-def _autocast_context(device: torch.device, dtype: torch.dtype):
-    """Return a device-appropriate autocast context or a no-op."""
-    if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16):
-        return torch.autocast(device_type=device.type, dtype=dtype)
-    return nullcontext()
-
-
-def _ensure_cuda() -> torch.device:
-    """Assert CUDA is available and return the default CUDA device."""
-    if not torch.cuda.is_available():
-        raise RuntimeError("CUDA GPU required for this benchmark.")
-    return torch.device("cuda")
-
-
-def _clone_config(cfg):
-    """Deep-copy a config so we can apply overrides without mutating the original."""
-    return copy.deepcopy(cfg)
-
-
-def _prepare_config(base_cfg, spec: ModelSpec, image_size: int):
-    """Clone the base config and apply model-size / resolution overrides."""
-    cfg = _clone_config(base_cfg)
-    overrides = [
-        f"net.hidden_dim={spec.hidden_dim}",
-        f"net.num_blocks={spec.num_layers}",
-        f"diffusion.time_embed_dim={spec.hidden_dim}",
-        f"diffusion.cosine_schedule_image_resolution={image_size}",
-        f"diffusion.cosine_schedule_noise_res_low={max(32, image_size // 2)}",
-        f"diffusion.cosine_schedule_noise_res_high={image_size}",
-    ]
-    overrides.append(f"dataset.image_size={image_size}")
-    overrides.append(f"dataset.final_image_size={image_size}")
-    return apply_config_overrides(cfg, overrides)
-
-
-def _instantiate_wrapper(
-    cfg,
-    device: torch.device,
-    dtype: torch.dtype,
-    compile_options: Optional[CompileOptions] = None,
-):
-    """Build a Lightning wrapper on *device*, optionally applying ``torch.compile``."""
-    network = instantiate(cfg.net, in_channels=3, out_channels=3)
-    wrapper = instantiate(cfg.lightning_wrapper_class, network=network, cfg=cfg)
-    wrapper = wrapper.to(device=device)
-    setattr(wrapper, "_compiled_shared_step", None)
-
-    if compile_options is not None:
-        if not hasattr(torch, "compile"):
-            raise RuntimeError("torch.compile is not available in this PyTorch build.")
-        compile_kwargs = {}
-        if compile_options.backend:
-            compile_kwargs["backend"] = compile_options.backend
-        if compile_options.mode:
-            compile_kwargs["mode"] = compile_options.mode
-        wrapper.network = torch.compile(wrapper.network, **compile_kwargs)
-        wrapper._compiled_shared_step = torch.compile(wrapper._shared_step, **compile_kwargs)
-    return wrapper
-
-
-def _make_images(resolution: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    """Generate random images in [-1, 1] for benchmarking."""
-    return torch.rand((BATCH_SIZE, resolution, resolution, 3), device=device, dtype=dtype) * 2.0 - 1.0
-
-
-def _make_labels(device: torch.device) -> torch.Tensor:
-    """Generate dummy class labels (all zeros)."""
-    return torch.zeros((BATCH_SIZE,), device=device, dtype=torch.long)
-
-
-def _measure(fn, device: torch.device, repeat: int) -> tuple[float, float]:
-    """Run *fn* ``repeat`` times and return (avg_seconds, peak_gpu_mb)."""
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats(device)
-    torch.cuda.synchronize(device)
-    start = time.perf_counter()
-    for _ in range(repeat):
-        fn()
-    torch.cuda.synchronize(device)
-    elapsed = (time.perf_counter() - start) / repeat
-    peak_mb = torch.cuda.max_memory_allocated(device) / (1024**2)
-    return elapsed, peak_mb
-
-
-def _inference_fn(wrapper, resolution: int, dtype: torch.dtype):
-    """Return a zero-arg closure that runs one inference step."""
-    device = next(wrapper.parameters()).device
-    images = _make_images(resolution, device, dtype)
-    labels = _make_labels(device)
-    timesteps = torch.randint(
-        0,
-        wrapper.scheduler.config.num_train_timesteps,
-        (BATCH_SIZE,),
-        device=device,
-        dtype=torch.long,
-    )
-
-    def _run():
-        wrapper.eval()
-        with torch.no_grad(), _autocast_context(device, dtype):
-            condition = wrapper._condition_from_timesteps(timesteps, labels=labels)
-            wrapper.network({"input": images, "condition": condition})
-
-    return _run
-
-
-def _training_fn(wrapper, resolution: int, dtype: torch.dtype):
-    """Return a zero-arg closure that runs one train step (fwd + bwd + optimizer)."""
-    device = next(wrapper.parameters()).device
-    wrapper.train()
-    images = _make_images(resolution, device, dtype)
-    labels = _make_labels(device)
-    batch = {"input": images, "label": labels, "condition": None}
-    optimizer = torch.optim.AdamW(wrapper.parameters(), lr=1e-4)
-
-    step_fn = getattr(wrapper, "_compiled_shared_step", None) or wrapper._shared_step
-
-    def _run():
-        optimizer.zero_grad(set_to_none=True)
-        with _autocast_context(device, dtype):
-            loss = step_fn(batch)
-        loss.backward()
-        optimizer.step()
-
-    return _run
-
-
-def _is_oom_error(exc: BaseException) -> bool:
-    """Check whether an exception is a CUDA out-of-memory error."""
-    message = str(exc).lower()
-    return "out of memory" in message
-
-
-def _run_mode(
-    *,
-    mode: str,
-    build_fn: Callable[[], Callable[[], None]],
-    spec: ModelSpec,
-    image_size: int,
-    device: torch.device,
-    repeat: int,
-    dtype_name: str,
-) -> dict:
-    """Run a single benchmark mode (inference or training), catching OOM gracefully."""
-    result = {
-        "mode": mode,
-        "model": spec.name,
-        "hidden_dim": spec.hidden_dim,
-        "num_layers": spec.num_layers,
-        "image_size": image_size,
-        "batch_size": BATCH_SIZE,
-        "dtype": dtype_name,
-        "num_params": spec.num_params,
-    }
-    try:
-        fn = build_fn()
-        fn()  # warmup
-        elapsed, peak = _measure(fn, device, repeat)
-        result["time_ms"] = elapsed * 1e3
-        result["peak_memory_mb"] = peak
-    except (RuntimeError, torch.cuda.OutOfMemoryError) as exc:
-        if _is_oom_error(exc):
-            torch.cuda.empty_cache()
-            result["error"] = "OOM"
-        else:
-            raise
-    return result
-
-
-def benchmark_spec(
-    base_cfg,
-    spec: ModelSpec,
-    image_size: int,
-    device: torch.device,
-    dtype: torch.dtype,
-    repeat: int,
-    dtype_name: str,
-    compile_options: Optional[CompileOptions] = None,
-) -> list[dict]:
-    """Run memory and latency benchmarks for a single model specification."""
-    cfg = _prepare_config(base_cfg, spec, image_size)
-    wrapper = _instantiate_wrapper(cfg, device, dtype, compile_options=compile_options)
-    if spec.num_params is None:
-        param_count = sum(p.numel() for p in wrapper.parameters())
-        spec = replace(spec, num_params=param_count)
-    results = []
-
-    results.append(
-        _run_mode(
-            mode="inference",
-            build_fn=partial(_inference_fn, wrapper, image_size, dtype),
-            spec=spec,
-            image_size=image_size,
-            device=device,
-            repeat=repeat,
-            dtype_name=dtype_name,
-        )
-    )
-
-    results.append(
-        _run_mode(
-            mode="training",
-            build_fn=partial(_training_fn, wrapper, image_size, dtype),
-            spec=spec,
-            image_size=image_size,
-            device=device,
-            repeat=repeat,
-            dtype_name=dtype_name,
-        )
-    )
-    del wrapper
-    torch.cuda.empty_cache()
-    return results
-
-
-def _print_table(rows: Iterable[dict]) -> None:
-    header = "{:<10} {:<8} {:>5} {:>5} {:<6} {:>10} {:>10} {:>6} {:>6} {:>10} {:<8}"
-    print(
-        header.format(
-            "mode", "model", "res", "bs", "dtype", "time_ms", "mem_mb", "hidden", "layers", "params", "status"
-        )
-    )
-    for row in rows:
-        time_val = row.get("time_ms")
-        mem_val = row.get("peak_memory_mb")
-        time_str = f"{time_val:10.2f}" if isinstance(time_val, (int, float)) else f"{'--':>10}"
-        mem_str = f"{mem_val:10.1f}" if isinstance(mem_val, (int, float)) else f"{'--':>10}"
-        status = row.get("error", "ok")
-        params = row.get("num_params")
-        params_str = f"{params / 1e6:10.2f}M" if isinstance(params, (int, float)) else f"{'--':>10}"
-        print(
-            f"{row['mode']:<10} {row['model']:<8} {row['image_size']:>5} {row['batch_size']:>5} {row['dtype']:<6} "
-            f"{time_str} {mem_str} {row['hidden_dim']:>6} {row['num_layers']:>6} {params_str} {status:<8}"
-        )
-
-
-def parse_args() -> argparse.Namespace:
-    """Parse command-line arguments for the ImageNet diffusion benchmark."""
-    parser = argparse.ArgumentParser(description="Benchmark ImageNet diffusion model memory/time.")
-    parser.add_argument(
-        "--config",
-        type=str,
-        default="examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py",
-        help="Path to the ImageNet diffusion config file.",
-    )
-    parser.add_argument(
-        "--repeat",
-        type=int,
-        default=3,
-        help="Number of iterations to average for each measurement.",
-    )
-    parser.add_argument(
-        "--dtypes",
-        type=str,
-        nargs="+",
-        default=["fp32", "bf16", "fp16"],
-        help="List of precisions to benchmark (choices: fp32 bf16 fp16).",
-    )
-    parser.add_argument(
-        "--output-json",
-        type=str,
-        default=None,
-        help="Optional path to save the raw metrics as JSON.",
-    )
-    parser.add_argument(
-        "--torch-compile",
-        action="store_true",
-        help="Wrap the network and training step with torch.compile before benchmarking.",
-    )
-    parser.add_argument(
-        "--compile-backend",
-        type=str,
-        default="inductor",
-        help="Backend to use when --torch-compile is enabled.",
-    )
-    parser.add_argument(
-        "--compile-mode",
-        type=str,
-        default="default",
-        help="Compilation mode to use when --torch-compile is enabled.",
-    )
-    return parser.parse_args()
-
-
-def _dtype_from_string(name: str) -> torch.dtype:
-    mapping = {
-        "fp32": torch.float32,
-        "bf16": torch.bfloat16,
-        "fp16": torch.float16,
-    }
-    return mapping[name.lower()]
-
-
-def main() -> None:
-    """Entry point: parse args, instantiate models, and run benchmarks."""
-    args = parse_args()
-    base_cfg = load_config_from_file(args.config)
-    device = _ensure_cuda()
-    compile_options = None
-    if args.torch_compile:
-        compile_options = CompileOptions(
-            backend=args.compile_backend,
-            mode=args.compile_mode,
-        )
-
-    results: List[dict] = []
-    repeat = max(1, args.repeat)
-
-    for dtype_name in args.dtypes:
-        dtype = _dtype_from_string(dtype_name)
-        for spec in MODEL_SPECS:
-            for res in RESOLUTIONS:
-                spec_results = benchmark_spec(
-                    base_cfg=base_cfg,
-                    spec=spec,
-                    image_size=res,
-                    device=device,
-                    dtype=dtype,
-                    repeat=repeat,
-                    dtype_name=dtype_name,
-                    compile_options=compile_options,
-                )
-                results.extend(spec_results)
-
-    _print_table(results)
-    if args.output_json:
-        with open(args.output_json, "w", encoding="utf-8") as fh:
-            json.dump(results, fh, indent=2)
-        print(f"\nSaved metrics to {args.output_json}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/benchmarks/ops/FP16_FFTCONV_RESULTS.md b/benchmarks/ops/FP16_FFTCONV_RESULTS.md
deleted file mode 100644
index 8b60d519..00000000
--- a/benchmarks/ops/FP16_FFTCONV_RESULTS.md
+++ /dev/null
@@ -1,101 +0,0 @@
-# FP16 FFT Convolution: Benchmark Results
-
-Benchmark results for the dual-centered FP16 circular FFT convolution
-(`use_fp16_fft=True`) compared to the default FP32 path.
-
-All measurements on **Gray-Scott Reaction Diffusion** (Hyena, 2D, 128x128
-spatial resolution, BF16-mixed precision training) on a single **H100 SXM
-80 GB** GPU.
-
-## End-to-End Training Throughput
-
-Measured with the `IterationSpeedCallback` (windowed steady-state after
-`torch.compile` warmup) and wall-clock tqdm parsing, 600 training steps.
-
-### patch_size=1 (full-resolution 128x128 sequences)
-
-| Configuration           | Steady it/s | fwd (ms) | bwd (ms)   | other (ms) | Peak GPU (MB) |
-| ----------------------- | ----------- | -------- | ---------- | ---------- | ------------- |
-| FP32 fftconv, bs=16     | 3.45        | 98.2     | 185.9      | 6.0        | 60,565        |
-| **FP16 fftconv, bs=16** | **3.58**    | 108.8    | 163.7      | 6.7        | **58,049**    |
-| Delta                   | **+3.8%**   | +10.8%   | **-12.0%** | —          | **-4.2%**     |
-
-At full resolution, the FFT convolution is a small fraction of total
-compute. The centering overhead slows down the forward pass (+10ms), but
-smaller `complex32` intermediates speed up the backward pass (-22ms).
-Net: a modest 3.8% throughput improvement and 2.5 GB memory savings —
-not enough to increase batch size (bs=24 OOMs at 78.6/79.3 GB).
-
-### patch_size=2 (downsampled 64x64 sequences)
-
-| Configuration       | Steady it/s | fwd (ms) | bwd (ms) | other (ms) | Peak GPU (MB) |
-| ------------------- | ----------- | -------- | -------- | ---------- | ------------- |
-| FP32 fftconv, bs=16 | 11.85       | 28.8     | 50.4     | 5.2        | 15,293        |
-| FP16 fftconv, bs=16 | 11.24       | 36.6     | 47.2     | 5.2        | 14,651        |
-| Delta               | **-5.1%**   | +27.1%   | -6.3%    | —          | -4.2%         |
-
-At lower resolution the FFTs are already fast (~29ms fwd), so the
-centering overhead dominates: the forward pass increases by 8ms (+27%)
-while the backward saves only 3ms. **FP16 fftconv is not beneficial
-at patch_size=2.**
-
-### Varying batch size at patch_size=2
-
-| Configuration | Steady it/s | fwd (ms) | bwd (ms) | Peak GPU (MB) |
-| ------------- | ----------- | -------- | -------- | ------------- |
-| ps=2, bs=16   | 11.85       | 28.8     | 50.4     | 15,293        |
-| ps=2, bs=32   | 6.56        | 51.5     | 95.1     | 30,225        |
-| ps=2, bs=64   | 3.54        | 96.7     | 180.1    | 60,113        |
-
-Memory and compute scale linearly with batch size.  At ps=2 the GPU is
-under-utilized at bs=16 (15 GB / 80 GB), leaving room for much larger
-batches.
-
-## Data Loading: Not the Bottleneck
-
-| Configuration      | Steady it/s | other (ms) | Peak GPU (MB) |
-| ------------------ | ----------- | ---------- | ------------- |
-| nw=12, no preload  | 3.44        | 6.3        | 60,565        |
-| nw=12, RAM preload | 3.46        | 5.4        | 60,565        |
-| nw=4, no preload   | 3.46        | 5.5        | 60,564        |
-| nw=0, RAM preload  | 3.37        | 12.5       | 60,565        |
-| nw=0, no preload   | 3.06        | 43.0       | 60,565        |
-
-With `num_workers >= 4`, data loading (`other_ms`) is 5-6ms regardless
-of RAM preloading. The GPU compute (~284ms/step) is the sole bottleneck.
-RAM preloading (126 GB for Gray-Scott) provides no steady-state benefit.
-
-## Correctness
-
-The dual-centered FP16 implementation was validated against the FP32
-reference on a trained Euler Hyena checkpoint (177M parameters):
-
-- **Relative error** (vs FP32 reference): \< 1e-3 mean absolute error
-- **Validation loss**: identical to 4 significant figures
-- **No NaNs or Infs** across all tested configurations (1D, 2D, 3D)
-
-See `tests/test_circular_fftconv_fp16.py` for the automated test suite.
-
-## Summary
-
-| Scenario                       | FP16 fftconv recommendation                          |
-| ------------------------------ | ---------------------------------------------------- |
-| ps=1 (full res), memory-tight  | Marginal: +3.8% speed, -4.2% memory                  |
-| ps=1, need to fit larger batch | Not enough savings to change batch size              |
-| ps=2 (lower res)               | **Do not use** — centering overhead hurts throughput |
-| Correctness-critical           | Safe — validated against FP32 reference              |
-
-The primary value of the FP16 fftconv work is **fixing the NaN bug** in
-the original implementation (which was unusable). The centering technique
-is mathematically sound and numerically stable, but the practical
-throughput gains are modest because FFT convolutions are a small fraction
-of total model compute.
-
-## Environment
-
-- GPU: NVIDIA H100 SXM 80 GB
-- PyTorch 2.6.0+cu129, CUDA 12.9
-- Precision: BF16-mixed (`bf16-mixed` Lightning trainer)
-- `torch.compile` enabled (default mode)
-- Dataset: Gray-Scott Reaction Diffusion (The Well)
-- Model: Hyena with Gaussian mask, 12 blocks
diff --git a/benchmarks/ops/README.md b/benchmarks/ops/README.md
index 061b35f0..aa2c57e1 100644
--- a/benchmarks/ops/README.md
+++ b/benchmarks/ops/README.md
@@ -46,12 +46,3 @@ PYTHONPATH=. conda run -n nv-subq python benchmarks/ops/bench_subquadratic_fftco
 ```
 
 Output: stdout.
-
-## Related
-
-- [`FP16_FFTCONV_RESULTS.md`](FP16_FFTCONV_RESULTS.md) — accuracy and
-  throughput of the FP16 FFT convolution path against the FP32
-  reference.  See
-  [`docs/ops/FP16_FFTCONV_DERIVATION.md`](../../docs/ops/FP16_FFTCONV_DERIVATION.md)
-  for the dual-mean-centering derivation that motivates the FP16
-  implementation.
diff --git a/docs-tracker.md b/docs-tracker.md
index cf4bfe01..4487bef5 100644
--- a/docs-tracker.md
+++ b/docs-tracker.md
@@ -36,18 +36,15 @@ Work bottom-up: primitive ops → modules → networks → experiments.
 
 ### `nvsubquadratic/ops/` — FFT convolution primitives
 
-| File                       | Status | Notes                                                                                   |
-| -------------------------- | ------ | --------------------------------------------------------------------------------------- |
-| `README.md`                | \[x\]  | Folder overview, decision tree, math primer (new file)                                  |
-| `fftconv.py`               | \[x\]  | Module + key per-fn docstrings rewritten with math                                      |
-| `circular_fftconv.py`      | \[x\]  | Already strong; left as-is                                                              |
-| `circular_fftconv_fp16.py` | \[x\]  | Already strong; relies on FP16_FFTCONV_DERIVATION.md                                    |
-| `fftconv_fp16.py`          | \[x\]  | Already adequate; left as-is                                                            |
-| `fftconv_multihead.py`     | \[x\]  | Module docstring expanded with multi-head/low-rank math                                 |
-| `fftconv_chunked.py`       | \[x\]  | Already strong; left as-is                                                              |
-| `fftconv_custom.py`        | \[x\]  | Module docstring expanded with motivation; 1D causal wrappers added in 1D PR            |
-| `causal_conv1d_custom.py`  | \[x\]  | New (1D PR): thin wrappers for direct `causal_conv1d` + fused `b2b_causal_conv1d`       |
-| `mixed_fftconv.py`         | \[x\]  | New (#120): per-axis mixed boundary-condition FFT conv; see `docs/ops/MIXED_BC_PLAN.md` |
+| File                      | Status | Notes                                                                                               |
+| ------------------------- | ------ | --------------------------------------------------------------------------------------------------- |
+| `README.md`               | \[x\]  | Folder overview, decision tree, math primer (new file)                                              |
+| `fftconv.py`              | \[x\]  | Module + key per-fn docstrings rewritten with math                                                  |
+| `circular_fftconv.py`     | \[x\]  | Already strong; left as-is                                                                          |
+| `fftconv_chunked.py`      | \[x\]  | Already strong; left as-is                                                                          |
+| `fftconv_custom.py`       | \[x\]  | Module docstring expanded with motivation; 1D causal wrappers added in 1D PR                        |
+| `causal_conv1d_custom.py` | \[x\]  | New (1D PR): thin wrappers for direct `causal_conv1d` + fused `b2b_causal_conv1d`                   |
+| `mixed_fftconv.py`        | \[x\]  | New (#120): per-axis mixed boundary-condition FFT conv; see `docs/ops/mixed_boundary_conditions.md` |
 
 ### `nvsubquadratic/modules/` — Building blocks
 
@@ -56,7 +53,6 @@ Work bottom-up: primitive ops → modules → networks → experiments.
 | `kernels_nd.py`                    | \[x\]  | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants                           |
 | `hyena_nd.py`                      | \[x\]  | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution                        |
 | `ckconv_nd.py`                     | \[x\]  | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes                  |
-| `ckconv_multihead_nd.py`           | \[x\]  | Multi-head CKConv — H heads, dense d×d kernel per head, low-rank U·V factorisation                |
 | `mamba_nd.py`                      | \[x\]  | Mamba SSM (ND) — selective SSM, ZOH discretisation, raster-scan ND, bidirectional mode            |
 | `attention.py`                     | \[x\]  | Scaled dot-product attention — multi-head, RoPE, ND spatial, O(L²) FLOP formula                   |
 | `vit5_attention.py`                | \[x\]  | ViT5 attention — register-aware 2D RoPE, QK-norm, CUDA-graph-safe buffers                         |
@@ -89,8 +85,6 @@ Work bottom-up: primitive ops → modules → networks → experiments.
 | `classification_resnet.py`            | \[x\]  | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork                                                         |
 | `vit5_classification.py`              | \[x\]  | ViT5 classification — token layout, hybrid blocks, CLS/GAP/register_concat readouts, FLOP count                                           |
 | `vit5_hierarchical_classification.py` | \[ \]  | Pending (#122 — `feat/patch-merging`): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout |
-| `huggingface_diffusers.py`            | \[x\]  | HF diffusers adapters — DiT and UVit wrappers, BHL↔BCHW translation, shared timestep state model                                          |
-| `jit.py` / `jit_utils.py`             | \[x\]  | JiT diffusion backbone port (LTH14/JiT) — patch embed, attention/SwiGLU blocks, RoPE helpers, RMSNorm                                     |
 | `baselines/unet_convnext.py`          | \[x\]  | UNet-ConvNeXt baseline ported from The Well — preserves upstream `skips[0]` bug for reproducibility                                       |
 | `baselines/unet_convnext_v2.py`       | \[x\]  | Fixed-skip variant of UNet-ConvNeXt — consumes the missing finest-resolution skip                                                         |
 
@@ -110,7 +104,6 @@ the `tests/` tree is out of scope for this tracker.)
 | File             | Status | Notes                                                                                 |
 | ---------------- | ------ | ------------------------------------------------------------------------------------- |
 | `lazy_config.py` | \[x\]  | LazyConfig + instantiate — deferred instantiation, nested configs, arithmetic strings |
-| `metrics/`       | \[x\]  | cleanfid.py — CleanFID wrapper, FID formula, usage context                            |
 | `utils/`         | \[x\]  | init.py (weight init factories), qk_norm.py (QK-norm, L2Norm module), quack_utils.py  |
 | `testing/`       | \[x\]  | utils.py — compute_relative_error, TTrace reference, already had good docstrings      |
 
@@ -124,7 +117,6 @@ the `tests/` tree is out of scope for this tracker.)
 | `lightning_wrappers/base_lightning_wrapper.py` | \[x\]  | LightningWrapperBase — param groups, scheduler, checkpoint resume, profiling                               |
 | `lightning_wrappers/classification_wrapper.py` | \[x\]  | ClassificationWrapper — cross_entropy / soft_target_ce / bce loss modes                                    |
 | `lightning_wrappers/regression_wrapper.py`     | \[x\]  | RegressionWrapper — MAE/MSE loss, base for WELLRegressionWrapper                                           |
-| `lightning_wrappers/diffusion_wrapper.py`      | \[x\]  | DiffusionWrapper — JiT continuous-time diffusion, ODE sampler                                              |
 | `lightning_wrappers/autoregressive_wrapper.py` | \[x\]  | Already had good module + class docstrings; left as-is                                                     |
 | `lightning_wrappers/arc_wrapper.py`            | \[ \]  | (untracked new file — out of scope until merged)                                                           |
 | `lightning_wrappers/well_lightning_wrapper.py` | \[x\]  | Already had good class docstring; left as-is                                                               |
@@ -132,7 +124,6 @@ the `tests/` tree is out of scope for this tracker.)
 | `datamodules/mnist.py`                         | \[x\]  | MNIST datamodule — channels-last reshape, train/val split                                                  |
 | `datamodules/emnist.py`                        | \[x\]  | EMNIST datamodule — all five splits (digits, letters, balanced, bymerge, byclass)                          |
 | `datamodules/tinyimagenet.py`                  | \[x\]  | TinyImageNet HF-backed datamodule — RandAugment, Mixup/CutMix, token access                                |
-| `datamodules/ucf101.py`                        | \[x\]  | UCF101 — video/sequence modes, frames_per_clip, deterministic workers                                      |
 | `datamodules/dali_imagenet_fused.py`           | \[x\]  | DALI ImageNet — fused GPU augmentation, MixupConfig/AugmentConfig, repeated aug                            |
 | `datamodules/spatial_recall_dataset.py`        | \[x\]  | Already had comprehensive module + class docstrings; left as-is                                            |
 | `datamodules/pde/well.py`                      | \[x\]  | WELL benchmark datamodule — persistent HDF5 handles, NVMe staging, RAM preload, val/test normalisation fix |
@@ -156,7 +147,6 @@ README per subdirectory.  No Sphinx API reference entry.
 | `ops/bench_fftconv2d.py`                         | \[x\]  | Already had a full docstring                                                                               |
 | `ops/bench_mlp.py`                               | \[x\]  | torch vs QuACK MLP correctness + timing                                                                    |
 | `ops/bench_subquadratic_fftconv.py`              | \[x\]  | Sanity gate for the CUDA fft_causal_conv1d kernel                                                          |
-| `ops/FP16_FFTCONV_RESULTS.md`                    | \[x\]  | FP16 FFT conv accuracy + throughput vs FP32 reference                                                      |
 | `vit5_imagenet/README.md`                        | \[x\]  | Per-script overview + pointer to the historical profiling report                                           |
 | `vit5_imagenet/bench_vit5_baseline.py`           | \[x\]  | Baseline (eager, torchvision dataloader) throughput                                                        |
 | `vit5_imagenet/bench_vit5_compile.py`            | \[x\]  | `torch.compile` configuration sweep                                                                        |
diff --git a/docs/api_reference/core.rst b/docs/api_reference/core.rst
index 9c6bc3c3..e1083472 100644
--- a/docs/api_reference/core.rst
+++ b/docs/api_reference/core.rst
@@ -7,8 +7,7 @@ Core
 
 Top-level utilities: the lazy-instantiation system that powers every
 config file, weight-init helpers, QK-norm and rotary embedding
-primitives, the QuACK-kernel capability probe, FID computation, and
-testing helpers.
+primitives, the QuACK-kernel capability probe, and testing helpers.
 
 Lazy configuration
 ------------------
@@ -83,15 +82,6 @@ QuACK capability probe
 
    ~utils.quack_utils.cuda_supports_quack
 
-Metrics
--------
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~metrics.cleanfid.compute_folder_fid
-
 Testing helpers
 ---------------
 
diff --git a/docs/api_reference/experiments.rst b/docs/api_reference/experiments.rst
index 6b58ce73..e9a3fde2 100644
--- a/docs/api_reference/experiments.rst
+++ b/docs/api_reference/experiments.rst
@@ -7,7 +7,7 @@ Experiments
 
 The ``experiments`` package wires nvSubquadratic modules into reproducible
 training pipelines built on PyTorch Lightning.  Each experiment is a
-``LazyConfig`` of dataclasses (``ExperimentConfig`` / ``DiffusionExperimentConfig``)
+``LazyConfig`` of dataclasses (``ExperimentConfig``)
 plus a wrapper subclass that defines the training step.
 
 Entry points
@@ -33,14 +33,12 @@ the per-experiment config files in ``examples/``.
    :template: class_template.rst
 
    ~default_cfg.ExperimentConfig
-   ~default_cfg.DiffusionExperimentConfig
    ~default_cfg.TrainConfig
    ~default_cfg.TrainerConfig
    ~default_cfg.SchedulerConfig
    ~default_cfg.WandbConfig
    ~default_cfg.AutoResumeConfig
    ~default_cfg.StartFromCheckpointConfig
-   ~default_cfg.DiffusionConfig
 
 Lightning wrappers
 ------------------
@@ -55,7 +53,6 @@ Task-specific wrappers around a common base.  Each wrapper defines
    ~lightning_wrappers.base_lightning_wrapper.LightningWrapperBase
    ~lightning_wrappers.classification_wrapper.ClassificationWrapper
    ~lightning_wrappers.classification_wrapper.SoftTargetCrossEntropy
-   ~lightning_wrappers.diffusion_wrapper.DiffusionWrapper
    ~lightning_wrappers.regression_wrapper.RegressionWrapper
    ~lightning_wrappers.well_lightning_wrapper.WELLRegressionWrapper
    ~lightning_wrappers.autoregressive_wrapper.AutoregressiveWrapper
@@ -98,7 +95,6 @@ that experiments target.
    ~datamodules.mnist
    ~datamodules.emnist
    ~datamodules.tinyimagenet
-   ~datamodules.ucf101
    ~datamodules.spatial_recall_dataset
    ~datamodules.dali_imagenet_fused
    ~datamodules.pde.well
diff --git a/docs/api_reference/modules.rst b/docs/api_reference/modules.rst
index 8b65ece0..6893a01c 100644
--- a/docs/api_reference/modules.rst
+++ b/docs/api_reference/modules.rst
@@ -38,7 +38,6 @@ context-parallel counterparts.
    ~modules.causal_conv1d.CausalConv1D
    ~modules.subq_ops_causal_conv1d.SubqOpsCausalConv1d
    ~modules.ckconv_nd.CKConvND
-   ~modules.ckconv_multihead_nd.CKConvMultiheadND
    ~modules.distributed_depthwise_conv_nd.DistributedDepthwiseConv1d
    ~modules.distributed_depthwise_conv_nd.DistributedDepthwiseConv2d
    ~modules.distributed_depthwise_conv_nd.DistributedDepthwiseConv3d
diff --git a/docs/api_reference/networks.rst b/docs/api_reference/networks.rst
index e1a92860..4761b18b 100644
--- a/docs/api_reference/networks.rst
+++ b/docs/api_reference/networks.rst
@@ -6,8 +6,7 @@ Networks
 ========
 
 End-to-end classification and general-purpose networks composing the
-modules above, plus the diffusion backbones and the UNet-ConvNeXt
-baselines used in benchmark comparisons.
+modules above, plus the UNet-ConvNeXt baselines used in benchmark comparisons.
 
 Classification & general-purpose
 --------------------------------
@@ -20,74 +19,6 @@ Classification & general-purpose
    ~networks.general_purpose_resnet.ResidualNetwork
    ~networks.vit5_classification.ViT5ClassificationNet
 
-Diffusion — Hugging Face adapters
----------------------------------
-
-Wrappers that expose :class:`diffusers.DiTTransformer2DModel` and
-:class:`diffusers.UVit2DModel` to the diffusion Lightning wrapper.
-
-.. autosummary::
-   :toctree: generated/
-   :template: class_template.rst
-
-   ~networks.huggingface_diffusers.HuggingFaceDiTConfig
-   ~networks.huggingface_diffusers.HuggingFaceUVitConfig
-   ~networks.huggingface_diffusers.DiffusersDiTWrapper
-   ~networks.huggingface_diffusers.DiffusersUVitWrapper
-
-Diffusion — JiT backbone
-------------------------
-
-Port of the JiT diffusion model (`LTH14/JiT <https://github.com/LTH14/JiT>`_) —
-patch-embedding, transformer blocks with RoPE and SwiGLU FFN, and the
-factory functions for the published model sizes.
-
-.. autosummary::
-   :toctree: generated/
-   :template: class_template.rst
-
-   ~networks.jit.JiT
-   ~networks.jit.JiTBlock
-   ~networks.jit.BottleneckPatchEmbed
-   ~networks.jit.TimestepEmbedder
-   ~networks.jit.LabelEmbedder
-   ~networks.jit.Attention
-   ~networks.jit.SwiGLUFFN
-   ~networks.jit.FinalLayer
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~networks.jit.modulate
-   ~networks.jit.JiT_B_4
-   ~networks.jit.JiT_B_16
-   ~networks.jit.JiT_B_32
-   ~networks.jit.JiT_L_16
-   ~networks.jit.JiT_L_32
-   ~networks.jit.JiT_H_16
-   ~networks.jit.JiT_H_32
-
-JiT helpers (rotary embeddings, RMSNorm, sin-cos position embeddings):
-
-.. autosummary::
-   :toctree: generated/
-   :template: class_template.rst
-
-   ~networks.jit_utils.VisionRotaryEmbedding
-   ~networks.jit_utils.VisionRotaryEmbeddingFast
-   ~networks.jit_utils.RMSNorm
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~networks.jit_utils.broadcat
-   ~networks.jit_utils.rotate_half
-   ~networks.jit_utils.get_1d_sincos_pos_embed_from_grid
-   ~networks.jit_utils.get_2d_sincos_pos_embed
-   ~networks.jit_utils.get_2d_sincos_pos_embed_from_grid
-
 Baselines
 ---------
 
diff --git a/docs/api_reference/ops.rst b/docs/api_reference/ops.rst
index 1f79c097..1ea8ab3a 100644
--- a/docs/api_reference/ops.rst
+++ b/docs/api_reference/ops.rst
@@ -9,10 +9,6 @@ Low-level convolution primitives.  Pure-PyTorch reference implementations
 double as the spec the CUDA kernels must match; the
 ``subquadratic_ops_torch`` wrappers are the production path on GPU.
 
-The fp16 variants require power-of-2 spatial dims (cuFFT constraint) and
-use dual mean-centering for numerical stability — see
-:doc:`../ops/FP16_FFTCONV_DERIVATION` for the derivation.
-
 FFT convolutions (reference fp32)
 ---------------------------------
 
@@ -76,21 +72,6 @@ Periodic-boundary FFT convolutions for global mixing without zero padding.
    ~ops.circular_fftconv.circular_fftconv2d_fp32_bhl
    ~ops.circular_fftconv.circular_fftconv3d_fp32_bhl
 
-Multi-head FFT convolutions
----------------------------
-
-Multi-head variants used by Hyena-style mixers, including low-rank
-factorizations.
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~ops.fftconv_multihead.fftconv2d_multihead_bhl
-   ~ops.fftconv_multihead.fftconv2d_multihead_lowrank_bhl
-   ~ops.fftconv_multihead.fftconv2d_multihead_circular_bhl
-   ~ops.fftconv_multihead.fftconv2d_multihead_lowrank_circular_bhl
-
 Chunking utilities
 ------------------
 
@@ -106,52 +87,13 @@ sequence axis in chunks.
    ~ops.fftconv_chunked.set_default_chunk_size
    ~ops.fftconv_chunked.get_default_chunk_size
 
-FFT convolutions (fp16)
------------------------
-
-Half-precision linear-convolution variants.  Internal compute is fp16
-via cuFFT; output dtype matches the caller's input.
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~ops.fftconv_fp16.fftconv1d_fp16_bhl
-   ~ops.fftconv_fp16.fftconv2d_fp16_bhl
-   ~ops.fftconv_fp16.fftconv3d_fp16_bhl
-   ~ops.fftconv_fp16.causal_fftconv1d_fp16_bhl
-   ~ops.fftconv_fp16.fftconv1d_fp16_bhl_w_reshape
-   ~ops.fftconv_fp16.fftconv2d_fp16_bhl_w_reshape
-   ~ops.fftconv_fp16.fftconv3d_fp16_bhl_w_reshape
-   ~ops.fftconv_fp16.causal_fftconv1d_fp16_bhl_w_reshape
-   ~ops.fftconv_fp16.fftconv1d_fp16_bhl_chunked
-   ~ops.fftconv_fp16.fftconv2d_fp16_bhl_chunked
-   ~ops.fftconv_fp16.fftconv3d_fp16_bhl_chunked
-   ~ops.fftconv_fp16.causal_fftconv1d_fp16_bhl_chunked
-
-Circular FFT convolutions (fp16)
---------------------------------
-
-Periodic-boundary half-precision variants.
-
-.. autosummary::
-   :toctree: generated/
-   :template: function_template.rst
-
-   ~ops.circular_fftconv_fp16.circular_fftconv1d_fp16_bhl
-   ~ops.circular_fftconv_fp16.circular_fftconv2d_fp16_bhl
-   ~ops.circular_fftconv_fp16.circular_fftconv3d_fp16_bhl
-   ~ops.circular_fftconv_fp16.circular_fftconv1d_fp16_bhl_w_reshape
-   ~ops.circular_fftconv_fp16.circular_fftconv2d_fp16_bhl_w_reshape
-   ~ops.circular_fftconv_fp16.circular_fftconv3d_fp16_bhl_w_reshape
-
-Mixed-precision FFT convolutions
---------------------------------
+Mixed boundary-condition FFT convolutions
+-----------------------------------------
 
-FFT convolutions that switch internal precision per-axis (e.g. fp16 on
-power-of-2 dims, fp32 on others).  See the
-`FP16 Circular FFT Convolution: Derivation <../ops/FP16_FFTCONV_DERIVATION.html>`_
-for the numerical-stability background.
+FFT convolutions with per-axis boundary conditions — periodic on some
+spatial axes, zero-padded on others.  See
+:doc:`../ops/mixed_boundary_conditions` for the per-axis algorithm and the
+``fft_padding`` API.
 
 .. autosummary::
    :toctree: generated/
diff --git a/docs/architecture.md b/docs/architecture.md
index a9b7e85a..c0730f8d 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -66,10 +66,6 @@ are documented in detail in `docs/ops/README.md`; the short version:
   to BHL, run the fast path, and reshape back.  Recommended entry point
   for channels-last callers.
 - **`_chunked`** — processes channels in groups to cap peak FFT memory.
-- **`fp32` vs `fp16`** — internal compute precision.  fp16 ops require
-  power-of-2 spatial dims (cuFFT constraint) and use dual mean-centering
-  for numerical stability — see
-  [FP16 Circular FFT Convolution: Derivation](ops/FP16_FFTCONV_DERIVATION.md).
 
 So `causal_fftconv1d_fp32_bhl_w_reshape` is a causal 1D FFT conv that
 accepts channels-last input, runs the fp32 channels-first kernel under
diff --git a/docs/benchmarks.md b/docs/benchmarks.md
index 7aec3b69..b0c0c817 100644
--- a/docs/benchmarks.md
+++ b/docs/benchmarks.md
@@ -1,6 +1,6 @@
 # Benchmarks
 
-Throughput numbers, FLOP scaling, and FP16 op-level results.  The
+Throughput numbers and FLOP scaling.  The
 tables below are included verbatim from the
 [`benchmarks/README.md`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/benchmarks/README.md)
 single-source — edits should land there, not here.
@@ -19,10 +19,3 @@ for the script that produced the plot.
 start-after: '# ViT-5-Small Throughput Benchmarks'
 ---
 ```
-
-## Op-level results
-
-- [FP16 FFT convolution results](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/benchmarks/ops/FP16_FFTCONV_RESULTS.md)
-  — accuracy and throughput of the FP16 path against the FP32 reference,
-  with the dual-mean-centering derivation summarised in
-  [FP16 Circular FFT Convolution: Derivation](ops/FP16_FFTCONV_DERIVATION.md).
diff --git a/docs/conf.py b/docs/conf.py
index a52961c3..8f02a59e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -92,8 +92,6 @@ def _read_version():
     "megatron",
     "megatron.core",
     "omegaconf",
-    "cleanfid",
-    "diffusers",
     "pytorch_lightning",
     "lightning",
     "matplotlib",
@@ -102,7 +100,6 @@ def _read_version():
     "h5py",
     "scipy",
     "the_well",
-    "torch_fidelity",
     "torchmetrics",
     "torchvision",
     "timm",
diff --git a/docs/examples/index.md b/docs/examples/index.md
index a3469a5d..09bc46ef 100644
--- a/docs/examples/index.md
+++ b/docs/examples/index.md
@@ -11,13 +11,6 @@ The active experimental roadmap (priorities, owners, status) lives at
 
 ## Classification
 
-### MNIST / SMNIST
-
-[`examples/mnist_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/mnist_classification)
-covers MNIST with both attention and Hyena baselines, plus a small CCNN
-backbone.  [`examples/smnist_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/smnist_classification)
-covers sequential MNIST (1D input).
-
 ### ImageNet
 
 [`examples/imagenet_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/imagenet_classification)
@@ -32,25 +25,6 @@ Representative entry points: `ccnn_7_512_hyena.py`,
 is the ViT-5 baseline suite (v1–v5) with its own
 [`TRACKER.md`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/vit5_imagenet/TRACKER.md).
 
-### UCF101
-
-[`examples/ucf101_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/ucf101_classification)
-covers video classification with both sequence- and clip-mode datamodules.
-
-## Diffusion
-
-### MNIST
-
-[`examples/mnist_diffusion/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/mnist_diffusion)
-is a small DDPM/JiT diffusion sanity-check.
-
-### ImageNet
-
-[`examples/imagenet_diffusion/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/imagenet_diffusion)
-is the full ImageNet diffusion setup.  See its
-[README](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/imagenet_diffusion/README.md)
-for the JiT vs Hyena-vs-attention comparison.
-
 ## Spatial recall
 
 [`examples/spatial_recall_1d/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/spatial_recall_1d),
diff --git a/docs/ops/FP16_FFTCONV_DERIVATION.md b/docs/ops/FP16_FFTCONV_DERIVATION.md
deleted file mode 100644
index afd745df..00000000
--- a/docs/ops/FP16_FFTCONV_DERIVATION.md
+++ /dev/null
@@ -1,283 +0,0 @@
-# FP16 Circular FFT Convolution: Derivation
-
-This document derives the numerically stable FP16 circular FFT convolution
-used in `circular_fftconv_fp16.py`.  The technique applies to 1D, 2D, and
-3D and is a drop-in replacement for the FP32 variants.
-
-## 1. Problem Statement
-
-We want to compute the circular (periodic) convolution of a signal
-**x** of length $N$ with a kernel **k** of length $K \le N$, using
-half-precision (FP16) FFTs for speed and memory savings.
-
-The standard approach is:
-
-$$
-y = \text{IFFT}\!\bigl( \text{FFT}(x) \odot \text{FFT}(k_{\text{pad}}) \bigr)
-$$
-
-where $k_{\text{pad}}$ is **k** zero-padded to length $N$.
-
-### Why naive FP16 fails
-
-cuFFT supports FP16 transforms for power-of-2 sizes, but `float16` can
-only represent values up to 65504.  Two overflow paths exist:
-
-1. **DC-bin overflow.**  The DC bin of an un-normalized FFT is
-   $\hat{x}[0] = \sum_n x[n]$.  For a signal of length $N$ with mean
-   $\mu_x$, this is $\mu_x N$.  Even for moderate means ($\mu_x = 5$,
-   $N = 16384$), $\mu_x N = 81920 > 65504$.  The frequency-domain
-   product at DC is $\mu_x \mu_k N^2$, which overflows even more
-   easily.
-
-1. **Internal accumulation overflow.**  cuFFT computes butterfly sums
-   in the working precision.  In FP16, partial sums during the FFT
-   reach $O(\mu N)$ before any normalization, causing intermediate
-   `Inf` or `NaN` even if the final result would be representable.
-
-### Using `norm="ortho"`
-
-PyTorch's `norm="ortho"` divides the forward FFT by $\sqrt{N}$ and the
-inverse by $\sqrt{N}$, so the round-trip gives
-$\text{ortho-IFFT}(\text{ortho-FFT}(x) \odot \text{ortho-FFT}(k)) = y / \sqrt{N}$.
-We multiply by $\sqrt{N}$ after the inverse to recover the correct
-scale:
-
-$$
-y = \sqrt{N} \cdot \text{IFFT}_{\text{ortho}}\!\bigl(
-\text{FFT}_{\text{ortho}}(x) \odot \text{FFT}_{\text{ortho}}(k_{\text{pad}})
-\bigr)
-$$
-
-This reduces the DC bin to $\mu_x \sqrt{N}$, and the DC product to
-$\mu_x \mu_k N$, which still overflows for large $N$.  More importantly,
-it does **not** fix internal accumulation overflow (#2 above), because
-cuFFT performs the $\sqrt{N}$ scaling *after* the butterfly, not during.
-
-## 2. Solution: Dual Mean-Centering
-
-### Core idea
-
-Remove the DC component from both signals before the FFT:
-
-$$
-x_c = x - \mu_x, \qquad k_c = k - \mu_k
-$$
-
-Both centered signals have zero mean, so:
-
-- Their DC bins are exactly 0 (fixing overflow path #1).
-- Internal FFT sums are $O(\sigma)$ instead of $O(\mu N)$ (fixing #2).
-
-We then recover the exact convolution result analytically.
-
-### 1D Derivation
-
-Let $L = N$ be the circular convolution length.  The kernel has $K$
-nonzero elements ($K \le L$, typically $K = L$ or $K = L-1$) and is
-zero-padded to length $L$.
-
-Define:
-
-- $\mu_x = \frac{1}{L}\sum_{n} x[n]$, and $x_c = x - \mu_x$
-- $\mu_k = \frac{1}{K}\sum_{m} k[m]$, and $k_c = k - \mu_k$
-- $\delta[m]$ = the "mean indicator": $\mu_k$ for $0 \le m < K$, zero
-  otherwise (i.e., $k_{\text{pad}} - k_{c,\text{pad}} = \delta$)
-
-Then the zero-padded kernel decomposes as:
-
-$$
-k_{\text{pad}} = k_{c,\text{pad}} + \delta
-$$
-
-Expanding the circular convolution:
-
-$$
-y[n] =
-\underbrace{(x_c * k_{c,\text{pad}})[n]}_{T_1}
-+ \underbrace{(x_c * \delta)[n]}_{T_2}
-+ \underbrace{\mu_x \sum_m k_{c,\text{pad}}[m]}_{T_3}
-+ \underbrace{\mu_x \cdot \mu_k \cdot K}_{T_4}
-$$
-
-where $*$ denotes circular convolution.
-
-**Term T1** — the "safe" convolution of two zero-mean signals.  Both DC
-bins are 0, internal magnitudes are $O(\sigma)$.  Computed via FP16 FFT.
-
-**Term T3** = 0, because $\sum_m k_c[m] = 0$ by construction.
-
-**Term T4** is a scalar constant, computed in FP32.
-
-**Term T2** — the centering correction — requires care.
-
-#### T2: Centering correction (1D)
-
-$$
-T_2[n] = \sum_m x_c[(n-m) \bmod L] \cdot \delta[m]
-= \mu_k \sum_{m=0}^{K-1} x_c[(n-m) \bmod L]
-$$
-
-**Case $K = L$** (kernel covers the full circle):
-
-$$
-T_2[n] = \mu_k \sum_{m=0}^{L-1} x_c[(n-m) \bmod L] = \mu_k \cdot 0 = 0
-$$
-
-because $\sum x_c = 0$.  **No correction needed.**
-
-**Case $K = L-1$** (one zero-padded position, at index $L-1$):
-
-The sum covers all positions except $m = L-1$:
-
-$$
-T_2[n] = \mu_k \!\!\sum_{m \ne L-1}\!\! x_c[(n-m) \bmod L]
-       = -\mu_k \cdot x_c[(n+1) \bmod L]
-$$
-
-This is just $-\mu_k$ times a circular shift of $x_c$ by $-1$.
-
-#### Phase-ramp absorption
-
-In frequency domain, a circular shift by $s$ corresponds to
-multiplication by the phase ramp $\phi_s[f] = e^{-2\pi i f s / L}$.
-
-When the kernel is centered with shift $s = -\lfloor(K-1)/2\rfloor$,
-the zero-padded position moves, and the correction becomes:
-
-$$
-T_2[n] = -\mu_k \cdot x_c[(n - s + 1) \bmod L]
-$$
-
-In frequency domain, the effective kernel spectrum (including T1 + T2)
-is:
-
-$$
-\hat{k}_{\text{eff}}[f] = \phi_s[f] \cdot
-\biggl(\hat{k}_c[f] - \tfrac{\mu_k}{\sqrt{L}} \cdot \phi_{-1}[f]\biggr)
-$$
-
-where $\hat{k}_c$ is the ortho-normalized FFT of $k_c$ (zero-padded),
-and $\phi_{-1}[f] = e^{2\pi i f / L}$ is the DFT of the delta at
-position $L-1$.
-
-The $1/\sqrt{L}$ factor keeps intermediate values small (matching the
-ortho scaling); the compensating $\sqrt{L}$ is applied after the inverse
-FFT in FP32.
-
-### Final 1D formula
-
-$$
-y = \sqrt{L}\,\text{IFFT}_{\text{ortho}}\!\bigl(
-\hat{x}_c \odot \hat{k}_{\text{eff}}
-\bigr) + \mu_x \mu_k K
-$$
-
-where $\hat{x}_c = \text{FFT}_{\text{ortho}}(x_c)$ is computed in FP16,
-$\hat{k}_{\text{eff}}$ is assembled in FP32 (small tensor, no batch dim)
-and cast to `complex32` before the element-wise multiply, and the
-$\sqrt{L}$ rescaling and DC correction are done in FP32.
-
-## 3. Extension to nD
-
-For $d$-dimensional signals of shape $N_1 \times \cdots \times N_d$
-with kernel shape $K_1 \times \cdots \times K_d$, the decomposition
-generalizes.  The T1 and T4 terms carry over directly:
-
-$$
-T_1 = \text{circ\_conv}(x_c, k_{c,\text{pad}}), \qquad
-T_4 = \mu_x \mu_k \prod_i K_i
-$$
-
-T3 is again zero.  **T2 becomes an inclusion-exclusion sum over
-corrected axes** — those axes where $K_i < N_i$ (i.e., there is at
-least one zero-padded position).
-
-### Geometric correction factor
-
-Define $\mathcal{C} = {i : K_i < N_i}$ as the set of corrected axes.
-For each nonempty subset $S \subseteq \mathcal{C}$, define the phase
-factor $p_i[f_i] = e^{2\pi i f_i / N_i}$ (the DFT of a delta at
-position $N_i - 1$ along axis $i$).
-
-The geometric correction in frequency domain is:
-
-$$
-\text{geo}[\mathbf{f}] = \sum_{\emptyset \ne S \subseteq \mathcal{C}}
-(-1)^{|S|} \Bigl(\prod_{i \in S} p_i[f_i]\Bigr)
-\Bigl(\prod_{j \notin S} N_j\Bigr)
-$$
-
-The corrected effective kernel spectrum is:
-
-$$
-\hat{k}_{\text{eff}}[\mathbf{f}] = \hat{k}_c[\mathbf{f}]
-+ \tfrac{\mu_k}{\sqrt{N}} \cdot \text{geo}[\mathbf{f}]
-$$
-
-followed by the phase-ramp shift for kernel centering.
-
-### 2D example
-
-For a 2D signal $X \times Y$ with kernel $K_x \times K_y$:
-
-**Both axes corrected** ($K_x < X$ and $K_y < Y$):
-
-$$
-\text{geo}[f_1, f_2] = -Y \cdot p_x[f_1]\,\delta_{f_2=0}
-- X \cdot \delta_{f_1=0}\,p_y[f_2]
-+ p_x[f_1]\,p_y[f_2]
-$$
-
-**One axis corrected** (e.g., $K_x < X$, $K_y = Y$):
-
-$$
-\text{geo}[f_1, f_2] = -Y \cdot p_x[f_1]\,\delta_{f_2=0}
-$$
-
-### Caching
-
-The geometric factor `geo` depends only on
-$(K_1, \ldots, K_d, N_1, \ldots, N_d, \text{device})$ and is
-**constant** during training.  It is computed once and cached with an
-LRU policy.  Only the scalar $\mu_k / \sqrt{N}$ changes per forward
-call.
-
-## 4. Implementation Details
-
-### Precision strategy
-
-| Operation                              | Precision        | Rationale                                                |
-| -------------------------------------- | ---------------- | -------------------------------------------------------- |
-| Mean computation ($\mu_x$, $\mu_k$)    | FP16             | Cheap, small values                                      |
-| Forward FFT                            | FP16 (complex32) | Centered signals, cuFFT half-precision                   |
-| $\hat{k}_{\text{eff}}$ assembly        | FP32 (complex64) | Small tensor (no batch), needs precision for phase ramps |
-| $\hat{x}_c \odot \hat{k}_{\text{eff}}$ | complex32        | Large tensor, cast $\hat{k}_{\text{eff}}$ down           |
-| Inverse FFT                            | FP16 (complex32) | Same cuFFT path as forward                               |
-| $\sqrt{N}$ rescaling + DC correction   | FP32             | Avoid overflow in final scaling                          |
-| Output cast                            | Original dtype   | Match caller expectations                                |
-
-### cuFFT power-of-2 constraint
-
-cuFFT only supports FP16 transforms for power-of-2 sizes.  Since we use
-same-size circular FFTs (no padding), the **input spatial dimensions
-must be powers of 2**.  This is asserted at runtime.
-
-### Phase ramps
-
-Alignment shifts are implemented as frequency-domain phase ramps
-(precomputed and cached in FP32, cast to complex32 before multiply).
-This avoids a spatial `torch.roll` which would prevent fusion under
-`torch.compile`.
-
-## 5. Correctness Guarantees
-
-- **Mathematically exact**: the dual-centering decomposition introduces
-  no approximation.  The only source of error is finite-precision
-  arithmetic in the FP16 FFT (vs FP32).
-- **Validated**: automated tests in `tests/test_circular_fftconv_fp16.py`
-  verify 1D, 2D, and 3D implementations against FP32 references across
-  multiple shapes and kernel sizes.
-- **Inference validated**: on a trained 177M-parameter Euler Hyena model,
-  the FP16 path produces identical validation loss (to 4 significant
-  figures) as the FP32 reference.
diff --git a/docs/ops/MIXED_BC_PLAN.md b/docs/ops/MIXED_BC_PLAN.md
deleted file mode 100644
index 7e97123d..00000000
--- a/docs/ops/MIXED_BC_PLAN.md
+++ /dev/null
@@ -1,345 +0,0 @@
-# Mixed Boundary-Condition FFT Convolution — Plan & Tracker
-
-**Status:** In progress (v1 ops + tests)
-**Branch:** `dwromero/mixed-bc-fftconv`
-**Owner:** dwromero
-**Started:** 2026-05-20
-
-This is the working plan and tracker for adding **per-axis boundary
-condition** support to the FFT-based convolution operators
-(`nvsubquadratic/ops/`) and the modules that consume them
-(`nvsubquadratic/modules/ckconv_nd.py`, etc.).
-
-If you pick this up later, read the [survey & decisions](#1-context--decisions)
-section, then check [§5 Tracked questions](#5-tracked-questions--revisit-items)
-for what still needs to be done.
-
-______________________________________________________________________
-
-## 1. Context & decisions
-
-### Motivation
-
-Several Well PDE datasets have boundaries that are **periodic on some axes
-and non-periodic (WALL or OPEN) on others**, e.g.
-
-| Dataset                        | x        | y        | z    |
-| ------------------------------ | -------- | -------- | ---- |
-| `rayleigh_benard`              | periodic | wall     | —    |
-| `viscoelastic_instability`     | periodic | wall     | —    |
-| `turbulent_radiative_layer_2D` | periodic | open     | —    |
-| `turbulent_radiative_layer_3D` | periodic | periodic | open |
-| `rayleigh_taylor_instability`  | periodic | periodic | wall |
-| `helmholtz_staircase`          | open     | per-face | —    |
-| `acoustic_scattering_maze`     | per-face | per-face | —    |
-
-Today the FFT-conv operators expose only two **global** modes
-(`fft_padding="zero"` or `"circular"`) selected per module. There is no way
-to say "periodic on x, zero-padded on y" for one and the same conv.
-
-### Decisions (locked in)
-
-1. **API** — extend `fft_padding` to accept either a **single mode
-   string** (applies to every axis) or a **list of mode strings** (one
-   per spatial axis):
-
-   ```python
-   fft_padding: str | Sequence[str] = "zero"
-   # "zero"                              -> all axes zero-padded.
-   # "circular"                          -> all axes periodic.
-   # ["circular", "zero"]                -> 2D, x periodic + y zero-padded.
-   # ["zero", "circular", "zero"]        -> 3D, etc.
-   # ("circular", "zero")                -> tuple form is equivalent.
-   ```
-
-   Internally everything normalises to a tuple `periodic: tuple[bool, ...]`
-   of length `data_dim`. Three inputs are deliberately rejected with an
-   error that redirects to the canonical form:
-
-   - **Booleans** (`(True, False)`, `True`): the per-axis intent is not
-     obvious from the boolean values.
-   - **Comma-separated strings** (`"circular, zero"`): redundant with the
-     list form and gives two ways to say the same thing; we keep one
-     canonical per-axis form.
-
-1. **WALL vs OPEN** — both treated as **zero-padded linear** at the
-   conv level. Physical distinctions are handled elsewhere (data
-   normalisation, loss, etc.). Per-face BC (different BC on opposite faces
-   of the same axis) is **out of scope** for v1.
-
-1. **Kernel size per axis** — auto-derived from the per-axis BC, **not** a
-   new knob:
-
-   | axis BC      | grid_lens per axis (in `CKConvND.forward`)   | SIREN kernel size on that axis |
-   | ------------ | -------------------------------------------- | ------------------------------ |
-   | periodic     | `(s+1)//2` (≡ today's `grid_type="single"`)  | `≈ s`                          |
-   | non-periodic | `s`         (≡ today's `grid_type="double"`) | `≈ 2s − 1`                     |
-
-   When `fft_padding` is a tuple, the legacy `grid_type` argument must be
-   `None` (or omitted) — raise on conflict, no silent overrides.
-
-1. **First-PR scope** — **ops + tests only.** Module wiring
-   (`CKConvND`/`CKConvMultiheadND`), Well config updates, fp16, multihead,
-   and the `subq_ops` CUDA path are explicitly deferred (see §4).
-
-1. **`subq_ops` CUDA kernel** — left at zero-only. Any
-   `fft_backend="subq_ops"` + mixed BC will raise in the future
-   `CKConvND` wiring PR.
-
-______________________________________________________________________
-
-## 2. Algorithm — per-axis recipe
-
-The mixed N-D FFT convolution applies, **independently per spatial axis**:
-
-| axis is      | FFT length `F_d`               | post-IFFT crop range             | phase ramp shift on that axis |
-| ------------ | ------------------------------ | -------------------------------- | ----------------------------- |
-| periodic     | `N_d` (no padding)             | `0 : N_d` (no crop)              | `−(K_d − 1) // 2`             |
-| non-periodic | `min(N_d + (K_d+1)//2, 2·N_d)` | `K_d//2 : K_d//2 + N_d` (center) | `0` (no shift)                |
-
-The whole conv is still **one** `rfftn` / `irfftn` over all spatial dims;
-the per-axis recipe just feeds different per-axis `F_d` values to `s=` and
-different per-axis slices to the post-IFFT crop. Phase ramps are the
-product of per-axis 1-D ramps (length-1 broadcast on non-periodic axes).
-
-Edge behaviour required by the tests:
-
-- All `periodic == False` → bit-identical to existing `fftconv*` linear op.
-- All `periodic == True`  → bit-identical to existing `circular_fftconv*` op.
-
-______________________________________________________________________
-
-## 3. v1 deliverables (this PR)
-
-### Code
-
-- [x] `nvsubquadratic/ops/mixed_fftconv.py` (fp32):
-  - Self-contained per-axis 1-D phase-ramp LRU cache; N-D ramp built on
-    demand by broadcasted multiplication so non-periodic axes contribute
-    nothing (skipped, not just length-1).
-  - 1D / 2D / 3D BHL variants.
-  - BHL `_w_reshape` wrappers (BLH inputs).
-  - Channel-chunked variants.
-  - Automatic dispatch to the existing linear / circular ops in the
-    all-False / all-True cases (no perf cost for legacy usage).
-- [x] No `nvsubquadratic/ops/__init__.py` exists — callers import from
-  submodules directly (matches existing convention).
-
-> **Note:** an earlier draft of this plan called for extracting
-> `_PhaseRampCache1D/2D/3D` from `circular_fftconv.py` into a shared
-> `_phase_ramp.py`. We **dropped** that refactor for v1 because the existing
-> caches hard-code `FFT_shape == input_shape` per axis, which is only true
-> for the all-circular case. There is nothing to share without first
-> generalising the API. We may revisit this as a unification refactor
-> (see §5 Q4).
-
-### Tests
-
-- [x] `tests/ops/test_mixed_fftconv.py` — 76 tests, all passing on H100:
-  - Reference comparison against
-    `F.pad(x, mode="circular"|"constant")` + `F.conv{1,2,3}d(padding=0)`
-    for every per-axis combo across 1D / 2D / 3D, including the K==N
-    edge case and even-K kernels (asymmetric "same" padding).
-  - Sanity: all-False matches existing `fftconv*`; all-True matches
-    existing `circular_fftconv*`.
-  - BHL ↔ BLH wrapper equivalence.
-  - Chunked vs non-chunked equivalence.
-  - Backward / gradient equivalence vs the spatial reference (1D, 2D, 3D).
-  - `use_phase_shift=False` (roll on periodic axes only) matches
-    `use_phase_shift=True` for every combo.
-  - Shortcut term equivalence and dtype preservation (fp32, bf16).
-  - Validation errors: wrong `periodic` length, oversized kernel,
-    mismatched shortcut dtype.
-- [x] Re-ran existing FFT-conv suites — **101 passed**, no regressions:
-  - `tests/ops/test_fftconv.py`
-  - `tests/ops/test_circular_fftconv.py`
-  - `tests/ops/test_fftconv_chunked.py`
-
-### Docs
-
-- [x] Updated `docs/ops/README.md` "File map" with the new
-  `mixed_fftconv.py` row.
-
-______________________________________________________________________
-
-## 4. Deferred — follow-up work
-
-Each item below is intentionally **not** part of this PR. They are
-listed so we don't lose track.
-
-### 4.1 Module wiring — `CKConvND` ✅ DONE (2026-05-20)
-
-- [x] `CKConvND` (`nvsubquadratic/modules/ckconv_nd.py`):
-  - Accepts `fft_padding: str | Sequence[str]` in two forms: a single
-    mode string (`"zero"` / `"circular"`) that applies to every axis, or
-    a list of mode strings (e.g. `["circular", "zero"]`) — one per axis.
-  - Resolves to a normalised per-axis `periodic` tuple via
-    `_resolve_periodic` (length checked against `data_dim`).
-  - When `fft_padding` is a per-axis list, **requires** `grid_type=None`
-    and raises `ValueError` otherwise. When it's a single mode string,
-    `grid_type` is required as before.
-  - Boolean inputs (e.g. `(True, False)`) and comma-separated strings
-    (`"circular, zero"`) are explicitly rejected with errors that
-    redirect to the list form.
-  - Per-axis `grid_lens` and per-axis `L_cache` halving auto-derived in
-    tuple mode (halve only on periodic axes); helper
-    `_grid_is_single_per_axis(grid_type, periodic)` is the single source
-    of truth used by both `__init__` (L_cache) and `forward`/`flop_count`.
-  - Dispatch: tuple mode routes through `MIXED_FFT_FUNCTIONS[_CHUNKED]`
-    (wrapped by `_wrap_mixed_op` to bind `periodic`). All-False / all-True
-    tuples internally fall back to the legacy linear / circular ops
-    bit-identically (verified by tests). String mode keeps the legacy
-    `FFT_FUNCTIONS` tables unchanged.
-  - Validation:
-    - `is_causal=True` + any periodic axis → `ValueError`.
-    - `fft_backend="subq_ops"` + tuple `fft_padding` → `ValueError`.
-    - `use_fp16_fft=True` + tuple `fft_padding` → `NotImplementedError`
-      (planned for v2; see §4.2).
-    - `use_chunked_fftconv` allowed with tuple `fft_padding` for **all**
-      per-axis combos (including all-True — new capability vs the legacy
-      string-mode where circular + chunked was an error).
-  - `flop_count` uses per-axis padded sizes: `s` on periodic axes,
-    `min(s + (k+1)//2, 2*s)` on non-periodic axes.
-- [x] `mixed_fftconv.py` op: `K <= N` assertion relaxed on non-periodic
-  axes to `K <= 2*N` to match the "double-grid" SIREN kernel size
-  (`2N - 1`) that `CKConvND` produces on non-periodic axes.
-- [x] `mixed_fftconv.py`: added `*_w_reshape_chunked` BLH wrappers so the
-  module dispatch table is symmetric with the legacy chunked path.
-- [x] Module-level tests in `tests/modules/test_ckconv_nd_mixed_bc.py`:
-  resolver / helper unit tests, validation errors, per-axis kernel
-  shape, tuple-vs-string bit-identical equivalence, mixed-mode
-  forward correctness (matches the underlying op called directly),
-  BHL/BLH layout equivalence, chunked-vs-non-chunked, and FLOP
-  accounting (mixed sits between all-zero and all-circular).
-
-### 4.2 fp16 variant (deferred)
-
-### 4.2 fp16 variant
-
-- [ ] `nvsubquadratic/ops/mixed_fftconv_fp16.py`:
-  - Per-axis cuFFT-fp16 constraints: pad-up to power-of-2 on linear axes;
-    require input dim power-of-2 on periodic axes (fallback to fp32 with a
-    warning otherwise).
-  - Reuse centering / DC-correction logic from `circular_fftconv_fp16`
-    on the periodic axes only.
-- [ ] `tests/ops/test_mixed_fftconv_fp16.py`.
-- [ ] `CKConvND` fp16 dispatch update.
-
-### 4.3 2D multi-head variant
-
-- [ ] `fftconv2d_multihead_mixed_bhl` (and `_bhi`) in
-  `nvsubquadratic/ops/fftconv_multihead.py`.
-- [ ] `CKConvMultiheadND` wiring + tests (same shape of changes as
-  `CKConvND` above; 2D only).
-
-### 4.4 Well experiment configs
-
-- [ ] Add Hyena variants with mixed BC for the datasets that need them.
-  Start with the ones we actually run:
-  - `rayleigh_benard`              — periodic on x → `(True, False)`
-  - `viscoelastic_instability`     — periodic on x → `(True, False)`
-  - `turbulent_radiative_layer_2D` — periodic on x → `(True, False)`
-  - `turbulent_radiative_layer_3D` — periodic on x,y → `(True, True, False)`
-  - `rayleigh_taylor_instability`  — periodic on x,y → `(True, True, False)`
-- [ ] Decide whether to fix `examples/well/v1/supernova_explosion_64/cfg_hyena.py`
-  which uses `FFT_PADDING="circular"` but the dataset is all-OPEN. v2 is
-  already corrected.
-- [ ] Longer-term: read Well HDF5 `boundary_conditions` from the
-  datamodule and auto-derive `periodic_axes` instead of hard-coding it
-  per config.
-
-### 4.5 Per-face BCs
-
-- [ ] Investigate whether `acoustic_scattering_maze` and
-  `helmholtz_staircase` benefit from a *per-face* BC treatment
-  (different BC on opposite faces of the same axis).
-- [ ] If yes: design a follow-up that goes beyond per-axis circular/linear.
-  v1 maps any non-periodic axis to symmetric zero-pad.
-
-### 4.6 Custom CUDA path (`subq_ops`)
-
-- [ ] If/when we want mixed-BC to use the
-  `subquadratic_ops_torch.fft_conv2d` fast path, the upstream kernel
-  must grow per-axis BC support. v1 leaves the kernel zero-only and
-  raises in `CKConvND`.
-
-### 4.7 SIREN kernel generator anisotropy
-
-- [ ] Re-verify that the SIREN kernel module
-  (`nvsubquadratic/modules/kernels_nd.py`) actually handles
-  anisotropic `grid_lens` (e.g. `(64, 128)`) cleanly end-to-end —
-  including positional embedding, masks, monitors, and FLOP accounting.
-  Used today for the all-isotropic case; we need it for the per-axis
-  grid in the mixed path.
-
-______________________________________________________________________
-
-## 5. Tracked questions & "revisit" items
-
-These are not bugs we plan to fix in this PR, just things we noticed
-along the way that may want attention later.
-
-### Q1. Silent mutation of user-provided `L_cache` in `CKConvND.__init__`
-
-`ckconv_nd.py` (~L266–289) does `copy.deepcopy(kernel_cfg)` and then
-silently halves `L_cache` when `grid_type=="single"` so the SIREN
-positional grid spans `[-1, 1]` over the actual kernel size. The
-behaviour is *intentional* (per the in-code comment) but the **silent
-mutation of a user-provided config** is mildly surprising. Cleaner
-pattern would be to compute the effective `L_cache` at construction
-time without round-tripping through the config object. Not a v1 bug —
-revisit in a separate refactor.
-
-For the mixed path, the same adjustment must become **per-axis** (halve
-only on periodic axes); design that in the module-wiring PR (§4.1).
-
-### Q2. `subq_ops` 2D linear kernel — make it BC-aware?
-
-Out of scope for v1, but worth a conversation with the kernel authors
-before we commit to a divergent fast path that only supports zero-pad.
-
-### Q3. Auto-wiring from Well HDF5 boundary_conditions
-
-`experiments/datamodules/pde/well.py` can return per-sample BC metadata
-from HDF5, but no model code consumes it. Long-term it would be cleaner
-to derive the `periodic_axes` tuple from the dataset rather than
-hard-coding in each config.
-
-### Q4. Unify the phase-ramp cache across `circular_fftconv` and `mixed_fftconv`
-
-After v1, the codebase will have two parallel phase-ramp caches:
-
-- The original `_PhaseRampCache1D/2D/3D` in `circular_fftconv.py`, which
-  hard-codes `FFT_shape == input_shape` per axis (only valid for the
-  all-circular case).
-- A new general N-D cache local to `mixed_fftconv.py`, which also handles
-  per-axis padded `F_d` and zero shifts on linear axes.
-
-These can be unified into a single shared helper once we are happy with
-the mixed op's API. Doing so is a pure refactor (no behaviour change for
-either op) and should be its own PR.
-
-______________________________________________________________________
-
-## 6. Changelog
-
-- **2026-05-20** — Plan written, feature branch created.
-- **2026-05-20** — v1 ops landed: `mixed_fftconv.py` (fp32, 1D/2D/3D,
-  BHL + BLH wrappers + channel-chunked), 76-test suite passing,
-  no regressions in existing FFT-conv tests.
-- **2026-05-20** — Tolerance tightening + analytical-truth tests
-  (impulse response, DC response) added; 102 op tests passing.
-- **2026-05-20** — `CKConvND` integration landed: per-axis
-  `fft_padding: Sequence[bool]` API, auto-derived per-axis grid +
-  `L_cache` halving, unified dispatch via `MIXED_FFT_FUNCTIONS*`,
-  module-level test suite (`tests/modules/test_ckconv_nd_mixed_bc.py`).
-  Full regression sweep on `tests/ops + tests/modules` → 802 passed.
-- **2026-05-21** — Public API revised on PR review: `fft_padding` now
-  accepts mode-name strings only. Two forms: a single mode string
-  (`"zero"` / `"circular"`) that applies to every axis, or a list of mode
-  strings (e.g. `["circular", "zero"]`) — one per axis. Comma-separated
-  strings and bool tuples are both rejected with redirecting errors.
-  Rationale: `(True, False)` did not convey which axis was periodic, and
-  having both a comma-string form and a list form was two ways to say
-  the same thing. The list form reads identically in Python and OmegaConf
-  / YAML overrides. `rayleigh_benard` config updated to the list form.
diff --git a/docs/ops/README.md b/docs/ops/README.md
index e6e06423..63fc706a 100644
--- a/docs/ops/README.md
+++ b/docs/ops/README.md
@@ -1,6 +1,6 @@
 # `nvsubquadratic.ops` — FFT convolution primitives
 
-This folder contains the **lowest-level building blocks** of the library: FFT-based convolution operators that turn an `O(N · K)` spatial convolution into an `O(N log N)` frequency-domain product. They are the workhorses behind every subquadratic mixer in the library (Hyena, CKConv, multi-head variants), and are kept here as plain functions — no `nn.Module` state, no learned parameters — so that higher-level modules can compose them freely.
+This folder contains the **lowest-level building blocks** of the library: FFT-based convolution operators that turn an `O(N · K)` spatial convolution into an `O(N log N)` frequency-domain product. They are the workhorses behind every subquadratic mixer in the library (Hyena, CKConv), and are kept here as plain functions — no `nn.Module` state, no learned parameters — so that higher-level modules can compose them freely.
 
 If you are reading the paper alongside this codebase, this is the file to start with.
 
@@ -35,19 +35,16 @@ ______________________________________________________________________
 
 ## File map
 
-| File                                                                          | Precision | Conv type         | Channel mixing                       | When you'd reach for it                                                                                                                                                                                                                           |
-| ----------------------------------------------------------------------------- | --------- | ----------------- | ------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| [fftconv.py](../../nvsubquadratic/ops/fftconv.py)                             | fp32      | linear            | depthwise                            | The default. 1D/2D/3D, causal & non-causal.                                                                                                                                                                                                       |
-| [circular_fftconv.py](../../nvsubquadratic/ops/circular_fftconv.py)           | fp32      | circular          | depthwise                            | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful.                                                                                                                                                               |
-| [mixed_fftconv.py](../../nvsubquadratic/ops/mixed_fftconv.py)                 | fp32      | per-axis BC       | depthwise                            | **Mixed boundaries** — periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. |
-| [fftconv_fp16.py](../../nvsubquadratic/ops/fftconv_fp16.py)                   | fp16      | linear            | depthwise                            | Memory/throughput savings on power-of-2 spatial dims. Drop-in for `fftconv.py`.                                                                                                                                                                   |
-| [circular_fftconv_fp16.py](../../nvsubquadratic/ops/circular_fftconv_fp16.py) | fp16      | circular          | depthwise                            | Same as above for the circular case. Uses **dual mean-centering** for fp16 stability — see [FP16_FFTCONV_DERIVATION.md](FP16_FFTCONV_DERIVATION.md).                                                                                              |
-| [fftconv_chunked.py](../../nvsubquadratic/ops/fftconv_chunked.py)             | fp32      | linear            | depthwise                            | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently.                                                                                                                                  |
-| [fftconv_multihead.py](../../nvsubquadratic/ops/fftconv_multihead.py)         | fp32      | linear & circular | dense within head, optional low-rank | Multi-head FFT conv — channel mixing inside each head, in the spirit of multi-head attention.                                                                                                                                                     |
-| [fftconv_custom.py](../../nvsubquadratic/ops/fftconv_custom.py)               | fp32      | linear            | depthwise                            | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`.                                                                                 |
-| [causal_conv1d_custom.py](../../nvsubquadratic/ops/causal_conv1d_custom.py)   | fp32      | direct causal     | depthwise                            | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block.                                             |
+| File                                                                        | Precision | Conv type     | Channel mixing | When you'd reach for it                                                                                                                                                                                                                           |
+| --------------------------------------------------------------------------- | --------- | ------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| [fftconv.py](../../nvsubquadratic/ops/fftconv.py)                           | fp32      | linear        | depthwise      | The default. 1D/2D/3D, causal & non-causal.                                                                                                                                                                                                       |
+| [circular_fftconv.py](../../nvsubquadratic/ops/circular_fftconv.py)         | fp32      | circular      | depthwise      | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful.                                                                                                                                                               |
+| [mixed_fftconv.py](../../nvsubquadratic/ops/mixed_fftconv.py)               | fp32      | per-axis BC   | depthwise      | **Mixed boundaries** — periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. |
+| [fftconv_chunked.py](../../nvsubquadratic/ops/fftconv_chunked.py)           | fp32      | linear        | depthwise      | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently.                                                                                                                                  |
+| [fftconv_custom.py](../../nvsubquadratic/ops/fftconv_custom.py)             | fp32      | linear        | depthwise      | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`.                                                                                 |
+| [causal_conv1d_custom.py](../../nvsubquadratic/ops/causal_conv1d_custom.py) | fp32      | direct causal | depthwise      | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block.                                             |
 
-`FP16_FFTCONV_DERIVATION.md` contains the full derivation of the numerically stable fp16 circular conv (dual mean-centering + inclusion-exclusion geometric correction). Read it if you are touching the fp16 path or want to understand the math behind those `T1, T2, T3, T4` terms in the code.
+`mixed_boundary_conditions.md` describes the per-axis boundary-condition support (periodic on some spatial axes, zero-padded on others) used by the Well PDE datasets.
 
 ______________________________________________________________________
 
@@ -56,21 +53,21 @@ ______________________________________________________________________
 Every function name encodes its contract:
 
 ```
-[causal_] fftconv {1d|2d|3d} _ {fp32|fp16} _ {bhl|blh} [_w_reshape] [_chunked]
+[causal_] fftconv {1d|2d|3d} _ fp32 _ {bhl|blh} [_w_reshape] [_chunked]
 ```
 
 | Part               | Meaning                                                                                                                                        |
 | ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- |
 | `causal_`          | Output at position `n` only sees inputs at positions `≤ n`. 1D only.                                                                           |
 | `1d` / `2d` / `3d` | Spatial rank.                                                                                                                                  |
-| `fp32` / `fp16`    | Internal compute precision. The output dtype always matches `x.dtype` regardless.                                                              |
+| `fp32`             | Internal compute precision. The output dtype always matches `x.dtype` regardless.                                                              |
 | `bhl` / `blh`      | Memory layout. `bhl` = channels-first (`[B, H, *spatial]`). `blh` = channels-last (`[B, *spatial, H]`).                                        |
 | `_w_reshape`       | Wrapper that accepts BLH input, internally reshapes to BHL (faster), and reshapes back. The recommended entry point for channels-last callers. |
 | `_chunked`         | Processes channels in groups to reduce peak GPU memory.                                                                                        |
 
 So `causal_fftconv1d_fp32_bhl_w_reshape` is: causal 1D FFT conv, fp32 internal, accepts channels-last input, internally uses the channels-first kernel.
 
-The CUDA-accelerated wrappers in `fftconv_custom.py` drop the `_fp32_` / `_fp16_` token because the underlying kernel manages its own precision internally — so the same name in `fftconv_custom` is `causal_fftconv1d_bhl_w_reshape`. The direct-conv wrappers in `causal_conv1d_custom.py` (`causal_conv1d`, `b2b_causal_conv1d`) do not follow this scheme because they are thin pass-throughs to the upstream API; see their docstrings for shapes.
+The CUDA-accelerated wrappers in `fftconv_custom.py` drop the `_fp32_` token because the underlying kernel manages its own precision internally — so the same name in `fftconv_custom` is `causal_fftconv1d_bhl_w_reshape`. The direct-conv wrappers in `causal_conv1d_custom.py` (`causal_conv1d`, `b2b_causal_conv1d`) do not follow this scheme because they are thin pass-throughs to the upstream API; see their docstrings for shapes.
 
 ______________________________________________________________________
 
@@ -112,21 +109,9 @@ ______________________________________________________________________
    - Channels-first (`[B, H, …]`) → use `_bhl` directly.
    - Channels-last (`[B, …, H]`) → use `_bhl_w_reshape`. Benchmarks show this is faster than a true `_blh` op because the FFT runs on contiguous spatial axes.
 
-1. **What's my precision budget?**
-
-   - fp32 is the default, always correct.
-   - For aggressive memory/throughput savings on **power-of-2 spatial dims**, use the `*_fp16_*` variant. The fp16 ops use `norm="ortho"` and (for circular) dual mean-centering to stay within fp16 dynamic range — see [FP16_FFTCONV_DERIVATION.md](FP16_FFTCONV_DERIVATION.md).
-   - If your spatial dims aren't powers of two, stay in fp32 (cuFFT requires power-of-2 for fp16 transforms).
-
 1. **Am I OOMing?**
 
    - Try `fftconv_chunked` — splits the channel dim into groups to cap peak memory. Default chunk size 128 gives ~26% memory savings for ~11% overhead.
-   - Or combine: `fftconv_fp16.py` already provides `_chunked` variants that stack both savings.
-
-1. **Do I need cross-channel mixing inside the conv itself?**
-
-   - Yes → `fftconv_multihead.py`. Splits the channel dim into heads and applies a *dense* (head_dim × head_dim) mixing in the frequency domain. Use the `_lowrank_*` factorisation when `head_dim` is large.
-   - No (the default) → depthwise variants — separate pointwise/MLP layer handles channel mixing.
 
 1. **Is there a fused CUDA kernel for my shape?**
 
@@ -138,8 +123,8 @@ ______________________________________________________________________
 
 ## Numerical notes
 
-- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`** — no need for a manual cast on the caller side.
-- The fp32 ops are correct for any input range. The fp16 ops impose two constraints: spatial dims must be powers of two (cuFFT), and the dynamic range is handled by mean-centering both `x` and `k` (see derivation doc).
+- All operators **accept any input dtype** but cast to fp32 before the FFT. The output is returned in the **original dtype of `x`** — no need for a manual cast on the caller side.
+- The fp32 ops are correct for any input range.
 - The non-causal linear ops match a standard `torch.nn.ConvNd(padding='same')` up to floating-point rounding. The circular ops match `torch.nn.functional.conv*d` after a circular pad. Both are exercised in `tests/`.
 
 ______________________________________________________________________
@@ -148,13 +133,12 @@ ______________________________________________________________________
 
 - **[`nvsubquadratic/modules/kernels_nd.py`](../../nvsubquadratic/modules/kernels_nd.py)** — learned kernel parametrisations that produce the kernels these ops consume.
 - **[`nvsubquadratic/modules/hyena_nd.py`](../../nvsubquadratic/modules/hyena_nd.py)** — the Hyena operator, the main consumer of these ops.
-- **[`nvsubquadratic/modules/ckconv_nd.py`](../../nvsubquadratic/modules/ckconv_nd.py)** / **[`ckconv_multihead_nd.py`](../../nvsubquadratic/modules/ckconv_multihead_nd.py)** — CKConv variants that compose these primitives.
+- **[`nvsubquadratic/modules/ckconv_nd.py`](../../nvsubquadratic/modules/ckconv_nd.py)** — the CKConv operator that composes these primitives.
 
 ```{toctree}
 ---
 maxdepth: 1
 caption: Further reading
 ---
-FP16_FFTCONV_DERIVATION
-MIXED_BC_PLAN
+mixed_boundary_conditions
 ```
diff --git a/docs/ops/mixed_boundary_conditions.md b/docs/ops/mixed_boundary_conditions.md
new file mode 100644
index 00000000..93653a6e
--- /dev/null
+++ b/docs/ops/mixed_boundary_conditions.md
@@ -0,0 +1,124 @@
+# Mixed Boundary-Condition FFT Convolution
+
+The FFT-convolution operators support **per-axis boundary conditions**: a
+single convolution can be periodic on some spatial axes and zero-padded
+(linear) on others. This is exposed through
+[`mixed_fftconv.py`](../../nvsubquadratic/ops/mixed_fftconv.py) at the op
+level and through the `fft_padding` argument of
+{class}`~nvsubquadratic.modules.ckconv_nd.CKConvND` at the module level.
+
+______________________________________________________________________
+
+## Motivation
+
+Several PDE datasets have boundaries that are **periodic on some axes and
+non-periodic (wall or open) on others**:
+
+| Dataset                        | x        | y        | z    |
+| ------------------------------ | -------- | -------- | ---- |
+| `rayleigh_benard`              | periodic | wall     | —    |
+| `viscoelastic_instability`     | periodic | wall     | —    |
+| `turbulent_radiative_layer_2D` | periodic | open     | —    |
+| `turbulent_radiative_layer_3D` | periodic | periodic | open |
+| `rayleigh_taylor_instability`  | periodic | periodic | wall |
+
+The two global modes (`fft_padding="zero"` or `"circular"`) cannot express
+"periodic on x, zero-padded on y" for one and the same convolution. The
+mixed path closes that gap.
+
+Wall and open boundaries are both treated as **zero-padded linear** at the
+convolution level; physical distinctions (if any) are handled elsewhere
+(data normalisation, loss). Per-face boundary conditions — a different BC
+on opposite faces of the same axis — are not supported (see
+[Limitations](#limitations)).
+
+______________________________________________________________________
+
+## API
+
+`fft_padding` accepts either a **single mode string** (applies to every
+axis) or a **list of mode strings** (one per spatial axis):
+
+```python
+fft_padding: str | Sequence[str] = "zero"
+# "zero"                       -> all axes zero-padded (linear "same" conv).
+# "circular"                   -> all axes periodic (wrap-around conv).
+# ["circular", "zero"]         -> 2D: x periodic, y zero-padded.
+# ["zero", "circular", "zero"] -> 3D mixed.
+```
+
+The list form reads identically in Python and in OmegaConf / YAML config
+overrides. Internally everything normalises to a per-axis boolean tuple
+`periodic: tuple[bool, ...]` of length `data_dim`. Two inputs are
+deliberately rejected with an error that redirects to the canonical form:
+
+- **Booleans** (`(True, False)`): the per-axis intent is not obvious from
+  the boolean values.
+- **Comma-separated strings** (`"circular, zero"`): redundant with the list
+  form.
+
+### Kernel size per axis
+
+The kernel grid size is auto-derived from the per-axis boundary condition —
+it is **not** a separate knob. When `fft_padding` is a list, the legacy
+`grid_type` argument must be `None` (a conflict raises rather than silently
+overriding):
+
+| axis BC      | grid length per axis (`CKConvND.forward`) | SIREN kernel size on that axis |
+| ------------ | ----------------------------------------- | ------------------------------ |
+| periodic     | `(s+1)//2` (≡ `grid_type="single"`)       | `≈ s`                          |
+| non-periodic | `s` (≡ `grid_type="double"`)              | `≈ 2s − 1`                     |
+
+______________________________________________________________________
+
+## Algorithm
+
+The mixed N-D FFT convolution is computed in **one** `rfftn` / `irfftn`
+call over all spatial dims. The per-axis variation is encoded entirely in
+the transform arguments:
+
+| axis is      | FFT length `F_d`               | post-IFFT crop range             | phase ramp shift on that axis |
+| ------------ | ------------------------------ | -------------------------------- | ----------------------------- |
+| periodic     | `N_d` (no padding)             | `0 : N_d` (no crop)              | `−(K_d − 1) // 2`             |
+| non-periodic | `min(N_d + (K_d+1)//2, 2·N_d)` | `K_d//2 : K_d//2 + N_d` (center) | `0` (no shift)                |
+
+Phase ramps are the product of per-axis 1-D ramps (length-1 broadcast on
+non-periodic axes, so they contribute nothing there).
+
+Two edge behaviours are guaranteed (and exercised in the tests):
+
+- All axes non-periodic → **bit-identical** to the linear
+  [`fftconv`](../../nvsubquadratic/ops/fftconv.py) op.
+- All axes periodic → **bit-identical** to the
+  [`circular_fftconv`](../../nvsubquadratic/ops/circular_fftconv.py) op.
+
+In those uniform corners the mixed op dispatches internally to the legacy
+linear / circular ops, so there is no performance cost for non-mixed usage.
+
+______________________________________________________________________
+
+## What's available
+
+- **Op level** ([`mixed_fftconv.py`](../../nvsubquadratic/ops/mixed_fftconv.py)):
+  fp32 1D / 2D / 3D, both BHL and BLH (`_w_reshape`) layouts, and
+  channel-chunked variants. Shortcut term and dtype preservation match the
+  other ops.
+- **Module level** ({class}`~nvsubquadratic.modules.ckconv_nd.CKConvND`):
+  pass `fft_padding` as a per-axis list. `use_chunked_fftconv` is supported
+  for every per-axis combination (including all-periodic, which the legacy
+  string-mode rejected). `flop_count` uses per-axis padded sizes.
+
+______________________________________________________________________
+
+## Limitations
+
+- **Per-face boundary conditions** are out of scope: any non-periodic axis
+  maps to symmetric zero-padding. Datasets with different BCs on opposite
+  faces of the same axis (e.g. `helmholtz_staircase`,
+  `acoustic_scattering_maze`) are approximated by the nearest per-axis BC.
+- **Custom CUDA path** (`fft_backend="subq_ops"`) supports zero-padding
+  only; combining it with a per-axis `fft_padding` raises. Use
+  `fft_backend="torch_fft"` (the default) for mixed boundaries.
+- **Auto-wiring** — the periodic axes are specified per config. The Well
+  datamodule can read `boundary_conditions` from the HDF5 metadata, but
+  model code does not yet consume it to derive `periodic` automatically.
diff --git a/docs/repository_overview.md b/docs/repository_overview.md
index a292001b..fc3316ea 100644
--- a/docs/repository_overview.md
+++ b/docs/repository_overview.md
@@ -20,12 +20,9 @@ nvSubquadratic/
 │   ├── callbacks/             FiLM monitor, image-grid viz, EMA, walltime checkpointer, …
 │   └── utils/                 cli + checkpointing helpers
 ├── examples/                LazyConfig recipes that feed experiments.run
-│   ├── mnist_classification/
 │   ├── imagenet_classification/
-│   ├── imagenet_diffusion/
 │   ├── vit5_imagenet/         ViT-5 baseline suite (v1–v5)
 │   ├── spatial_recall_{1,2,3}d/ and spatial_recall_v2/
-│   ├── ucf101_classification/
 │   ├── well/                  The Well PDE benchmark suite
 │   └── overview_tracker.md    active experimental roadmap
 ├── benchmarks/              performance measurement (the canonical home)
@@ -43,8 +40,7 @@ nvSubquadratic/
 │   └── vit5_imagenet_dataloader_profiling/
 ├── scripts/                 utilities (data prep, sanity, SLURM, viz)
 │   ├── slurm/                 SLURM submit scripts (portable wrapper + per-experiment)
-│   ├── data/                  data prep (ImageNet folder extraction, FID stats, …)
-│   ├── evaluation/            eval helpers (FID, CMMD)
+│   ├── data/                  data prep (ImageNet folder extraction, normalization stats, …)
 │   ├── visualization/         kernel viewers + throughput plot
 │   └── check_gpu_availability.py, license_check.py, …
 ├── tests/                   correctness tests
@@ -131,8 +127,6 @@ nvsubquadratic/
 │   ├── qk_norm.py                QK normalization (apply + L2Norm)
 │   ├── rope.py                   rotary position embedding (1D / 2D / 3D)
 │   └── quack_utils.py            QuACK capability probe
-├── metrics/
-│   └── cleanfid.py               FID via cleanfid
 └── testing/
     └── utils.py                  compute_relative_error
 ```
@@ -167,8 +161,7 @@ docs page.
 Indexed at {doc}`reports`.
 
 **`scripts/`** — Utility / glue scripts.  SLURM submit drivers
-(`scripts/slurm/`), data prep (`scripts/data/`), evaluation helpers
-(`scripts/evaluation/`), kernel viewers (`scripts/visualization/`),
+(`scripts/slurm/`), data prep (`scripts/data/`), kernel viewers (`scripts/visualization/`),
 and standalone sanity scripts.  No benchmarks live here — those moved
 to `benchmarks/`.
 
diff --git a/examples/imagenet_diffusion/README.md b/examples/imagenet_diffusion/README.md
deleted file mode 100644
index bb790c08..00000000
--- a/examples/imagenet_diffusion/README.md
+++ /dev/null
@@ -1,92 +0,0 @@
-# ImageNet Diffusion Experiments
-
-Class-conditional ImageNet generation using CCNN (Hyena + SIREN kernels) with the JiT flow-matching framework.
-
-## Experiment matrix
-
-| Config                        | Resolution | Model size | hidden | blocks | patch | omega_0 | Gated | FiLM | Notes                         |
-| ----------------------------- | ---------- | ---------- | ------ | ------ | ----- | ------- | ----- | ---- | ----------------------------- |
-| `ccnn_jit_baseline`           | 64x64      | JiT-B      | 768    | 12     | 4     | 10.0    | no    | no   | Existing baseline             |
-| `jit_baseline`                | 64x64      | JiT-B      | 768    | 12     | 4     | —       | no    | no   | Reference JiT transformer     |
-| `ccnn_jit_128`                | 128x128    | JiT-B      | 768    | 12     | 8     | 10.0    | no    | no   | Higher res, same token count  |
-| `ccnn_jit_256`                | 256x256    | JiT-B      | 768    | 12     | 16    | 10.0    | no    | no   | Higher res, same token count  |
-| `ccnn_jit_128_low_omega`      | 128x128    | JiT-B      | 768    | 12     | 8     | **3.0** | no    | no   | Reduced omega_0               |
-| `ccnn_jit_256_low_omega`      | 256x256    | JiT-B      | 768    | 12     | 16    | **3.0** | no    | no   | Reduced omega_0               |
-| `ccnn_jit_L_256`              | 256x256    | JiT-L      | 1024   | 24     | 16    | 10.0    | no    | no   | Larger model                  |
-| `ccnn_jit_128_gated_film_ema` | 128x128    | JiT-B      | 768    | 12     | 8     | 10.0    | yes   | yes  | Gated Hyena + FiLM (untested) |
-
-## Design decisions
-
-### Resolution scaling
-
-All configs keep the token count at 16x16 = 256 tokens to match the 64px baseline:
-
-- 64px / patch 4 = 16x16
-- 128px / patch 8 = 16x16
-- 256px / patch 16 = 16x16
-
-This means the patch embedding carries more information at higher resolutions (larger patches) while the sequence model operates on the same length. Batch sizes and gradient accumulation are adjusted to maintain ~1024 effective batch size.
-
-### Low omega_0 experiments
-
-The `_low_omega` variants reduce `omega_0` from 10.0 to 3.0 on the first SIREN layer. This controls the initial frequency range of the implicit convolution kernel — a lower value biases the kernel toward lower spatial frequencies, which should help suppress the high-frequency artifacts observed in prior runs. `hidden_omega_0` remains at 1.0.
-
-### JiT-L model size
-
-`ccnn_jit_L_256` matches the JiT-L transformer architecture:
-
-- 1024 hidden dim (vs 768 for JiT-B)
-- 24 blocks (vs 12)
-- Kernel MLP scaled to 64 hidden / 64 embedding (vs 32/32)
-- Lower learning rate (1e-4 vs 2e-4)
-- Higher EMA decay (0.9999 vs 0.9998)
-
-### Gated FiLM experiment
-
-`ccnn_jit_128_gated_film_ema` adapts the classification architecture from `vit5_small_pretrain_hyena_cls_row_gated_film_ema` to diffusion:
-
-- **Gated Hyena**: dual nonlinearity (SiLU + Sigmoid gates)
-- **FiLM-conditioned SIREN**: the timestep condition vector (collapsed to `[B, hidden_dim]` by AdaLN) is forwarded through the mixer chain to modulate SIREN hidden layers via learned (gamma, beta) pairs
-- **EMA**: decay 0.9998
-
-In the classification setup, FiLM conditioning comes from register tokens pooled via `RegisterPooling`. In diffusion there are no registers — instead the AdaLN timestep condition serves as the FiLM input. A one-line change in `AdaLNZeroResidualBlock.forward` (passing `conditioning=cond` to the sequence mixer) wires this up. This change is safe for all non-FiLM configs since the kwarg flows through `**mixer_kwargs` and is only consumed when a `film_generator` exists on the SIREN kernel. **This path is untested end-to-end.**
-
-## Shared settings (all configs)
-
-| Parameter                    | Value                   |
-| ---------------------------- | ----------------------- |
-| Training timesteps           | 1000                    |
-| Inference steps (Heun)       | 50                      |
-| CFG guidance scale           | 2.9                     |
-| CFG interval                 | \[0.1, 1.0\]            |
-| Condition dropout            | 0.1                     |
-| Flow-matching p_mean / p_std | -0.8 / 0.8              |
-| Optimizer                    | Adam (betas 0.9, 0.95)  |
-| Scheduler                    | Constant with 2% warmup |
-| Gradient clip                | 1.0                     |
-| Classes                      | 1000 (ImageNet)         |
-| Training iterations          | 250k                    |
-
-## Running
-
-```bash
-# Snellius (SLURM)
-sbatch scripts/slurm/diffusion/ccnn_jit_128_snellius.sh
-
-# Local / interactive (single node, 4 GPUs)
-PYTHONPATH=. python experiments/run.py \
-    --config examples/imagenet_diffusion/ccnn_jit_128.py \
-    experiment_dir=runs/ccnn_jit_128
-
-# Override omega_0 at launch time (any config)
-PYTHONPATH=. python experiments/run.py \
-    --config examples/imagenet_diffusion/ccnn_jit_128.py \
-    experiment_dir=runs/ccnn_jit_128_omega5 \
-    net.block_cfg.sequence_mixer_cfg.mixer_cfg.global_conv_cfg.kernel_cfg.omega_0=5.0
-```
-
-## SLURM scripts
-
-- `scripts/slurm/diffusion/ccnn_jit_128_snellius.sh` — template for 128px on Snellius H100 nodes. Copy and edit `CONFIG_FILE` / `EXPERIMENT_NAME` for other configs.
-- `scripts/slurm/diffusion/ccnn_jit_diff_baseline_snellius.sh` — existing 64px CCNN baseline.
-- `scripts/slurm/diffusion/jit_diff_baseline_snellius.sh` — existing 64px JiT transformer baseline.
diff --git a/examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py b/examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py
deleted file mode 100644
index 3fc16983..00000000
--- a/examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py
+++ /dev/null
@@ -1,240 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for ImageNet diffusion using the shared ResNet backbone."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import RandomFourierKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 3  # RGB images
-OUTPUT_CHANNELS = 3  # Reconstruct RGB
-NUM_CLASSES = 1_000
-DATA_DIM = 2
-
-# Training parameters
-BATCH_SIZE = 42
-IMAGENET_PATH = os.environ.get("IMAGENET_CACHE", "/home/dknigge/project_dir/huggingface/imagenet")
-HF_DATASET_NAME = "imagenet-1k"
-HF_DATASET_CONFIG = None
-IMAGE_SIZE = 256
-FINAL_IMAGE_SIZE = 64
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# Optimisation parameters
-TRAINING_ITERATIONS = 800_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-WEIGHT_DECAY = 1e-3
-LEARNING_RATE = 2e-4
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 1
-
-# Diffusion parameters
-PREDICTION_TYPE = "v_prediction"
-NUM_TRAIN_TIMESTEPS = 1_000
-BETA_START = 1e-4
-BETA_END = 2e-2
-BETA_SCHEDULE = "cosine_interpolated"
-TIME_EMBED_DIM = NUM_HIDDEN_CHANNELS
-MAX_PERIOD = 10_000.0
-LOG_SAMPLES = True
-
-# Classifier-free guidance
-CFG_ENABLED = True
-GUIDANCE_SCALE = 3.5
-CONDITION_DROPOUT_PROB = 0.25
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Return the ImageNet diffusion configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-    config.compile = True
-
-    hf_token = os.environ.get("HF_TOKEN")
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=IMAGENET_PATH,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        seed=config.seed,
-        image_size=IMAGE_SIZE,
-        final_image_size=FINAL_IMAGE_SIZE,
-        center_crop=True,
-        drop_labels=False,
-        hf_dataset_name="imagenet-1k",
-        hf_dataset_config=None,
-        hf_auth_token=hf_token,
-        task="generation",
-    )
-
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(RandomFourierKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=64,
-                            num_layers=3,
-                            embedding_dim=64,
-                            omega_0=100.0,
-                            L_cache=32,
-                            use_bias=True,
-                            nonlinear_cfg=LazyConfig(torch.nn.SiLU)(),
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-        precision="bf16-mixed",
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        prediction_type=PREDICTION_TYPE,
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        beta_start=BETA_START,
-        beta_end=BETA_END,
-        beta_schedule=BETA_SCHEDULE,
-        cosine_schedule_image_resolution=FINAL_IMAGE_SIZE,
-        cosine_schedule_noise_res_high=FINAL_IMAGE_SIZE,
-        cosine_schedule_noise_res_low=max(32, FINAL_IMAGE_SIZE // 2),
-        time_embed_dim=TIME_EMBED_DIM,
-        max_period=MAX_PERIOD,
-        num_classes=1_000,
-        use_classifier_free_guidance=CFG_ENABLED,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        fid_enabled=False,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet-diffusion",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_6_256_hyena_rope_qknorm.py b/examples/imagenet_diffusion/ccnn_6_256_hyena_rope_qknorm.py
deleted file mode 100644
index 0915c902..00000000
--- a/examples/imagenet_diffusion/ccnn_6_256_hyena_rope_qknorm.py
+++ /dev/null
@@ -1,265 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for ImageNet diffusion using the shared ResNet backbone."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import RandomFourierKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 3  # RGB images
-OUTPUT_CHANNELS = 3  # Reconstruct RGB
-NUM_CLASSES = 1_000
-DATA_DIM = 2
-
-# Training parameters
-BATCH_SIZE = 64
-IMAGENET_PATH = os.environ.get("IMAGENET_CACHE", "/projects/0/prjs1161/imagenet")
-HF_DATASET_NAME = "imagenet-1k"
-HF_DATASET_CONFIG = None
-IMAGE_SIZE = 256
-FINAL_IMAGE_SIZE = 64
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 256
-NUM_BLOCKS = 6
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# Optimisation parameters
-TRAINING_ITERATIONS = 800_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-WEIGHT_DECAY = 1e-3
-LEARNING_RATE = 2e-4
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 1
-
-# Diffusion parameters
-NUM_TRAIN_TIMESTEPS = 1_000
-CLASSIFIER_FREE_GUIDANCE = True
-BETA_START = 1e-4
-BETA_END = 2e-2
-BETA_SCHEDULE = "linear"
-TIME_EMBED_DIM = NUM_HIDDEN_CHANNELS
-MAX_PERIOD = 10_000.0
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 8
-LOG_SAMPLES = True
-COSINE_SCHEDULE_LOGSNR_MIN = -10.0
-COSINE_SCHEDULE_LOGSNR_MAX = 10.0
-PREDICTION_TYPE = "v_prediction"
-DDIM_ETA = 0.0
-EMA_ENABLED = True
-EMA_DECAY = 0.999
-EMA_WARMUP_STEPS = 1_000
-EMA_UPDATE_EVERY = 1
-GUIDANCE_SCALE = 3.5
-CONDITION_DROPOUT_PROB = 0.1
-USE_SIGMOID_LOSS_WEIGHTING = True
-SIGMOID_LOSS_BIAS = -3.0
-FID_ENABLED = True
-FID_NUM_BATCHES = 8
-FID_NUM_INFERENCE_STEPS = NUM_INFERENCE_STEPS
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Return the ImageNet diffusion configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-    hf_token = os.environ.get("HF_TOKEN")
-    config.classifier_free_guidance = CLASSIFIER_FREE_GUIDANCE
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=IMAGENET_PATH,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        seed=config.seed,
-        image_size=IMAGE_SIZE,
-        final_image_size=FINAL_IMAGE_SIZE,
-        center_crop=True,
-        num_classes=NUM_CLASSES,
-        drop_labels=False,
-        hf_dataset_name="imagenet-1k",
-        hf_dataset_config=None,
-        hf_auth_token=hf_token,
-        task="generation",
-    )
-
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(RandomFourierKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=64,
-                            num_layers=3,
-                            embedding_dim=64,
-                            omega_0=100.0,
-                            L_cache=32,
-                            use_bias=True,
-                            nonlinear_cfg=LazyConfig(torch.nn.SiLU)(),
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    # Compose diffusion config with explicit schedule, sampling, and EMA parameters.
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        beta_start=BETA_START,
-        beta_end=BETA_END,
-        beta_schedule=BETA_SCHEDULE,
-        cosine_schedule_logsnr_min=COSINE_SCHEDULE_LOGSNR_MIN,
-        cosine_schedule_logsnr_max=COSINE_SCHEDULE_LOGSNR_MAX,
-        cosine_schedule_image_resolution=FINAL_IMAGE_SIZE,
-        cosine_schedule_noise_res_high=FINAL_IMAGE_SIZE,
-        cosine_schedule_noise_res_low=max(32, FINAL_IMAGE_SIZE // 2),
-        prediction_type=PREDICTION_TYPE,
-        time_embed_dim=TIME_EMBED_DIM,
-        max_period=MAX_PERIOD,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ddim_eta=DDIM_ETA,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_sigmoid_loss_weighting=USE_SIGMOID_LOSS_WEIGHTING,
-        sigmoid_loss_bias=SIGMOID_LOSS_BIAS,
-        use_classifier_free_guidance=CLASSIFIER_FREE_GUIDANCE,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        fid_enabled=FID_ENABLED,
-        fid_num_batches=FID_NUM_BATCHES,
-        fid_num_inference_steps=FID_NUM_INFERENCE_STEPS,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet-diffusion",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_128.py b/examples/imagenet_diffusion/ccnn_jit_128.py
deleted file mode 100644
index 29e574f4..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_128.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 128x128, JiT-B matched (768 hidden, 12 blocks, patch 8)."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 128  # Per GPU — halved vs 64px to fit memory
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 128
-PATCH_SIZE = 8  # 128 / 8 = 16x16 tokens (same count as 64px / 4)
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 2  # 128 * 4 GPUs * 2 = 1024 effective batch size
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    # CCNN Model — JiT-B matched at 128x128
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,  # 16
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_128",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet128", "jit-B-matched"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_128_gated_film_ema.py b/examples/imagenet_diffusion/ccnn_jit_128_gated_film_ema.py
deleted file mode 100644
index cc10fe31..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_128_gated_film_ema.py
+++ /dev/null
@@ -1,278 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 128x128, JiT-B matched, gated Hyena + FiLM-conditioned SIREN + EMA.
-
-This config mirrors the gated FiLM architecture from the classification config
-``vit5_small_pretrain_hyena_cls_row_gated_film_ema`` adapted for diffusion:
-  - Gated dual nonlinearity: SiLU gate + Sigmoid gate
-  - FiLM-conditioned SIREN kernels (input-dependent convolution kernels)
-  - EMA (decay 0.9998)
-
-NOTE — untested change:
-  ``AdaLNZeroResidualBlock.forward`` was modified to pass ``conditioning=cond``
-  (the collapsed timestep vector) through to the sequence mixer so that
-  FiLM-enabled SIREN kernels receive it.  This is a one-line change in
-  ``nvsubquadratic/modules/residual_block.py`` (line ~194).  The kwarg is
-  harmless for non-FiLM mixers (absorbed by ``**mixer_kwargs``), but the
-  end-to-end FiLM diffusion path has not been validated yet.
-"""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.film import KernelFiLMGenerator
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.rms_norm import RMSNorm
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 128
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 128
-PATCH_SIZE = 8  # 128 / 8 = 16x16 tokens
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel
-KERNEL_MLP_HIDDEN_DIM = 32
-KERNEL_NUM_LAYERS = 3
-KERNEL_EMBEDDING_DIM = 32
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# FiLM conditioning
-FILM_HIDDEN_DIM = 64
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 2
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    # FiLM generator: timestep condition -> per-layer (gamma, beta) for SIREN
-    film_cfg = LazyConfig(KernelFiLMGenerator)(
-        cond_dim=NUM_HIDDEN_CHANNELS,
-        kernel_hidden_dim=KERNEL_MLP_HIDDEN_DIM,
-        num_film_layers=KERNEL_NUM_LAYERS - 1,  # One (gamma, beta) per hidden SIREN layer
-        film_hidden_dim=FILM_HIDDEN_DIM,
-    )
-
-    # CCNN Model — JiT-B matched, gated Hyena + FiLM SIREN
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM,
-                            num_layers=KERNEL_NUM_LAYERS,
-                            embedding_dim=KERNEL_EMBEDDING_DIM,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,
-                            use_bias=True,
-                            film_cfg=film_cfg,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    # Gated dual nonlinearity (matches classification config)
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(),
-                    pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS, eps=1e-6),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                    output_norm_cfg=LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS, eps=1e-6),
-                    gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_128_gated_film_ema",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet128", "jit-B-matched", "gated-film", "ema"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_128_low_omega.py b/examples/imagenet_diffusion/ccnn_jit_128_low_omega.py
deleted file mode 100644
index e62809d8..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_128_low_omega.py
+++ /dev/null
@@ -1,249 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 128x128, JiT-B matched, low omega_0 (3.0).
-
-Identical to ccnn_jit_128 except omega_0 is reduced from 10.0 to 3.0 to
-suppress high-frequency artifacts observed in earlier runs.
-"""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 128
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 128
-PATCH_SIZE = 8
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel — reduced omega_0 to suppress high-frequency artifacts
-KERNEL_OMEGA_0 = 3.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 2
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_128_low_omega",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet128", "jit-B-matched", "low-omega"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_256.py b/examples/imagenet_diffusion/ccnn_jit_256.py
deleted file mode 100644
index 98465aeb..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_256.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 256x256, JiT-B matched (768 hidden, 12 blocks, patch 16)."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 64  # Per GPU — reduced for 256px memory
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 256
-PATCH_SIZE = 16  # 256 / 16 = 16x16 tokens (same count as 64px / 4)
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 4  # 64 * 4 GPUs * 4 = 1024 effective batch size
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    # CCNN Model — JiT-B matched at 256x256
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,  # 16
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_256",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet256", "jit-B-matched"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_256_low_omega.py b/examples/imagenet_diffusion/ccnn_jit_256_low_omega.py
deleted file mode 100644
index 2576511f..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_256_low_omega.py
+++ /dev/null
@@ -1,249 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 256x256, JiT-B matched, low omega_0 (3.0).
-
-Identical to ccnn_jit_256 except omega_0 is reduced from 10.0 to 3.0 to
-suppress high-frequency artifacts observed in earlier runs.
-"""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 64
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 256
-PATCH_SIZE = 16
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel — reduced omega_0 to suppress high-frequency artifacts
-KERNEL_OMEGA_0 = 3.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 4
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_256_low_omega",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet256", "jit-B-matched", "low-omega"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_L_256.py b/examples/imagenet_diffusion/ccnn_jit_L_256.py
deleted file mode 100644
index 4b939be6..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_L_256.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN diffusion on ImageNet 256x256, JiT-L matched (1024 hidden, 24 blocks, patch 16)."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 32  # Per GPU — reduced for large model at 256px
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 256
-PATCH_SIZE = 16  # 256 / 16 = 16x16 tokens
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-L)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 1024
-NUM_BLOCKS = 24
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# SIREN kernel
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 1e-4  # Lower LR for larger model
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 8  # 32 * 4 GPUs * 8 = 1024 effective batch size
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9999
-EMA_WARMUP_STEPS = 2000
-EMA_UPDATE_EVERY = 1
-
-# CFG
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",
-        drop_labels=False,
-    )
-
-    # CCNN Model — JiT-L matched at 256x256
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=64,  # Larger kernel MLP for L-size
-                            num_layers=3,
-                            embedding_dim=64,
-                            omega_0=KERNEL_OMEGA_0,
-                            hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                            L_cache=FINAL_IMAGE_SIZE // PATCH_SIZE,  # 16
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        fid_online_jit=False,
-        fid_stats_file="",
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_L_256",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet256", "jit-L-matched"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/ccnn_jit_baseline.py b/examples/imagenet_diffusion/ccnn_jit_baseline.py
deleted file mode 100644
index c47820e1..00000000
--- a/examples/imagenet_diffusion/ccnn_jit_baseline.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CCNN-Hyena baseline config for class-conditional ImageNet 64×64 flow-matching diffusion, size-matched to JiT-B."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.patchify import Patchify, Unpatchify
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 256  # Per GPU (Total 1024 with accumulation)
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 64
-PATCH_SIZE = int(FINAL_IMAGE_SIZE / 16)  # 4 -> patch_size 4
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Network params (CCNN matching JiT-B)
-INPUT_CHANNELS = 3
-OUTPUT_CHANNELS = 3
-DATA_DIM = 2
-NUM_HIDDEN_CHANNELS = 768
-NUM_BLOCKS = 12
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# Optimisation -----------------------------------------------------------------
-# 200 epochs * 1.28M images / 1024 batch size = 250,000 iterations
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 1  # 256 * 4 = 1024 effective batch size
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50  # Heun steps
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG, per JiT
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",  # normalizes to [-1, 1]
-        drop_labels=False,
-    )
-
-    # CCNN Model
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        # Patchification matching JiT
-        in_proj_cfg=LazyConfig(Patchify)(
-            in_features="${net.in_channels}",
-            out_features="${net.hidden_dim}",
-            data_dim="${net.data_dim}",
-            patch_size=PATCH_SIZE,
-            stride=PATCH_SIZE,
-        ),
-        out_proj_cfg=LazyConfig(Unpatchify)(
-            in_features="${net.hidden_dim}",
-            out_features="${net.out_channels}",
-            data_dim="${net.data_dim}",
-            patch_size="${net.in_proj_cfg.patch_size}",
-            stride="${net.in_proj_cfg.stride}",
-            weight_init="zeros",
-        ),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=10.0,
-                            hidden_omega_0=1.0,
-                            L_cache=16,
-                            use_bias=True,
-                        ),
-                        mask_cfg=LazyConfig(torch.nn.Identity)(),  # No mask
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                    use_rope=False,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            condition_norm_cfg="${net.norm_cfg}",
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-        betas=(0.9, 0.95),
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        # CFG
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        # JiT flow-matching params
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        # Online FID
-        fid_online_jit=True,
-        fid_stats_file="examples/imagenet_diffusion/fid_stats/jit_in64_train_stats_full.npz",
-        fid_interval=100,
-        fid_num_samples=50_000,
-        fid_batch_size=1024,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_ccnn_jit_baseline",
-        entity=WANDB_ENTITY,
-        tags=["ccnn", "diffusion", "imagenet64", "jit-matched"],
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/fid_stats/jit_in128_train_stats_full.npz b/examples/imagenet_diffusion/fid_stats/jit_in128_train_stats_full.npz
deleted file mode 100644
index 78c46085..00000000
Binary files a/examples/imagenet_diffusion/fid_stats/jit_in128_train_stats_full.npz and /dev/null differ
diff --git a/examples/imagenet_diffusion/fid_stats/jit_in64_train_stats_full.npz b/examples/imagenet_diffusion/fid_stats/jit_in64_train_stats_full.npz
deleted file mode 100644
index a969ea93..00000000
Binary files a/examples/imagenet_diffusion/fid_stats/jit_in64_train_stats_full.npz and /dev/null differ
diff --git a/examples/imagenet_diffusion/hf_uvit_baseline.py b/examples/imagenet_diffusion/hf_uvit_baseline.py
deleted file mode 100644
index 0815e030..00000000
--- a/examples/imagenet_diffusion/hf_uvit_baseline.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""HuggingFace UViT baseline config for class-conditional ImageNet 64×64 diffusion."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.networks.huggingface_diffusers import DiffusersUVitWrapper, HuggingFaceUVitConfig
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 32
-MAX_WORKERS = 8
-NUM_WORKERS = min(MAX_WORKERS, os.cpu_count() or MAX_WORKERS)
-IMAGE_SIZE = 256
-FINAL_IMAGE_SIZE = 64
-IMAGENET_CACHE_DIR = os.environ.get("IMAGENET_CACHE", "/projects/0/prjs1161/imagenet")
-HF_DATASET_NAME = "imagenet-1k"
-HF_DATASET_CONFIG = None
-
-# UVit architecture ------------------------------------------------------------
-UVIT_SAMPLE_SIZE = FINAL_IMAGE_SIZE
-UVIT_IN_CHANNELS = 3
-UVIT_OUT_CHANNELS = 3
-UVIT_HIDDEN_SIZE = 256
-UVIT_COND_EMBED_DIM = 128
-UVIT_ENCODER_HIDDEN_SIZE = 128
-UVIT_BLOCK_OUT_CHANNELS = 256
-UVIT_NUM_HIDDEN_LAYERS = 8
-UVIT_NUM_ATTENTION_HEADS = 8
-UVIT_INTERMEDIATE_SIZE = 512
-UVIT_LAYER_NORM_EPS = 1e-5
-UVIT_MICRO_COND_ENCODE_DIM = None
-UVIT_MICRO_COND_EMBED_DIM = None
-UVIT_CODEBOOK_SIZE = None
-UVIT_VOCAB_SIZE = None
-
-# Optimisation -----------------------------------------------------------------
-TRAINING_ITERATIONS = 800_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.02
-LEARNING_RATE = 1e-4
-WEIGHT_DECAY = 0.01
-GRAD_CLIP = 1.0
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1_000
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 8
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.999
-EMA_WARMUP_STEPS = 1_000
-EMA_UPDATE_EVERY = 1
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    hf_token = os.environ.get("HF_TOKEN")
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=IMAGENET_CACHE_DIR,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        seed=config.seed,
-        image_size=IMAGE_SIZE,
-        final_image_size=FINAL_IMAGE_SIZE,
-        center_crop=True,
-        # Keep labels so classifier-free guidance has valid class indices.
-        drop_labels=False,
-        hf_dataset_name=HF_DATASET_NAME,
-        hf_dataset_config=HF_DATASET_CONFIG,
-        hf_auth_token=hf_token,
-        task="generation",
-    )
-
-    hf_cfg = HuggingFaceUVitConfig(
-        sample_size=UVIT_SAMPLE_SIZE,
-        in_channels=UVIT_IN_CHANNELS,
-        out_channels=UVIT_OUT_CHANNELS,
-        hidden_size=UVIT_HIDDEN_SIZE,
-        cond_embed_dim=UVIT_COND_EMBED_DIM,
-        encoder_hidden_size=UVIT_ENCODER_HIDDEN_SIZE,
-        block_out_channels=UVIT_BLOCK_OUT_CHANNELS,
-        num_hidden_layers=UVIT_NUM_HIDDEN_LAYERS,
-        num_attention_heads=UVIT_NUM_ATTENTION_HEADS,
-        intermediate_size=UVIT_INTERMEDIATE_SIZE,
-        layer_norm_eps=UVIT_LAYER_NORM_EPS,
-        micro_cond_encode_dim=UVIT_MICRO_COND_ENCODE_DIM,
-        micro_cond_embed_dim=UVIT_MICRO_COND_EMBED_DIM,
-        codebook_size=UVIT_CODEBOOK_SIZE,
-        vocab_size=UVIT_VOCAB_SIZE,
-    )
-
-    config.net = LazyConfig(DiffusersUVitWrapper)(hf_config=hf_cfg)
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_hf_uvit_baseline",
-        entity=WANDB_ENTITY,
-    )
-
-    return config
diff --git a/examples/imagenet_diffusion/jit_baseline.py b/examples/imagenet_diffusion/jit_baseline.py
deleted file mode 100644
index 28df617d..00000000
--- a/examples/imagenet_diffusion/jit_baseline.py
+++ /dev/null
@@ -1,151 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""JiT-B baseline config for class-conditional ImageNet 64×64 flow-matching diffusion."""
-
-import os
-
-import torch
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import LazyConfig
-from nvsubquadratic.networks.jit import JiT_models
-
-
-PLACEHOLDER = None
-WANDB_ENTITY = "dafidofff"
-
-# Dataset ----------------------------------------------------------------------
-BATCH_SIZE = 256  # Per GPU (Total 1024 with accumulation)
-NUM_WORKERS = min(12, os.cpu_count() - 2 or 4)
-FINAL_IMAGE_SIZE = 64
-PATCH_SIZE = FINAL_IMAGE_SIZE // 16  # 4 -> JiT-B/4
-HF_DATASET = os.environ.get("IMAGENET_HF_DATASET", "imagenet-1k")
-HF_CACHE = os.environ.get("IMAGENET_PATH", "/scratch-shared/dknigge/hf_cache")
-
-# Optimisation -----------------------------------------------------------------
-# 200 epochs * 1.28M images / 1024 batch size = 250,000 iterations
-TRAINING_ITERATIONS = 250_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.025
-LEARNING_RATE = 2e-4
-WEIGHT_DECAY = 0.0
-GRAD_CLIP = 1.0
-ACCUMULATE_GRAD_STEPS = 1  # 256 * 4 = 1024 effective batch size
-
-# Diffusion --------------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1000
-NUM_INFERENCE_STEPS = 50  # Heun steps
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.9998
-EMA_WARMUP_STEPS = 1000
-EMA_UPDATE_EVERY = 1
-
-# CFG, per JiT
-USE_CFG = True
-GUIDANCE_SCALE = 2.9
-CONDITION_DROPOUT_PROB = 0.1
-NUM_CLASSES = 1000
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(ImageNetDataModule)(
-        data_dir=HF_CACHE,
-        hf_dataset_name=HF_DATASET,
-        image_size=FINAL_IMAGE_SIZE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=42,
-        task="generation",  # normalizes to [-1, 1]
-        drop_labels=False,
-    )
-
-    # JiT Model
-    config.net = LazyConfig(JiT_models[f"JiT-B/{PATCH_SIZE}"])(
-        input_size=FINAL_IMAGE_SIZE,
-        num_classes=NUM_CLASSES,
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.Adam)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        accumulate_grad_steps=ACCUMULATE_GRAD_STEPS,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="constant",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        # CFG
-        use_classifier_free_guidance=USE_CFG,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        num_classes=NUM_CLASSES,
-        # JiT flow-matching params
-        p_mean=-0.8,
-        p_std=0.8,
-        cfg_interval_start=0.1,
-        cfg_interval_end=1.0,
-        # Online FID
-        fid_online_jit=True,
-        fid_stats_file="examples/imagenet_diffusion/fid_stats/jit_in64_train_stats_full.npz",
-        fid_interval=100,
-        fid_num_samples=50_000,
-        fid_batch_size=512,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="imagenet_diffusion_jit_baseline",
-        entity=WANDB_ENTITY,
-        tags=["jit", "diffusion", "imagenet64"],
-    )
-
-    return config
diff --git a/examples/mnist_classification/ccnn_4_160_attn.py b/examples/mnist_classification/ccnn_4_160_attn.py
deleted file mode 100644
index ebedc4b1..00000000
--- a/examples/mnist_classification/ccnn_4_160_attn.py
+++ /dev/null
@@ -1,159 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for MNIST classification."""
-
-import os
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.attention import Attention
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import ResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.classification_resnet import ClassificationResNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-
-
-# Dataset parameters
-INPUT_CHANNELS = 1  # MNIST grayscale
-OUTPUT_CHANNELS = 10  # 10 digit classes
-DATA_TYPE = "image"
-DATA_DIM = 2
-
-# Training parameters
-BATCH_SIZE = 128
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 160
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-
-# TRAINING parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-GRAD_CLIP = 10.0
-
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 0.001
-
-
-def get_config() -> ExperimentConfig:
-    """Get the configuration for the MNIST classification experiment.
-
-    Returns:
-        ExperimentConfig: The configuration for the MNIST classification experiment.
-    """
-    # Sratr with default config
-    config = ExperimentConfig()
-
-    # Add dataset config
-    # Update dataset with LazyConfig directly referencing the MNISTDataModule class
-    # and providing all parameters directly (no nested params dataclass)
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=True,  # Flag to use deterministic worker initialization
-        seed=config.seed,  # Pass the seed value instead of a Generator object
-        task="classification",
-    )
-
-    # Add net config
-    config.net = LazyConfig(ClassificationResNet)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Attention)(
-                    hidden_dim="${net.hidden_dim}",
-                    num_heads=8,
-                    apply_qk_norm=True,
-                    use_rope=True,
-                    is_causal=False,
-                    rope_base=10000.0,
-                    attn_dropout=DROPOUT_RATE,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            # Condition mixer
-            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
-            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
-            # MLP
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=1.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            # Dropout
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-    )
-
-    # Add lightning wrapper config
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
-
-    # Add optimizer config
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # Modify the train config - only set what is different from the default
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    # Modify the scheduler config - only set what's different from default
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-    )
-
-    # Add wandb group
-    config.wandb = WandbConfig(
-        job_group="mnist_classification",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py b/examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py
deleted file mode 100644
index 6dd15a77..00000000
--- a/examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py
+++ /dev/null
@@ -1,194 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for MNIST classification."""
-
-import os
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import ResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.classification_resnet import ClassificationResNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 1  # MNIST grayscale
-OUTPUT_CHANNELS = 10  # 10 digit classes
-DATA_TYPE = "image"
-DATA_DIM = 2
-
-# Training parameters
-BATCH_SIZE = 128
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 160
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-
-# TRAINING parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-GRAD_CLIP = 10.0
-
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 0.001
-
-
-def get_config() -> ExperimentConfig:
-    """Get the configuration for the MNIST classification experiment.
-
-    Returns:
-        ExperimentConfig: The configuration for the MNIST classification experiment.
-    """
-    # Sratr with default config
-    config = ExperimentConfig()
-
-    # Add dataset config
-    # Update dataset with LazyConfig directly referencing the MNISTDataModule class
-    # and providing all parameters directly (no nested params dataclass)
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=True,  # Flag to use deterministic worker initialization
-        seed=config.seed,  # Pass the seed value instead of a Generator object
-        task="classification",
-    )
-
-    # Add net config
-    config.net = LazyConfig(ClassificationResNet)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=100.0,
-                            L_cache=32,
-                            use_bias=True,
-                            hidden_omega_0=1.0,
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels="${net.hidden_dim}"),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            # Condition mixer
-            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
-            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
-            # MLP
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=1.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            # Dropout
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-    )
-
-    # Add lightning wrapper config
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
-
-    # Add optimizer config
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # Modify the train config - only set what is different from the default
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    # Modify the scheduler config - only set what's different from default
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-    )
-
-    # Add wandb group
-    config.wandb = WandbConfig(
-        job_group="mnist_classification",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/mnist_diffusion/ccnn_4_160_hyena_rope_qknorm.py b/examples/mnist_diffusion/ccnn_4_160_hyena_rope_qknorm.py
deleted file mode 100644
index 3cf647bd..00000000
--- a/examples/mnist_diffusion/ccnn_4_160_hyena_rope_qknorm.py
+++ /dev/null
@@ -1,230 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for MNIST diffusion using the shared ResNet backbone."""
-
-import os
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import AdaLNZeroResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 1  # MNIST grayscale
-OUTPUT_CHANNELS = 1  # Reconstruct grayscale
-DATA_TYPE = "image"
-DATA_DIM = 2
-
-# Training parameters
-BATCH_SIZE = 128
-MAX_WORKERS = 16
-MNIST_PATH = ".data/mnist"
-PRECISION = "bf16-mixed"
-NUM_WORKERS = min(MAX_WORKERS, os.cpu_count() or MAX_WORKERS)
-IMAGE_SIZE = 28
-
-# Model parameters
-HIDDEN_DIM = 160
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "circular"
-NUM_CLASSES = 10
-
-# Optimisation parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-GRAD_CLIP = 10.0
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 1e-3
-
-# Diffusion parameters ---------------------------------------------------------------
-NUM_TRAIN_TIMESTEPS = 1_000
-TIME_EMBED_DIM = HIDDEN_DIM
-MAX_PERIOD = 10_000.0
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-EMA_ENABLED = True
-EMA_DECAY = 0.999
-EMA_WARMUP_STEPS = 1_000
-EMA_UPDATE_EVERY = 1
-FID_NUM_INFERENCE_STEPS = NUM_INFERENCE_STEPS
-GUIDANCE_SCALE = 3.0
-CONDITION_DROPOUT_PROB = 0.1
-CFG_ENABLED = True
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Return the MNIST diffusion configuration."""
-    config = DiffusionExperimentConfig()
-    config.debug = False
-    config.seed = 42
-
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=MNIST_PATH,
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=config.deterministic,
-        seed=config.seed,
-        task="generation",
-    )
-
-    config.net = LazyConfig(ResidualNetwork)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=HIDDEN_DIM,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(AdaLNZeroResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=10.0,
-                            L_cache=32,
-                            use_bias=True,
-                            hidden_omega_0=1.0,
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(
-                        num_groups=1,
-                        num_channels="${net.hidden_dim}",
-                    ),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=2.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            condition_norm_cfg="${net.norm_cfg}",
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-            hidden_dim="${net.hidden_dim}",
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-        condition_in_proj_cfg=LazyConfig(torch.nn.Linear)(
-            in_features="${net.hidden_dim}", out_features="${net.hidden_dim}"
-        ),
-    )
-
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        precision=PRECISION,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        time_embed_dim=TIME_EMBED_DIM,
-        max_period=MAX_PERIOD,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        num_classes=NUM_CLASSES,
-        use_classifier_free_guidance=CFG_ENABLED,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        fid_num_inference_steps=FID_NUM_INFERENCE_STEPS,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="mnist-diffusion",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/mnist_diffusion/hf_dit_baseline.py b/examples/mnist_diffusion/hf_dit_baseline.py
deleted file mode 100644
index 41051135..00000000
--- a/examples/mnist_diffusion/hf_dit_baseline.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Baseline diffusion configuration leveraging a Hugging Face DiT transformer."""
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.networks.huggingface_diffusers import DiffusersDiTWrapper, HuggingFaceDiTConfig
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset
-BATCH_SIZE = 32
-NUM_WORKERS = 16
-IMAGE_SIZE = 28
-NUM_CLASSES = 10
-
-# Optimisation
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-LEARNING_RATE = 1e-4
-WEIGHT_DECAY = 0.01
-GRAD_CLIP = 1.0
-
-# Diffusion
-NUM_TRAIN_TIMESTEPS = 1_000
-BETA_START = 1e-4
-BETA_END = 0.02
-BETA_SCHEDULE = "squaredcos_cap_v2"
-COSINE_SCHEDULE_LOGSNR_MIN = -10.0
-COSINE_SCHEDULE_LOGSNR_MAX = 10.0
-COSINE_SCHEDULE_IMAGE_RESOLUTION = IMAGE_SIZE
-COSINE_SCHEDULE_NOISE_RES_HIGH = IMAGE_SIZE
-COSINE_SCHEDULE_NOISE_RES_LOW = max(16, IMAGE_SIZE // 2)
-PREDICTION_TYPE = "v_prediction"
-TIME_EMBED_DIM = None
-MAX_PERIOD = 10_000.0
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-DDIM_ETA = 0.0
-EMA_ENABLED = False
-EMA_DECAY = 0.999
-EMA_WARMUP_STEPS = 0
-EMA_UPDATE_EVERY = 1
-CFG_ENABLED = True
-GUIDANCE_SCALE = 3.0
-CONDITION_DROPOUT_PROB = 0.1
-USE_SIGMOID_LOSS_WEIGHTING = True
-SIGMOID_LOSS_BIAS = 0.0
-FID_ENABLED = False
-FID_NUM_BATCHES = 0
-FID_NUM_INFERENCE_STEPS = None
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type="image",
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=config.deterministic,
-        seed=config.seed,
-        task="generation",
-    )
-
-    hf_cfg = HuggingFaceDiTConfig(
-        sample_size=28,
-        patch_size=2,
-        in_channels=1,
-        out_channels=1,
-        num_layers=6,
-        num_attention_heads=4,
-        attention_head_dim=64,
-        dropout=0.0,
-        num_embeds_ada_norm=NUM_TRAIN_TIMESTEPS,
-        activation_fn="gelu-approximate",
-        norm_type="ada_norm_zero",
-        norm_num_groups=32,
-    )
-
-    config.net = LazyConfig(DiffusersDiTWrapper)(hf_config=hf_cfg)
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        beta_start=BETA_START,
-        beta_end=BETA_END,
-        beta_schedule=BETA_SCHEDULE,
-        cosine_schedule_logsnr_min=COSINE_SCHEDULE_LOGSNR_MIN,
-        cosine_schedule_logsnr_max=COSINE_SCHEDULE_LOGSNR_MAX,
-        cosine_schedule_image_resolution=COSINE_SCHEDULE_IMAGE_RESOLUTION,
-        cosine_schedule_noise_res_high=COSINE_SCHEDULE_NOISE_RES_HIGH,
-        cosine_schedule_noise_res_low=COSINE_SCHEDULE_NOISE_RES_LOW,
-        prediction_type=PREDICTION_TYPE,
-        time_embed_dim=TIME_EMBED_DIM,
-        max_period=MAX_PERIOD,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ddim_eta=DDIM_ETA,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_sigmoid_loss_weighting=USE_SIGMOID_LOSS_WEIGHTING,
-        sigmoid_loss_bias=SIGMOID_LOSS_BIAS,
-        num_classes=NUM_CLASSES,
-        use_classifier_free_guidance=CFG_ENABLED,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        fid_enabled=FID_ENABLED,
-        fid_num_batches=FID_NUM_BATCHES,
-        fid_num_inference_steps=FID_NUM_INFERENCE_STEPS,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="mnist_diffusion_hf_baseline",
-        entity=WANDB_ENTITY,
-    )
-
-    return config
diff --git a/examples/mnist_diffusion/hf_uvit_baseline.py b/examples/mnist_diffusion/hf_uvit_baseline.py
deleted file mode 100644
index 949bad67..00000000
--- a/examples/mnist_diffusion/hf_uvit_baseline.py
+++ /dev/null
@@ -1,188 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Baseline diffusion configuration leveraging Hugging Face UVit transformer."""
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import (
-    DiffusionConfig,
-    DiffusionExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.networks.huggingface_diffusers import DiffusersUVitWrapper, HuggingFaceUVitConfig
-
-
-WANDB_ENTITY = "dafidofff"
-
-# Dataset
-BATCH_SIZE = 16
-NUM_WORKERS = 16
-IMAGE_SIZE = 28
-NUM_CLASSES = 10
-
-# UVit architecture ------------------------------------------------------------
-UVIT_SAMPLE_SIZE = 28
-UVIT_IN_CHANNELS = 1
-UVIT_OUT_CHANNELS = 1
-UVIT_HIDDEN_SIZE = 256
-UVIT_COND_EMBED_DIM = 128
-UVIT_ENCODER_HIDDEN_SIZE = 128
-UVIT_BLOCK_OUT_CHANNELS = 256
-UVIT_NUM_HIDDEN_LAYERS = 8
-UVIT_NUM_ATTENTION_HEADS = 8
-UVIT_INTERMEDIATE_SIZE = 512
-UVIT_LAYER_NORM_EPS = 1e-5
-UVIT_MICRO_COND_ENCODE_DIM = None
-UVIT_MICRO_COND_EMBED_DIM = None
-UVIT_CODEBOOK_SIZE = None
-UVIT_VOCAB_SIZE = None
-
-# Optimisation
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-LEARNING_RATE = 1e-4
-WEIGHT_DECAY = 0.01
-GRAD_CLIP = 1.0
-
-# Diffusion
-NUM_TRAIN_TIMESTEPS = 1_000
-BETA_START = 1e-4
-BETA_END = 0.02
-BETA_SCHEDULE = "squaredcos_cap_v2"
-COSINE_SCHEDULE_LOGSNR_MIN = -10.0
-COSINE_SCHEDULE_LOGSNR_MAX = 10.0
-COSINE_SCHEDULE_IMAGE_RESOLUTION = IMAGE_SIZE
-COSINE_SCHEDULE_NOISE_RES_HIGH = IMAGE_SIZE
-COSINE_SCHEDULE_NOISE_RES_LOW = max(16, IMAGE_SIZE // 2)
-PREDICTION_TYPE = "v_prediction"
-TIME_EMBED_DIM = None
-MAX_PERIOD = 10_000.0
-NUM_INFERENCE_STEPS = 50
-NUM_SAMPLES = 16
-LOG_SAMPLES = True
-DDIM_ETA = 0.0
-EMA_ENABLED = False
-EMA_DECAY = 0.999
-EMA_WARMUP_STEPS = 0
-EMA_UPDATE_EVERY = 1
-CFG_ENABLED = True
-GUIDANCE_SCALE = 3.0
-CONDITION_DROPOUT_PROB = 0.1
-USE_SIGMOID_LOSS_WEIGHTING = True
-SIGMOID_LOSS_BIAS = 0.0
-FID_ENABLED = False
-FID_NUM_BATCHES = 0
-FID_NUM_INFERENCE_STEPS = None
-
-
-def get_config() -> DiffusionExperimentConfig:
-    """Build the experiment configuration."""
-    config = DiffusionExperimentConfig()
-
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type="image",
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=config.deterministic,
-        seed=config.seed,
-        task="generation",
-    )
-
-    hf_cfg = HuggingFaceUVitConfig(
-        sample_size=UVIT_SAMPLE_SIZE,
-        in_channels=UVIT_IN_CHANNELS,
-        out_channels=UVIT_OUT_CHANNELS,
-        hidden_size=UVIT_HIDDEN_SIZE,
-        cond_embed_dim=UVIT_COND_EMBED_DIM,
-        encoder_hidden_size=UVIT_ENCODER_HIDDEN_SIZE,
-        block_out_channels=UVIT_BLOCK_OUT_CHANNELS,
-        num_hidden_layers=UVIT_NUM_HIDDEN_LAYERS,
-        num_attention_heads=UVIT_NUM_ATTENTION_HEADS,
-        intermediate_size=UVIT_INTERMEDIATE_SIZE,
-        layer_norm_eps=UVIT_LAYER_NORM_EPS,
-        micro_cond_encode_dim=UVIT_MICRO_COND_ENCODE_DIM,
-        micro_cond_embed_dim=UVIT_MICRO_COND_EMBED_DIM,
-        codebook_size=UVIT_CODEBOOK_SIZE,
-        vocab_size=UVIT_VOCAB_SIZE,
-    )
-
-    config.net = LazyConfig(DiffusersUVitWrapper)(hf_config=hf_cfg)
-    config.lightning_wrapper_class = LazyConfig(DiffusionWrapper)()
-
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="min",
-    )
-
-    config.diffusion = DiffusionConfig(
-        num_train_timesteps=NUM_TRAIN_TIMESTEPS,
-        beta_start=BETA_START,
-        beta_end=BETA_END,
-        beta_schedule=BETA_SCHEDULE,
-        cosine_schedule_logsnr_min=COSINE_SCHEDULE_LOGSNR_MIN,
-        cosine_schedule_logsnr_max=COSINE_SCHEDULE_LOGSNR_MAX,
-        cosine_schedule_image_resolution=COSINE_SCHEDULE_IMAGE_RESOLUTION,
-        cosine_schedule_noise_res_high=COSINE_SCHEDULE_NOISE_RES_HIGH,
-        cosine_schedule_noise_res_low=COSINE_SCHEDULE_NOISE_RES_LOW,
-        prediction_type=PREDICTION_TYPE,
-        time_embed_dim=TIME_EMBED_DIM,
-        max_period=MAX_PERIOD,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        num_samples=NUM_SAMPLES,
-        log_samples=LOG_SAMPLES,
-        ddim_eta=DDIM_ETA,
-        ema_enabled=EMA_ENABLED,
-        ema_decay=EMA_DECAY,
-        ema_update_every=EMA_UPDATE_EVERY,
-        ema_warmup_steps=EMA_WARMUP_STEPS,
-        use_sigmoid_loss_weighting=USE_SIGMOID_LOSS_WEIGHTING,
-        sigmoid_loss_bias=SIGMOID_LOSS_BIAS,
-        num_classes=NUM_CLASSES,
-        use_classifier_free_guidance=CFG_ENABLED,
-        guidance_scale=GUIDANCE_SCALE,
-        condition_dropout_prob=CONDITION_DROPOUT_PROB,
-        fid_enabled=FID_ENABLED,
-        fid_num_batches=FID_NUM_BATCHES,
-        fid_num_inference_steps=FID_NUM_INFERENCE_STEPS,
-    )
-
-    config.wandb = WandbConfig(
-        job_group="mnist_diffusion_hf_uvit_baseline",
-        entity=WANDB_ENTITY,
-    )
-
-    return config
diff --git a/examples/smnist_classification/ccnn_4_160_attn.py b/examples/smnist_classification/ccnn_4_160_attn.py
deleted file mode 100644
index ecb64458..00000000
--- a/examples/smnist_classification/ccnn_4_160_attn.py
+++ /dev/null
@@ -1,159 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for MNIST classification."""
-
-import os
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.attention import Attention
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import ResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.classification_resnet import ClassificationResNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-
-
-# Dataset parameters
-INPUT_CHANNELS = 1  # MNIST grayscale
-OUTPUT_CHANNELS = 10  # 10 digit classes
-DATA_TYPE = "sequence"
-DATA_DIM = 1
-
-# Training parameters
-BATCH_SIZE = 128
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 160
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "double"
-
-# TRAINING parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-GRAD_CLIP = 10.0
-
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 0.001
-
-
-def get_config() -> ExperimentConfig:
-    """Get the configuration for the MNIST classification experiment.
-
-    Returns:
-        ExperimentConfig: The configuration for the MNIST classification experiment.
-    """
-    # Sratr with default config
-    config = ExperimentConfig()
-
-    # Add dataset config
-    # Update dataset with LazyConfig directly referencing the MNISTDataModule class
-    # and providing all parameters directly (no nested params dataclass)
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=True,  # Flag to use deterministic worker initialization
-        seed=config.seed,  # Pass the seed value instead of a Generator object
-    )
-
-    # Add net config
-    config.net = LazyConfig(ClassificationResNet)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Attention)(
-                    hidden_dim="${net.hidden_dim}",
-                    num_heads=8,
-                    apply_qk_norm=True,
-                    use_rope=True,
-                    is_causal=False,
-                    rope_base=10000.0,
-                    attn_dropout=DROPOUT_RATE,
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            # Condition mixer
-            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
-            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
-            # MLP
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=1.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            # Dropout
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-    )
-
-    # Add lightning wrapper config
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
-
-    # Add optimizer config
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # Modify the train config - only set what is different from the default
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    # Modify the scheduler config - only set what's different from default
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-    )
-
-    # Add wandb group
-    config.wandb = WandbConfig(
-        job_group="smnist_classification",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/smnist_classification/ccnn_4_160_hyena_rope_qknorm.py b/examples/smnist_classification/ccnn_4_160_hyena_rope_qknorm.py
deleted file mode 100644
index f04dbc54..00000000
--- a/examples/smnist_classification/ccnn_4_160_hyena_rope_qknorm.py
+++ /dev/null
@@ -1,193 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for MNIST classification."""
-
-import os
-
-import torch
-
-from experiments.datamodules.mnist import MNISTDataModule
-from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import ResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.classification_resnet import ClassificationResNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 1  # MNIST grayscale
-OUTPUT_CHANNELS = 10  # 10 digit classes
-DATA_TYPE = "sequence"
-DATA_DIM = 1
-
-# Training parameters
-BATCH_SIZE = 128
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 160
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "double"
-FFT_PADDING = "zero"
-
-# TRAINING parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-GRAD_CLIP = 10.0
-
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 0.001
-
-
-def get_config() -> ExperimentConfig:
-    """Get the configuration for the MNIST classification experiment.
-
-    Returns:
-        ExperimentConfig: The configuration for the MNIST classification experiment.
-    """
-    # Sratr with default config
-    config = ExperimentConfig()
-
-    # Add dataset config
-    # Update dataset with LazyConfig directly referencing the MNISTDataModule class
-    # and providing all parameters directly (no nested params dataclass)
-    config.dataset = LazyConfig(MNISTDataModule)(
-        data_dir=".data/mnist",
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=True,  # Flag to use deterministic worker initialization
-        seed=config.seed,  # Pass the seed value instead of a Generator object
-    )
-
-    # Add net config
-    config.net = LazyConfig(ClassificationResNet)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=100.0,
-                            L_cache=32,
-                            use_bias=True,
-                            hidden_omega_0=1.0,
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv1d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels="${net.hidden_dim}"),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            # Condition mixer
-            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
-            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
-            # MLP
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=1.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            # Dropout
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-    )
-
-    # Add lightning wrapper config
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
-
-    # Add optimizer config
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # Modify the train config - only set what is different from the default
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    # Modify the scheduler config - only set what's different from default
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-    )
-
-    # Add wandb group
-    config.wandb = WandbConfig(
-        job_group="smnist_classification",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/ucf101_classification/ccnn_4_160_hyena_rope_qknorm.py b/examples/ucf101_classification/ccnn_4_160_hyena_rope_qknorm.py
deleted file mode 100644
index d35f6125..00000000
--- a/examples/ucf101_classification/ccnn_4_160_hyena_rope_qknorm.py
+++ /dev/null
@@ -1,204 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Config file for UCF101 classification."""
-
-import os
-
-import torch
-
-from experiments.datamodules.ucf101 import UCF101DataModule
-from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_nd import CKConvND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.masks_nd import GaussianModulationND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.residual_block import ResidualBlock
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.networks.classification_resnet import ClassificationResNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-from nvsubquadratic.utils.qk_norm import L2Norm
-
-
-# Dataset parameters
-INPUT_CHANNELS = 3  # RGB video frames
-OUTPUT_CHANNELS = 101  # UCF101 action classes
-DATA_TYPE = "video"
-DATA_DIM = 3
-
-# Training parameters
-BATCH_SIZE = 1
-PRECISION = "bf16-mixed"  # Tested options: "32-true", "bf16-mixed"
-
-# Model parameters
-NUM_HIDDEN_CHANNELS = 156  # Must be divisible by 6
-NUM_BLOCKS = 4
-DROPOUT_IN_RATE = 0.0
-DROPOUT_RATE = 0.1
-GRID_TYPE = "single"
-FFT_PADDING = "zero"
-
-# TRAINING parameters
-TRAINING_ITERATIONS = 100_000
-WARMUP_ITERATIONS_PERCENTAGE = 0.05
-NUM_WORKERS = os.cpu_count() // torch.cuda.device_count() if torch.cuda.is_available() else os.cpu_count()
-GRAD_CLIP = 10.0
-
-WEIGHT_DECAY = 0.01
-LEARNING_RATE = 0.001
-
-# Dataset parameters
-FRAME_SIZE = (128, 128)
-FRAMES_PER_CLIP = 16
-STEP_BETWEEN_CLIPS = 1
-VAL_SPLIT_FRACTION = 0.1
-
-
-def get_config() -> ExperimentConfig:
-    """Get the configuration for the UCF101 classification experiment.
-
-    Returns:
-        ExperimentConfig: The configuration for the UCF101 classification experiment.
-    """
-    # Sratr with default config
-    config = ExperimentConfig()
-
-    # Add dataset config
-    # Update dataset with LazyConfig directly referencing the UCF101DataModule class
-    # and providing all parameters directly (no nested params dataclass)
-    config.dataset = LazyConfig(UCF101DataModule)(
-        data_dir=".data/ucf101",
-        split_fold=1,
-        data_type=DATA_TYPE,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=torch.cuda.is_available() and config.device == "cuda",
-        use_deterministic_worker_init=True,  # Flag to use deterministic worker initialization
-        seed=config.seed,  # Pass the seed value instead of a Generator object
-        frames_per_clip=FRAMES_PER_CLIP,
-        step_between_clips=STEP_BETWEEN_CLIPS,
-        frame_size=FRAME_SIZE,
-        val_split_fraction=VAL_SPLIT_FRACTION,
-    )
-
-    # Add net config
-    config.net = LazyConfig(ClassificationResNet)(
-        in_channels=INPUT_CHANNELS,
-        out_channels=OUTPUT_CHANNELS,
-        num_blocks=NUM_BLOCKS,
-        hidden_dim=NUM_HIDDEN_CHANNELS,
-        data_dim=DATA_DIM,
-        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
-        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
-        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
-        block_cfg=LazyConfig(ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
-                hidden_dim="${net.hidden_dim}",
-                mixer_cfg=LazyConfig(Hyena)(
-                    global_conv_cfg=LazyConfig(CKConvND)(
-                        data_dim="${net.data_dim}",
-                        hidden_dim="${net.hidden_dim}",
-                        kernel_cfg=LazyConfig(SIRENKernelND)(
-                            data_dim="${net.data_dim}",
-                            out_dim="${net.hidden_dim}",
-                            mlp_hidden_dim=32,
-                            num_layers=3,
-                            embedding_dim=32,
-                            omega_0=100.0,
-                            L_cache=32,
-                            use_bias=True,
-                            hidden_omega_0=1.0,
-                        ),
-                        mask_cfg=LazyConfig(GaussianModulationND)(
-                            data_dim="${net.data_dim}",
-                            num_channels="${net.hidden_dim}",
-                            min_attenuation_at_step=0.1,
-                            max_attenuation_at_limit=0.95,
-                            init_extent=1.0,
-                            parametrization="direct",
-                        ),
-                        grid_type=GRID_TYPE,
-                        fft_padding=FFT_PADDING,
-                    ),
-                    short_conv_cfg=LazyConfig(torch.nn.Conv3d)(
-                        in_channels="3 * ${net.hidden_dim}",
-                        out_channels="3 * ${net.hidden_dim}",
-                        kernel_size=3,
-                        groups="3 * ${net.hidden_dim}",
-                        padding=1,
-                        bias=False,
-                    ),
-                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
-                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels="${net.hidden_dim}"),
-                    qk_norm_cfg=LazyConfig(L2Norm)(),
-                ),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            sequence_mixer_norm_cfg="${net.norm_cfg}",
-            # Condition mixer
-            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
-            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
-            # MLP
-            mlp_cfg=LazyConfig(MLP)(
-                dim="${net.hidden_dim}",
-                activation="glu",
-                expansion_factor=1.0,
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
-            ),
-            mlp_norm_cfg="${net.norm_cfg}",
-            # Dropout
-            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE),
-        ),
-        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE),
-    )
-
-    # Add lightning wrapper config
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
-
-    # Add optimizer config
-    config.optimizer = LazyConfig(torch.optim.AdamW)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # Modify the train config - only set what is different from the default
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TRAINING_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-    )
-
-    # Modify the scheduler config - only set what's different from default
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-    )
-
-    # Add wandb group
-    config.wandb = WandbConfig(
-        job_group="ucf101_classification",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    return config
diff --git a/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_cls_row_apex.py b/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_cls_row_apex.py
deleted file mode 100644
index b91b29fb..00000000
--- a/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_cls_row_apex.py
+++ /dev/null
@@ -1,269 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""ViT-5-Small + Multi-Head Hyena ImageNet-1k — CLS-row variant, Apex FusedLAMB.
-
-Based on v2/vit5_small_pretrain_hyena_cls_row_apex.py but replaces the depthwise
-CKConvND with CKConvMultiheadND (dense within-head channel mixing).
-
-Key differences from v2/vit5_small_pretrain_hyena_cls_row_apex.py:
-- CKConvMultiheadND replaces CKConvND: each head performs dense [head_dim x
-  head_dim] channel mixing within heads while keeping heads isolated.
-- SIREN kernel out_dim = NUM_HEADS * HEAD_DIM * HEAD_DIM (dense kernel per head)
-  instead of HIDDEN_DIM (depthwise).
-- NUM_HEADS=6, HEAD_DIM=64 (consistent with the ViT-5-Small attention variant).
-- PerHeadRMSNorm for QK normalization (each head normalized independently).
-- SiLU gate nonlinearity + output RMSNorm (same as other v2 configs).
-- CLS-row architecture: CLS + 13 registers as extra row → 15×14 grid.
-- DALI fused data pipeline with local NVMe staging.
-"""
-
-import os
-
-import torch
-from apex.optimizers import FusedLAMB as Lamb
-
-from experiments.datamodules.dali_imagenet_fused import (
-    AugmentConfig,
-    DALIImageNetFusedDataModule,
-    MixupConfig,
-)
-from experiments.default_cfg import (
-    AutoResumeConfig,
-    ExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    TrainerConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_multihead_nd import CKConvMultiheadND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.rms_norm import PerHeadRMSNorm, RMSNorm
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter
-from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock
-from nvsubquadratic.networks.vit5_classification import ViT5ClassificationNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-
-
-# ─── Dataset ────────────────────────────────────────────────────────────────────
-INPUT_CHANNELS = 3
-NUM_CLASSES = 1000
-IMAGE_SIZE = 224
-FINAL_IMAGE_SIZE = 224
-IMAGENET_PATH = os.environ.get("IMAGENET_PATH", "/shared/data/image_datasets/imagenet")
-IMAGENET_FOLDER_PATH = os.environ.get("IMAGENET_FOLDER_PATH", "/shared/data/image_datasets/imagenet_folder")
-
-# ─── Model (ViT-5-Small + Multi-Head Hyena, CLS-row) ────────────────────────────
-HIDDEN_DIM = 384
-NUM_BLOCKS = 12
-NUM_HEADS = 6
-HEAD_DIM = HIDDEN_DIM // NUM_HEADS  # 64
-PATCH_SIZE = 16
-LAYER_SCALE_INIT = 1e-4
-DROP_PATH_RATE = 0.05
-MLP_RATIO = 4
-NUM_PATCHES_H = FINAL_IMAGE_SIZE // PATCH_SIZE  # 14
-NUM_PATCHES_W = FINAL_IMAGE_SIZE // PATCH_SIZE  # 14
-NUM_REGISTERS = NUM_PATCHES_W - 1  # 13 — fills the extra row: [CLS, regs, patches] → (H'+1)×W' grid
-
-# ─── Multi-Head Hyena / SIREN kernel hyperparameters ─────────────────────────────
-KERNEL_MLP_HIDDEN_DIM = 32
-KERNEL_NUM_LAYERS = 3
-KERNEL_EMBEDDING_DIM = 32
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-KERNEL_OUT_DIM = NUM_HEADS * HEAD_DIM * HEAD_DIM  # dense kernel per head
-
-# ─── Training recipe ────────────────────────────────────────────────────────────
-BATCH_SIZE = 256
-EPOCHS = 800
-IMAGENET_TRAIN_SIZE = 1_281_167
-EFFECTIVE_BATCH_SIZE = 2048
-ITERS_PER_EPOCH = IMAGENET_TRAIN_SIZE // EFFECTIVE_BATCH_SIZE
-TOTAL_ITERATIONS = EPOCHS * ITERS_PER_EPOCH
-WARMUP_EPOCHS = 5
-WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / EPOCHS
-
-LEARNING_RATE = 4e-3
-WEIGHT_DECAY = 0.05
-GRAD_CLIP = 1.0
-PRECISION = "bf16-mixed"
-
-NUM_WORKERS = 12
-
-
-def get_config() -> ExperimentConfig:
-    """Return the ViT-5-Small + Multi-Head Hyena CLS-row config with Apex FusedLAMB."""
-    config = ExperimentConfig()
-    config.debug = False
-    config.seed = 42
-    config.compile = True
-    config.compile_mode = "max-autotune"
-    # ─── Dataset (fused DALI + local NVMe staging) ────────────────────────
-    config.dataset = LazyConfig(DALIImageNetFusedDataModule)(
-        data_dir=IMAGENET_PATH,
-        imagefolder_dir=IMAGENET_FOLDER_PATH,
-        prefetch_factor=3,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=config.seed,
-        image_size=IMAGE_SIZE,
-        final_image_size=FINAL_IMAGE_SIZE,
-        num_classes=NUM_CLASSES,
-        drop_labels=False,
-        task="classification",
-        mixup_cfg=LazyConfig(MixupConfig)(
-            mixup=0.8,
-            cutmix=1.0,
-            mixup_prob=1.0,
-            mixup_switch_prob=0.5,
-            mixup_mode="batch",
-            smoothing=0.0,
-        ),
-        augment_cfg=LazyConfig(AugmentConfig)(
-            use_three_augment=True,
-            color_jitter=0.3,
-        ),
-        device_id=0,
-        local_staging_dir=f"/scratch/{os.environ.get('USER', 'unknown')}/imagenet_dataset",
-    )
-
-    # ─── Network ────────────────────────────────────────────────────────────
-    hyena_mixer_cfg = LazyConfig(QKVSequenceMixer)(
-        hidden_dim=HIDDEN_DIM,
-        mixer_cfg=LazyConfig(Hyena)(
-            global_conv_cfg=LazyConfig(CKConvMultiheadND)(
-                data_dim=2,
-                hidden_dim=HIDDEN_DIM,
-                num_heads=NUM_HEADS,
-                kernel_cfg=LazyConfig(SIRENKernelND)(
-                    data_dim=2,
-                    out_dim=KERNEL_OUT_DIM,
-                    mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM,
-                    num_layers=KERNEL_NUM_LAYERS,
-                    embedding_dim=KERNEL_EMBEDDING_DIM,
-                    omega_0=KERNEL_OMEGA_0,
-                    L_cache=NUM_PATCHES_H + 1,  # 15: grid is (H'+1)×W' due to the extra CLS row.
-                    use_bias=True,
-                    hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                ),
-                mask_cfg=LazyConfig(torch.nn.Identity)(),
-                grid_type="double",
-                fft_padding="zero",
-            ),
-            short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                in_channels=3 * HIDDEN_DIM,
-                out_channels=3 * HIDDEN_DIM,
-                kernel_size=3,
-                groups=3 * HIDDEN_DIM,
-                padding=1,
-                bias=False,
-            ),
-            gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(),
-            pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            qk_norm_cfg=LazyConfig(PerHeadRMSNorm)(num_heads=NUM_HEADS, head_dim=HEAD_DIM, eps=1e-6),
-            output_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-        ),
-        init_method_in=small_init,
-        init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS),
-    )
-
-    config.net = LazyConfig(ViT5ClassificationNet)(
-        in_channels=INPUT_CHANNELS,
-        num_classes=NUM_CLASSES,
-        hidden_dim=HIDDEN_DIM,
-        num_blocks=NUM_BLOCKS,
-        patch_size=PATCH_SIZE,
-        image_size=FINAL_IMAGE_SIZE,
-        num_registers=NUM_REGISTERS,
-        dropout_rate=0.0,
-        readout="cls",
-        prepend_registers=True,
-        norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-        block_cfg=LazyConfig(ViT5ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)(
-                inner_mixer_cfg=hyena_mixer_cfg,
-                grid_w=NUM_PATCHES_W,
-            ),
-            sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            mlp_cfg=LazyConfig(MLP)(
-                dim=HIDDEN_DIM,
-                activation="gelu",
-                expansion_factor=float(MLP_RATIO),
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS),
-            ),
-            mlp_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            hidden_dim=HIDDEN_DIM,
-            layer_scale_init=LAYER_SCALE_INIT,
-            drop_path_rate=DROP_PATH_RATE,
-        ),
-    )
-
-    # ─── Lightning wrapper ──────────────────────────────────────────────────
-    # NOTE: The ViT-5 reference uses BCE for pretraining, but we observed that
-    # pretraining with BCE leads to significantly lower finetuning accuracy
-    # (~76%) compared to SoftTargetCE (~82%).
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce")
-
-    # ─── Optimizer (Apex FusedLAMB) ─────────────────────────────────────────
-    config.optimizer = LazyConfig(Lamb)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # ─── Training ───────────────────────────────────────────────────────────
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TOTAL_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        precision=PRECISION,
-    )
-
-    config.trainer = TrainerConfig(
-        check_val_every_n_epoch=4,
-        checkpoint_every_n_steps=5000,
-        find_unused_parameters=True,
-    )
-
-    # ─── Scheduler ──────────────────────────────────────────────────────────
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="max",
-    )
-
-    # ─── Wandb ──────────────────────────────────────────────────────────────
-    config.wandb = WandbConfig(
-        job_group="vit5_imagenet_pretrain",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    # ─── Auto-resume ────────────────────────────────────────────────────────
-    config.autoresume = AutoResumeConfig(
-        enabled=False,
-    )
-
-    return config
diff --git a/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_gap_apex.py b/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_gap_apex.py
deleted file mode 100644
index 0b79be92..00000000
--- a/examples/vit5_imagenet/v2/vit5_small_pretrain_multihead_hyena_gap_apex.py
+++ /dev/null
@@ -1,268 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""ViT-5-Small + Multi-Head Hyena + GAP ImageNet-1k — Apex FusedLAMB variant.
-
-Based on v2/vit5_small_pretrain_hyena_gap_apex.py but replaces the depthwise
-CKConvND with CKConvMultiheadND (dense within-head channel mixing).
-
-Key differences from v2/vit5_small_pretrain_hyena_gap_apex.py:
-- CKConvMultiheadND replaces CKConvND: each head performs dense [head_dim x
-  head_dim] channel mixing, enabling cross-channel feature learning within
-  heads while keeping heads isolated.
-- SIREN kernel out_dim = NUM_HEADS * HEAD_DIM * HEAD_DIM (dense kernel per head)
-  instead of HIDDEN_DIM (depthwise).
-- NUM_HEADS=6, HEAD_DIM=64 (consistent with the ViT-5-Small attention variant).
-- PerHeadRMSNorm for QK normalization (each head normalized independently).
-- SiLU gate nonlinearity + output RMSNorm (same as other v2 configs).
-- DALI fused data pipeline with local NVMe staging.
-"""
-
-import os
-
-import torch
-from apex.optimizers import FusedLAMB as Lamb
-
-from experiments.datamodules.dali_imagenet_fused import (
-    AugmentConfig,
-    DALIImageNetFusedDataModule,
-    MixupConfig,
-)
-from experiments.default_cfg import (
-    AutoResumeConfig,
-    ExperimentConfig,
-    SchedulerConfig,
-    TrainConfig,
-    TrainerConfig,
-    WandbConfig,
-)
-from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
-from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig
-from nvsubquadratic.modules.ckconv_multihead_nd import CKConvMultiheadND
-from nvsubquadratic.modules.hyena_nd import Hyena
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.modules.mlp import MLP
-from nvsubquadratic.modules.rms_norm import PerHeadRMSNorm, RMSNorm
-from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
-from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter
-from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock
-from nvsubquadratic.networks.vit5_classification import ViT5ClassificationNet
-from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
-
-
-# ─── Dataset ────────────────────────────────────────────────────────────────────
-INPUT_CHANNELS = 3
-NUM_CLASSES = 1000
-IMAGE_SIZE = 224
-FINAL_IMAGE_SIZE = 224
-IMAGENET_PATH = os.environ.get("IMAGENET_PATH", "/shared/data/image_datasets/imagenet")
-IMAGENET_FOLDER_PATH = os.environ.get("IMAGENET_FOLDER_PATH", "/shared/data/image_datasets/imagenet_folder")
-
-# ─── Model (ViT-5-Small + Multi-Head Hyena, no CLS) ─────────────────────────────
-HIDDEN_DIM = 384
-NUM_BLOCKS = 12
-NUM_HEADS = 6
-HEAD_DIM = HIDDEN_DIM // NUM_HEADS  # 64
-PATCH_SIZE = 16
-NUM_REGISTERS = 0
-LAYER_SCALE_INIT = 1e-4
-DROP_PATH_RATE = 0.05
-MLP_RATIO = 4
-NUM_PATCHES_H = FINAL_IMAGE_SIZE // PATCH_SIZE  # 14
-NUM_PATCHES_W = FINAL_IMAGE_SIZE // PATCH_SIZE  # 14
-
-# ─── Multi-Head Hyena / SIREN kernel hyperparameters ─────────────────────────────
-KERNEL_MLP_HIDDEN_DIM = 32
-KERNEL_NUM_LAYERS = 3
-KERNEL_EMBEDDING_DIM = 32
-KERNEL_OMEGA_0 = 10.0
-KERNEL_HIDDEN_OMEGA_0 = 1.0
-KERNEL_OUT_DIM = NUM_HEADS * HEAD_DIM * HEAD_DIM  # dense kernel per head
-
-# ─── Training recipe ────────────────────────────────────────────────────────────
-BATCH_SIZE = 256
-EPOCHS = 800
-IMAGENET_TRAIN_SIZE = 1_281_167
-EFFECTIVE_BATCH_SIZE = 2048
-ITERS_PER_EPOCH = IMAGENET_TRAIN_SIZE // EFFECTIVE_BATCH_SIZE
-TOTAL_ITERATIONS = EPOCHS * ITERS_PER_EPOCH
-WARMUP_EPOCHS = 5
-WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / EPOCHS
-
-LEARNING_RATE = 4e-3
-WEIGHT_DECAY = 0.05
-GRAD_CLIP = 1.0
-PRECISION = "bf16-mixed"
-
-NUM_WORKERS = 12
-
-
-def get_config() -> ExperimentConfig:
-    """Return the ViT-5-Small + Multi-Head Hyena (GAP) config with Apex FusedLAMB."""
-    config = ExperimentConfig()
-    config.debug = False
-    config.seed = 42
-    config.compile = True
-    config.compile_mode = "max-autotune"
-    # ─── Dataset (fused DALI + local NVMe staging) ────────────────────────
-    config.dataset = LazyConfig(DALIImageNetFusedDataModule)(
-        data_dir=IMAGENET_PATH,
-        imagefolder_dir=IMAGENET_FOLDER_PATH,
-        prefetch_factor=3,
-        batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        seed=config.seed,
-        image_size=IMAGE_SIZE,
-        final_image_size=FINAL_IMAGE_SIZE,
-        num_classes=NUM_CLASSES,
-        drop_labels=False,
-        task="classification",
-        mixup_cfg=LazyConfig(MixupConfig)(
-            mixup=0.8,
-            cutmix=1.0,
-            mixup_prob=1.0,
-            mixup_switch_prob=0.5,
-            mixup_mode="batch",
-            smoothing=0.0,
-        ),
-        augment_cfg=LazyConfig(AugmentConfig)(
-            use_three_augment=True,
-            color_jitter=0.3,
-        ),
-        device_id=0,
-        local_staging_dir=f"/scratch/{os.environ.get('USER', 'unknown')}/imagenet_dataset",
-    )
-
-    # ─── Network ────────────────────────────────────────────────────────────
-    hyena_mixer_cfg = LazyConfig(QKVSequenceMixer)(
-        hidden_dim=HIDDEN_DIM,
-        mixer_cfg=LazyConfig(Hyena)(
-            global_conv_cfg=LazyConfig(CKConvMultiheadND)(
-                data_dim=2,
-                hidden_dim=HIDDEN_DIM,
-                num_heads=NUM_HEADS,
-                kernel_cfg=LazyConfig(SIRENKernelND)(
-                    data_dim=2,
-                    out_dim=KERNEL_OUT_DIM,
-                    mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM,
-                    num_layers=KERNEL_NUM_LAYERS,
-                    embedding_dim=KERNEL_EMBEDDING_DIM,
-                    omega_0=KERNEL_OMEGA_0,
-                    L_cache=NUM_PATCHES_H,
-                    use_bias=True,
-                    hidden_omega_0=KERNEL_HIDDEN_OMEGA_0,
-                ),
-                mask_cfg=LazyConfig(torch.nn.Identity)(),
-                grid_type="double",
-                fft_padding="zero",
-            ),
-            short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
-                in_channels=3 * HIDDEN_DIM,
-                out_channels=3 * HIDDEN_DIM,
-                kernel_size=3,
-                groups=3 * HIDDEN_DIM,
-                padding=1,
-                bias=False,
-            ),
-            gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(),
-            pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            qk_norm_cfg=LazyConfig(PerHeadRMSNorm)(num_heads=NUM_HEADS, head_dim=HEAD_DIM, eps=1e-6),
-            output_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-        ),
-        init_method_in=small_init,
-        init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS),
-    )
-
-    config.net = LazyConfig(ViT5ClassificationNet)(
-        in_channels=INPUT_CHANNELS,
-        num_classes=NUM_CLASSES,
-        hidden_dim=HIDDEN_DIM,
-        num_blocks=NUM_BLOCKS,
-        patch_size=PATCH_SIZE,
-        image_size=FINAL_IMAGE_SIZE,
-        num_registers=NUM_REGISTERS,
-        dropout_rate=0.0,
-        readout="gap",
-        norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-        block_cfg=LazyConfig(ViT5ResidualBlock)(
-            sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)(
-                inner_mixer_cfg=hyena_mixer_cfg,
-                grid_w=NUM_PATCHES_W,
-            ),
-            sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            mlp_cfg=LazyConfig(MLP)(
-                dim=HIDDEN_DIM,
-                activation="gelu",
-                expansion_factor=float(MLP_RATIO),
-                dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0),
-                init_method_in=small_init,
-                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS),
-            ),
-            mlp_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6),
-            hidden_dim=HIDDEN_DIM,
-            layer_scale_init=LAYER_SCALE_INIT,
-            drop_path_rate=DROP_PATH_RATE,
-        ),
-    )
-
-    # ─── Lightning wrapper ──────────────────────────────────────────────────
-    # NOTE: The ViT-5 reference uses BCE for pretraining, but we observed that
-    # pretraining with BCE leads to significantly lower finetuning accuracy
-    # (~76%) compared to SoftTargetCE (~82%).
-    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce")
-
-    # ─── Optimizer (Apex FusedLAMB) ─────────────────────────────────────────
-    config.optimizer = LazyConfig(Lamb)(
-        params=PLACEHOLDER,
-        lr=LEARNING_RATE,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    # ─── Training ───────────────────────────────────────────────────────────
-    config.train = TrainConfig(
-        batch_size="${dataset.batch_size}",
-        iterations=TOTAL_ITERATIONS,
-        grad_clip=GRAD_CLIP,
-        precision=PRECISION,
-    )
-
-    config.trainer = TrainerConfig(
-        check_val_every_n_epoch=4,
-        checkpoint_every_n_steps=5000,
-        find_unused_parameters=True,
-    )
-
-    # ─── Scheduler ──────────────────────────────────────────────────────────
-    config.scheduler = SchedulerConfig(
-        name="cosine",
-        warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE,
-        total_iterations="${train.iterations}",
-        mode="max",
-    )
-
-    # ─── Wandb ──────────────────────────────────────────────────────────────
-    config.wandb = WandbConfig(
-        job_group="vit5_imagenet_pretrain",
-        entity="implicit-long-convs",
-        project="nvsubquadratic",
-    )
-
-    # ─── Auto-resume ────────────────────────────────────────────────────────
-    config.autoresume = AutoResumeConfig(
-        enabled=False,
-    )
-
-    return config
diff --git a/examples/well/v1/euler_multi_quadrants_periodicBC/cfg_hyena.py b/examples/well/v1/euler_multi_quadrants_periodicBC/cfg_hyena.py
index 0a96c797..1f325c09 100644
--- a/examples/well/v1/euler_multi_quadrants_periodicBC/cfg_hyena.py
+++ b/examples/well/v1/euler_multi_quadrants_periodicBC/cfg_hyena.py
@@ -106,7 +106,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v1/gray_scott_reaction_diffusion/cfg_hyena_gaussian_mask.py b/examples/well/v1/gray_scott_reaction_diffusion/cfg_hyena_gaussian_mask.py
index 1b6b8d4d..7fccacf6 100644
--- a/examples/well/v1/gray_scott_reaction_diffusion/cfg_hyena_gaussian_mask.py
+++ b/examples/well/v1/gray_scott_reaction_diffusion/cfg_hyena_gaussian_mask.py
@@ -113,7 +113,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/MHD_64/hyena.py b/examples/well/v2/MHD_64/hyena.py
index a59bb3cf..c356cf17 100644
--- a/examples/well/v2/MHD_64/hyena.py
+++ b/examples/well/v2/MHD_64/hyena.py
@@ -110,7 +110,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/acoustic_scattering_maze/hyena.py b/examples/well/v2/acoustic_scattering_maze/hyena.py
index 27229676..04566551 100644
--- a/examples/well/v2/acoustic_scattering_maze/hyena.py
+++ b/examples/well/v2/acoustic_scattering_maze/hyena.py
@@ -99,7 +99,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/active_matter/compile_norm_benchmarks/hyena_cf_norm.py b/examples/well/v2/active_matter/compile_norm_benchmarks/hyena_cf_norm.py
index 1c626406..075be481 100644
--- a/examples/well/v2/active_matter/compile_norm_benchmarks/hyena_cf_norm.py
+++ b/examples/well/v2/active_matter/compile_norm_benchmarks/hyena_cf_norm.py
@@ -104,7 +104,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/active_matter/hyena.py b/examples/well/v2/active_matter/hyena.py
index 19de0677..70494987 100644
--- a/examples/well/v2/active_matter/hyena.py
+++ b/examples/well/v2/active_matter/hyena.py
@@ -99,7 +99,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/gray_scott_reaction_diffusion/hyena.py b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py
index 1f4283ef..5dfd2665 100644
--- a/examples/well/v2/gray_scott_reaction_diffusion/hyena.py
+++ b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py
@@ -99,7 +99,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/rayleigh_benard/hyena.py b/examples/well/v2/rayleigh_benard/hyena.py
index 725d612e..e1179714 100644
--- a/examples/well/v2/rayleigh_benard/hyena.py
+++ b/examples/well/v2/rayleigh_benard/hyena.py
@@ -131,7 +131,6 @@ def get_config() -> ExperimentConfig:
                         # Mixed boundary conditions: list of per-axis padding
                         # modes (one per spatial axis, in order).
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/examples/well/v2/supernova_explosion_64/hyena.py b/examples/well/v2/supernova_explosion_64/hyena.py
index 101e559b..a19d8ed4 100644
--- a/examples/well/v2/supernova_explosion_64/hyena.py
+++ b/examples/well/v2/supernova_explosion_64/hyena.py
@@ -99,7 +99,6 @@ def get_config() -> ExperimentConfig:
                         data_dim=DATA_DIM,
                         hidden_dim=NUM_HIDDEN_CHANNELS,
                         fft_padding=FFT_PADDING,
-                        use_fp16_fft=False,
                         kernel_cfg=LazyConfig(SIRENKernelND)(
                             data_dim=DATA_DIM,
                             out_dim=NUM_HIDDEN_CHANNELS,
diff --git a/experiments/callbacks/wandb_cache_cleanup.py b/experiments/callbacks/wandb_cache_cleanup.py
index 7949d6d7..45670090 100644
--- a/experiments/callbacks/wandb_cache_cleanup.py
+++ b/experiments/callbacks/wandb_cache_cleanup.py
@@ -13,10 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: Add licence header
-
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Callback to periodically run `wandb artifact cache cleanup ${X}GB` to cap local cache size."""
 
 import subprocess
diff --git a/experiments/datamodules/ucf101.py b/experiments/datamodules/ucf101.py
deleted file mode 100644
index 03f9255c..00000000
--- a/experiments/datamodules/ucf101.py
+++ /dev/null
@@ -1,446 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""UCF101 video-classification datamodule for PyTorch Lightning.
-
-Provides :class:`UCF101DataModule`, which wraps
-:class:`torchvision.datasets.UCF101`.  :meth:`prepare_data` automatically
-downloads and extracts both the video archive and the train/test annotation
-split files if they are not already present under ``data_dir`` /
-``annotation_dir``.
-
-**Two output modes** (selected via ``data_type``):
-
-- ``"video"`` — returns clips shaped ``[B, T, H, W, C]`` (channels-last,
-  float32).  Each clip spans ``frames_per_clip`` frames sampled at
-  ``frame_rate`` fps with ``step_between_clips`` temporal stride.
-- ``"sequence"`` — flattens the temporal and spatial axes into a 1-D token
-  sequence: ``[B, T*H*W, C]``.  Suitable for models that treat video as a
-  flat sequence of patches.
-
-**Deterministic workers**
-
-When ``use_deterministic_worker_init=True`` each DataLoader worker is seeded
-with ``base_seed + worker_id`` via :func:`deterministic_worker_init_fn`,
-ensuring reproducible random clip sampling across runs and restarts.
-"""
-
-import os
-import random
-from pathlib import Path
-from typing import Literal, Optional
-
-import numpy as np
-import pytorch_lightning as pl
-import torch
-import torchvision.transforms._transforms_video as TV
-from einops import rearrange
-from torch.utils.data import DataLoader, random_split
-from torchvision import datasets, transforms
-
-
-_BASE_SEED = 0
-
-
-def set_base_seed(seed):
-    """Set the base seed for worker initialization."""
-    global _BASE_SEED
-    _BASE_SEED = seed
-
-
-def deterministic_worker_init_fn(worker_id: int):
-    """Initialize the worker with a deterministic seed derived from base_seed and worker_id.
-
-    Each worker gets a unique but deterministic seed: base_seed + worker_id
-    """
-    global _BASE_SEED
-    seed = _BASE_SEED + worker_id
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-class UCF101DataModule(pl.LightningDataModule):
-    """UCF101 Lightning data module.
-
-    Expects videos under data_dir and annotation split files under annotation_dir.
-    prepare_data() will attempt to download/extract both if missing.
-    """
-
-    def __init__(
-        self,
-        data_dir: str,
-        batch_size: int,
-        data_type: Literal["sequence", "video"],
-        num_workers: int,
-        pin_memory: bool,
-        use_deterministic_worker_init: bool,
-        seed: int,
-        frames_per_clip: int = 16,
-        step_between_clips: int = 1,
-        frame_size: Optional[tuple[int, int]] = (256, 256),
-        val_split_fraction: float = 0.1,
-        split_fold: int = 1,
-    ):
-        """Initialize the UCF101DataModule.
-
-        Args:
-            data_dir: Directory to save the data
-            batch_size: Batch size
-            data_type: Type of data. Can be "sequence" or "video".
-            num_workers: Number of workers
-            pin_memory: Whether to pin memory
-            use_deterministic_worker_init: Whether to use deterministic worker initialization
-            seed: Seed for the data
-            frames_per_clip: Number of frames per clip
-            step_between_clips: Step between clips
-            frame_size: Size of the frames
-            val_split_fraction: Fraction of the data to use for validation
-            split_fold: Fold to use for the data. Can be 1, 2 or 3.
-        """
-        assert data_type in ["video"], f"data_type must be 'video', got {data_type}"
-        assert split_fold in [1, 2, 3], f"split_fold must be 1, 2 or 3, got {split_fold}"
-
-        super().__init__()
-
-        self.data_dir = data_dir
-        self.annotation_dir = Path(data_dir) / "annotations"
-        self.videos_dir = Path(data_dir) / "videos"
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-        self.seed = seed
-        self.frames_per_clip = frames_per_clip
-        self.step_between_clips = step_between_clips
-        self.frame_size = frame_size
-        self.val_split_fraction = val_split_fraction
-        self.split_fold = split_fold
-
-        self.generator = torch.Generator().manual_seed(seed)
-        self.worker_init_fn = deterministic_worker_init_fn if use_deterministic_worker_init else None
-
-        self.input_channels = 3
-        self.output_channels = 101
-
-        self.data_type = data_type
-
-        if self.frame_size is not None:
-            self.frame_transform = transforms.Compose(
-                [
-                    transforms.Resize(self.frame_size),
-                ]
-            )
-        else:
-            self.frame_transform = None
-
-        self.video_transform = transforms.Compose(
-            [
-                TV.ToTensorVideo(),
-                TV.NormalizeVideo(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            ]
-        )
-
-        self.train_dataset = None
-        self.val_dataset = None
-        self.test_dataset = None
-
-    def prepare_data(self):
-        """Download and prepare UCF101 videos and annotation splits in-place if missing."""
-        import shutil
-        import tempfile
-        import urllib.request
-        import zipfile
-        from pathlib import Path
-
-        try:
-            from tqdm import tqdm as _tqdm  # type: ignore
-        except Exception:
-            _tqdm = None
-
-        def _download_with_progress(urls, filename: str, desc: str):
-            import ssl
-
-            if isinstance(urls, str):
-                urls = [urls]
-
-            def _stream(url: str, verify: bool):
-                ctx = None if verify else ssl._create_unverified_context()
-                opener = urllib.request.build_opener(
-                    urllib.request.HTTPSHandler(context=ctx) if ctx is not None else urllib.request.HTTPSHandler()
-                )
-                req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
-                with opener.open(req) as resp, open(filename, "wb") as out_f:
-                    total = resp.length
-                    if total is None:
-                        total = int(resp.headers.get("Content-Length", 0)) or None
-                    if _tqdm is None:
-                        while True:
-                            chunk = resp.read(1024 * 1024)
-                            if not chunk:
-                                break
-                            out_f.write(chunk)
-                    else:
-                        bar = _tqdm(total=total, unit="B", unit_scale=True, desc=desc)
-                        while True:
-                            chunk = resp.read(1024 * 1024)
-                            if not chunk:
-                                break
-                            out_f.write(chunk)
-                            bar.update(len(chunk))
-                        bar.close()
-
-            last_error = None
-            for url in urls:
-                try:
-                    _stream(url, verify=True)
-                    return
-                except Exception as e1:
-                    last_error = e1
-                    try:
-                        _stream(url, verify=False)
-                        return
-                    except Exception as e2:
-                        last_error = e2
-                        continue
-            raise RuntimeError(f"Failed downloading from provided URLs. Last error: {last_error}")
-
-        def _extract_zip_with_progress(zip_path: Path, dest_dir: Path, desc: str):
-            if _tqdm is None:
-                with zipfile.ZipFile(zip_path, "r") as zf:
-                    zf.extractall(dest_dir)
-                return
-            with zipfile.ZipFile(zip_path, "r") as zf:
-                members = zf.infolist()
-                bar = _tqdm(total=len(members), unit="file", desc=desc)
-                for m in members:
-                    zf.extract(m, dest_dir)
-                    bar.update(1)
-                bar.close()
-
-        videos_root = Path(self.videos_dir)
-        ann_root = Path(self.annotation_dir)
-        videos_root.mkdir(parents=True, exist_ok=True)
-        ann_root.mkdir(parents=True, exist_ok=True)
-
-        def has_videos(root: Path) -> bool:
-            video_exts = (".avi", ".mp4", ".webm", ".mkv")
-            # Any video file anywhere under root
-            for ext in video_exts:
-                if next(root.rglob(f"*{ext}"), None) is not None:
-                    return True
-            return False
-
-        def organize_ucf101_structure(root: Path):
-            import re
-
-            video_exts = (".avi", ".mp4", ".webm", ".mkv")
-            # Move top-level videos like v_ApplyEyeMakeup_g01_c01.avi into class dirs ApplyEyeMakeup/
-            top_level_files = [p for p in root.iterdir() if p.is_file() and p.suffix.lower() in video_exts]
-            iterator = (
-                _tqdm(top_level_files, desc="Organizing videos by class") if _tqdm is not None else top_level_files
-            )
-            for p in iterator:
-                m = re.match(r"^v_([^_]+)_", p.name)
-                if not m:
-                    # Skip files that don't follow the expected naming
-                    continue
-                class_name = m.group(1)
-                class_dir = root / class_name
-                class_dir.mkdir(parents=True, exist_ok=True)
-                shutil.move(p.as_posix(), (class_dir / p.name).as_posix())
-
-        def _pair_exists_for_fold(root: Path, fold: int) -> bool:
-            tfile = f"trainlist{fold:02d}.txt"
-            vfile = f"testlist{fold:02d}.txt"
-            return (root / tfile).is_file() and (root / vfile).is_file()
-
-        def _detect_available_fold(root: Path) -> Optional[int]:
-            for f in (1, 2, 3):
-                if _pair_exists_for_fold(root, f):
-                    return f
-            return None
-
-        def has_annotations(root: Path) -> bool:
-            # Need classInd and at least one fold pair
-            if not (root / "classInd.txt").is_file():
-                return False
-            return _detect_available_fold(root) is not None
-
-        if not has_videos(videos_root):
-            mirror_zip = "https://storage.googleapis.com/thumos14_files/UCF101_videos.zip"
-            with tempfile.TemporaryDirectory() as td:
-                tmp_zip = Path(td) / "UCF101_videos.zip"
-                _download_with_progress(mirror_zip, tmp_zip.as_posix(), desc="Downloading UCF101 videos")
-                _extract_zip_with_progress(tmp_zip, videos_root, desc=f"Extracting videos in {videos_root}")
-            # Flatten common nested folders created by the archive
-            for nested in [videos_root / "UCF-101", videos_root / "UCF101"]:
-                if nested.exists() and nested.is_dir():
-                    items = list(nested.iterdir())
-                    iterator = _tqdm(items, desc=f"Flattening {nested.name}") if _tqdm is not None else items
-                    for p in iterator:
-                        shutil.move(p.as_posix(), (videos_root / p.name).as_posix())
-                    try:
-                        nested.rmdir()
-                    except OSError:
-                        pass
-            # Ensure class-separated structure
-            organize_ucf101_structure(videos_root)
-            if not has_videos(videos_root):
-                raise RuntimeError("UCF101 videos not found after extraction.")
-
-        if not has_annotations(ann_root):
-            splits_zip = "https://www.crcv.ucf.edu/data/UCF101/UCF101TrainTestSplits-RecognitionTask.zip"
-            with tempfile.TemporaryDirectory() as td:
-                tmp_zip = Path(td) / "UCF101TrainTestSplits-RecognitionTask.zip"
-                _download_with_progress(splits_zip, tmp_zip.as_posix(), desc="Downloading UCF101 splits")
-                _extract_zip_with_progress(tmp_zip, ann_root, desc=f"Extracting annotations in {ann_root}")
-
-            nested = ann_root / "ucfTrainTestlist"
-            if nested.exists() and nested.is_dir():
-                items = list(nested.iterdir())
-                iterator = _tqdm(items, desc="Organizing annotations") if _tqdm is not None else items
-                for p in iterator:
-                    if p.is_file():
-                        shutil.move(p.as_posix(), (ann_root / p.name).as_posix())
-                try:
-                    nested.rmdir()
-                except OSError:
-                    pass
-            if not has_annotations(ann_root):
-                raise RuntimeError("UCF101 annotation files not found after extraction.")
-
-        # Decide on fold
-        if self.split_fold is None:
-            detected = _detect_available_fold(ann_root)
-            if detected is None:
-                raise RuntimeError(
-                    "Could not detect any available fold (trainlistXX/testlistXX). Please provide fold files."
-                )
-            self.split_fold = detected
-        else:
-            # Validate explicit fold exists
-            if not _pair_exists_for_fold(ann_root, self.split_fold):
-                raise RuntimeError(
-                    f"Requested split_fold={self.split_fold} but corresponding files are missing in {ann_root}."
-                )
-
-        self.annotation_dir = str(ann_root)
-
-    def setup(self, stage=None):
-        """Function to setup the datamodule."""
-        # we set up only relevant datamodules when stage is specified
-        if stage == "fit" or stage is None:
-            full_train = datasets.UCF101(
-                root=self.videos_dir,
-                annotation_path=self.annotation_dir,
-                frames_per_clip=self.frames_per_clip,
-                step_between_clips=self.step_between_clips,
-                fold=self.split_fold,
-                train=True,
-                num_workers=self.num_workers,
-                transform=self.video_transform,
-                frame_rate=None,
-            )
-            num_full = len(full_train)
-            num_val = max(1, int(num_full * self.val_split_fraction))
-            num_train = num_full - num_val
-            self.train_dataset, self.val_dataset = random_split(
-                full_train, [num_train, num_val], generator=self.generator
-            )
-
-        if stage == "test" or stage is None:
-            self.test_dataset = datasets.UCF101(
-                root=self.videos_dir,
-                annotation_path=self.annotation_dir,
-                frames_per_clip=self.frames_per_clip,
-                step_between_clips=self.step_between_clips,
-                fold=self.split_fold,
-                train=False,
-                num_workers=self.num_workers,
-                transform=self.video_transform,
-                frame_rate=None,
-            )
-
-    def _build_loader(self, dataset, shuffle: bool, drop_last: bool = False):
-        """Function to create dataloaders given a dataset and a few arguments.
-
-        Reused for train, val and test dataloaders.
-
-        Args:
-            dataset: Dataset to create a dataloader for.
-            shuffle: Whether to shuffle the dataset.
-            drop_last: Whether to drop the last batch if it's not complete.
-
-        Returns:
-            DataLoader: DataLoader for the dataset.
-        """
-        return DataLoader(
-            dataset,
-            batch_size=self.batch_size,
-            shuffle=shuffle,
-            num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
-            drop_last=drop_last,
-            # worker_init_fn=self.worker_init_fn,  # No longer needed with pl.seed_everything(workers=True)
-            generator=self.generator,
-            persistent_workers=self.num_workers > 0,
-            collate_fn=self._ignore_audio_collate_fn,
-        )
-
-    def train_dataloader(self):
-        """Function to create the train dataloader."""
-        return self._build_loader(self.train_dataset, shuffle=True, drop_last=True)
-
-    def val_dataloader(self):
-        """Function to create the validation dataloader."""
-        return self._build_loader(self.val_dataset, shuffle=False)
-
-    def test_dataloader(self):
-        """Function to create the test dataloader."""
-        return self._build_loader(self.test_dataset, shuffle=False)
-
-    def _ignore_audio_collate_fn(self, samples):
-        """Collate that ignores audio and stacks videos and labels.
-
-        Samples are shaped as (video, audio, label).
-
-        Returns (videos [batch_size, ...], labels [batch_size]).
-        """
-        # Keep only valid (video, label)
-        videos, _audios, labels = zip(*samples)
-        return torch.stack(list(videos), dim=0), torch.as_tensor(labels, dtype=torch.long)
-
-    @torch.no_grad()
-    def on_before_batch_transfer(self, batch, dataloader_idx) -> dict[str, torch.Tensor]:
-        """Function to reshape (if needed) the frames before batch transfer.
-
-        Returns:
-            dict[str, torch.Tensor]: A dictionary containing the input, label and condition.
-                Keys: "input", "label" and "condition".
-        """
-        video, label = batch
-        # Must reshape the frames
-        if self.frame_transform is not None:
-            batch_size, channels, temporal_length, _height, _width = video.shape
-            video = rearrange(video, "b c t h w -> (b c t) h w")
-            # Apply frame transform
-            video = self.frame_transform(video)
-            # Reconstruct to original shape
-            video = rearrange(video, "(b c t) h w -> b t h w c", b=batch_size, t=temporal_length, c=channels)
-        # Return the dictionary
-        return {"input": video, "label": label, "condition": None}
diff --git a/experiments/default_cfg.py b/experiments/default_cfg.py
index d5d103c0..b6416aa5 100644
--- a/experiments/default_cfg.py
+++ b/experiments/default_cfg.py
@@ -13,12 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Typed configuration dataclasses for nvSubquadratic experiments.
 
-Every training run is fully described by an :class:`ExperimentConfig` (or one
-of its task-specific subclasses such as :class:`DiffusionExperimentConfig`).
+Every training run is fully described by an :class:`ExperimentConfig`.
 These are plain Python :mod:`dataclasses` so they can be instantiated directly
 in a Python config file, serialised via OmegaConf, and overridden at the CLI.
 
@@ -40,8 +37,6 @@
 
 :data:`PLACEHOLDER` is re-exported from :mod:`nvsubquadratic.lazy_config` for
 convenience in config files.
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 from dataclasses import dataclass, field
@@ -216,56 +211,3 @@ class ExperimentConfig:
     start_from_checkpoint: StartFromCheckpointConfig = field(default_factory=StartFromCheckpointConfig)
     autoresume: AutoResumeConfig = field(default_factory=AutoResumeConfig)
     callbacks: list[LazyConfig] = field(default_factory=list)
-
-
-@dataclass
-class DiffusionConfig:
-    """Diffusion configuration for JiT-style continuous-time flow matching."""
-
-    num_train_timesteps: int = 1_000
-    time_embed_dim: Optional[int] = None
-    max_period: float = 10_000.0
-
-    # Noise scale for initial sample (1.0 for 256px, 2.0 for 512px per JiT).
-    noise_scale: float = 1.0
-
-    # Logit-normal time sampling parameters (JiT defaults).
-    p_mean: float = -0.8
-    p_std: float = 0.8
-
-    num_inference_steps: int = 50
-    num_samples: int = 25
-    log_samples: bool = True
-
-    ema_enabled: bool = True
-    ema_decay: float = 0.9995
-    ema_update_every: int = 1
-    ema_warmup_steps: int = 5_000
-
-    # Classifier-free guidance settings, enabled by default.
-    use_classifier_free_guidance: bool = True
-    guidance_scale: float = 3.5
-    condition_dropout_prob: float = 0.1
-    num_classes: Optional[int] = 1000
-
-    # CFG time interval: apply guidance only within [start, end].
-    cfg_interval_start: float = 0.1
-    cfg_interval_end: float = 1.0
-
-    # Online FID evaluation (JiT-style).
-    fid_online_jit: bool = False
-    fid_stats_file: str = ""
-    fid_num_samples: int = 50_000
-    fid_interval: int = 100
-    fid_batch_size: int = 512
-    fid_num_inference_steps: Optional[int] = None
-
-
-@dataclass
-class DiffusionExperimentConfig(ExperimentConfig):
-    """Experiment configuration for diffusion runs."""
-
-    # Override debug mode.
-    debug: bool = False
-
-    diffusion: DiffusionConfig = field(default_factory=DiffusionConfig)
diff --git a/experiments/lightning_wrappers/base_lightning_wrapper.py b/experiments/lightning_wrappers/base_lightning_wrapper.py
index 49bd5139..76440602 100644
--- a/experiments/lightning_wrappers/base_lightning_wrapper.py
+++ b/experiments/lightning_wrappers/base_lightning_wrapper.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Base PyTorch Lightning wrapper for all nvSubquadratic experiments.
 
 Provides :class:`LightningWrapperBase`, the shared superclass for all task-specific
@@ -36,8 +34,6 @@
 Task-specific forward passes, losses, and metrics live in the subclasses
 (:class:`~experiments.lightning_wrappers.classification_wrapper.ClassificationWrapper`,
 :class:`~experiments.lightning_wrappers.regression_wrapper.RegressionWrapper`, etc.).
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 import warnings
diff --git a/experiments/lightning_wrappers/classification_wrapper.py b/experiments/lightning_wrappers/classification_wrapper.py
index 808e543c..d0c6f81b 100644
--- a/experiments/lightning_wrappers/classification_wrapper.py
+++ b/experiments/lightning_wrappers/classification_wrapper.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Lightning wrapper for image classification tasks.
 
 Provides :class:`ClassificationWrapper` and :class:`SoftTargetCrossEntropy`.
@@ -26,8 +24,6 @@
   ``-sum(target * log_softmax(logits))``.  Use with Mixup/CutMix (DeiT III recipe).
 - ``"bce"`` — :class:`torch.nn.BCEWithLogitsLoss` with binarised multi-hot targets.
   Matches the ViT-5 / DeiT III pre-training recipe (``--bce-loss``).
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 import torch
diff --git a/experiments/lightning_wrappers/diffusion_wrapper.py b/experiments/lightning_wrappers/diffusion_wrapper.py
deleted file mode 100644
index 9dcdebf6..00000000
--- a/experiments/lightning_wrappers/diffusion_wrapper.py
+++ /dev/null
@@ -1,752 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
-"""Lightning wrapper for continuous-time diffusion (JiT-style) experiments.
-
-Provides :class:`DiffusionWrapper`, which implements the flow-matching / JiT
-training loop: noises inputs according to a time-dependent schedule, trains a
-denoiser network, and generates samples via an ODE integrator at inference time.
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
-"""
-
-import copy
-import math
-import os
-import shutil
-from typing import Optional
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-import wandb
-from torchvision.utils import make_grid, save_image
-from tqdm.auto import tqdm
-
-from experiments.default_cfg import DiffusionExperimentConfig
-from experiments.lightning_wrappers.base_lightning_wrapper import LightningWrapperBase
-
-
-class DiffusionWrapper(LightningWrapperBase):
-    """Lightning module for JiT-style continuous-time diffusion."""
-
-    def __init__(
-        self,
-        network: torch.nn.Module,
-        cfg: DiffusionExperimentConfig,
-    ) -> None:
-        """Initialize the DiffusionWrapper.
-
-        Args:
-            network: The neural network to be used as the denoiser model.
-            cfg: The diffusion experiment configuration.
-        """
-        super().__init__(network=network, cfg=cfg)
-
-        if not isinstance(cfg, DiffusionExperimentConfig):
-            raise TypeError("DiffusionWrapper requires cfg to be a DiffusionExperimentConfig instance.")
-        diffusion_cfg = cfg.diffusion
-        if diffusion_cfg is None:
-            raise ValueError("DiffusionWrapper requires cfg.diffusion to be provided.")
-
-        # JiT continuous-time diffusion setup.
-        self.num_train_timesteps = int(diffusion_cfg.num_train_timesteps)
-        self.noise_scale = float(diffusion_cfg.noise_scale)
-        self.p_mean = float(diffusion_cfg.p_mean)
-        self.p_std = float(diffusion_cfg.p_std)
-        self.cfg_interval_start = float(diffusion_cfg.cfg_interval_start)
-        self.cfg_interval_end = float(diffusion_cfg.cfg_interval_end)
-
-        hidden_dim = getattr(network, "hidden_dim", None)
-        if hidden_dim is None:
-            raise AttributeError("DiffusionWrapper requires the network to expose a 'hidden_dim' attribute.")
-
-        schedule_time_embed = diffusion_cfg.time_embed_dim
-        timestep_dim = int(schedule_time_embed) if schedule_time_embed is not None else hidden_dim * 2
-        if timestep_dim % 2 != 0:
-            timestep_dim += 1
-        self.timestep_dim = timestep_dim
-        self.max_period = float(diffusion_cfg.max_period)
-
-        # Time conditioning pipeline: sinusoidal embedding followed by an MLP so the backbone always
-        # receives conditioning in its native hidden dimension.
-        self.time_mlp = torch.nn.Sequential(
-            torch.nn.Linear(self.timestep_dim, hidden_dim * 2),
-            torch.nn.SiLU(),
-            torch.nn.Linear(hidden_dim * 2, hidden_dim),
-        )
-
-        # Sampling and logging configuration.
-        self.example_input_shape: Optional[torch.Size] = None
-        self.default_inference_steps = int(diffusion_cfg.num_inference_steps)
-        self.log_samples = bool(diffusion_cfg.log_samples)
-        self.num_generated_samples = int(diffusion_cfg.num_samples)
-
-        # Optional online FID evaluation (JiT Style).
-        self.fid_online_jit = bool(getattr(diffusion_cfg, "fid_online_jit", False))
-        self.fid_stats_file = getattr(diffusion_cfg, "fid_stats_file", None)
-        self.fid_num_samples = int(getattr(diffusion_cfg, "fid_num_samples", 50000))
-        self.fid_interval = int(getattr(diffusion_cfg, "fid_interval", 50))
-        self.fid_batch_size = int(getattr(diffusion_cfg, "fid_batch_size", 64))
-        fid_steps_cfg = getattr(diffusion_cfg, "fid_num_inference_steps", None)
-        self.fid_num_inference_steps = (
-            int(fid_steps_cfg) if fid_steps_cfg is not None else self.default_inference_steps
-        )
-
-        # Classifier-free guidance (CFG) settings.
-        # Configurable to support both conditional and unconditional training within the same wrapper.
-        self.class_conditioning = diffusion_cfg.num_classes is not None
-        self.cfg_enabled = bool(diffusion_cfg.use_classifier_free_guidance) and self.class_conditioning
-        self.guidance_scale = float(diffusion_cfg.guidance_scale)
-
-        # Dropout probability for the conditioning signal during training.
-        self.condition_dropout_prob = float(diffusion_cfg.condition_dropout_prob) if self.class_conditioning else 0.0
-        self.num_classes: Optional[int] = (
-            int(diffusion_cfg.num_classes) if diffusion_cfg.num_classes is not None else None
-        )
-
-        if diffusion_cfg.use_classifier_free_guidance and not self.class_conditioning:
-            raise ValueError(
-                "Classifier-free guidance requires 'diffusion.num_classes' to be set so labels can be embedded."
-            )
-
-        if self.class_conditioning:
-            if self.num_classes is None or self.num_classes <= 0:
-                raise ValueError("diffusion.num_classes must be a positive integer when enabling class conditioning.")
-            # We dedicate one additional embedding slot to represent the unconditional branch.
-            # This embedding is learned during training (initialized to zero).
-            self.null_label_index = self.num_classes
-            self.label_embed = torch.nn.Embedding(self.num_classes + 1, hidden_dim)
-            torch.nn.init.normal_(self.label_embed.weight, mean=0.0, std=0.02)
-            with torch.no_grad():
-                # Initialize the unconditional embedding to zero.
-                self.label_embed.weight[self.null_label_index].zero_()
-        else:
-            self.null_label_index = None
-            self.label_embed = None
-
-        # Exponential Moving Average (EMA) tracking.
-        self.ema_enabled = bool(diffusion_cfg.ema_enabled)
-        self.ema_decay = float(diffusion_cfg.ema_decay)
-        self.ema_update_every = int(diffusion_cfg.ema_update_every)
-        self.ema_warmup_steps = int(diffusion_cfg.ema_warmup_steps)
-        self._ema_model: Optional[torch.nn.Module] = None
-        self._ema_has_been_updated = False
-        if self.ema_enabled:
-            # Create a shadow copy of the network that does not receive gradients.
-            self._ema_model = copy.deepcopy(self.network)
-            for p in self._ema_model.parameters():
-                p.detach_()
-                p.requires_grad_(False)
-
-        # Allow Hugging Face-backed models to register themselves for callbacks.
-        register_fn = getattr(network, "hf_register_diffusion_wrapper", None)
-        if callable(register_fn):
-            register_fn(self)
-
-    @staticmethod
-    def _channels_last_to_first(tensor: torch.Tensor) -> torch.Tensor:
-        """Convert an image tensor from channels-last to channels-first layout."""
-        return torch.moveaxis(tensor, -1, 1).contiguous()
-
-    @staticmethod
-    def _channels_first_to_last(tensor: torch.Tensor) -> torch.Tensor:
-        """Convert an image tensor from channels-first to channels-last layout."""
-        return torch.moveaxis(tensor, 1, -1).contiguous()
-
-    def _timestep_embedding(self, timesteps: torch.Tensor) -> torch.Tensor:
-        """Create sinusoidal timestep embeddings.
-
-        Standard sinusoidal embedding with configurable dimensionality, identical to the previous
-        implementation so we preserve conditioning behaviour.
-
-        Args:
-            timesteps: a 1-D Tensor of N indices, one per batch element.
-        """
-        device = timesteps.device
-        if timesteps.dtype in (torch.int8, torch.int16, torch.int32, torch.int64):
-            working = timesteps.to(torch.float32)
-            embed_dtype = torch.float32
-        else:
-            working = timesteps
-            embed_dtype = timesteps.dtype
-
-        half_dim = self.timestep_dim // 2
-        exponent = torch.arange(half_dim, device=device, dtype=torch.float32)
-        exponent = -math.log(self.max_period) * exponent / max(half_dim - 1, 1)
-        freqs = torch.exp(exponent).to(working.dtype)
-        args = working.view(-1, 1) * freqs.view(1, -1)
-        embedding = torch.cat([torch.sin(args), torch.cos(args)], dim=-1).to(embed_dtype)
-        if embedding.shape[-1] < self.timestep_dim:
-            embedding = F.pad(embedding, (0, self.timestep_dim - embedding.shape[-1]))
-        return embedding
-
-    def _noiselevel_embedding(self, timesteps: torch.Tensor) -> torch.Tensor:
-        """Embed continuous timesteps t ∈ [0, 1] into the model hidden dimension."""
-        target_dtype = self.time_mlp[0].weight.dtype
-        # Scale t into the sinusoidal range expected by time_mlp (matching JiT continuous-time).
-        scaled_t = timesteps.to(target_dtype).clamp_(0.0, 1.0) * float(self.num_train_timesteps)
-        return self.time_mlp(self._timestep_embedding(scaled_t))
-
-    def _condition_from_timesteps(
-        self,
-        timesteps: torch.Tensor,
-        *,
-        labels: Optional[torch.Tensor] = None,
-        unconditional: bool = False,
-        dropout_mask: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """Return the conditioning vector and raw label embedding for a batch of timesteps.
-
-        Returns a tuple ``(combined_condition, label_emb)`` where ``label_emb`` is the raw class
-        embedding before adding the time embedding. Pass ``label_emb`` to the network as
-        ``"class_emb"`` so that in-context tokens receive a pure class signal (no time information),
-        matching the JiT reference. ``label_emb`` is ``None`` when class conditioning is disabled.
-
-        Args:
-            timesteps: Diffusion timestep indices sampled for each element in the batch.
-            labels: Optional class labels associated with the batch. Required when class conditioning
-                is enabled because each label switches us to a different guidance direction.
-            unconditional: When ``True`` we force the method to emit the unconditional embedding by
-                routing all samples to the extra "null" slot in ``self.label_embed``.
-            dropout_mask: Boolean mask selecting which labels should be dropped for classifier-free
-                guidance training. Elements set to ``True`` fall back to the unconditional embedding.
-        """
-        # Step 1: obtain the base time embedding as usual.
-        time_emb = self._noiselevel_embedding(timesteps)
-
-        # Without class conditioning the timestep embedding is the entire conditioning signal.
-        if self.label_embed is None:
-            return time_emb, None
-
-        # If we do expect labels, make sure the caller provided them.
-        if labels is None:
-            if unconditional:
-                # During sampling we sometimes request unconditional guidance without providing the
-                # original labels, so we create a tensor filled with the null label index on demand.
-                labels_to_embed = torch.full_like(timesteps, self.null_label_index, dtype=torch.long)
-            else:
-                raise ValueError(
-                    "Class conditioning requested but no labels were provided. "
-                    "Ensure the datamodule keeps labels (drop_labels=False) and the caller forwards them."
-                )
-        else:
-            labels_to_embed = labels.to(timesteps.device, dtype=torch.long).view(-1)
-
-        # Clone to avoid in-place edits that would leak outwards.
-        labels_to_embed = labels_to_embed.clone()
-
-        if unconditional:
-            labels_to_embed.fill_(self.null_label_index)
-
-        if dropout_mask is not None:
-            if dropout_mask.shape != labels_to_embed.shape:
-                raise ValueError("dropout_mask must match the shape of the labels tensor.")
-            labels_to_embed[dropout_mask] = self.null_label_index
-
-        if (labels_to_embed < 0).any():
-            raise ValueError("Encountered negative labels while class conditioning; check datamodule configuration.")
-        if (labels_to_embed > self.null_label_index).any():
-            raise ValueError("Label index out of range for classifier-free guidance.")
-
-        label_emb = self.label_embed(labels_to_embed)
-        return time_emb + label_emb, label_emb
-
-    def _shared_step(
-        self,
-        batch: dict[str, torch.Tensor],
-        *,
-        return_clean_images: bool = False,
-    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, Optional[torch.Tensor]]]:
-        """Shared logic for training and validation steps.
-
-        Args:
-            batch: A datamodule batch containing at least the "input" key with clean images.
-            return_clean_images: When ``True``, the method returns a tuple containing the loss
-                and a dict with clean images for FID computation. Used during validation only.
-
-        Returns:
-            - During training: The computed loss tensor.
-            - During validation with ``return_clean_images=True``: A tuple of the loss tensor and
-              a dict with clean images under the "clean_images_bchw" key and optional labels.
-        """
-        # Inputs arrive in channels-last format from the datamodule; we keep that convention for the
-        # network but convert to channels-first whenever diffusers expects it.
-        images = batch["input"].to(self.device)
-        if self.example_input_shape is None:
-            # Cache the tensor shape so we can initialise random noise for sampling later on.
-            self.example_input_shape = images.shape[1:]
-
-        images_bchw = self._channels_last_to_first(images)
-        batch_size = images_bchw.shape[0]
-
-        labels_tensor: Optional[torch.Tensor] = None
-        if self.class_conditioning:
-            # Retrieve labels if the datamodule provided them. For class-conditioned runs the labels are
-            # essential; otherwise we quietly fall back to the unconditional path.
-            if "label" not in batch:
-                raise RuntimeError(
-                    "Class conditioning requires datamodule batches to include 'label'. "
-                    "Set drop_labels=False on the datamodule to keep them."
-                )
-            labels_tensor = batch["label"].to(self.device, non_blocking=True).long().view(-1)
-
-        # 1. Sample timestep t in [0, 1] from the JiT logit-normal distribution.
-        t_logit = torch.randn(batch_size, device=self.device) * self.p_std + self.p_mean
-        timesteps = torch.sigmoid(t_logit)
-
-        # 2. Sample noise scaled by the configured noise_scale (matches JiT denoiser.py).
-        eps_bchw = torch.randn_like(images_bchw) * self.noise_scale
-
-        # 3. Mix clean image and noise to get z_t.
-        t_b = timesteps.view(batch_size, 1, 1, 1)
-        z_bchw = t_b * images_bchw + (1.0 - t_b) * eps_bchw
-        target_v = images_bchw - eps_bchw
-
-        # 4. Predict x from z_t and conditioning.
-        z_cl = self._channels_first_to_last(z_bchw)
-        dropout_mask = None
-        if self.class_conditioning and self.condition_dropout_prob > 0.0:
-            dropout_mask = torch.rand(batch_size, device=self.device) < self.condition_dropout_prob
-
-        condition, class_emb = self._condition_from_timesteps(
-            timesteps,
-            labels=labels_tensor,
-            dropout_mask=dropout_mask,
-        )
-        net_input = {"input": z_cl, "condition": condition}
-        if class_emb is not None:
-            net_input["class_emb"] = class_emb
-        prediction = self.network(net_input)["logits"]
-        prediction_bchw = self._channels_last_to_first(prediction)
-
-        # 5. JiT objective: network predicts x, loss is applied in v-space.
-        denominator = torch.clamp(1.0 - t_b, min=0.05)
-        predicted_v = (prediction_bchw - z_bchw) / denominator
-        loss = F.mse_loss(predicted_v, target_v)
-
-        if return_clean_images:
-            aux = {
-                "clean_images_bchw": images_bchw.detach(),
-                "labels": labels_tensor.detach() if labels_tensor is not None else None,
-            }
-            return loss, aux
-
-        return loss
-
-    def training_step(self, batch, batch_idx):
-        """Run one training step and log the loss."""
-        loss = self._shared_step(batch)
-        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=self.distributed)
-        self.log(
-            "global_step", float(self.trainer.global_step), on_step=True, on_epoch=False, prog_bar=True, logger=False
-        )
-        self.log("current_step", float(self.trainer.global_step), on_step=True, on_epoch=False, prog_bar=False)
-        return loss
-
-    def on_validation_epoch_start(self) -> None:
-        """Reset validation metrics."""
-        super().on_validation_epoch_start()
-
-    def validation_step(self, batch, batch_idx):
-        """Compute validation loss."""
-        loss = self._shared_step(batch)
-        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=self.distributed)
-        return loss
-
-    def test_step(self, batch, batch_idx):
-        """Lightning test loop placeholder (no dedicated metric)."""
-        # Generation experiments do not have a dedicated test metric.
-        pass
-
-    @torch.no_grad()
-    def sample(
-        self,
-        num_samples: int,
-        num_inference_steps: Optional[int] = None,
-        labels: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Generate samples with the JiT continuous-time sampler."""
-        return self._sample_continuous(num_samples, num_inference_steps, labels)
-
-    def on_fit_start(self) -> None:
-        """Move EMA weights to device before training begins."""
-        super().on_fit_start()
-        if self.ema_enabled and self._ema_model is not None:
-            self._ema_model.to(self.device)
-            self._ema_model.eval()
-
-    def on_train_batch_end(self, outputs, batch, batch_idx) -> None:
-        """Update EMA parameters at the configured interval."""
-        super().on_train_batch_end(outputs, batch, batch_idx)
-
-        if (
-            self.ema_enabled
-            and self._ema_model is not None
-            and self.global_step >= self.ema_warmup_steps
-            and (self.global_step % self.ema_update_every == 0)
-        ):
-            with torch.no_grad():
-                decay = self.ema_decay
-                for ema_param, param in zip(self._ema_model.parameters(), self.network.parameters()):
-                    ema_param.mul_(decay).add_(param, alpha=1.0 - decay)
-                for ema_buffer, buffer in zip(self._ema_model.buffers(), self.network.buffers()):
-                    if ema_buffer.shape != buffer.shape:
-                        ema_buffer.resize_as_(buffer)
-                    ema_buffer.copy_(buffer)
-                self._ema_has_been_updated = True
-
-    def on_validation_epoch_end(self, outputs=None):
-        """Compute and log validation summary metrics and sample grids."""
-        if self.trainer.sanity_checking:
-            return
-
-        if self.fid_online_jit and self.current_epoch % self.fid_interval == 0:
-            if self.global_rank == 0:
-                print(f"Starting FID evaluation for epoch {self.current_epoch}...")
-            self._run_jit_online_eval()
-
-        if not self.log_samples:
-            return
-        if self.example_input_shape is None:
-            return
-        if self.logger is None or not hasattr(self.logger, "experiment"):
-            return
-
-        num_samples = int(self.num_generated_samples)
-
-        # When class conditioning is enabled we draw random labels so validation grids vary each epoch.
-        # Guidance scale determines whether the conditional branch is actually used during sampling.
-        labels_for_sampling = None
-        if self.class_conditioning:
-            assert self.num_classes is not None
-            labels_for_sampling = torch.randint(
-                low=0,
-                high=self.num_classes,
-                size=(num_samples,),
-                device=self.device,
-                dtype=torch.long,
-            )
-
-        samples = self.sample(num_samples=num_samples, labels=labels_for_sampling)
-        value_range = (-1.0, 1.0)
-        normalize_grid = True
-
-        datamodule = getattr(self.trainer, "datamodule", None)
-        if datamodule is not None:
-            unnormalize_fn = getattr(datamodule, "unnormalize", None)
-            if callable(unnormalize_fn):
-                try:
-                    samples = unnormalize_fn(samples)
-                except (TypeError, ValueError):
-                    pass
-                else:
-                    samples = torch.clamp(samples, 0.0, 1.0)
-                    normalize_grid = False
-                    value_range = (0.0, 1.0)
-
-        samples_bchw = self._channels_last_to_first(samples)
-        grid = make_grid(
-            samples_bchw.detach().cpu(),
-            nrow=max(1, int(math.sqrt(num_samples))),
-            normalize=normalize_grid,
-            value_range=value_range,
-        )
-        self.logger.experiment.log(
-            {
-                "val/samples": wandb.Image(grid.cpu()),
-                "global_step": self.global_step,
-            }
-        )
-
-    # -------------------------------------------------------------------------
-    # Continuous Time / Flow Matching Logic (JiT Port)
-    # -------------------------------------------------------------------------
-
-    def _sample_continuous(
-        self,
-        num_samples: int,
-        num_inference_steps: Optional[int] = None,
-        labels: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Sampling loop for continuous time flow matching."""
-        if self.example_input_shape is None:
-            raise RuntimeError("Cannot sample before observing at least one training batch.")
-
-        num_inference_steps = num_inference_steps or self.default_inference_steps
-        device = self.device
-        height, width, channels = self.example_input_shape
-
-        # Prepare labels
-        labels_tensor: Optional[torch.Tensor]
-        if self.class_conditioning:
-            if labels is None:
-                labels_tensor = torch.randint(0, self.num_classes, (num_samples,), device=device, dtype=torch.long)
-            else:
-                labels_tensor = torch.as_tensor(labels, device=device, dtype=torch.long).view(-1)
-                if labels_tensor.numel() == 1 and num_samples > 1:
-                    labels_tensor = labels_tensor.expand(num_samples)
-                if labels_tensor.shape[0] != num_samples:
-                    raise ValueError("labels must either be a scalar or have the same length as num_samples.")
-                if (labels_tensor < 0).any():
-                    raise ValueError("labels must contain non-negative class indices.")
-        else:
-            if labels is not None:
-                raise ValueError("labels were provided but the model was configured without class conditioning.")
-            labels_tensor = None
-
-        # 1. Initialize with noise scaled by the configured noise_scale (matches JiT denoiser.py).
-        z = torch.randn((num_samples, channels, height, width), device=device) * self.noise_scale
-
-        # 2. Integrate from t=0 (noise) to t=1 (data) using a Heun solver.
-        dt = 1.0 / num_inference_steps
-
-        # Use EMA model if enabled
-        use_ema = self.ema_enabled and self._ema_model is not None and self._ema_has_been_updated
-        inference_model = self._ema_model if use_ema else self.network
-        was_training = inference_model.training
-        inference_model.eval()
-        inference_model = inference_model.to(device)
-
-        for i in range(num_inference_steps - 1):
-            t = i * dt
-            z = self._heun_step(inference_model, z, t, dt, labels_tensor)
-
-        # Last step: Euler only (matches JiT reference — avoids an extra model call at t≈1
-        # where the ODE denominator is near its minimum clamp value).
-        t_last = (num_inference_steps - 1) * dt
-        v_last = self._model_forward_continuous(inference_model, z, t_last, labels_tensor)
-        z = z + v_last * dt
-
-        if was_training:
-            inference_model.train()
-
-        return torch.clamp(self._channels_first_to_last(z), -1.0, 1.0)
-
-    def _model_forward_continuous(self, model, x, t_scalar, labels):
-        # Broadcast t
-        t_batch = torch.full((x.shape[0],), t_scalar, device=x.device, dtype=torch.float32)
-
-        # 1. Condition
-        do_cfg = self.cfg_enabled and (
-            min(self.cfg_interval_start, self.cfg_interval_end)
-            <= t_scalar
-            <= max(self.cfg_interval_start, self.cfg_interval_end)
-        )
-
-        pred_bchw = None
-
-        if do_cfg:
-            # CFG
-            cond_uncond, class_emb_uncond = self._condition_from_timesteps(
-                t_batch,
-                labels=labels,
-                unconditional=True,
-            )
-            cond_cond, class_emb_cond = self._condition_from_timesteps(
-                t_batch,
-                labels=labels,
-            )
-
-            # Note: Model expects channels-last in 'input'
-            x_cl = self._channels_first_to_last(x)
-
-            inp_uncond = {"input": x_cl, "condition": cond_uncond}
-            if class_emb_uncond is not None:
-                inp_uncond["class_emb"] = class_emb_uncond
-            inp_cond = {"input": x_cl, "condition": cond_cond}
-            if class_emb_cond is not None:
-                inp_cond["class_emb"] = class_emb_cond
-
-            out_uncond = model(inp_uncond)["logits"]
-            out_cond = model(inp_cond)["logits"]
-
-            pred_uncond = self._channels_last_to_first(out_uncond)
-            pred_cond = self._channels_last_to_first(out_cond)
-
-            pred_bchw = pred_uncond + self.guidance_scale * (pred_cond - pred_uncond)
-        else:
-            # Standard forward
-            condition, class_emb = self._condition_from_timesteps(t_batch, labels=labels)
-            x_cl = self._channels_first_to_last(x)
-            inp = {"input": x_cl, "condition": condition}
-            if class_emb is not None:
-                inp["class_emb"] = class_emb
-            out = model(inp)["logits"]
-            pred_bchw = self._channels_last_to_first(out)
-
-        # JiT mode predicts x and converts to velocity with denominator clipping.
-        denominator = max(1.0 - t_scalar, 0.05)
-        return (pred_bchw - x) / denominator
-
-    def _heun_step(self, model, x, t, dt, labels):
-        v = self._model_forward_continuous(model, x, t, labels)
-        x_euler = x + v * dt
-        v_next = self._model_forward_continuous(model, x_euler, t + dt, labels)
-        x_next = x + 0.5 * dt * (v + v_next)
-        return x_next
-
-    def _run_jit_online_eval(self) -> None:
-        """Run standard FID evaluation matching the JiT repository's methodology."""
-        if self.fid_stats_file is None:
-            if self.global_rank == 0:
-                print("FID stats file not configured. Skipping online evaluation.")
-            return
-
-        if not os.path.exists(self.fid_stats_file) and not os.environ.get("SKIP_FID_STATS_CHECK"):
-            if self.global_rank == 0:
-                print(f"FID stats file {self.fid_stats_file} not found. Skipping online evaluation.")
-            return
-
-        fid_run_dir = os.path.join(self.trainer.default_root_dir, f"fid_eval_{self.global_step}")
-
-        if self.global_rank == 0:
-            os.makedirs(fid_run_dir, exist_ok=True)
-
-        if torch.distributed.is_initialized():
-            torch.distributed.barrier()
-
-        world_size = self.trainer.world_size
-        global_rank = self.global_rank
-
-        total_samples = self.fid_num_samples
-        base_samples = total_samples // world_size
-        remainder = total_samples % world_size
-
-        # Non-overlapping slice [start_idx, start_idx + my_count) for this rank.
-        start_idx = base_samples * global_rank + min(global_rank, remainder)
-        my_count = base_samples + (1 if global_rank < remainder else 0)
-
-        if self.num_classes is not None:
-            samples_per_class = total_samples // self.num_classes
-            remainder_classes = total_samples % self.num_classes
-
-            class_labels = np.arange(self.num_classes).repeat(samples_per_class)
-
-            if remainder_classes > 0:
-                class_labels = np.concatenate([class_labels, np.zeros(remainder_classes, dtype=int)])
-
-            end_idx = start_idx + my_count
-            my_labels = class_labels[start_idx:end_idx]
-            my_labels = torch.tensor(my_labels, device=self.device, dtype=torch.long)
-
-        else:
-            my_labels = None
-
-        batches = (
-            math.ceil(len(my_labels) / self.fid_batch_size)
-            if my_labels is not None
-            else math.ceil(my_count / self.fid_batch_size)
-        )
-
-        self.eval()
-
-        samples_generated = 0
-        batch_iterator = range(batches)
-        if global_rank == 0:
-            batch_iterator = tqdm(batch_iterator, total=batches, desc="FID sample generation", leave=True)
-
-        for i in batch_iterator:
-            current_batch_size = (
-                min(self.fid_batch_size, len(my_labels) - samples_generated)
-                if my_labels is not None
-                else min(self.fid_batch_size, my_count - samples_generated)
-            )
-
-            batch_labels = (
-                my_labels[samples_generated : samples_generated + current_batch_size]
-                if my_labels is not None
-                else None
-            )
-
-            with torch.no_grad():
-                samples = self.sample(
-                    num_samples=current_batch_size,
-                    num_inference_steps=self.fid_num_inference_steps,
-                    labels=batch_labels,
-                )
-                # sample() returns (B, H, W, C), but save_image expects (C, H, W).
-                samples = samples.permute(0, 3, 1, 2)  # (B, C, H, W)
-
-            # Denormalize
-            samples = (samples + 1.0) / 2.0
-            samples = torch.clamp(samples, 0.0, 1.0)
-
-            for b in range(current_batch_size):
-                filename = f"sample_{start_idx + samples_generated + b:08d}.png"
-                save_path = os.path.join(fid_run_dir, filename)
-                save_image(samples[b], save_path)
-
-            samples_generated += current_batch_size
-
-        if torch.distributed.is_initialized():
-            torch.distributed.barrier()
-
-        if global_rank == 0:
-            try:
-                from torch_fidelity.metric_fid import fid_featuresdict_to_statistics, fid_statistics_to_metric
-                from torch_fidelity.metric_isc import isc_featuresdict_to_metric
-                from torch_fidelity.utils import (
-                    create_feature_extractor,
-                    extract_featuresdict_from_input_id,
-                    resolve_feature_layer_for_metric,
-                )
-
-                feat_layer_fid = resolve_feature_layer_for_metric("fid", fid=True)
-                feat_layer_isc = resolve_feature_layer_for_metric("isc", isc=True)
-                feat_layers = list({feat_layer_fid, feat_layer_isc})
-
-                feat_extractor = create_feature_extractor("inception-v3-compat", feat_layers, cuda=True, verbose=False)
-
-                featuresdict = extract_featuresdict_from_input_id(
-                    input_id=1, feat_extractor=feat_extractor, input1=fid_run_dir, cuda=True, verbose=False
-                )
-
-                stats_1 = fid_featuresdict_to_statistics(featuresdict, feat_layer_fid)
-                f = np.load(self.fid_stats_file)
-                stats_2 = {"mu": f["mu"], "sigma": f["sigma"]}
-                f.close()
-
-                stats_1["mu"] = stats_1["mu"].astype(np.float64)
-                stats_1["sigma"] = stats_1["sigma"].astype(np.float64)
-                stats_2["mu"] = stats_2["mu"].astype(np.float64)
-                stats_2["sigma"] = stats_2["sigma"].astype(np.float64)
-
-                metrics_fid = fid_statistics_to_metric(stats_1, stats_2, verbose=False)
-                metrics_isc = isc_featuresdict_to_metric(featuresdict, feat_layer_isc, isc_splits=10)
-
-                fid = metrics_fid["frechet_inception_distance"]
-                isc = metrics_isc["inception_score_mean"]
-
-                print(f"FID: {fid:.4f}, IS: {isc:.4f}")
-
-                self.log("metrics/fid_online", fid, sync_dist=False)
-                self.log("metrics/is_online", isc, sync_dist=False)
-
-                if self.logger is not None and hasattr(self.logger, "experiment"):
-                    self.logger.experiment.log(
-                        {"metrics/fid_online": fid, "metrics/is_online": isc, "global_step": self.global_step}
-                    )
-
-            except Exception as e:
-                print(f"Error calculating FID: {e}")
-            finally:
-                print(f"Removing temporary FID directory: {fid_run_dir}")
-                shutil.rmtree(fid_run_dir, ignore_errors=True)
-
-        if torch.distributed.is_initialized():
-            torch.distributed.barrier()
diff --git a/experiments/lightning_wrappers/regression_wrapper.py b/experiments/lightning_wrappers/regression_wrapper.py
index ba3f23f7..efcae94d 100644
--- a/experiments/lightning_wrappers/regression_wrapper.py
+++ b/experiments/lightning_wrappers/regression_wrapper.py
@@ -13,15 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Lightning wrapper for regression tasks (MAE / MSE loss).
 
 Provides :class:`RegressionWrapper`, which supports both Mean Absolute Error
 (L1) and Mean Squared Error (L2) regression objectives.  It is also the base
 class for :class:`~experiments.lightning_wrappers.well_lightning_wrapper.WELLRegressionWrapper`.
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 from typing import Literal
diff --git a/experiments/run.py b/experiments/run.py
index 7dd140a3..658704e6 100644
--- a/experiments/run.py
+++ b/experiments/run.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 r"""Main entry point for launching nvSubquadratic training runs.
 
 **Usage**::
@@ -45,8 +43,6 @@
 6. Instantiates the Lightning wrapper from ``cfg.lightning_wrapper_cfg``.
 7. Calls :func:`~experiments.trainer.construct_trainer` and
    :meth:`pytorch_lightning.Trainer.fit`.
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 import argparse
@@ -105,7 +101,7 @@ def parse_args() -> argparse.Namespace:
         "--config",
         type=str,
         required=True,
-        help="Path to the configuration file, e.g., config/experiments/mnist/mnist_classification_cfg.py",
+        help="Path to the configuration file, e.g., examples/imagenet_classification/ccnn_7_512_hyena.py",
     )
 
     # Add a catch-all for arbitrary config overrides
diff --git a/experiments/trainer.py b/experiments/trainer.py
index 4ee3e798..59cd3043 100644
--- a/experiments/trainer.py
+++ b/experiments/trainer.py
@@ -13,10 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: Add licence header
-
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """PyTorch Lightning trainer factory for nvSubquadratic experiments.
 
 :func:`construct_trainer` is the single entry point for building a
@@ -40,8 +36,6 @@
 7. Returns the fully configured :class:`pytorch_lightning.Trainer` together with
    the :class:`~pytorch_lightning.callbacks.ModelCheckpoint` callback (so the
    caller can inspect ``checkpoint_callback.best_model_path``).
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 from pathlib import Path
diff --git a/experiments/utils/checkpointing.py b/experiments/utils/checkpointing.py
index 58e3d962..d4b3a2cf 100644
--- a/experiments/utils/checkpointing.py
+++ b/experiments/utils/checkpointing.py
@@ -13,10 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: Add licence header
-
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Utility helpers for checkpointing operations (loading, downloading, uploading, pruning, etc.)."""
 
 from __future__ import annotations
diff --git a/nvsubquadratic/lazy_config.py b/nvsubquadratic/lazy_config.py
index 61a21c4f..eb418d88 100644
--- a/nvsubquadratic/lazy_config.py
+++ b/nvsubquadratic/lazy_config.py
@@ -51,8 +51,6 @@
 :func:`save_config` / :func:`load_config` serialise configs to YAML via
 OmegaConf.  :func:`to_config` reverse-engineers a ``__target__`` dict from an
 already-instantiated object (best-effort; works for simple cases).
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 import ast
diff --git a/nvsubquadratic/metrics/cleanfid.py b/nvsubquadratic/metrics/cleanfid.py
deleted file mode 100644
index e0284656..00000000
--- a/nvsubquadratic/metrics/cleanfid.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Thin wrapper around CleanFID image generation quality metrics.
-
-Fréchet Inception Distance (FID) measures the distributional similarity between
-a set of generated images and a reference dataset by comparing the mean and
-covariance of Inception-v3 feature vectors (Heusel et al., "GANs Trained by a
-Two Time-Scale Update Rule Converge to a Local Nash Equilibrium", NeurIPS 2017).
-
-CleanFID (Parmar et al., "On Aliased Resizing and Surprising Subtleties in GAN
-Evaluation", CVPR 2022) corrects for common pre-processing inconsistencies
-(e.g. JPEG re-compression, bilinear vs Lanczos resizing) that cause FID scores
-to be non-reproducible across libraries.  This module delegates to the
-``cleanfid`` package, which ships pre-computed reference statistics for standard
-benchmarks (FFHQ, CIFAR-10, ImageNet, etc.).
-
-Usage::
-
-    score = compute_folder_fid(
-        sample_dir="outputs/samples/",
-        dataset_name="imagenet",
-        dataset_resolution=256,
-        dataset_split="train",
-    )
-    print(f"FID: {score:.2f}")
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-from cleanfid import fid
-
-
-def compute_folder_fid(
-    sample_dir: str | Path,
-    *,
-    dataset_name: str,
-    dataset_resolution: int,
-    dataset_split: str = "train",
-) -> float:
-    """Compute CleanFID between a folder of generated images and a reference dataset.
-
-    Calls ``cleanfid.fid.compute_fid`` with pre-computed reference statistics
-    for ``dataset_name`` at ``dataset_resolution``, so no reference images need
-    to be stored locally.
-
-    Args:
-        sample_dir: Path to the directory containing generated images (PNG/JPEG).
-            Expanded and resolved to an absolute path before use.
-        dataset_name: Name of the CleanFID reference dataset, e.g.
-            ``"imagenet"``, ``"ffhq"``, ``"cifar10"``.  Must match a dataset
-            whose statistics are bundled with or downloaded by ``cleanfid``.
-        dataset_resolution: Reference image resolution in pixels, e.g. ``256``
-            for 256×256 ImageNet.
-        dataset_split: Which split of the reference dataset to compare against.
-            Default ``"train"``.
-
-    Returns:
-        FID score as a Python ``float``.  Lower is better; 0 means the
-        generated and reference distributions are identical under the Inception
-        feature extractor.
-
-    Raises:
-        FileNotFoundError: If ``sample_dir`` does not exist.
-    """
-    sample_path = Path(sample_dir).expanduser().resolve()
-    if not sample_path.exists():
-        raise FileNotFoundError(f"Sample directory not found: {sample_path}")
-
-    score = fid.compute_fid(
-        fdir1=sample_path.as_posix(),
-        dataset_name=dataset_name,
-        dataset_res=dataset_resolution,
-        dataset_split=dataset_split,
-    )
-    return float(score)
diff --git a/nvsubquadratic/modules/ckconv_multihead_nd.py b/nvsubquadratic/modules/ckconv_multihead_nd.py
deleted file mode 100644
index 524f6dfe..00000000
--- a/nvsubquadratic/modules/ckconv_multihead_nd.py
+++ /dev/null
@@ -1,856 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-r"""Multi-head Continuous Kernel Convolution (CKConv) for 2D signals.
-
-Background
-----------
-This module implements a **multi-head** extension of the CKConv operator
-(Romero et al., "CKConv: Continuous Kernel Convolution With Arbitrary
-Resolution", ICLR 2022, arXiv:2102.02611).  For the single-head (depthwise)
-variant see :mod:`nvsubquadratic.modules.ckconv_nd`.
-
-In the single-head variant, every channel is convolved independently with its
-own implicit kernel — the kernel has shape ``[C, K_x, K_y]`` and there is no
-cross-channel mixing in the convolution layer.  The multi-head variant enables
-*dense channel mixing within each head*, analogous to multi-head attention but
-for convolutions:
-
-- The hidden dimension ``C`` is split into ``H`` heads of ``d = C / H``
-  channels each.
-- Within head ``h``, the convolution kernel ``K^h`` has shape
-  ``[d, d, K_x, K_y]``, so every output channel sees all ``d`` input channels
-  of that head through a spatially-varying weight.
-- Across heads, channels remain isolated (no cross-head mixing).
-
-Formally, for head ``h`` and a 2D input :math:`x^h \in \mathbb{R}^{d \times H \times W}`:
-
-.. code-block:: none
-
-    y^h = K^h * x^h + shortcut^h ⊙ x^h
-
-where ``*`` denotes a dense 2D convolution (not depthwise), and the
-concatenated output across all heads is:
-
-.. code-block:: none
-
-    y = [y^0 | y^1 | ... | y^{H-1}]   (concatenated along channel axis)
-
-Each per-head kernel :math:`K^h_\theta` is produced by a shared SIREN network
-evaluated on a continuous positional grid normalised to ``[-1, 1]^2``:
-
-.. code-block:: none
-
-    K^h_\theta(p) = MLP_\theta(pos_enc(p))
-
-Because the MLP is small and the same for all heads, the number of parameters
-scales with ``H * d^2 = C^2 / H`` (dense) rather than ``C^2`` (full
-unstructured mixing), giving a parameter count between the depthwise and fully
-dense extremes.
-
-Low-rank / factored kernel
---------------------------
-For large ``d``, the dense kernel ``[d, d, K_x, K_y]`` per head can be
-expensive.  An optional low-rank factorisation (``kernel_rank=r``) decomposes
-the kernel as:
-
-.. code-block:: none
-
-    K^h ≈ U^h · V^h,    U^h ∈ R^{d × r × K_x × K_y},
-                          V^h ∈ R^{r × d × K_x × K_y}
-
-The SIREN outputs ``num_heads * 2 * r * d`` values per position rather than
-``num_heads * d^2``, and the frequency-domain contraction is split into two
-cheaper steps:
-
-.. code-block:: none
-
-    z = V^h x     (rank-to-d_in contraction)
-    y = U^h z     (d_out-to-rank contraction)
-
-This reduces the per-head compute from :math:`O(d^2)` to :math:`O(2 d r)` and
-the kernel parameter count by the same factor, while preserving the ``d × d``
-mixing capacity at rank ``r``.
-
-Boundary conditions
--------------------
-Two boundary conditions are supported:
-
-* **Zero-padding** (``fft_padding="zero"``): standard linear convolution with
-  "same" output size.  Kernel size is ``2*N``, covering the full receptive
-  field without wrap-around.
-* **Circular / periodic** (``fft_padding="circular"``): wrap-around
-  convolution.  The kernel size equals the input size (requires
-  ``grid_type="single"``).
-
-Shortcut (skip connection)
---------------------------
-Every forward pass adds a per-channel learnable shortcut term:
-
-.. code-block:: none
-
-    y ← y + shortcut ⊙ x
-
-where ``shortcut`` is a ``[hidden_dim]`` parameter vector initialised with
-Kaiming-uniform scale, identical in design to the shortcut in
-:class:`nvsubquadratic.modules.ckconv_nd.CKConvND`.
-
-Current limitations
--------------------
-* Only 2D inputs are supported (``data_dim`` must equal 2).
-* Context parallelism (``cp_group``) is not yet implemented.
-* Only ``"torch_fft"`` backend (no ``"subq_ops"`` backend support).
-
-Related modules
----------------
-* :mod:`nvsubquadratic.modules.ckconv_nd` — single-head (depthwise) CKConv,
-  supporting 1D / 2D / 3D inputs, mixed boundary conditions, and causal mode.
-* :mod:`nvsubquadratic.modules.kernels_nd` — implicit kernel parametrisation
-  (``SIRENKernelND``, FiLM-conditioned variants) consumed by both variants.
-* :mod:`nvsubquadratic.ops.fftconv_multihead` — the FFT convolution primitives
-  that implement the dense per-head spatial mixing called from ``forward``.
-
-References:
-----------
-* Romero et al. (2022). *CKConv: Continuous Kernel Convolution With Arbitrary
-  Resolution*. ICLR 2022. https://arxiv.org/abs/2102.02611
-* Sitzmann et al. (2020). *Implicit Neural Representations with Periodic
-  Activation Functions*. NeurIPS 2020. (SIREN kernel used for ``MLP_θ``)
-"""
-
-import math
-from typing import Literal
-
-import torch
-from einops import rearrange
-
-from nvsubquadratic.lazy_config import LazyConfig, instantiate
-from nvsubquadratic.ops.fftconv_multihead import (
-    fftconv2d_multihead_bhl,
-    fftconv2d_multihead_circular_bhl,
-    fftconv2d_multihead_lowrank_bhl,
-    fftconv2d_multihead_lowrank_circular_bhl,
-)
-
-
-class CKConvMultiheadND(torch.nn.Module):
-    r"""Multi-head CKConv for 2D signals with dense within-head channel mixing.
-
-    Extends :class:`nvsubquadratic.modules.ckconv_nd.CKConvND` from depthwise
-    convolution to *dense within-head* convolution by splitting the channel
-    dimension into ``num_heads`` groups and applying a separate
-    ``[head_dim × head_dim]`` implicit kernel to each group.
-
-    **Mathematical description**
-
-    Let :math:`x \\in \\mathbb{R}^{B \times C \times H \times W}` with
-    :math:`C = H_\text{heads} \\cdot d` (``num_heads × head_dim``).  The input
-    is partitioned into heads:
-
-    .. code-block:: none
-
-        x^h ∈ R^{B × d × H × W},   h = 0, ..., H_heads - 1
-
-    Per head ``h``, the implicit kernel network produces
-    :math:`K^h_\theta \\in \\mathbb{R}^{d \times d \times K_x \times K_y}`
-    (full-rank) or its low-rank factorisation
-    :math:`U^h \\in \\mathbb{R}^{d \times r \times K_x \times K_y}`,
-    :math:`V^h \\in \\mathbb{R}^{r \times d \times K_x \times K_y}` (low-rank,
-    when ``kernel_rank`` is set).  The output for each head is:
-
-    .. code-block:: none
-
-        y^h = K^h_θ * x^h + shortcut^h ⊙ x^h
-
-    where ``*`` is a dense 2D linear (or circular) convolution computed via
-    FFT.  The final output concatenates all head outputs:
-
-    .. code-block:: none
-
-        y = concat([y^0, y^1, ..., y^{H-1}], dim=channel)
-
-    **Key differences from CKConvND (single-head)**
-
-    * *Dense within-head mixing*: whereas :class:`CKConvND` uses a depthwise
-      kernel of shape ``[C, K_x, K_y]`` (one scalar kernel per channel),
-      ``CKConvMultiheadND`` uses a dense kernel of shape
-      ``[H, d, d, K_x, K_y]`` (a :math:`d \times d` mixing matrix per head
-      per spatial frequency).
-    * *No context parallelism*: CP support is not yet implemented; passing
-      ``cp_group`` to ``forward`` raises ``NotImplementedError``.
-    * *2D only*: the current implementation only supports ``data_dim=2``
-      (images with two spatial axes).  The single-head variant supports 1D,
-      2D, and 3D inputs.
-    * *No causal mode, no fp16 FFT, no chunked FFT*: these features from the
-      single-head variant are not carried over here.
-    * *Optional low-rank kernel*: the ``kernel_rank`` parameter (absent in
-      ``CKConvND``) enables a factored ``U · V`` kernel decomposition that
-      reduces SIREN output size and FFT cost by approximately ``2r / d``.
-
-    Attributes:
-        data_dim (int): Spatial rank of the input.  Always 2 for this class.
-        hidden_dim (int): Total number of channels ``C = num_heads * head_dim``.
-        num_heads (int): Number of independent heads ``H``.
-        head_dim (int): Channels per head ``d = hidden_dim // num_heads``.
-        fft_padding (str): Boundary condition — ``"zero"`` (linear conv) or
-            ``"circular"`` (periodic conv).
-        grid_type (str): Kernel grid size mode — ``"single"`` (kernel size
-            equals input size, for circular conv) or ``"double"`` (kernel
-            size is ``2N``, for zero-padded conv).
-        kernel_rank (int or None): Rank of the low-rank kernel factorisation.
-            ``None`` means full-rank ``[d, d, K_x, K_y]`` kernels are used.
-        kernel (nn.Module): Implicit kernel generator (SIREN or similar).
-            Called as ``kernel(grid_lens, conditioning=...)`` and returns
-            ``(kernel_values, grid)`` where ``kernel_values`` has shape
-            ``[1_or_B, K_x, K_y, num_heads * d * d]`` (full-rank) or
-            ``[1_or_B, K_x, K_y, num_heads * 2 * r * d]`` (low-rank).
-        mask (nn.Module): Attenuation mask applied to kernel values after
-            generation.  ``nn.Identity`` when no mask is configured.
-        shortcut (nn.Parameter): Learnable per-channel skip-connection scale
-            of shape ``(hidden_dim,)``.  Added as ``shortcut ⊙ x`` after each
-            convolution.  Initialised with Kaiming-uniform scale
-            ``uniform(-1/√hidden_dim, 1/√hidden_dim)``.
-        fftconv_fn (callable): Selected FFT convolution function, one of the
-            four functions from :mod:`nvsubquadratic.ops.fftconv_multihead`.
-            Signature varies by full-rank vs. low-rank path:
-            full-rank: ``(x, kernel, shortcut) → output``;
-            low-rank: ``(x, kernel_u, kernel_v, shortcut) → output``.
-    """
-
-    def __init__(
-        self,
-        data_dim: int,
-        hidden_dim: int,
-        num_heads: int,
-        kernel_cfg: LazyConfig,
-        mask_cfg: LazyConfig,
-        grid_type: Literal["double", "single"],
-        fft_padding: Literal["zero", "circular"],
-        kernel_rank: int | None = None,
-    ):
-        """Construct a CKConvMultiheadND operator.
-
-        Validates parameter combinations, derives ``head_dim``, adjusts the
-        kernel output-scale for variance control, initialises the ``shortcut``
-        parameter, and selects the appropriate FFT convolution function from
-        :mod:`nvsubquadratic.ops.fftconv_multihead`.
-
-        Args:
-            data_dim: Spatial rank of the input signal.  Must be ``2``; a
-                value other than 2 raises ``AssertionError``.
-            hidden_dim: Total number of channels ``C``.  Must be divisible by
-                ``num_heads``; a violation raises ``AssertionError``.
-            num_heads: Number of independent convolution heads ``H``.  The
-                channels are split evenly: ``head_dim = hidden_dim // num_heads``.
-            kernel_cfg: ``LazyConfig`` that instantiates the implicit kernel
-                generator (e.g. ``SIRENKernelND``).  The generator's output
-                dimension must be set externally to match the expected flat
-                kernel size:
-
-                * Full-rank: ``out_dim = num_heads * head_dim * head_dim``
-                * Low-rank: ``out_dim = num_heads * 2 * kernel_rank * head_dim``
-
-                The output-scale weight ``kernel.out_linear.weight`` is
-                multiplied in-place at construction time for variance control
-                (see Notes).
-            mask_cfg: ``LazyConfig`` for an optional attenuation mask applied
-                to the generated kernel values.  Use ``torch.nn.Identity``
-                for no masking.
-            grid_type: Relationship between the SIREN coordinate grid and the
-                input spatial size:
-
-                * ``"single"``: grid spans ``(N+1)//2`` points per axis,
-                  producing a kernel of size ``≈ N`` (for circular conv).
-                * ``"double"``: grid spans ``N`` points per axis, producing a
-                  kernel of size ``2N - 1 ≈ 2N`` (for zero-padded conv).
-
-            fft_padding: Boundary condition for the convolution:
-
-                * ``"zero"``: zero-padded linear convolution, "same" output
-                  size.
-                * ``"circular"``: periodic (wrap-around) convolution.  Requires
-                  ``grid_type="single"`` and ``K_x == H``, ``K_y == W`` at
-                  runtime (enforced by an ``AssertionError``).
-
-            kernel_rank: Rank ``r`` for the low-rank kernel factorisation.
-                When ``None`` (default), full-rank ``[d, d, K_x, K_y]``
-                kernels are used and the SIREN outputs
-                ``num_heads * d^2`` values per spatial position.  When set to
-                an integer ``r < d``, the kernel is factored as
-                :math:`K = U V` with ``U`` of shape ``[d, r, K_x, K_y]`` and
-                ``V`` of shape ``[r, d, K_x, K_y]``, and the SIREN outputs
-                ``num_heads * 2 * r * d`` values instead.
-
-        Raises:
-            AssertionError: If ``data_dim != 2``, ``hidden_dim % num_heads != 0``,
-                ``grid_type`` is not ``"double"`` or ``"single"``,
-                ``fft_padding`` is not ``"zero"`` or ``"circular"``, or
-                ``fft_padding="circular"`` is combined with
-                ``grid_type != "single"``.
-
-        Notes:
-            **Output-scale initialisation for variance control.**  The SIREN's
-            final linear layer weight is rescaled in-place (via
-            ``torch.no_grad()``) so that the convolution output has unit
-            variance at initialisation:
-
-            * Full-rank: each output channel sums over ``head_dim`` input
-              channels weighted by the kernel, so the weight is multiplied by
-              ``1 / √head_dim``.
-            * Low-rank: the two-step ``U @ V`` contraction has variance that
-              depends on both ``head_dim`` and ``kernel_rank``, so the weight
-              is multiplied by ``(1 / (head_dim * kernel_rank))^{1/4}`` — the
-              geometric mean of the per-factor scales.
-        """
-        assert data_dim == 2, f"CKConvMultiheadND currently only supports 2D. Got {data_dim}D."
-        assert hidden_dim % num_heads == 0, f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({num_heads})"
-        assert grid_type in ["double", "single"], f"Invalid grid_type: {grid_type}"
-        assert fft_padding in ["zero", "circular"], f"Invalid fft_padding: {fft_padding}"
-
-        if fft_padding == "circular":
-            assert grid_type == "single", (
-                "fft_padding='circular' requires grid_type='single' (kernel size equals input size)."
-            )
-
-        super().__init__()
-        self.data_dim = data_dim
-        self.hidden_dim = hidden_dim
-        self.num_heads = num_heads
-        self.head_dim = hidden_dim // num_heads
-        self.fft_padding = fft_padding
-        self.grid_type = grid_type
-        self.kernel_rank = kernel_rank
-
-        # Construct kernel (SIREN) and mask
-        self.kernel = instantiate(kernel_cfg)
-        self.mask = instantiate(mask_cfg)
-
-        # Apply output scaling for variance control.
-        # Full-rank: conv sums over head_dim terms → scale by 1/sqrt(head_dim).
-        # Low-rank: conv sums over head_dim then rank → scale by 1/(head_dim * rank)^{1/4}
-        #   so that each factor (U, V) has the right variance independently.
-        with torch.no_grad():
-            if kernel_rank is not None:
-                self.kernel.out_linear.weight.data *= (1.0 / (self.head_dim * kernel_rank)) ** 0.25
-            else:
-                self.kernel.out_linear.weight.data *= math.sqrt(1.0 / self.head_dim)
-
-        # Shortcut parameter for residual connection
-        self.shortcut = torch.nn.Parameter(torch.empty(hidden_dim, dtype=torch.float32))
-        bounds = math.sqrt(1.0 / hidden_dim)
-        self.shortcut.data.uniform_(-bounds, bounds)
-
-        # Select FFT convolution function
-        if kernel_rank is not None:
-            if fft_padding == "circular":
-                self.fftconv_fn = fftconv2d_multihead_lowrank_circular_bhl
-            else:
-                self.fftconv_fn = fftconv2d_multihead_lowrank_bhl
-        else:
-            if fft_padding == "circular":
-                self.fftconv_fn = fftconv2d_multihead_circular_bhl
-            else:
-                self.fftconv_fn = fftconv2d_multihead_bhl
-
-    def extra_repr(self) -> str:
-        """Return a concise summary string for ``print(module)`` and ``repr(module)``.
-
-        Returns:
-            A human-readable string listing ``data_dim``, ``hidden_dim``,
-            ``num_heads``, ``head_dim``, ``fft_padding``, ``grid_type``, and
-            (if set) ``kernel_rank``.
-        """
-        parts = (
-            f"data_dim={self.data_dim}, hidden_dim={self.hidden_dim}, "
-            f"num_heads={self.num_heads}, head_dim={self.head_dim}, "
-            f"fft_padding={self.fft_padding!r}, grid_type={self.grid_type!r}"
-        )
-        if self.kernel_rank is not None:
-            parts += f", kernel_rank={self.kernel_rank}"
-        return parts
-
-    def _reshape_lowrank_kernel(self, conv_kernel_flat: torch.Tensor, K_x: int, K_y: int, B: int | None = None):
-        """Reshape the flat SIREN output into the low-rank ``U`` and ``V`` factors.
-
-        The SIREN network outputs a flat tensor whose last dimension encodes
-        all heads and both low-rank factors interleaved.  This method splits
-        and permutes the tensor into the shapes expected by the low-rank FFT
-        convolution ops in :mod:`nvsubquadratic.ops.fftconv_multihead`.
-
-        The flat layout (last dim of ``conv_kernel_flat``) is:
-
-        .. code-block:: none
-
-            [head_0_factor_U | head_0_factor_V | head_1_factor_U | ... ]
-
-        More precisely, for each head the SIREN outputs
-        ``2 * rank * head_dim`` values arranged as
-        ``[rank, head_dim]`` for U followed by ``[rank, head_dim]`` for V,
-        with a factor index (``0`` = U, ``1`` = V) in the second-to-last
-        logical dimension after reshaping to
-        ``[..., num_heads, 2, rank, head_dim]``.
-
-        Args:
-            conv_kernel_flat: Flat SIREN output tensor.
-
-                * Unbatched (``B=None``): shape
-                  ``[K_x, K_y, num_heads * 2 * rank * head_dim]``.
-                * FiLM-batched (``B`` is not ``None``): shape
-                  ``[B, K_x, K_y, num_heads * 2 * rank * head_dim]``.
-
-            K_x: Kernel spatial height (first spatial axis of the SIREN grid).
-            K_y: Kernel spatial width (second spatial axis of the SIREN grid).
-            B: Batch size when each sample has its own kernel (FiLM
-                conditioning).  ``None`` for the standard unbatched path where
-                one kernel is shared across all samples in the batch.
-
-        Returns:
-            A ``(kernel_u, kernel_v)`` tuple of contiguous tensors:
-
-            * Unbatched (``B=None``):
-
-              * ``kernel_u``: shape ``[num_heads, head_dim, rank, K_x, K_y]``
-                — the "output projection" factor.
-              * ``kernel_v``: shape ``[num_heads, rank, head_dim, K_x, K_y]``
-                — the "input projection" factor.
-
-            * FiLM-batched (``B`` is not ``None``):
-
-              * ``kernel_u``: shape ``[B, num_heads, head_dim, rank, K_x, K_y]``.
-              * ``kernel_v``: shape ``[B, num_heads, rank, head_dim, K_x, K_y]``.
-
-        Notes:
-            The contraction order in the FFT convolution is
-            ``z = V x`` (input projection) followed by ``y = U z``
-            (output projection), i.e. the einsum chain is
-            ``(n, r, d_in) × (n, d_in) → (n, r)`` then
-            ``(n, d_out, r) × (n, r) → (n, d_out)``.  The shape convention
-            for ``kernel_u`` and ``kernel_v`` matches the einsum indices used
-            in :func:`~CKConvMultiheadND.apply_convolution_batched_lowrank`.
-        """
-        rank = self.kernel_rank
-        head_dim = self.head_dim
-        num_heads = self.num_heads
-
-        if B is not None:
-            # FiLM-batched: [B, K_x, K_y, num_heads * 2 * rank * head_dim]
-            # -> [B, K_x, K_y, num_heads, 2, rank, head_dim]
-            reshaped = conv_kernel_flat.view(B, K_x, K_y, num_heads, 2, rank, head_dim)
-            # After [:, :, :, :, idx, :, :] -> [B, K_x, K_y, num_heads, rank, head_dim]
-            # U target: [B, num_heads, head_dim, rank, K_x, K_y]
-            kernel_u = reshaped[:, :, :, :, 0, :, :].permute(0, 3, 5, 4, 1, 2).contiguous()
-            # V target: [B, num_heads, rank, head_dim, K_x, K_y]
-            kernel_v = reshaped[:, :, :, :, 1, :, :].permute(0, 3, 4, 5, 1, 2).contiguous()
-        else:
-            # Unbatched: [1, K_x, K_y, num_heads * 2 * rank * head_dim]
-            # -> [K_x, K_y, num_heads, 2, rank, head_dim]
-            reshaped = conv_kernel_flat.view(K_x, K_y, num_heads, 2, rank, head_dim)
-            # After [:, :, :, idx, :, :] -> [K_x, K_y, num_heads, rank, head_dim]
-            # U target: [num_heads, head_dim, rank, K_x, K_y]
-            kernel_u = reshaped[:, :, :, 0, :, :].permute(2, 4, 3, 0, 1).contiguous()
-            # V target: [num_heads, rank, head_dim, K_x, K_y]
-            kernel_v = reshaped[:, :, :, 1, :, :].permute(2, 3, 4, 0, 1).contiguous()
-
-        return kernel_u, kernel_v
-
-    def apply_convolution(self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut: torch.Tensor) -> torch.Tensor:
-        """Apply the full-rank multi-head FFT convolution (shared kernel across batch).
-
-        Calls the pre-selected ``fftconv_fn`` from
-        :mod:`nvsubquadratic.ops.fftconv_multihead`.  The convolution is
-        performed in float32 regardless of the input dtype to avoid numerical
-        instability; the output is cast back to the input dtype before
-        returning.
-
-        The operation per head ``h`` is:
-
-        .. code-block:: none
-
-            y^h = K^h * x^h + shortcut^h ⊙ x^h
-
-        implemented via rFFT in the frequency domain as a dense
-        ``[head_dim × head_dim]`` matrix multiply at each spatial frequency.
-
-        Args:
-            x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``,
-                any floating-point dtype.
-            conv_kernel: Full-rank kernel tensor of shape
-                ``[num_heads, head_dim, head_dim, K_x, K_y]``, float32.
-                ``head_dim`` appears twice — first as ``d_out``, second as
-                ``d_in``.
-            shortcut: Learnable skip-connection scale of shape
-                ``[hidden_dim]``, float32.  Fused into the FFT op.
-
-        Returns:
-            Output tensor of shape ``[B, num_heads, head_dim, H, W]`` in the
-            same dtype as the input ``x``.
-        """
-        x_dtype = x.dtype
-        out = self.fftconv_fn(
-            x.to(torch.float32),
-            conv_kernel.to(torch.float32),
-            shortcut.to(torch.float32),
-        )
-        return out.to(x_dtype)
-
-    def apply_convolution_lowrank(
-        self, x: torch.Tensor, kernel_u: torch.Tensor, kernel_v: torch.Tensor, shortcut: torch.Tensor
-    ) -> torch.Tensor:
-        """Apply the low-rank multi-head FFT convolution (shared kernel across batch).
-
-        Calls the pre-selected low-rank ``fftconv_fn`` from
-        :mod:`nvsubquadratic.ops.fftconv_multihead`.  The two-step contraction
-        avoids materialising the full ``[head_dim × head_dim]`` kernel spectrum
-        per spatial frequency:
-
-        .. code-block:: none
-
-            z = V x          (shape: [B, num_heads, rank, H, W])
-            y = U z          (shape: [B, num_heads, head_dim, H, W])
-            y += shortcut ⊙ x
-
-        The convolution is performed in float32; the output is cast back to
-        the input dtype before returning.
-
-        Args:
-            x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``,
-                any floating-point dtype.
-            kernel_u: Output-projection factor of shape
-                ``[num_heads, head_dim, rank, K_x, K_y]``, float32.
-            kernel_v: Input-projection factor of shape
-                ``[num_heads, rank, head_dim, K_x, K_y]``, float32.
-            shortcut: Learnable skip-connection scale of shape
-                ``[hidden_dim]``, float32.  Fused into the FFT op.
-
-        Returns:
-            Output tensor of shape ``[B, num_heads, head_dim, H, W]`` in the
-            same dtype as the input ``x``.
-        """
-        x_dtype = x.dtype
-        out = self.fftconv_fn(
-            x.to(torch.float32),
-            kernel_u.to(torch.float32),
-            kernel_v.to(torch.float32),
-            shortcut.to(torch.float32),
-        )
-        return out.to(x_dtype)
-
-    def apply_convolution_batched(
-        self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut: torch.Tensor
-    ) -> torch.Tensor:
-        """Apply the full-rank multi-head FFT convolution with per-sample kernels.
-
-        Used when FiLM conditioning is active and each sample in the batch has
-        its own kernel (``conv_kernel.shape[0] == B``).  The FFT convolution is
-        implemented directly via ``torch.fft.rfft2`` / ``irfft2`` and
-        ``torch.einsum``, bypassing the ``fftconv_fn`` (which expects a
-        shared-kernel layout).
-
-        The zero-padded FFT size is computed per-axis as:
-
-        .. code-block:: none
-
-            fft_h = min(H + (K_x + 1) // 2, 2 * H)
-            fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-        which matches the padding convention of the standard (non-batched)
-        FFT convolution ops in :mod:`nvsubquadratic.ops.fftconv_multihead`.
-
-        The frequency-domain operation per sample is:
-
-        .. code-block:: none
-
-            ŷ_{b,n,o,fx,fy} = Σ_i  K̂_{b,n,o,i,fx,fy} · x̂_{b,n,i,fx,fy}
-
-        implemented as a single einsum ``"bnihw,bnoihw->bnohw"``.
-
-        After the inverse FFT, the output is cropped to ``[H, W]`` using a
-        centered crop starting at ``(K_x // 2, K_y // 2)``, and the shortcut
-        term is added.
-
-        Args:
-            x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``,
-                any floating-point dtype.  Cast to float32 internally.
-            conv_kernel: Per-sample full-rank kernel of shape
-                ``[B, num_heads, head_dim, head_dim, K_x, K_y]``, float32.
-                ``head_dim`` appears twice (``d_out``, ``d_in``).
-            shortcut: Learnable skip-connection scale of shape
-                ``[hidden_dim]``, float32.  Reshaped to
-                ``[1, num_heads, head_dim, 1, 1]`` and added as
-                ``shortcut ⊙ x`` after the inverse FFT.
-
-        Returns:
-            Output tensor of shape ``[B, num_heads, head_dim, H, W]``, float32.
-        """
-        x = x.to(torch.float32)
-        conv_kernel = conv_kernel.to(torch.float32)
-        shortcut = shortcut.to(torch.float32)
-
-        _B, num_heads, head_dim, H, W = x.shape
-        K_x, K_y = conv_kernel.shape[-2], conv_kernel.shape[-1]
-
-        fft_h = min(H + (K_x + 1) // 2, 2 * H)
-        fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-        x_fft = torch.fft.rfft2(x, s=(fft_h, fft_w))
-        k_fft = torch.fft.rfft2(conv_kernel, s=(fft_h, fft_w))
-
-        # Batched dense conv: both x and kernel have batch dim
-        out_fft = torch.einsum("bnihw,bnoihw->bnohw", x_fft, k_fft)
-
-        crop_h = K_x // 2
-        crop_w = K_y // 2
-        out_full = torch.fft.irfft2(out_fft, s=(fft_h, fft_w))
-        out = out_full[..., crop_h : crop_h + H, crop_w : crop_w + W]
-
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-        return out
-
-    def apply_convolution_batched_lowrank(
-        self,
-        x: torch.Tensor,
-        kernel_u: torch.Tensor,
-        kernel_v: torch.Tensor,
-        shortcut: torch.Tensor,
-    ) -> torch.Tensor:
-        """Apply the low-rank multi-head FFT convolution with per-sample kernels.
-
-        Used when FiLM conditioning is active and each sample in the batch has
-        its own low-rank kernel pair ``(U, V)``.  The two-step frequency-domain
-        contraction avoids materialising the full ``[head_dim × head_dim]``
-        kernel spectrum per frequency bin:
-
-        .. code-block:: none
-
-            ẑ_{b,n,r,fx,fy} = Σ_i  V̂_{b,n,r,i,fx,fy} · x̂_{b,n,i,fx,fy}   ("bnihw,bnrihw->bnrhw")
-            ŷ_{b,n,o,fx,fy} = Σ_r  Û_{b,n,o,r,fx,fy} · ẑ_{b,n,r,fx,fy}   ("bnrhw,bnorhw->bnohw")
-
-        The same zero-padded FFT size and centered-crop convention as
-        :meth:`apply_convolution_batched` are used.  After the inverse FFT the
-        shortcut term is added element-wise.
-
-        Args:
-            x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``,
-                any floating-point dtype.  Cast to float32 internally.
-            kernel_u: Per-sample output-projection factor of shape
-                ``[B, num_heads, head_dim, rank, K_x, K_y]``, float32.
-            kernel_v: Per-sample input-projection factor of shape
-                ``[B, num_heads, rank, head_dim, K_x, K_y]``, float32.
-            shortcut: Learnable skip-connection scale of shape
-                ``[hidden_dim]``, float32.  Reshaped to
-                ``[1, num_heads, head_dim, 1, 1]`` and added after the inverse
-                FFT.
-
-        Returns:
-            Output tensor of shape ``[B, num_heads, head_dim, H, W]``, float32.
-        """
-        x = x.to(torch.float32)
-        kernel_u = kernel_u.to(torch.float32)
-        kernel_v = kernel_v.to(torch.float32)
-        shortcut = shortcut.to(torch.float32)
-
-        _B, num_heads, head_dim, H, W = x.shape
-        K_x, K_y = kernel_u.shape[-2], kernel_u.shape[-1]
-
-        fft_h = min(H + (K_x + 1) // 2, 2 * H)
-        fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-        x_fft = torch.fft.rfft2(x, s=(fft_h, fft_w))
-        u_fft = torch.fft.rfft2(kernel_u, s=(fft_h, fft_w))
-        v_fft = torch.fft.rfft2(kernel_v, s=(fft_h, fft_w))
-
-        # Two-step low-rank conv (avoids materializing full [head_dim x head_dim] K_fft)
-        # Step 1: z = V @ x — contract over head_dim_in
-        z_fft = torch.einsum("bnihw,bnrihw->bnrhw", x_fft, v_fft)
-        # Step 2: y = U @ z — contract over rank
-        out_fft = torch.einsum("bnrhw,bnorhw->bnohw", z_fft, u_fft)
-
-        crop_h = K_x // 2
-        crop_w = K_y // 2
-        out_full = torch.fft.irfft2(out_fft, s=(fft_h, fft_w))
-        out = out_full[..., crop_h : crop_h + H, crop_w : crop_w + W]
-
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-        return out
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        is_bhl_input: bool = False,
-        cp_group: torch.distributed.ProcessGroup = None,
-        **mixer_kwargs,
-    ) -> torch.Tensor:
-        """Run the CKConvMultiheadND forward pass.
-
-        Generates the per-head implicit kernel from the SIREN network,
-        optionally applies the attenuation mask, reshapes the flat SIREN output
-        into the per-head kernel layout, and applies the FFT convolution with
-        the shortcut term.
-
-        The computation (non-FiLM path) is:
-
-        .. code-block:: none
-
-            # 1. Determine grid size
-            grid_lens = [(s + 1) // 2 for s in (H, W)]   # if grid_type == "single"
-                      = [H, W]                             # if grid_type == "double"
-
-            # 2. Generate kernel from SIREN
-            k_flat, grid = self.kernel(grid_lens)          # [1, K_x, K_y, C_flat]
-
-            # 3. Apply mask (if not Identity)
-            k_flat = self.mask(grid=grid, x=k_flat)
-
-            # 4. Reshape flat output into per-head kernel
-            # Full-rank:  conv_kernel [num_heads, head_dim, head_dim, K_x, K_y]
-            # Low-rank:   kernel_u [num_heads, head_dim, rank, K_x, K_y],
-            #             kernel_v [num_heads, rank, head_dim, K_x, K_y]
-
-            # 5. Apply FFT convolution + shortcut
-            out_heads = apply_convolution*(x_heads, kernel*, shortcut)
-
-        When the kernel is FiLM-conditioned (``batched_kernel=True``), each
-        sample receives its own kernel and the batched convolution methods
-        (:meth:`apply_convolution_batched` or
-        :meth:`apply_convolution_batched_lowrank`) are used instead.
-
-        Args:
-            x: Input signal tensor.  Two supported layouts:
-
-                * **Channels-last** (``is_bhl_input=False``, default): shape
-                  ``[B, H, W, hidden_dim]``.  Rearranged internally to
-                  ``[B, num_heads, head_dim, H, W]`` via einops.
-                * **Channels-first** (``is_bhl_input=True``): shape
-                  ``[B, hidden_dim, H, W]``.  Reshaped internally to
-                  ``[B, num_heads, head_dim, H, W]`` via ``view``.
-
-            is_bhl_input: If ``True``, treat ``x`` as channels-first
-                (BHL) layout.  Default: ``False`` (channels-last / BLH).
-            cp_group: Context-parallel process group.  **Not supported** — if
-                provided and ``cp_group.size() > 1`` a ``NotImplementedError``
-                is raised immediately.  Accepted as a keyword argument for
-                interface compatibility with :class:`CKConvND`.
-            **mixer_kwargs: Additional keyword arguments forwarded to the
-                kernel generator.  Recognised key:
-
-                * ``conditioning`` (``torch.Tensor``, shape ``[B, cond_dim]``):
-                  conditioning vector for FiLM-enabled kernels such as
-                  ``SIRENKernelND`` with a ``film_cfg``.  When supplied, the
-                  SIREN returns a per-sample kernel (batch dimension == B);
-                  otherwise the kernel is shared across the batch (batch
-                  dimension == 1).
-
-        Returns:
-            Output tensor in the same memory layout as the input ``x``:
-
-            * Channels-last: shape ``[B, H, W, hidden_dim]``
-              (when ``is_bhl_input=False``).
-            * Channels-first: shape ``[B, hidden_dim, H, W]``
-              (when ``is_bhl_input=True``).
-
-        Raises:
-            NotImplementedError: If ``cp_group`` is provided with
-                ``cp_group.size() > 1``.  Context parallelism is not yet
-                implemented for ``CKConvMultiheadND``.
-        """
-        if cp_group is not None and cp_group.size() > 1:
-            raise NotImplementedError("Context parallelism not yet supported for CKConvMultiheadND")
-
-        # Get spatial dimensions
-        if is_bhl_input:
-            B, _C, H, W = x.shape
-            # Reshape to [B, num_heads, head_dim, H, W]
-            x_heads = x.view(B, self.num_heads, self.head_dim, H, W)
-        else:
-            B, H, W, _C = x.shape
-            # Reshape to [B, num_heads, head_dim, H, W]
-            x_heads = rearrange(x, "b h w (n d) -> b n d h w", n=self.num_heads, d=self.head_dim)
-
-        spatial_dims = (H, W)
-
-        # Determine grid lengths for kernel generation
-        if self.grid_type == "single":
-            grid_lens = [(s + 1) // 2 for s in spatial_dims]
-        else:  # "double"
-            grid_lens = spatial_dims
-
-        # Generate kernel from SIREN (pass conditioning if available for FiLM-enabled kernels)
-        # Full-rank output: [1, *spatial, num_heads * head_dim * head_dim]
-        # Low-rank output: [1, *spatial, num_heads * 2 * rank * head_dim]
-        conditioning = mixer_kwargs.get("conditioning", None)
-        conv_kernel_flat, grid = self.kernel(grid_lens, conditioning=conditioning)
-
-        # Apply mask if not identity
-        if not isinstance(self.mask, torch.nn.Identity):
-            conv_kernel_flat = self.mask(grid=grid, x=conv_kernel_flat)
-
-        K_x, K_y = conv_kernel_flat.shape[-3], conv_kernel_flat.shape[-2]
-        batched_kernel = conv_kernel_flat.shape[0] > 1
-
-        if self.kernel_rank is not None:
-            # Low-rank path: reshape SIREN output into U and V factors
-            kernel_u, kernel_v = self._reshape_lowrank_kernel(
-                conv_kernel_flat, K_x, K_y, B=B if batched_kernel else None
-            )
-
-            # Cache kernel stats for debugging
-            if getattr(self, "_cache_debug_stats", False):
-                with torch.no_grad():
-                    self._debug_stats = {
-                        "norm_u": kernel_u.norm().item(),
-                        "norm_v": kernel_v.norm().item(),
-                        "max_abs_u": kernel_u.abs().max().item(),
-                        "max_abs_v": kernel_v.abs().max().item(),
-                    }
-
-            if batched_kernel:
-                out_heads = self.apply_convolution_batched_lowrank(x_heads, kernel_u, kernel_v, self.shortcut)
-            else:
-                out_heads = self.apply_convolution_lowrank(x_heads, kernel_u, kernel_v, self.shortcut)
-        else:
-            # Full-rank path: reshape SIREN output into [num_heads, head_dim, head_dim, K_x, K_y]
-            if batched_kernel:
-                conv_kernel = conv_kernel_flat.view(B, K_x, K_y, self.num_heads, self.head_dim, self.head_dim)
-                conv_kernel = conv_kernel.permute(0, 3, 4, 5, 1, 2).contiguous()
-            else:
-                conv_kernel = conv_kernel_flat.view(K_x, K_y, self.num_heads, self.head_dim, self.head_dim)
-                conv_kernel = conv_kernel.permute(2, 3, 4, 0, 1).contiguous()
-
-            # Cache kernel stats for debugging
-            if getattr(self, "_cache_debug_stats", False):
-                with torch.no_grad():
-                    self._debug_stats = {
-                        "norm": conv_kernel.norm().item(),
-                        "max_abs": conv_kernel.abs().max().item(),
-                    }
-
-            if batched_kernel:
-                out_heads = self.apply_convolution_batched(x_heads, conv_kernel, self.shortcut)
-            else:
-                out_heads = self.apply_convolution(x_heads, conv_kernel, self.shortcut)
-
-        # Reshape back to original format
-        if is_bhl_input:
-            # [B, num_heads, head_dim, H, W] -> [B, hidden_dim, H, W]
-            out = out_heads.reshape(B, self.hidden_dim, H, W)
-        else:
-            # [B, num_heads, head_dim, H, W] -> [B, H, W, hidden_dim]
-            out = rearrange(out_heads, "b n d h w -> b h w (n d)")
-
-        return out
diff --git a/nvsubquadratic/modules/ckconv_nd.py b/nvsubquadratic/modules/ckconv_nd.py
index c13a9ac9..55d5ef59 100644
--- a/nvsubquadratic/modules/ckconv_nd.py
+++ b/nvsubquadratic/modules/ckconv_nd.py
@@ -120,7 +120,6 @@
 import copy
 import inspect
 import math
-import warnings
 from collections.abc import Sequence
 from typing import Literal
 
@@ -139,16 +138,6 @@
     circular_fftconv3d_fp32_bhl,
     circular_fftconv3d_fp32_bhl_w_reshape,
 )
-
-# FP16 circular FFT convolutions (requires power-of-2 spatial dimensions)
-from nvsubquadratic.ops.circular_fftconv_fp16 import (
-    circular_fftconv1d_fp16_bhl,
-    circular_fftconv1d_fp16_bhl_w_reshape,
-    circular_fftconv2d_fp16_bhl,
-    circular_fftconv2d_fp16_bhl_w_reshape,
-    circular_fftconv3d_fp16_bhl,
-    circular_fftconv3d_fp16_bhl_w_reshape,
-)
 from nvsubquadratic.ops.fftconv import (
     causal_fftconv1d_fp32_bhl,
     causal_fftconv1d_fp32_bhl_w_reshape,
@@ -187,26 +176,6 @@
     fftconv3d_fp32_bhl_w_reshape as fftconv3d_fp32_bhl_w_reshape_chunked,
 )
 
-# FP16 FFT convolutions (power-of-2 padding + ortho normalization)
-from nvsubquadratic.ops.fftconv_fp16 import (
-    causal_fftconv1d_fp16_bhl,
-    causal_fftconv1d_fp16_bhl_chunked,
-    causal_fftconv1d_fp16_bhl_w_reshape,
-    causal_fftconv1d_fp16_bhl_w_reshape_chunked,
-    fftconv1d_fp16_bhl,
-    fftconv1d_fp16_bhl_chunked,
-    fftconv1d_fp16_bhl_w_reshape,
-    fftconv1d_fp16_bhl_w_reshape_chunked,
-    fftconv2d_fp16_bhl,
-    fftconv2d_fp16_bhl_chunked,
-    fftconv2d_fp16_bhl_w_reshape,
-    fftconv2d_fp16_bhl_w_reshape_chunked,
-    fftconv3d_fp16_bhl,
-    fftconv3d_fp16_bhl_chunked,
-    fftconv3d_fp16_bhl_w_reshape,
-    fftconv3d_fp16_bhl_w_reshape_chunked,
-)
-
 # Mixed boundary-condition FFT convolutions (per-axis periodic / non-periodic).
 # Used when ``fft_padding`` is given as a list of mode strings (e.g.
 # ``["circular", "zero"]``) rather than a single mode; the all-zero /
@@ -262,39 +231,7 @@
     },
 }
 
-# FP16 versions (power-of-2 padding + ortho normalization to prevent overflow)
-# Note: circular fp16 requires power-of-2 spatial dimensions (cuFFT constraint).
-FFT_FUNCTIONS_FP16 = {
-    "circular": {
-        1: (circular_fftconv1d_fp16_bhl_w_reshape, circular_fftconv1d_fp16_bhl),
-        2: (circular_fftconv2d_fp16_bhl_w_reshape, circular_fftconv2d_fp16_bhl),
-        3: (circular_fftconv3d_fp16_bhl_w_reshape, circular_fftconv3d_fp16_bhl),
-    },
-    "zero": {
-        1: (fftconv1d_fp16_bhl_w_reshape, fftconv1d_fp16_bhl),
-        2: (fftconv2d_fp16_bhl_w_reshape, fftconv2d_fp16_bhl),
-        3: (fftconv3d_fp16_bhl_w_reshape, fftconv3d_fp16_bhl),
-    },
-    "causal": {
-        1: (causal_fftconv1d_fp16_bhl_w_reshape, causal_fftconv1d_fp16_bhl),
-        # Causal is only supported for 1D (sequences)
-    },
-}
-
-# FP16 + chunked: combines fp16 memory savings with channel-chunking savings
-FFT_FUNCTIONS_FP16_CHUNKED = {
-    "zero": {
-        1: (fftconv1d_fp16_bhl_w_reshape_chunked, fftconv1d_fp16_bhl_chunked),
-        2: (fftconv2d_fp16_bhl_w_reshape_chunked, fftconv2d_fp16_bhl_chunked),
-        3: (fftconv3d_fp16_bhl_w_reshape_chunked, fftconv3d_fp16_bhl_chunked),
-    },
-    "causal": {
-        1: (causal_fftconv1d_fp16_bhl_w_reshape_chunked, causal_fftconv1d_fp16_bhl_chunked),
-        # Causal is only supported for 1D (sequences)
-    },
-}
-
-# Mixed-BC FFT convolutions: only fp32 in v1 (see docs/ops/MIXED_BC_PLAN.md).
+# Mixed-BC FFT convolutions: fp32 (see docs/ops/mixed_boundary_conditions.md).
 # Each entry is ``(fn_for_BLH_input (bhl_w_reshape), fn_for_BHL_input)`` and
 # takes an additional ``periodic`` argument compared to the legacy ops; the
 # wrapper ``_wrap_mixed_op`` below adapts the call signature.
@@ -560,7 +497,6 @@ class CKConvND(torch.nn.Module):
             convolution.  Only valid when ``data_dim=1``.
         use_chunked_fftconv (bool): Whether to process channels in chunks to
             reduce peak GPU memory.
-        use_fp16_fft (bool): Whether to use fp16 FFT convolution ops.
         fft_backend (str): FFT backend identifier, ``"torch_fft"`` or
             ``"subq_ops"``.
         grid_type (str or None): Kernel grid size mode (``"single"``,
@@ -595,13 +531,12 @@ def __init__(
         fft_padding: "Literal['zero', 'circular'] | str | Sequence[str]",
         is_causal: bool = False,
         use_chunked_fftconv: bool = False,
-        use_fp16_fft: bool = False,
         fft_backend: Literal["torch_fft", "subq_ops"] = "torch_fft",
     ):
         """Construct a CKConvND operator.
 
         Validates the combination of ``fft_padding``, ``grid_type``,
-        ``is_causal``, ``use_fp16_fft``, and ``fft_backend``, normalises the
+        ``is_causal``, and ``fft_backend``, normalises the
         per-axis boundary-condition representation, adjusts ``kernel_cfg``
         and ``mask_cfg`` to match the resolved kernel grid geometry, and
         selects the appropriate FFT convolution function pair.
@@ -647,7 +582,7 @@ def __init__(
                 * ``["circular", "zero"]`` (list/tuple of mode strings, one
                   per spatial axis, length must equal ``data_dim``): per-axis
                   mixed boundary conditions.  Requires ``grid_type=None`` and
-                  ``fft_backend="torch_fft"`` and ``use_fp16_fft=False``.
+                  ``fft_backend="torch_fft"``.
                   Mode names are case-insensitive and whitespace-stripped.
 
                 Must be ``"zero"`` (or an all-``"zero"`` list) when
@@ -662,16 +597,6 @@ def __init__(
                 Typical savings: ~26% memory at ~11% compute overhead.
                 Not supported with ``fft_padding="circular"``.
                 Default: ``False``.
-            use_fp16_fft: If ``True``, use fp16 FFT convolution ops.
-                Uses ``norm="ortho"`` internally to prevent overflow.  Saves
-                ~36% peak memory per convolution at ~0.8% mean relative error
-                vs fp32.  For zero/causal padding, spatial dims are
-                auto-padded to the next power of two.  For circular padding,
-                the input dims must already be powers of two (a runtime
-                assertion fires otherwise).  Not supported with a per-axis
-                ``fft_padding`` list (see ``docs/ops/MIXED_BC_PLAN.md``).
-                Not supported with ``fft_backend="subq_ops"``.
-                Default: ``False``.
             fft_backend: Which FFT convolution backend to use.
 
                 * ``"torch_fft"`` (default): torch.fft-based implementations
@@ -692,16 +617,13 @@ def __init__(
         Raises:
             AssertionError: If ``fft_backend`` is not one of the recognised
                 values, or if a constraint between ``grid_type``,
-                ``fft_padding``, ``is_causal``, ``use_fp16_fft``, and
-                ``fft_backend`` is violated.
+                ``fft_padding``, ``is_causal``, and ``fft_backend`` is
+                violated.
             ValueError: If ``fft_padding`` is invalid (wrong type, wrong
                 length, comma-separated string, boolean), if ``is_causal``
                 is combined with a per-axis padding list or periodic padding,
                 or if the resolved ``(fft_padding, data_dim)`` combination
                 has no registered FFT function.
-            NotImplementedError: If ``use_fp16_fft=True`` is requested
-                together with a per-axis ``fft_padding`` list (fp16 mixed
-                ops are not yet implemented).
         """
         assert fft_backend in ["torch_fft", "subq_ops"], (
             f"Invalid fft_backend: {fft_backend!r}. Must be 'torch_fft' or 'subq_ops'."
@@ -775,22 +697,6 @@ def __init__(
                 "Circular convolutions already have lower memory overhead due to no padding."
             )
 
-        # ---- fp16 + mixed-BC: not supported in v1 -----------------------------
-        if use_fp16_fft and _is_tuple_mode:
-            raise NotImplementedError(
-                "use_fp16_fft is not supported with a per-axis fft_padding in v1. "
-                "Either drop the fp16 flag or use a uniform 'zero'/'circular' fft_padding. "
-                "See docs/ops/MIXED_BC_PLAN.md (§4.2) for the planned fp16 mixed op."
-            )
-
-        if use_fp16_fft and not _is_tuple_mode and fft_padding == "circular":
-            warnings.warn(
-                "use_fp16_fft with circular padding requires power-of-2 spatial "
-                "dimensions (cuFFT fp16 constraint). A runtime assertion will fire "
-                "if the input is not power-of-2.",
-                stacklevel=2,
-            )
-
         # subq_ops backend constraints
         if fft_backend == "subq_ops":
             if _is_tuple_mode:
@@ -816,18 +722,12 @@ def __init__(
                 raise AssertionError(
                     f"fft_backend='subq_ops' only supports data_dim in (1, 2). Got data_dim={data_dim}."
                 )
-            assert not use_fp16_fft, (
-                "fft_backend='subq_ops' does not support fp16 FFT — the CUDA kernel "
-                "manages its own precision internally. Use use_fp16_fft=False."
-            )
-
         super().__init__()
         self.data_dim = data_dim
         self.hidden_dim = hidden_dim
         self.fft_padding = fft_padding
         self.is_causal = is_causal
         self.use_chunked_fftconv = use_chunked_fftconv
-        self.use_fp16_fft = use_fp16_fft
         self.fft_backend = fft_backend
         # Per-axis BC: single source of truth used by forward() and flop_count().
         # Always present (length == data_dim), even in legacy single-mode form.
@@ -960,15 +860,7 @@ def __init__(
             else:
                 effective_padding = "zero"
 
-            # Choose FFT functions: fp16+chunked > fp16 > chunked > standard
-            if use_fp16_fft and use_chunked_fftconv:
-                fft_fn_table = FFT_FUNCTIONS_FP16_CHUNKED
-            elif use_fp16_fft:
-                fft_fn_table = FFT_FUNCTIONS_FP16
-            elif use_chunked_fftconv:
-                fft_fn_table = FFT_FUNCTIONS_CHUNKED
-            else:
-                fft_fn_table = FFT_FUNCTIONS
+            fft_fn_table = FFT_FUNCTIONS_CHUNKED if use_chunked_fftconv else FFT_FUNCTIONS
             try:
                 self.fftconv_fn, self.fftconv_fn_bhl_input = fft_fn_table[effective_padding][self.data_dim]
             except KeyError:
@@ -991,7 +883,7 @@ def extra_repr(self) -> str:
             ``data_dim``, ``hidden_dim``, ``fft_padding``,
             ``periodic_per_axis`` (only when in per-axis list mode),
             ``grid_type``, ``is_causal``, ``use_chunked_fftconv``,
-            ``use_fp16_fft``, and ``fft_backend``.
+            and ``fft_backend``.
         """
         bc_repr = f"fft_padding={self.fft_padding!r}"
         if self._is_tuple_mode:
@@ -999,7 +891,7 @@ def extra_repr(self) -> str:
         return (
             f"data_dim={self.data_dim}, hidden_dim={self.hidden_dim}, "
             f"{bc_repr}, grid_type={self.grid_type!r}, is_causal={self.is_causal}, "
-            f"use_chunked_fftconv={self.use_chunked_fftconv}, use_fp16_fft={self.use_fp16_fft}, "
+            f"use_chunked_fftconv={self.use_chunked_fftconv}, "
             f"fft_backend={self.fft_backend!r}"
         )
 
diff --git a/nvsubquadratic/modules/condition_mixer.py b/nvsubquadratic/modules/condition_mixer.py
index 6424cc41..278f18df 100644
--- a/nvsubquadratic/modules/condition_mixer.py
+++ b/nvsubquadratic/modules/condition_mixer.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """QKV cross-attention condition mixer for injecting a conditioning signal into a feature map.
 
 **Role in the residual block**
diff --git a/nvsubquadratic/networks/baselines/unet_convnext.py b/nvsubquadratic/networks/baselines/unet_convnext.py
index eec442c5..5a447f99 100644
--- a/nvsubquadratic/networks/baselines/unet_convnext.py
+++ b/nvsubquadratic/networks/baselines/unet_convnext.py
@@ -1,4 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024 Polymathic AI.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,6 +13,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
+# Portions ported from PolymathicAI/the_well (BSD-3-Clause), itself adapted from
+# facebookresearch/ConvNeXt (MIT; Copyright (c) Meta Platforms, Inc. and affiliates).
+# See THIRD_PARTY_NOTICES.txt for the full BSD-3-Clause and MIT license texts.
 
 """UNet-ConvNeXt baseline from The Well benchmark.
 
diff --git a/nvsubquadratic/networks/baselines/unet_convnext_v2.py b/nvsubquadratic/networks/baselines/unet_convnext_v2.py
index 672c6d18..4583cf9c 100644
--- a/nvsubquadratic/networks/baselines/unet_convnext_v2.py
+++ b/nvsubquadratic/networks/baselines/unet_convnext_v2.py
@@ -1,4 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024 Polymathic AI.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,6 +13,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
+# Builds on unet_convnext.py, which is ported from PolymathicAI/the_well
+# (BSD-3-Clause), itself adapted from facebookresearch/ConvNeXt (MIT; Copyright
+# (c) Meta Platforms, Inc. and affiliates). See THIRD_PARTY_NOTICES.txt.
 
 """UNet-ConvNeXt V2 — fixed skip connections.
 
diff --git a/nvsubquadratic/networks/classification_resnet.py b/nvsubquadratic/networks/classification_resnet.py
index e021009c..36c8798e 100644
--- a/nvsubquadratic/networks/classification_resnet.py
+++ b/nvsubquadratic/networks/classification_resnet.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """Classification residual network — global-average-pool readout.
 
 :class:`ClassificationResNet` subclasses :class:`ResidualNetwork` and overrides
@@ -38,8 +36,6 @@
 :class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork`;
 ``target_size`` is not used by this subclass (the GAP step replaces the
 readout-crop mechanism).
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 import torch
diff --git a/nvsubquadratic/networks/general_purpose_resnet.py b/nvsubquadratic/networks/general_purpose_resnet.py
index 2543917c..089f4c59 100644
--- a/nvsubquadratic/networks/general_purpose_resnet.py
+++ b/nvsubquadratic/networks/general_purpose_resnet.py
@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Adapted from https://github.com/implicit-long-convs/ccnn_v2
-
 """General-purpose residual network backbone.
 
 :class:`ResidualNetwork` is a flexible, task-agnostic sequence-of-blocks backbone
@@ -57,8 +55,6 @@
 
 Set ``gradient_checkpointing=True`` to recompute activations during the backward
 pass instead of storing them, trading compute for memory at large scale.
-
-Adapted from https://github.com/implicit-long-convs/ccnn_v2.
 """
 
 from typing import Sequence
diff --git a/nvsubquadratic/networks/huggingface_diffusers.py b/nvsubquadratic/networks/huggingface_diffusers.py
deleted file mode 100644
index 5c58f3b9..00000000
--- a/nvsubquadratic/networks/huggingface_diffusers.py
+++ /dev/null
@@ -1,471 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Hugging Face ``diffusers`` adapters for the nvSubquadratic diffusion pipeline.
-
-This module bridges the two-sided contract between the
-``experiments.lightning_wrappers.diffusion_wrapper.DiffusionWrapper`` (which
-drives noise schedules, timestep conditioning, and the optimisation loop)
-and stock HF backbones (:class:`diffusers.DiTTransformer2DModel`,
-:class:`diffusers.UVit2DModel`) that expect their own input layout and
-timestep API.
-
-Two thin ``nn.Module`` wrappers do the translation:
-
-- :class:`DiffusersDiTWrapper` — for DiT-style denoisers that take a
-  ``(B, C, H, W)`` sample and a ``[B]`` discrete-timestep tensor.
-- :class:`DiffusersUVitWrapper` — for the UVit codebook predictor that
-  takes a ``(B, C, H, W)`` sample plus an explicit conditioning tensor.
-
-Both wrappers expose the same ``forward({"input": ...})`` -> ``{"logits": ...}``
-contract that the diffusion wrapper expects, hide the HF model's preferred
-channel layout (BLH inside the library; BCHW for the HF backbones), and
-share a small ``_SharedTimestepState`` object so that the ``DiffusionWrapper``
-can push the current timestep into the HF model just before each forward
-without polluting the wrapper's own constructor signature.  Registration
-happens via :meth:`DiffusersDiTWrapper.hf_register_diffusion_wrapper`, which
-monkey-patches the wrapper's ``_condition_from_timesteps`` to fan out to a
-list of timestep callbacks — this is intentional and reversible
-(``_hf_timestep_callbacks`` is created lazily and only patched once per
-wrapper instance).
-
-Dtype handling: the configs carry a ``dtype`` string (``"float32"``,
-``"bfloat16"``, …); the wrapper casts the sample to that dtype before the
-HF model and casts the output back to the caller's dtype on the way out.
-"""
-
-from __future__ import annotations
-
-import types
-import weakref
-from dataclasses import asdict, dataclass, is_dataclass
-from typing import TYPE_CHECKING, Callable, Iterable
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.models import DiTTransformer2DModel
-from diffusers.models.unets.uvit_2d import UVit2DModel
-
-
-if TYPE_CHECKING:
-    from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-
-
-class _SharedTimestepState:
-    __slots__ = ("latest",)
-
-    def __init__(self) -> None:
-        self.latest: torch.Tensor | None = None
-
-    def __deepcopy__(self, memo) -> "_SharedTimestepState":
-        memo[id(self)] = self
-        return self
-
-
-@dataclass
-class HuggingFaceDiTConfig:
-    """Construction config for a :class:`diffusers.DiTTransformer2DModel` backbone.
-
-    Every field maps 1:1 to the corresponding HF constructor argument.  Fields
-    set to ``None`` are filtered out before being passed to HF so that the
-    library's own defaults apply.  ``dtype`` is a string (e.g. ``"float32"``,
-    ``"bfloat16"``) used at forward time to cast the sample tensor before the
-    HF model runs; weights themselves are not re-cast here.
-    """
-
-    sample_size: int = 28
-    patch_size: int | None = 2
-    in_channels: int = 1
-    out_channels: int | None = None
-
-    num_layers: int = 12
-    num_attention_heads: int = 16
-    attention_head_dim: int = 64
-    dropout: float = 0.0
-    norm_num_groups: int = 32
-    attention_bias: bool = False
-    activation_fn: str = "gelu-approximate"
-    num_embeds_ada_norm: int | None = 1_000
-    upcast_attention: bool = False
-    norm_type: str = "ada_norm_zero"
-    norm_elementwise_affine: bool = True
-    norm_eps: float = 1e-6
-
-    dtype: str = "float32"
-
-
-@dataclass
-class HuggingFaceUVitConfig:
-    """Construction config for a :class:`diffusers.UVit2DModel` backbone.
-
-    Fields map 1:1 to the HF constructor; ``None`` values are filtered out
-    so the library defaults apply.  Note that UVit predicts logits over a
-    discrete codebook rather than continuous samples, so ``out_channels`` is
-    accepted for interface parity but is *not* propagated to the underlying
-    model.
-    """
-
-    sample_size: int = 32
-    in_channels: int = 3
-    out_channels: int | None = None
-
-    hidden_size: int = 256
-    cond_embed_dim: int = 128
-    encoder_hidden_size: int = 128
-    block_out_channels: int = 256
-    num_hidden_layers: int = 8
-    num_attention_heads: int = 8
-    intermediate_size: int = 512
-    layer_norm_eps: float = 1e-5
-    micro_cond_encode_dim: int | None = None
-    micro_cond_embed_dim: int | None = None
-    codebook_size: int | None = None
-    vocab_size: int | None = None
-
-    dtype: str = "float32"
-
-
-class DiffusersDiTWrapper(nn.Module):
-    """Adapter that exposes a :class:`diffusers.DiTTransformer2DModel` as an nvSubquadratic denoiser.
-
-    The HF DiT expects ``(B, C, H, W)`` samples and a ``[B]`` long-dtype
-    timestep tensor; the diffusion wrapper hands the model
-    ``{"input": (B, H, W, C)}`` and pushes timesteps through a side channel.
-    This adapter:
-
-    - moveaxes ``(B, H, W, C) -> (B, C, H, W)`` on input and back on output;
-    - casts the sample to ``hf_config.dtype`` and restores the caller's dtype;
-    - reads the latest timesteps from a shared state object that the diffusion
-      wrapper populates via the callback installed by
-      :meth:`hf_register_diffusion_wrapper`;
-    - synthesises a zero class-label tensor (the configs we use don't
-      condition on labels) so the HF API doesn't complain.
-
-    Call :meth:`hf_register_diffusion_wrapper` once after construction —
-    typically by the diffusion wrapper itself — to wire the timestep
-    callback before the first forward.
-    """
-
-    def __init__(
-        self,
-        hf_config: HuggingFaceDiTConfig,
-        in_channels: int | None = None,
-        out_channels: int | None = None,
-    ) -> None:
-        """Instantiate the wrapped DiT model with optional channel overrides."""
-        super().__init__()
-        self.hf_config = hf_config
-        if in_channels is not None:
-            self.hf_config.in_channels = in_channels
-        if out_channels is not None:
-            self.hf_config.out_channels = out_channels
-
-        self.transformer = self._build_transformer(self.hf_config)
-        self.hidden_dim = self.transformer.config.attention_head_dim * self.transformer.config.num_attention_heads
-
-        self._timestep_state = _SharedTimestepState()
-        self._registered_wrapper_ref: weakref.ReferenceType | None = None
-
-    def _get_latest_timesteps(self) -> torch.Tensor | None:
-        return self._timestep_state.latest
-
-    def _set_latest_timesteps(self, value: torch.Tensor | None) -> None:
-        self._timestep_state.latest = value
-
-    _latest_timesteps = property(_get_latest_timesteps, _set_latest_timesteps)
-
-    def _build_transformer(self, cfg: HuggingFaceDiTConfig) -> nn.Module:
-        if DiTTransformer2DModel is None:
-            raise ImportError("diffusers>=0.27 with DiTTransformer2DModel support is required for DiffusersDiTWrapper")
-
-        def _filtered_kwargs(pairs: Iterable[tuple[str, object | None]]) -> dict[str, object]:
-            return {k: v for k, v in pairs if v is not None}
-
-        shared_args = (
-            ("sample_size", cfg.sample_size),
-            ("patch_size", cfg.patch_size),
-            ("in_channels", cfg.in_channels),
-            ("out_channels", cfg.out_channels),
-            ("num_layers", cfg.num_layers),
-            ("num_attention_heads", cfg.num_attention_heads),
-            ("attention_head_dim", cfg.attention_head_dim),
-            ("dropout", cfg.dropout),
-            ("norm_num_groups", cfg.norm_num_groups),
-            ("attention_bias", cfg.attention_bias),
-            ("activation_fn", cfg.activation_fn),
-            ("num_embeds_ada_norm", cfg.num_embeds_ada_norm),
-            ("upcast_attention", cfg.upcast_attention),
-            ("norm_type", cfg.norm_type),
-            ("norm_elementwise_affine", cfg.norm_elementwise_affine),
-            ("norm_eps", cfg.norm_eps),
-        )
-
-        kwargs = _filtered_kwargs(shared_args)
-        return DiTTransformer2DModel(**kwargs)
-
-    def hf_register_diffusion_wrapper(self, wrapper: "DiffusionWrapper") -> None:
-        """Attach a diffusion wrapper so timestep hooks can be updated."""
-        if self._registered_wrapper_ref is not None and self._registered_wrapper_ref() is wrapper:
-            return
-
-        self._registered_wrapper_ref = weakref.ref(wrapper)
-
-        if not hasattr(wrapper, "_hf_timestep_callbacks"):
-            original_condition = wrapper._condition_from_timesteps
-
-            def patched_condition(
-                self_wrapper,
-                timesteps: torch.LongTensor,
-                *args,
-                **kwargs,
-            ) -> torch.Tensor:
-                conditioned = original_condition(timesteps, *args, **kwargs)
-                for callback in getattr(self_wrapper, "_hf_timestep_callbacks", []):
-                    callback(timesteps)
-                return conditioned
-
-            wrapper._hf_timestep_callbacks: list[Callable[[torch.LongTensor], None]] = []
-            wrapper._condition_from_timesteps = types.MethodType(patched_condition, wrapper)
-
-        def _update_timesteps(timesteps: torch.LongTensor) -> None:
-            self._latest_timesteps = timesteps
-
-        wrapper._hf_timestep_callbacks.append(_update_timesteps)
-
-    def forward(self, input_and_condition: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        """Forward pass that mirrors the diffusers DiT interface."""
-        if self._latest_timesteps is None:
-            raise RuntimeError(
-                "DiffusersDiTWrapper.forward() was called before timesteps were populated. "
-                "Ensure that hf_register_diffusion_wrapper has been called."
-            )
-
-        noisy_sample = input_and_condition["input"]
-        sample_bchw = torch.moveaxis(noisy_sample, -1, 1).contiguous()
-
-        try:
-            model_dtype = next(self.transformer.parameters()).dtype
-        except StopIteration:  # pragma: no cover
-            model_dtype = sample_bchw.dtype
-        sample_bchw = sample_bchw.to(dtype=getattr(torch, self.hf_config.dtype, model_dtype))
-
-        timesteps = self._latest_timesteps.to(device=sample_bchw.device, dtype=torch.long)
-
-        if isinstance(self.transformer, DiTTransformer2DModel):
-            class_labels = torch.zeros(sample_bchw.shape[0], device=sample_bchw.device, dtype=torch.long)
-            output = self.transformer(sample_bchw, timestep=timesteps, class_labels=class_labels, return_dict=True)
-        else:
-            output = self.transformer(sample_bchw, timestep=timesteps, return_dict=True)
-
-        prediction = torch.moveaxis(output.sample, 1, -1).contiguous().to(dtype=noisy_sample.dtype)
-        return {"logits": prediction}
-
-    def extra_repr(self) -> str:  # pragma: no cover - debugging helper
-        """Return a compact repr of the underlying transformer and config."""
-        if hasattr(self.hf_config, "items"):
-            items = self.hf_config.items()
-        elif is_dataclass(self.hf_config):
-            items = asdict(self.hf_config).items()
-        else:
-            items = vars(self.hf_config).items()
-        trimmed = {k: v for k, v in items if v is not None}
-        return f"transformer={self.transformer.__class__.__name__}, config={trimmed}"
-
-
-class DiffusersUVitWrapper(nn.Module):
-    """Adapter that exposes a :class:`diffusers.UVit2DModel` as an nvSubquadratic discrete-codebook denoiser.
-
-    UVit predicts logits over a discrete codebook rather than continuous
-    samples, so the forward returns ``{"logits": [B, codebook_size, H, W]}``
-    rather than a reconstructed sample.  Layout/dtype handling mirrors
-    :class:`DiffusersDiTWrapper`.  Conditioning tensors that the HF UVit
-    requires (text embeddings or similar) are routed through the registered
-    diffusion wrapper.
-    """
-
-    def __init__(
-        self,
-        hf_config: HuggingFaceUVitConfig,
-        in_channels: int | None = None,
-        out_channels: int | None = None,
-    ) -> None:
-        """Instantiate the wrapped UVit model with optional channel overrides."""
-        super().__init__()
-        if UVit2DModel is None:
-            raise ImportError("diffusers>=0.35 with UVit2DModel support is required for DiffusersUVitWrapper")
-
-        self.hf_config = hf_config
-        if in_channels is not None:
-            self.hf_config.in_channels = in_channels
-        # UVit outputs logits over the latent codebook; out_channels unused but accepted for interface parity
-        self.hf_config.out_channels = out_channels
-
-        self.transformer = self._build_uvit(self.hf_config)
-        self.hidden_dim = getattr(self.transformer.config, "block_out_channels", self.hf_config.in_channels)
-
-        self._registered_wrapper_ref: weakref.ReferenceType | None = None
-        self._timestep_state = _SharedTimestepState()
-
-    def _get_latest_timesteps(self) -> torch.Tensor | None:
-        return self._timestep_state.latest
-
-    def _set_latest_timesteps(self, value: torch.Tensor | None) -> None:
-        self._timestep_state.latest = value
-
-    _latest_timesteps = property(_get_latest_timesteps, _set_latest_timesteps)
-
-    def _build_uvit(self, cfg: HuggingFaceUVitConfig) -> UVit2DModel:
-        kwargs = {
-            "sample_size": cfg.sample_size,
-            "in_channels": cfg.in_channels,
-            "hidden_size": cfg.hidden_size,
-            "cond_embed_dim": cfg.cond_embed_dim,
-            "encoder_hidden_size": cfg.encoder_hidden_size,
-            "block_out_channels": cfg.block_out_channels,
-            "num_hidden_layers": cfg.num_hidden_layers,
-            "num_attention_heads": cfg.num_attention_heads,
-            "intermediate_size": cfg.intermediate_size,
-            "layer_norm_eps": cfg.layer_norm_eps,
-        }
-        if cfg.micro_cond_encode_dim is not None:
-            kwargs["micro_cond_encode_dim"] = cfg.micro_cond_encode_dim
-        if cfg.micro_cond_embed_dim is not None:
-            kwargs["micro_cond_embed_dim"] = cfg.micro_cond_embed_dim
-        if cfg.codebook_size is not None:
-            kwargs["codebook_size"] = cfg.codebook_size
-        if cfg.vocab_size is not None:
-            kwargs["vocab_size"] = cfg.vocab_size
-        return UVit2DModel(**kwargs)
-
-    def hf_register_diffusion_wrapper(self, wrapper: "DiffusionWrapper") -> None:
-        """Attach a diffusion wrapper so timestep callbacks receive updates."""
-        if self._registered_wrapper_ref is not None and self._registered_wrapper_ref() is wrapper:
-            return
-
-        self._registered_wrapper_ref = weakref.ref(wrapper)
-
-        if not hasattr(wrapper, "_hf_timestep_callbacks"):
-            original_condition = wrapper._condition_from_timesteps
-
-            def patched_condition(
-                self_wrapper,
-                timesteps: torch.LongTensor,
-                *args,
-                **kwargs,
-            ) -> torch.Tensor:
-                conditioned = original_condition(timesteps, *args, **kwargs)
-                for callback in getattr(self_wrapper, "_hf_timestep_callbacks", []):
-                    callback(timesteps)
-                return conditioned
-
-            wrapper._hf_timestep_callbacks: list[Callable[[torch.LongTensor], None]] = []
-            wrapper._condition_from_timesteps = types.MethodType(patched_condition, wrapper)
-
-        def _update_timesteps(timesteps: torch.LongTensor) -> None:
-            self._latest_timesteps = timesteps
-
-        wrapper._hf_timestep_callbacks.append(_update_timesteps)
-
-    def forward(self, input_and_condition: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        """Forward pass of the UVit wrapper."""
-        batch = input_and_condition["input"]
-        sample_bchw = torch.moveaxis(batch, -1, 1).contiguous()
-
-        conditioning = input_and_condition.get("condition")
-        config = self.transformer.config
-        batch_size = sample_bchw.shape[0]
-        device = sample_bchw.device
-        dtype = sample_bchw.dtype
-
-        if isinstance(conditioning, dict):
-            encoder_hidden_states = conditioning.get("encoder_hidden_states")
-            pooled_text_emb = conditioning.get("pooled_text_emb")
-            micro_conds = conditioning.get("micro_conds")
-            input_ids = conditioning.get("input_ids", sample_bchw)
-        else:
-            encoder_hidden_states = None
-            pooled_text_emb = conditioning
-            micro_conds = None
-            input_ids = sample_bchw
-
-        if encoder_hidden_states is None:
-            enc_dim = getattr(config, "encoder_hidden_size", self.hidden_dim)
-            encoder_hidden_states = torch.zeros(batch_size, 1, enc_dim, device=device, dtype=dtype)
-
-        cond_dim = getattr(
-            config,
-            "cond_embed_dim",
-            pooled_text_emb.shape[-1] if isinstance(pooled_text_emb, torch.Tensor) else self.hidden_dim,
-        )
-        if pooled_text_emb is None:
-            pooled_text_emb = torch.zeros(batch_size, cond_dim, device=device, dtype=dtype)
-        else:
-            pooled_text_emb = pooled_text_emb.to(device=device, dtype=dtype)
-            if pooled_text_emb.shape[-1] != cond_dim:
-                if pooled_text_emb.shape[-1] < cond_dim:
-                    pad = cond_dim - pooled_text_emb.shape[-1]
-                    pooled_text_emb = F.pad(pooled_text_emb, (0, pad))
-                else:
-                    pooled_text_emb = pooled_text_emb[..., :cond_dim]
-
-        if micro_conds is None:
-            if self._latest_timesteps is None:
-                raise RuntimeError(
-                    "UVit wrapper requires timestep information; ensure hf_register_diffusion_wrapper was invoked"
-                )
-            timesteps = self._latest_timesteps.to(device=device, dtype=torch.float32)
-            micro_encode_dim = getattr(config, "micro_cond_encode_dim", 1) or 1
-            micro_embed_dim = getattr(config, "micro_cond_embed_dim", micro_encode_dim)
-            repeat = max(1, micro_embed_dim // micro_encode_dim)
-            micro_conds = timesteps.unsqueeze(1).repeat(1, repeat)
-        else:
-            micro_conds = micro_conds.to(device=device)
-
-        input_ids = self._ensure_token_ids(input_ids)
-
-        logits = self.transformer(
-            input_ids=input_ids,
-            encoder_hidden_states=encoder_hidden_states,
-            pooled_text_emb=pooled_text_emb,
-            micro_conds=micro_conds,
-            cross_attention_kwargs=None,
-        )
-
-        target_channels = self.hf_config.out_channels or self.hf_config.in_channels
-        if logits.shape[1] != target_channels:
-            logits = logits[:, :target_channels]
-
-        prediction = torch.moveaxis(logits, 1, -1).contiguous().to(dtype=batch.dtype)
-        return {"logits": prediction}
-
-    def _ensure_token_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        if not torch.is_floating_point(input_ids):
-            tokens = input_ids.long()
-        else:
-            if input_ids.dim() == 4:
-                values = input_ids.mean(dim=1)
-            elif input_ids.dim() == 3:
-                values = input_ids
-            else:
-                raise ValueError("UVit wrapper expects conditioning images shaped [B, C, H, W] or [B, H, W]")
-
-            vocab_size = getattr(self.transformer.config, "vocab_size", 8192)
-            tokens = ((values + 1.0) * (vocab_size - 1) / 2.0).round()
-            tokens = tokens.clamp_(0, vocab_size - 1).to(dtype=torch.long)
-
-        if tokens.dim() == 4 and tokens.shape[1] == 1:
-            tokens = tokens.squeeze(1)
-        return tokens
diff --git a/nvsubquadratic/networks/jit.py b/nvsubquadratic/networks/jit.py
deleted file mode 100644
index e569f949..00000000
--- a/nvsubquadratic/networks/jit.py
+++ /dev/null
@@ -1,521 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-FileCopyrightText: Copyright (c) 2023 The JiT authors and contributors.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# Portions of this file are derived from https://github.com/LTH14/JiT
-# (MIT License). See LICENSE/third_party.txt for the full MIT notice.
-
-"""JiT model implementation.
-
-Ported from https://github.com/LTH14/JiT.
-"""
-
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from nvsubquadratic.networks.jit_utils import (
-    RMSNorm,
-    VisionRotaryEmbeddingFast,
-    get_2d_sincos_pos_embed,
-)
-
-
-def modulate(x, shift, scale):
-    """Apply adaptive scale/shift modulation to token features."""
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-class BottleneckPatchEmbed(nn.Module):
-    """Image-to-patch embedding with a bottleneck projection."""
-
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
-        """Initialize the bottleneck patch embedding module."""
-        super().__init__()
-        img_size = (img_size, img_size)
-        patch_size = (patch_size, patch_size)
-        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = num_patches
-
-        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
-        self.proj2 = nn.Conv2d(pca_dim, embed_dim, kernel_size=1, stride=1, bias=bias)
-
-    def forward(self, x):
-        """Project an input image batch into patch tokens."""
-        _B, _C, H, W = x.shape
-        assert H == self.img_size[0] and W == self.img_size[1], (
-            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        )
-        x = self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
-        return x
-
-
-class TimestepEmbedder(nn.Module):
-    """Embeds scalar timesteps into vector representations."""
-
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        """Initialize the timestep embedding MLP."""
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        """Create sinusoidal timestep embeddings.
-
-        :param t: a 1-D Tensor of N indices, one per batch element.
-                          These may be fractional.
-        :param dim: the dimension of the output.
-        :param max_period: controls the minimum frequency of the embeddings.
-        :return: an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        half = dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-            device=t.device
-        )
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t):
-        """Embed scalar timesteps into the model hidden space."""
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class LabelEmbedder(nn.Module):
-    """Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance."""
-
-    def __init__(self, num_classes, hidden_size):
-        """Initialize the class-label embedding table."""
-        super().__init__()
-        self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)
-        self.num_classes = num_classes
-
-    def forward(self, labels):
-        """Embed integer class labels."""
-        embeddings = self.embedding_table(labels)
-        return embeddings
-
-
-def scaled_dot_product_attention(query, key, value, dropout_p=0.0) -> torch.Tensor:
-    """Compute scaled dot-product attention via the native PyTorch op."""
-    # Use native flash attention!
-    # No need for manual attn_bias of zeros, it just disables flash attention.
-    return F.scaled_dot_product_attention(query, key, value, dropout_p=dropout_p)
-
-
-class Attention(nn.Module):
-    """Multi-head self-attention with optional Q/K normalization and RoPE."""
-
-    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True, attn_drop=0.0, proj_drop=0.0):
-        """Initialize the self-attention module."""
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-
-        self.q_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
-        self.k_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-    def forward(self, x, rope):
-        """Apply attention to token sequence ``x`` using the provided RoPE module."""
-        B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = (
-            qkv[0],
-            qkv[1],
-            qkv[2],
-        )  # make torchscript happy (cannot use tensor as tuple)
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        q = rope(q)
-        k = rope(k)
-
-        x = scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
-
-        x = x.transpose(1, 2).reshape(B, N, C)
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class SwiGLUFFN(nn.Module):
-    """SwiGLU feed-forward network used inside transformer blocks."""
-
-    def __init__(self, dim: int, hidden_dim: int, drop=0.0, bias=True) -> None:
-        """Initialize the SwiGLU feed-forward network."""
-        super().__init__()
-        hidden_dim = int(hidden_dim * 2 / 3)
-        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=bias)
-        self.w3 = nn.Linear(hidden_dim, dim, bias=bias)
-        self.ffn_dropout = nn.Dropout(drop)
-
-    def forward(self, x):
-        """Apply the SwiGLU feed-forward transformation."""
-        x12 = self.w12(x)
-        x1, x2 = x12.chunk(2, dim=-1)
-        hidden = F.silu(x1) * x2
-        return self.w3(self.ffn_dropout(hidden))
-
-
-class FinalLayer(nn.Module):
-    """The final layer of JiT."""
-
-    def __init__(self, hidden_size, patch_size, out_channels):
-        """Initialize the JiT output projection layer."""
-        super().__init__()
-        self.norm_final = RMSNorm(hidden_size)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
-
-    def forward(self, x, c):
-        """Project conditioned hidden states back to patch predictions."""
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-class JiTBlock(nn.Module):
-    """Transformer block with adaptive layer norm modulation."""
-
-    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, attn_drop=0.0, proj_drop=0.0):
-        """Initialize one JiT transformer block."""
-        super().__init__()
-        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
-        self.attn = Attention(
-            hidden_size,
-            num_heads=num_heads,
-            qkv_bias=True,
-            qk_norm=True,
-            attn_drop=attn_drop,
-            proj_drop=proj_drop,
-        )
-        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
-        mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        self.mlp = SwiGLUFFN(hidden_size, mlp_hidden_dim, drop=proj_drop)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
-
-    def forward(self, x, c, feat_rope=None):
-        """Run one conditioned transformer block pass."""
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
-        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), rope=feat_rope)
-        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
-        return x
-
-
-class JiT(nn.Module):
-    """Just image Transformer."""
-
-    def __init__(
-        self,
-        input_size=256,
-        patch_size=16,
-        in_channels=3,
-        hidden_size=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4.0,
-        attn_drop=0.0,
-        proj_drop=0.0,
-        num_classes=1000,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=8,
-    ):
-        """Initialize the JiT backbone."""
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.num_heads = num_heads
-        self.hidden_size = hidden_size
-        self.input_size = input_size
-        self.in_context_len = in_context_len
-        self.in_context_start = in_context_start
-        self.num_classes = num_classes
-
-        # time and class embed (unused in DiffusionWrapper forward, DDP will crash if requires_grad=True)
-        self.t_embedder = TimestepEmbedder(hidden_size)
-        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
-        self.t_embedder.requires_grad_(False)
-        self.y_embedder.requires_grad_(False)
-
-        # linear embed
-        self.x_embedder = BottleneckPatchEmbed(
-            input_size, patch_size, in_channels, bottleneck_dim, hidden_size, bias=True
-        )
-
-        # use fixed sin-cos embedding
-        num_patches = self.x_embedder.num_patches
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
-
-        # in-context cls token
-        if self.in_context_len > 0:
-            self.in_context_posemb = nn.Parameter(torch.zeros(1, self.in_context_len, hidden_size), requires_grad=True)
-            torch.nn.init.normal_(self.in_context_posemb, std=0.02)
-
-        # rope
-        half_head_dim = hidden_size // num_heads // 2
-        hw_seq_len = input_size // patch_size
-        self.feat_rope = VisionRotaryEmbeddingFast(dim=half_head_dim, pt_seq_len=hw_seq_len, num_cls_token=0)
-        self.feat_rope_incontext = VisionRotaryEmbeddingFast(
-            dim=half_head_dim,
-            pt_seq_len=hw_seq_len,
-            num_cls_token=self.in_context_len,
-        )
-
-        # transformer
-        self.blocks = nn.ModuleList(
-            [
-                JiTBlock(
-                    hidden_size,
-                    num_heads,
-                    mlp_ratio=mlp_ratio,
-                    attn_drop=attn_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
-                    proj_drop=proj_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
-                )
-                for i in range(depth)
-            ]
-        )
-
-        # linear predict
-        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
-
-        self.initialize_weights()
-
-    # Required for DiffusionWrapper compatibility
-    @property
-    def hidden_dim(self):
-        """Return the model hidden dimension."""
-        return self.hidden_size
-
-    def initialize_weights(self):
-        """Initialize trainable parameters following JiT defaults."""
-
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-
-        self.apply(_basic_init)
-
-        # Initialize (and freeze) pos_embed by sin-cos embedding:
-        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches**0.5))
-        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
-
-        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
-        w1 = self.x_embedder.proj1.weight.data
-        nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
-        w2 = self.x_embedder.proj2.weight.data
-        nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj2.bias, 0)
-
-        # Initialize label embedding table:
-        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
-
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-
-        # Zero-out adaLN modulation layers:
-        for block in self.blocks:
-            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-
-        # Zero-out output layers:
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
-
-        nn.init.constant_(self.final_layer.linear.weight, 0)
-        nn.init.constant_(self.final_layer.linear.bias, 0)
-
-    def unpatchify(self, x, p):
-        """Reconstruct spatial images from patch tokens."""
-        c = self.out_channels
-        h = w = int(x.shape[1] ** 0.5)
-        assert h * w == x.shape[1]
-
-        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
-        x = torch.einsum("nhwpqc->nchpwq", x)
-        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
-        return imgs
-
-    def forward(self, input_and_condition):
-        """Run a forward pass from a DiffusionWrapper-style input dictionary."""
-        x = input_and_condition["input"]
-        c = input_and_condition["condition"]
-
-        # DiffusionWrapper passes BHWC, but JiT expects BCHW (Conv2d).
-        x = x.permute(0, 3, 1, 2)
-
-        # forward JiT
-        x = self.x_embedder(x)
-        x += self.pos_embed
-
-        added_context = False
-        for i, block in enumerate(self.blocks):
-            if self.in_context_len > 0 and i == self.in_context_start:
-                # DiffusionWrapper passes the raw label embedding under "class_emb" so that
-                # in-context tokens carry only class identity, matching the JiT reference.
-                y_emb = input_and_condition["class_emb"]
-
-                in_context_tokens = y_emb.unsqueeze(1).repeat(1, self.in_context_len, 1)
-                in_context_tokens += self.in_context_posemb
-                x = torch.cat([in_context_tokens, x], dim=1)
-                added_context = True
-
-            x = block(x, c, self.feat_rope if i < self.in_context_start else self.feat_rope_incontext)
-
-        if added_context:
-            x = x[:, self.in_context_len :]
-
-        x = self.final_layer(x, c)
-        output = self.unpatchify(x, self.patch_size)
-
-        # DiffusionWrapper expects BHWC logits, unpatchify returns BCHW
-        output = output.permute(0, 2, 3, 1)
-
-        return {"logits": output}
-
-
-def JiT_B_4(**kwargs):
-    """Build the JiT-B/4 variant."""
-    return JiT(
-        depth=12,
-        hidden_size=768,
-        num_heads=12,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=4,
-        patch_size=4,
-        **kwargs,
-    )
-
-
-def JiT_B_16(**kwargs):
-    """Build the JiT-B/16 variant."""
-    return JiT(
-        depth=12,
-        hidden_size=768,
-        num_heads=12,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=4,
-        patch_size=16,
-        **kwargs,
-    )
-
-
-def JiT_B_32(**kwargs):
-    """Build the JiT-B/32 variant."""
-    return JiT(
-        depth=12,
-        hidden_size=768,
-        num_heads=12,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=4,
-        patch_size=32,
-        **kwargs,
-    )
-
-
-def JiT_L_16(**kwargs):
-    """Build the JiT-L/16 variant."""
-    return JiT(
-        depth=24,
-        hidden_size=1024,
-        num_heads=16,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=8,
-        patch_size=16,
-        **kwargs,
-    )
-
-
-def JiT_L_32(**kwargs):
-    """Build the JiT-L/32 variant."""
-    return JiT(
-        depth=24,
-        hidden_size=1024,
-        num_heads=16,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=8,
-        patch_size=32,
-        **kwargs,
-    )
-
-
-def JiT_H_16(**kwargs):
-    """Build the JiT-H/16 variant."""
-    return JiT(
-        depth=32,
-        hidden_size=1280,
-        num_heads=16,
-        bottleneck_dim=256,
-        in_context_len=32,
-        in_context_start=10,
-        patch_size=16,
-        **kwargs,
-    )
-
-
-def JiT_H_32(**kwargs):
-    """Build the JiT-H/32 variant."""
-    return JiT(
-        depth=32,
-        hidden_size=1280,
-        num_heads=16,
-        bottleneck_dim=256,
-        in_context_len=32,
-        in_context_start=10,
-        patch_size=32,
-        **kwargs,
-    )
-
-
-JiT_models = {
-    "JiT-B/4": JiT_B_4,
-    "JiT-B/16": JiT_B_16,
-    "JiT-B/32": JiT_B_32,
-    "JiT-L/16": JiT_L_16,
-    "JiT-L/32": JiT_L_32,
-    "JiT-H/16": JiT_H_16,
-    "JiT-H/32": JiT_H_32,
-}
diff --git a/nvsubquadratic/networks/jit_utils.py b/nvsubquadratic/networks/jit_utils.py
deleted file mode 100644
index 716d9fbc..00000000
--- a/nvsubquadratic/networks/jit_utils.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Helpers used by :mod:`nvsubquadratic.networks.jit`.
-
-Ported from the reference JiT implementation at
-https://github.com/LTH14/JiT.  Contains the rotary-embedding primitives
-(:class:`VisionRotaryEmbedding`, :class:`VisionRotaryEmbeddingFast`),
-the local :class:`RMSNorm`, and the 1D/2D sin-cos positional-embedding
-helpers used to seed the patch tokens.
-"""
-
-from math import pi
-
-import numpy as np
-import torch
-from einops import rearrange, repeat
-from torch import nn
-
-
-def broadcat(tensors, dim=-1):
-    """Concatenate tensors with broadcasting on non-concatenated dimensions."""
-    num_tensors = len(tensors)
-    shape_lens = {len(t.shape) for t in tensors}
-    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
-    shape_len = next(iter(shape_lens))
-    dim = (dim + shape_len) if dim < 0 else dim
-    dims = list(zip(*(list(t.shape) for t in tensors)))
-    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
-    assert all(len(set(values)) <= 2 for _, values in expandable_dims), (
-        "invalid dimensions for broadcastable concatenation"
-    )
-    max_dims = [(axis, max(values)) for axis, values in expandable_dims]
-    expanded_dims = [(axis, (axis_size,) * num_tensors) for axis, axis_size in max_dims]
-    expanded_dims.insert(dim, (dim, dims[dim]))
-    expandable_shapes = list(zip(*(sizes for _, sizes in expanded_dims)))
-    expanded_tensors = [tensor.expand(*shape) for tensor, shape in zip(tensors, expandable_shapes)]
-    return torch.cat(expanded_tensors, dim=dim)
-
-
-def rotate_half(x):
-    """Rotate tensor pairs in the last dimension for rotary embeddings."""
-    x = rearrange(x, "... (d r) -> ... d r", r=2)
-    x1, x2 = x.unbind(dim=-1)
-    x = torch.stack((-x2, x1), dim=-1)
-    return rearrange(x, "... d r -> ... (d r)")
-
-
-class VisionRotaryEmbedding(nn.Module):
-    """2D rotary embedding helper with configurable frequency construction."""
-
-    def __init__(
-        self,
-        dim,
-        pt_seq_len,
-        ft_seq_len=None,
-        custom_freqs=None,
-        freqs_for="lang",
-        theta=10000,
-        max_freq=10,
-        num_freqs=1,
-    ):
-        """Initialize a full 2D rotary embedding table."""
-        super().__init__()
-        if custom_freqs:
-            freqs = custom_freqs
-        elif freqs_for == "lang":
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-        elif freqs_for == "pixel":
-            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
-        elif freqs_for == "constant":
-            freqs = torch.ones(num_freqs).float()
-        else:
-            raise ValueError(f"unknown modality {freqs_for}")
-
-        if ft_seq_len is None:
-            ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs_h = torch.einsum("..., f -> ... f", t, freqs)
-        freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
-
-        freqs_w = torch.einsum("..., f -> ... f", t, freqs)
-        freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)
-
-        freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
-
-        self.register_buffer("freqs_cos", freqs.cos())
-        self.register_buffer("freqs_sin", freqs.sin())
-
-    def forward(self, t, start_index=0):
-        """Apply rotary embedding to the selected slice of ``t``."""
-        rot_dim = self.freqs_cos.shape[-1]
-        end_index = start_index + rot_dim
-        assert rot_dim <= t.shape[-1], (
-            f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
-        )
-        t_left, t, t_right = (
-            t[..., :start_index],
-            t[..., start_index:end_index],
-            t[..., end_index:],
-        )
-        t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
-        return torch.cat((t_left, t, t_right), dim=-1)
-
-
-class VisionRotaryEmbeddingFast(nn.Module):
-    """Flattened 2D rotary embedding helper optimized for sequence inputs."""
-
-    def __init__(
-        self,
-        dim,
-        pt_seq_len=16,
-        ft_seq_len=None,
-        custom_freqs=None,
-        freqs_for="lang",
-        theta=10000,
-        max_freq=10,
-        num_freqs=1,
-        num_cls_token=0,
-    ):
-        """Initialize flattened cosine and sine tables for rotary embedding."""
-        super().__init__()
-        if custom_freqs:
-            freqs = custom_freqs
-        elif freqs_for == "lang":
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-        elif freqs_for == "pixel":
-            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
-        elif freqs_for == "constant":
-            freqs = torch.ones(num_freqs).float()
-        else:
-            raise ValueError(f"unknown modality {freqs_for}")
-
-        if ft_seq_len is None:
-            ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _N_img, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype, device=cos_img.device)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype, device=sin_img.device)
-
-            freqs_cos = torch.cat([cos_pad, cos_img], dim=0).contiguous()  # [N_cls+N_img, D]
-            freqs_sin = torch.cat([sin_pad, sin_img], dim=0).contiguous()
-        else:
-            freqs_cos = freqs.cos().view(-1, freqs.shape[-1]).contiguous()
-            freqs_sin = freqs.sin().view(-1, freqs.shape[-1]).contiguous()
-
-        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
-        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
-
-    def forward(self, t):
-        """Apply precomputed flattened rotary embedding values to ``t``."""
-        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
-
-
-class RMSNorm(nn.Module):
-    """Root-mean-square normalization used by JiT blocks."""
-
-    def __init__(self, hidden_size, eps=1e-6):
-        """Initialize RMSNorm parameters."""
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        """Normalize ``hidden_states`` with RMS statistics on the last dimension."""
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return (self.weight * hidden_states).to(input_dtype)
-
-
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
-    """Build 2D sine-cosine positional embeddings for a square grid."""
-    grid_h = np.arange(grid_size, dtype=np.float32)
-    grid_w = np.arange(grid_size, dtype=np.float32)
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-
-    grid = grid.reshape([2, 1, grid_size, grid_size])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    """Build 2D sine-cosine positional embeddings from a provided grid."""
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    return emb
-
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """Build 1D sine-cosine positional embeddings from a position array."""
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.0
-    omega = 1.0 / 10000**omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out)  # (M, D/2)
-    emb_cos = np.cos(out)  # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb
diff --git a/nvsubquadratic/ops/circular_fftconv_fp16.py b/nvsubquadratic/ops/circular_fftconv_fp16.py
deleted file mode 100644
index 562d8646..00000000
--- a/nvsubquadratic/ops/circular_fftconv_fp16.py
+++ /dev/null
@@ -1,664 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""FP16 circular (periodic) FFT-based convolution for 1D, 2D, and 3D signals.
-
-Leverages PyTorch's native fp16 FFT support (cuFFT C2C half-precision).
-Two constraints must be satisfied:
-
-1. cuFFT requires **power-of-2 sizes** for half-precision transforms.
-   Since circular convolution uses same-size FFTs (FFT length == input length),
-   the *input* spatial dimensions must themselves be powers of 2.
-2. The frequency-domain products can exceed fp16's 65504 limit with real model
-   weights, so we use ``norm="ortho"`` which divides both the forward and
-   inverse FFT by ``sqrt(N)``, keeping intermediates in range.  The ortho pair
-   computes ``circular_conv / sqrt(N)`` instead of ``circular_conv``, so we
-   multiply by ``sqrt(N)`` after the inverse FFT to restore the correct scale.
-
-Numerical stability (dual mean-centering)
-------------------------------------------
-Ortho normalization alone is insufficient: the DC bin and internal FFT
-accumulations can still overflow in fp16 for signals with nonzero mean.
-To fix this, both the input ``x`` and kernel ``k`` are mean-centered before
-the forward FFT, which zeros the DC bins and reduces element magnitudes to
-``O(std)``.  The exact convolution result is recovered analytically via:
-
-- **DC correction** (T4): ``mu_x * mu_k * K`` added after the inverse FFT.
-- **Centering correction** (T2, nD): an inclusion-exclusion geometric factor
-  ``geo`` (cached per kernel/spatial shape) that accounts for zero-padded
-  kernel positions.  Only the scalar ``k_mean / sqrt_N`` changes per call;
-  ``geo`` is recomputed only when kernel or spatial dimensions change.
-
-See ``circular_fftconv1d_fp16_bhl`` docstring for the full derivation.
-
-Alignment
----------
-Alignment is done via a frequency-domain phase ramp, the same approach used by
-the fp32 variants.  The ramp is computed in float32 for precision and cast to
-complex32 before the multiply so all FFT intermediates stay in half-precision.
-
-Families provided
------------------
-- 1D: ``circular_fftconv1d_fp16_bhl`` (+``_w_reshape``)
-- 2D: ``circular_fftconv2d_fp16_bhl`` (+``_w_reshape``)
-- 3D: ``circular_fftconv3d_fp16_bhl`` (+``_w_reshape``)
-"""
-
-from __future__ import annotations
-
-import math
-from collections import OrderedDict
-
-import torch
-from einops import rearrange
-
-from nvsubquadratic.ops.circular_fftconv import (
-    _phase_ramp_cache_1d,
-    _phase_ramp_cache_2d,
-    _phase_ramp_cache_3d,
-)
-
-
-__all__ = [
-    "circular_fftconv1d_fp16_bhl",
-    "circular_fftconv1d_fp16_bhl_w_reshape",
-    "circular_fftconv2d_fp16_bhl",
-    "circular_fftconv2d_fp16_bhl_w_reshape",
-    "circular_fftconv3d_fp16_bhl",
-    "circular_fftconv3d_fp16_bhl_w_reshape",
-]
-
-
-def _is_power_of_2(n: int) -> bool:
-    return n > 0 and (n & (n - 1)) == 0
-
-
-# Phase ramps are computed in float32 (full-precision trig) and cast to
-# complex32 so the multiply with fp16 FFT data stays in half-precision.
-_PHASE_RAMP_COMPUTE_DTYPE = torch.float32
-_PHASE_RAMP_TARGET_DTYPE = torch.complex32
-
-
-###############################################################################
-# Helpers: centering correction geometry caches (2D / 3D)
-###############################################################################
-
-
-def _phase_m1(N: int, rfft: bool, device: torch.device) -> torch.Tensor:
-    r"""Phase factor for a unit impulse at position :math:`N-1` (i.e. :math:`-1 \bmod N`).
-
-    Returns the 1-D DFT :math:`p[f] = e^{2\pi i f / N}` evaluated on the
-    appropriate frequency grid.  This is used to build the inclusion-exclusion
-    geometric correction in the nD centering formulas.
-
-    Args:
-        N: Transform length along this axis.
-        rfft: If True, use ``rfftfreq`` (length ``N//2+1``, for the last axis
-            of ``rfftn``); otherwise use ``fftfreq`` (length ``N``).
-        device: Torch device for the output tensor.
-
-    Returns:
-        Complex64 tensor of shape ``[N//2+1]`` (rfft) or ``[N]`` (fft).
-    """
-    f = (torch.fft.rfftfreq if rfft else torch.fft.fftfreq)(
-        N,
-        device=device,
-        dtype=torch.float32,
-    )
-    a = 2.0 * math.pi * f
-    return torch.complex(a.cos(), a.sin())
-
-
-class _CenteringCorrectionCache2D:
-    r"""LRU cache for the 2D geometric correction factor ``geo``.
-
-    When the kernel is smaller than the spatial dimensions along one or
-    both axes (:math:`K_d < N_d`), the mean-centering decomposition
-    produces a T2 correction term that depends on an inclusion-exclusion
-    geometric factor.  This factor depends only on the kernel/spatial
-    shape and device, so it is computed once and reused.
-
-    The full frequency-domain correction added to :math:`\hat{k}_c` is:
-
-    .. math::
-        \frac{\mu_k}{\sqrt{N}} \cdot \text{geo}
-
-    Only the scalar :math:`\mu_k` changes per forward call.
-    """
-
-    def __init__(self, maxsize: int = 32):
-        self.maxsize = maxsize
-        self._cache: OrderedDict[tuple, torch.Tensor] = OrderedDict()
-
-    @staticmethod
-    def _key(K_x, K_y, X, Y, device):
-        return (K_x, K_y, X, Y, device.type, device.index if device.index is not None else -1)
-
-    def get(self, K_x: int, K_y: int, X: int, Y: int, device: torch.device) -> torch.Tensor | None:
-        r"""Return cached ``geo`` of shape ``[X, Y//2+1]`` (complex64), or None.
-
-        Returns ``None`` when no correction is needed (K_x == X and K_y == Y).
-        Otherwise computes the inclusion-exclusion geometric factor over the
-        corrected axes (those where ``K_d < N_d``).
-        """
-        cx, cy = int(K_x < X), int(K_y < Y)
-        if not (cx or cy):
-            return None
-
-        key = self._key(K_x, K_y, X, Y, device)
-        geo = self._cache.get(key)
-        if geo is not None:
-            self._cache.move_to_end(key)
-            return geo
-
-        Yf = Y // 2 + 1
-        p_x = _phase_m1(X, rfft=False, device=device) if cx else None
-        p_y = _phase_m1(Y, rfft=True, device=device) if cy else None
-
-        geo = torch.zeros(X, Yf, device=device, dtype=torch.complex64)
-        if cx:
-            geo[:, 0] -= Y * p_x
-        if cy:
-            geo[0, :] -= X * p_y
-        if cx and cy:
-            geo += p_x[:, None] * p_y[None, :]
-
-        self._cache[key] = geo
-        if len(self._cache) > self.maxsize:
-            self._cache.popitem(last=False)
-        return geo
-
-
-class _CenteringCorrectionCache3D:
-    r"""LRU cache for the 3D geometric correction factor ``geo``.
-
-    Same structure as :class:`_CenteringCorrectionCache2D` but with
-    three-axis inclusion-exclusion (signs: single terms ``-1``,
-    pair terms ``+1``, triple term ``-1``).
-    """
-
-    def __init__(self, maxsize: int = 32):
-        self.maxsize = maxsize
-        self._cache: OrderedDict[tuple, torch.Tensor] = OrderedDict()
-
-    @staticmethod
-    def _key(Kx, Ky, Kz, X, Y, Z, device):
-        return (Kx, Ky, Kz, X, Y, Z, device.type, device.index if device.index is not None else -1)
-
-    def get(self, Kx: int, Ky: int, Kz: int, X: int, Y: int, Z: int, device: torch.device) -> torch.Tensor | None:
-        """Return cached ``geo`` of shape ``[X, Y, Z//2+1]`` (complex64), or None."""
-        cx, cy, cz = int(Kx < X), int(Ky < Y), int(Kz < Z)
-        if not (cx or cy or cz):
-            return None
-
-        key = self._key(Kx, Ky, Kz, X, Y, Z, device)
-        geo = self._cache.get(key)
-        if geo is not None:
-            self._cache.move_to_end(key)
-            return geo
-
-        Zf = Z // 2 + 1
-        p_x = _phase_m1(X, rfft=False, device=device) if cx else None
-        p_y = _phase_m1(Y, rfft=False, device=device) if cy else None
-        p_z = _phase_m1(Z, rfft=True, device=device) if cz else None
-
-        geo = torch.zeros(X, Y, Zf, device=device, dtype=torch.complex64)
-
-        # Single-axis terms (sign: -1, coeff: product of OTHER axis sizes)
-        if cx:
-            geo[:, 0, 0] -= (Y * Z) * p_x
-        if cy:
-            geo[0, :, 0] -= (X * Z) * p_y
-        if cz:
-            geo[0, 0, :] -= (X * Y) * p_z
-
-        # Pair terms (sign: +1, coeff: remaining axis size)
-        if cx and cy:
-            geo[:, :, 0] += Z * (p_x[:, None] * p_y[None, :])
-        if cx and cz:
-            geo[:, 0, :] += Y * (p_x[:, None] * p_z[None, :])
-        if cy and cz:
-            geo[0, :, :] += X * (p_y[:, None] * p_z[None, :])
-
-        # Triple term (sign: -1)
-        if cx and cy and cz:
-            geo -= p_x[:, None, None] * p_y[None, :, None] * p_z[None, None, :]
-
-        self._cache[key] = geo
-        if len(self._cache) > self.maxsize:
-            self._cache.popitem(last=False)
-        return geo
-
-
-_centering_geo_cache_2d = _CenteringCorrectionCache2D()
-_centering_geo_cache_3d = _CenteringCorrectionCache3D()
-
-
-###############################################################################
-# 1D
-###############################################################################
-
-
-def circular_fftconv1d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    r"""1D circular FFT convolution in fp16 with dual mean-centering.
-
-    Numerically stable drop-in replacement for ``circular_fftconv1d_fp32_bhl``.
-    Casts *x* and *kernel* to fp16 internally; shortcut is never cast.
-    Returns in the **original** dtype of *x*.
-
-    Requires ``L`` to be a power of 2 (cuFFT fp16 constraint).
-
-    Root cause of NaN
-    -----------------
-    With ``norm="ortho"``, the FFT scales by ``1/sqrt(N)``.  Two overflow
-    sources exist in fp16 (max representable value 65504):
-
-    1. **DC-bin product overflow**: the DC bin of ``rfft(x, ortho)`` is
-       ``sum(x)/sqrt(N) = mean(x)*sqrt(N)``.  The product of two DC bins
-       is ``mean(x)*mean(k)*N``, which exceeds 65504 for moderate means.
-    2. **Internal FFT accumulation overflow**: cuFFT accumulates partial
-       sums in fp16 before dividing by ``sqrt(N)``.  For a signal with
-       mean ``mu`` and ``N`` elements, internal sums reach ``~mu*N``,
-       which overflows for ``mu*N > 65504``.
-
-    Fix: dual mean-centering
-    ------------------------
-    Subtract the spatial mean from both ``x`` and ``k`` before the FFT.
-    This zeros both DC bins (fixing #1) and reduces element magnitudes to
-    ``O(std)`` (fixing #2).  The exact result is recovered analytically.
-
-    Decomposition (circular conv of length ``L``, kernel zero-padded to ``L``)::
-
-        y[n] = sum_m x[n-m] * k_padded[m]
-
-    Let ``x = x_c + mu_x``, ``k_padded = k_c_padded + delta``, where
-    ``delta = [mu_k, ..., mu_k, 0, ..., 0]`` (``K`` copies of ``mu_k``,
-    then ``L-K`` zeros).  Expanding gives four terms::
-
-        T1 = circ_conv(x_c, k_c_padded)    [both DC=0, safe in fp16]
-        T2 = mu_k * circ_conv(x_c, delta)  [correction for centering k]
-        T3 = mu_x * sum(k_c_padded) = 0    [k_c sums to zero]
-        T4 = mu_x * k_sum                  [correction for centering x]
-
-    K=L-1 correction (T2)
-    ~~~~~~~~~~~~~~~~~~~~~~
-    ``delta`` is all-``mu_k`` except one zero at position ``L-1``.
-    ``circ_conv(x_c, delta)[n] = mu_k * sum_{m != L-1} x_c[n-m]``.
-    Since ``sum(x_c) = 0``, this equals ``-mu_k * x_c[(n+1) mod L]``,
-    i.e., a roll of ``x_c`` by ``-1`` scaled by ``-mu_k``.
-
-    When a kernel-centering shift ``s`` is applied, the zero in ``delta``
-    moves, and the identity becomes ``-mu_k * roll(x_c, s-1)``, or
-    equivalently ``-mu_k * fft_x * phase(s-1)`` in frequency domain.
-
-    Phase-ramp absorption
-    ~~~~~~~~~~~~~~~~~~~~~
-    Factor ``phase(s-1) = phase(s) * phase(-1)`` so the full effective
-    kernel spectrum is::
-
-        fft_k_eff = phase(s) * [fft_k_c - (mu_k / sqrt_N) * phase(-1)]
-
-    The ``1/sqrt_N`` keeps frequency values small (preventing irfft
-    internal overflow); ``sqrt_N`` is applied after irfft in float32.
-
-    K=L case
-    ~~~~~~~~
-    No zero-padding => ``delta`` covers all ``L`` positions =>
-    ``circ_conv(x_c, delta) = mu_k * sum(x_c) = 0``.  No correction.
-
-    Args:
-        x: Input tensor ``[B, H, L]`` (any dtype).
-        kernel: Kernel tensor ``[1|B, H, K]`` with ``K in {L, L-1}`` (any dtype).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, L]`` in the original dtype of *x*.
-    """
-    B, H, L = x.shape
-    assert _is_power_of_2(L), f"L must be a power of 2 for fp16 circular FFT. Got L={L}."
-    assert len(kernel.shape) == 3, f"Unexpected kernel shape: {kernel.shape}."
-    assert kernel.shape[0] in (1, B), (
-        f"Leading dimension must be 1 or batch_size ({B}). Got kernel.shape={kernel.shape}."
-    )
-    _, Hk, K = kernel.shape
-    assert H == Hk, "Input and kernel must have the same number of channels (H)."
-    assert K in (L, L - 1), f"FP16 centering requires K in {{L, L-1}}, got K={K}, L={L}"
-
-    x_fp16 = x.to(torch.float16, copy=True)
-    k_fp16 = kernel.to(torch.float16, copy=True)
-    sqrt_N = math.sqrt(L)
-
-    # ── Center both x and k in-place (zeros DC bins, no extra allocation) ──
-    # copy=True above ensures we never mutate the caller's tensors.
-    x_mean = x_fp16.mean(dim=-1, keepdim=True)  # [B, H, 1]  fp16
-    x_fp16.sub_(x_mean)  # x_fp16 is now x_c
-
-    k_mean = k_fp16.mean(dim=-1, keepdim=True)  # [1|B, H, 1]  fp16
-    k_fp16.sub_(k_mean)  # k_fp16 is now k_c
-
-    # DC correction: mu_x * mu_k * K, precomputed in fp32
-    dc_corr = x_mean.float() * (k_mean.float() * K)  # [B, H, 1]  fp32
-
-    # ── Forward FFTs in fp16 (both DC bins = 0) ──
-    fft_x = torch.fft.rfft(x_fp16, n=L, dim=2, norm="ortho")
-    fft_k = torch.fft.rfft(k_fp16, n=L, dim=2, norm="ortho")
-    del k_fp16
-
-    # ── Build effective kernel spectrum ──
-    # fft_k is small (no batch dim), so we work in complex64 for precision
-    # and cast back to complex32 before the large multiply with fft_x.
-    shift = -((K - 1) // 2)
-
-    del x_fp16  # x_c no longer needed (correction folded into fft_k)
-
-    fft_k_eff = fft_k.to(torch.complex64)
-    del fft_k
-
-    # Build effective kernel spectrum in one expression per case:
-    #   K=L  :  fft_k_eff = fft_k_c * phase(s)
-    #   K=L-1:  fft_k_eff = fft_k_c * phase(s) - (mu_k / sqrt_N) * phase(s-1)
-    # The K=L-1 form comes from expanding:
-    #   phase(s) * [fft_k_c - (mu_k/sqrt_N) * phase(-1)]
-    if K == L - 1:
-        phase_s = _phase_ramp_cache_1d.get(L, shift, x.device, _PHASE_RAMP_COMPUTE_DTYPE)
-        phase_sm1 = _phase_ramp_cache_1d.get(L, shift - 1, x.device, _PHASE_RAMP_COMPUTE_DTYPE)
-        fft_k_eff = fft_k_eff * phase_s - (k_mean.float() / sqrt_N) * phase_sm1
-    elif K == L:
-        if shift != 0:
-            phase_s = _phase_ramp_cache_1d.get(L, shift, x.device, _PHASE_RAMP_COMPUTE_DTYPE)
-            fft_k_eff = fft_k_eff * phase_s
-    else:
-        raise ValueError(f"Dual-centering correction requires K=L or K=L-1, got K={K}, L={L}")
-
-    fft_x.mul_(fft_k_eff.to(torch.complex32))
-
-    # ── Inverse FFT (fp16) ──
-    y = torch.fft.irfft(fft_x, n=L, dim=2, norm="ortho")
-    del fft_x
-
-    # ── Undo ortho scaling + DC correction in one expression (float32) ──
-    y = (y.float() * sqrt_N + dc_corr).to(x.dtype)
-
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        assert shortcut.shape == (H,)
-        y = y + rearrange(shortcut, "h -> 1 h 1") * x
-    return y
-
-
-def circular_fftconv1d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """1D circular FFT convolution in fp16, for inputs with layout ``[B, L, H]``.
-
-    Wrapper around :func:`circular_fftconv1d_fp16_bhl` that handles the
-    BLH <-> BHL reshape.
-    """
-    x_bhl = rearrange(x, "b l h -> b h l")
-    k_bhl = rearrange(kernel, "b k h -> b h k")
-    y = circular_fftconv1d_fp16_bhl(x_bhl, k_bhl, shortcut)
-    return rearrange(y, "b h l -> b l h")
-
-
-###############################################################################
-# 2D
-###############################################################################
-
-
-def circular_fftconv2d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    r"""2D circular FFT convolution in fp16 with dual mean-centering.
-
-    Numerically stable drop-in replacement for ``circular_fftconv2d_fp32_bhl``.
-    See :func:`circular_fftconv1d_fp16_bhl` for the derivation;  the 2D
-    centering correction uses a cached inclusion-exclusion geometric factor
-    ``geo`` (see :class:`_CenteringCorrectionCache2D`).
-
-    Requires ``X`` and ``Y`` to be powers of 2 (cuFFT fp16 constraint).
-
-    Args:
-        x: Input tensor ``[B, H, X, Y]`` (any dtype).
-        kernel: Kernel tensor ``[1|B, H, K_x, K_y]`` (any dtype).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, X, Y]`` in the original dtype of *x*.
-    """
-    B, H, X, Y = x.shape
-    assert _is_power_of_2(X) and _is_power_of_2(Y), (
-        f"Spatial dims must be powers of 2 for fp16 circular FFT. Got X={X}, Y={Y}."
-    )
-    assert len(kernel.shape) == 4, f"Unexpected kernel shape: {kernel.shape}."
-    assert kernel.shape[0] in (1, B), (
-        f"Leading dimension must be 1 or batch_size ({B}). Got kernel.shape={kernel.shape}."
-    )
-    _, Hk, K_x, K_y = kernel.shape
-    assert H == Hk, "Input and kernel must have the same number of channels (H)."
-    assert K_x in (X, X - 1), f"FP16 centering requires K_x in {{X, X-1}}, got K_x={K_x}, X={X}"
-    assert K_y in (Y, Y - 1), f"FP16 centering requires K_y in {{Y, Y-1}}, got K_y={K_y}, Y={Y}"
-
-    N = X * Y
-    sqrt_N = math.sqrt(N)
-
-    x_fp16 = x.to(torch.float16, copy=True)
-    k_fp16 = kernel.to(torch.float16, copy=True)
-
-    # ── Center both x and k in-place ──
-    # copy=True above ensures we never mutate the caller's tensors.
-    x_mean = x_fp16.mean(dim=(-2, -1), keepdim=True)  # [B, H, 1, 1]
-    x_fp16.sub_(x_mean)
-
-    k_mean = k_fp16.mean(dim=(-2, -1), keepdim=True)  # [1|B, H, 1, 1]
-    k_fp16.sub_(k_mean)
-
-    dc_corr = x_mean.float() * (k_mean.float() * (K_x * K_y))
-
-    # ── Forward FFTs (both DC bins = 0) ──
-    fft_x = torch.fft.rfft2(x_fp16, s=(X, Y), dim=(2, 3), norm="ortho")
-    fft_k = torch.fft.rfft2(k_fp16, s=(X, Y), dim=(2, 3), norm="ortho")
-    del k_fp16
-
-    shift_x = -((K_x - 1) // 2)
-    shift_y = -((K_y - 1) // 2)
-
-    del x_fp16  # x_c no longer needed (correction folded into fft_k)
-
-    # Build effective kernel spectrum in complex64 (small tensor, no batch dim).
-    fft_k_eff = fft_k.to(torch.complex64)
-    del fft_k
-
-    # T2: centering correction for zero-padded kernel positions.
-    # The geometric factor is cached and only the scalar mu_k/sqrt_N changes.
-    geo = _centering_geo_cache_2d.get(K_x, K_y, X, Y, x.device)
-    if geo is not None:
-        fft_k_eff = fft_k_eff + (k_mean.float() / sqrt_N) * geo
-
-    # Apply phase ramp for kernel alignment (spatial centering).
-    if shift_x != 0 or shift_y != 0:
-        phase = _phase_ramp_cache_2d.get(
-            X,
-            Y,
-            shift_x,
-            shift_y,
-            x.device,
-            _PHASE_RAMP_COMPUTE_DTYPE,
-        )
-        fft_k_eff = fft_k_eff * phase
-
-    # Multiply in complex32 (large batched tensor stays in half precision).
-    fft_x.mul_(fft_k_eff.to(torch.complex32))
-
-    # ── Inverse FFT (fp16) ──
-    y = torch.fft.irfft2(fft_x, s=(X, Y), dim=(2, 3), norm="ortho")
-    del fft_x
-
-    # ── Undo ortho scaling (×√N) and add DC correction (T4), in float32 ──
-    y = (y.float() * sqrt_N + dc_corr).to(x.dtype)
-
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        assert shortcut.shape == (H,)
-        y = y + rearrange(shortcut, "h -> 1 h 1 1") * x
-    return y
-
-
-def circular_fftconv2d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D circular FFT convolution in fp16, for inputs with layout ``[B, X, Y, H]``.
-
-    Wrapper around :func:`circular_fftconv2d_fp16_bhl` that handles the
-    BLH <-> BHL reshape.
-    """
-    x_bhl = rearrange(x, "b x y h -> b h x y")
-    k_bhl = rearrange(kernel, "b kx ky h -> b h kx ky")
-    y = circular_fftconv2d_fp16_bhl(x_bhl, k_bhl, shortcut)
-    return rearrange(y, "b h x y -> b x y h")
-
-
-###############################################################################
-# 3D
-###############################################################################
-
-
-def circular_fftconv3d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    r"""3D circular FFT convolution in fp16 with dual mean-centering.
-
-    Numerically stable drop-in replacement for ``circular_fftconv3d_fp32_bhl``.
-    See :func:`circular_fftconv1d_fp16_bhl` for the derivation;  the 3D
-    centering correction uses a cached inclusion-exclusion geometric factor
-    ``geo`` (see :class:`_CenteringCorrectionCache3D`).
-
-    Requires ``X``, ``Y``, and ``Z`` to be powers of 2 (cuFFT fp16 constraint).
-
-    Args:
-        x: Input tensor ``[B, H, X, Y, Z]`` (any dtype).
-        kernel: Kernel tensor ``[1|B, H, Kx, Ky, Kz]`` (any dtype).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, X, Y, Z]`` in the original dtype of *x*.
-    """
-    B, H, X, Y, Z = x.shape
-    assert _is_power_of_2(X) and _is_power_of_2(Y) and _is_power_of_2(Z), (
-        f"Spatial dims must be powers of 2 for fp16 circular FFT. Got X={X}, Y={Y}, Z={Z}."
-    )
-    assert len(kernel.shape) == 5, f"Unexpected kernel shape: {kernel.shape}."
-    assert kernel.shape[0] in (1, B), (
-        f"Leading dimension must be 1 or batch_size ({B}). Got kernel.shape={kernel.shape}."
-    )
-    _, Hk, Kx, Ky, Kz = kernel.shape
-    assert H == Hk, "Input and kernel must have the same number of channels (H)."
-    assert Kx in (X, X - 1), f"FP16 centering requires Kx in {{X, X-1}}, got Kx={Kx}, X={X}"
-    assert Ky in (Y, Y - 1), f"FP16 centering requires Ky in {{Y, Y-1}}, got Ky={Ky}, Y={Y}"
-    assert Kz in (Z, Z - 1), f"FP16 centering requires Kz in {{Z, Z-1}}, got Kz={Kz}, Z={Z}"
-
-    N = X * Y * Z
-    sqrt_N = math.sqrt(N)
-
-    x_fp16 = x.to(torch.float16, copy=True)
-    k_fp16 = kernel.to(torch.float16, copy=True)
-
-    # ── Center both x and k in-place ──
-    # copy=True above ensures we never mutate the caller's tensors.
-    x_mean = x_fp16.mean(dim=(-3, -2, -1), keepdim=True)  # [B, H, 1, 1, 1]
-    x_fp16.sub_(x_mean)
-
-    k_mean = k_fp16.mean(dim=(-3, -2, -1), keepdim=True)  # [1|B, H, 1, 1, 1]
-    k_fp16.sub_(k_mean)
-
-    dc_corr = x_mean.float() * (k_mean.float() * (Kx * Ky * Kz))
-
-    # ── Forward FFTs (both DC bins = 0) ──
-    fft_x = torch.fft.rfftn(x_fp16, s=(X, Y, Z), dim=(2, 3, 4), norm="ortho")
-    fft_k = torch.fft.rfftn(k_fp16, s=(X, Y, Z), dim=(2, 3, 4), norm="ortho")
-    del k_fp16
-
-    shift_x = -((Kx - 1) // 2)
-    shift_y = -((Ky - 1) // 2)
-    shift_z = -((Kz - 1) // 2)
-
-    del x_fp16  # x_c no longer needed (correction folded into fft_k)
-
-    # Build effective kernel spectrum in complex64 (small tensor, no batch dim).
-    fft_k_eff = fft_k.to(torch.complex64)
-    del fft_k
-
-    # T2: centering correction for zero-padded kernel positions.
-    # The geometric factor is cached and only the scalar mu_k/sqrt_N changes.
-    geo = _centering_geo_cache_3d.get(Kx, Ky, Kz, X, Y, Z, x.device)
-    if geo is not None:
-        fft_k_eff = fft_k_eff + (k_mean.float() / sqrt_N) * geo
-
-    # Apply phase ramp for kernel alignment (spatial centering).
-    if shift_x != 0 or shift_y != 0 or shift_z != 0:
-        phase = _phase_ramp_cache_3d.get(
-            X,
-            Y,
-            Z,
-            shift_x,
-            shift_y,
-            shift_z,
-            x.device,
-            _PHASE_RAMP_COMPUTE_DTYPE,
-        )
-        fft_k_eff = fft_k_eff * phase
-
-    # Multiply in complex32 (large batched tensor stays in half precision).
-    fft_x.mul_(fft_k_eff.to(torch.complex32))
-
-    # ── Inverse FFT (fp16) ──
-    y = torch.fft.irfftn(fft_x, s=(X, Y, Z), dim=(2, 3, 4), norm="ortho")
-    del fft_x
-
-    # ── Undo ortho scaling (×√N) and add DC correction (T4), in float32 ──
-    y = (y.float() * sqrt_N + dc_corr).to(x.dtype)
-
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        assert shortcut.shape == (H,)
-        y = y + rearrange(shortcut, "h -> 1 h 1 1 1") * x
-    return y
-
-
-def circular_fftconv3d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """3D circular FFT convolution in fp16, for inputs with layout ``[B, X, Y, Z, H]``.
-
-    Wrapper around :func:`circular_fftconv3d_fp16_bhl` that handles the
-    BLH <-> BHL reshape.
-    """
-    x_bhl = rearrange(x, "b x y z h -> b h x y z")
-    k_bhl = rearrange(kernel, "b kx ky kz h -> b h kx ky kz")
-    y = circular_fftconv3d_fp16_bhl(x_bhl, k_bhl, shortcut)
-    return rearrange(y, "b h x y z -> b x y z h")
diff --git a/nvsubquadratic/ops/fftconv.py b/nvsubquadratic/ops/fftconv.py
index bf02eac4..df3cb75c 100644
--- a/nvsubquadratic/ops/fftconv.py
+++ b/nvsubquadratic/ops/fftconv.py
@@ -92,18 +92,14 @@
     y \leftarrow y + \text{shortcut} \odot x
 
 broadcast along the spatial dimensions. This fuses the residual into the same
-kernel launch and matches the algebra used by the multi-head FFT conv (see
-:mod:`nvsubquadratic.ops.fftconv_multihead`) and by Hyena gating.
+kernel launch and matches the algebra used by Hyena gating.
 
 Precision
 ---------
 All operators accept any input dtype. Internally ``x`` and ``kernel`` are
 cast to ``float32`` for numerical stability (the frequency-domain product
 amplifies the dynamic range of intermediate values); the output is returned
-in the original dtype of ``x``. For aggressive memory/compute savings on
-power-of-two spatial dims, see the fp16 counterparts in
-:mod:`nvsubquadratic.ops.fftconv_fp16` and
-:mod:`nvsubquadratic.ops.circular_fftconv_fp16`.
+in the original dtype of ``x``.
 
 Performance
 -----------
diff --git a/nvsubquadratic/ops/fftconv_fp16.py b/nvsubquadratic/ops/fftconv_fp16.py
deleted file mode 100644
index d14a35a1..00000000
--- a/nvsubquadratic/ops/fftconv_fp16.py
+++ /dev/null
@@ -1,536 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""FP16 FFT-based convolution for 1D, 2D, and 3D signals.
-
-Leverages PyTorch's native fp16 FFT support (cuFFT C2C half-precision under the hood).
-Two constraints must be satisfied:
-
-1. cuFFT requires **power-of-2 sizes** for half-precision transforms.
-2. The frequency-domain products can exceed fp16's 65504 limit with real model weights.
-
-To avoid overflow we use ``norm="ortho"`` which divides both the forward and inverse
-FFT by ``sqrt(N)``.  This keeps all intermediate complex values within fp16 range.
-The ortho pair computes ``conv / sqrt(N)`` instead of ``conv``, so we multiply by
-``sqrt(N)`` after the inverse FFT to restore the correct scale.
-
-Families provided
------------------
-Standard:
-  - 1D: ``fftconv1d_fp16_bhl``, ``causal_fftconv1d_fp16_bhl`` (+``_w_reshape`` variants)
-  - 2D: ``fftconv2d_fp16_bhl`` (+``_w_reshape``)
-  - 3D: ``fftconv3d_fp16_bhl`` (+``_w_reshape``)
-
-Chunked (memory-efficient, processes channels in chunks):
-  - 1D: ``fftconv1d_fp16_bhl_chunked``, ``causal_fftconv1d_fp16_bhl_chunked``
-        (+``_w_reshape`` variants)
-  - 2D: ``fftconv2d_fp16_bhl_chunked`` (+``_w_reshape``)
-  - 3D: ``fftconv3d_fp16_bhl_chunked`` (+``_w_reshape``)
-"""
-
-import math
-
-import torch
-from einops import rearrange
-
-
-def _next_power_of_2(n: int) -> int:
-    """Return the smallest power of 2 >= n."""
-    if n <= 0:
-        return 1
-    return 1 << (n - 1).bit_length()
-
-
-###############################################################################
-# 1D — non-causal
-###############################################################################
-
-
-def fftconv1d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """1D FFT convolution in fp16 with power-of-2 padding and ortho normalization.
-
-    Drop-in replacement for ``fftconv1d_fp32_bhl``.  Casts *x* and *kernel* to fp16
-    internally; shortcut is never cast.  Returns in the **original** dtype of *x*.
-
-    Args:
-        x: Input tensor ``[B, H, L]`` (any dtype).
-        kernel: Kernel tensor ``[1, H, K]`` or ``[B, H, K]`` (same dtype as *x*).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, L]`` in the original dtype of *x*.
-    """
-    x_fp16 = x.to(torch.float16)
-    k_fp16 = kernel.to(torch.float16)
-
-    _, _hidden_dim, seq_len = x.shape
-    _, _, kernel_len = kernel.shape
-
-    fft_len = _next_power_of_2(min(seq_len + (kernel_len + 1) // 2, 2 * seq_len))
-    sqrt_N = math.sqrt(fft_len)
-
-    fft_x = torch.fft.rfft(x_fp16, n=fft_len, dim=2, norm="ortho")
-    fft_k = torch.fft.rfft(k_fp16, n=fft_len, dim=2, norm="ortho")
-    fft_x.mul_(fft_k)
-
-    crop_start = kernel_len // 2
-    y = torch.fft.irfft(fft_x, n=fft_len, dim=2, norm="ortho")[..., crop_start : crop_start + seq_len]
-
-    # Upcast before scaling to avoid fp16 overflow (sqrt_N can be large)
-    y = y.to(x.dtype)
-    y = y * sqrt_N
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        y = y + rearrange(shortcut, "h -> 1 h 1") * x
-
-    return y
-
-
-def fftconv1d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """1D FFT convolution in fp16, for inputs with layout ``[B, L, H]``.
-
-    Wrapper around :func:`fftconv1d_fp16_bhl` that handles the BLH <-> BHL reshape.
-    """
-    x = rearrange(x, "b l h -> b h l")
-    kernel = rearrange(kernel, "b l h -> b h l")
-    y = fftconv1d_fp16_bhl(x, kernel, shortcut)
-    return rearrange(y, "b h l -> b l h")
-
-
-###############################################################################
-# 1D — causal
-###############################################################################
-
-
-def causal_fftconv1d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Causal 1D FFT convolution in fp16 with power-of-2 padding and ortho normalization.
-
-    Drop-in replacement for ``causal_fftconv1d_fp32_bhl``.  Casts *x* and *kernel* to
-    fp16 internally; shortcut is never cast.  Returns in the **original** dtype of *x*.
-
-    Args:
-        x: Input tensor ``[B, H, L]`` (any dtype).
-        kernel: Kernel tensor ``[1, H, K]`` or ``[B, H, K]`` (same dtype as *x*).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, L]`` in the original dtype of *x*.
-    """
-    x_fp16 = x.to(torch.float16)
-    k_fp16 = kernel.to(torch.float16)
-
-    _, _hidden_dim, seq_len = x.shape
-    _, _, kernel_len = kernel.shape
-
-    # Causal: need seq_len + kernel_len to avoid wraparound
-    fft_len = _next_power_of_2(min(seq_len + kernel_len, 2 * seq_len))
-    sqrt_N = math.sqrt(fft_len)
-
-    fft_x = torch.fft.rfft(x_fp16, n=fft_len, dim=2, norm="ortho")
-    fft_k = torch.fft.rfft(k_fp16, n=fft_len, dim=2, norm="ortho")
-    fft_x.mul_(fft_k)
-
-    y = torch.fft.irfft(fft_x, n=fft_len, dim=2, norm="ortho")[..., :seq_len]
-
-    # Upcast before scaling to avoid fp16 overflow (sqrt_N can be large)
-    y = y.to(x.dtype)
-    y = y * sqrt_N
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        y = y + rearrange(shortcut, "h -> 1 h 1") * x
-
-    return y
-
-
-def causal_fftconv1d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Causal 1D FFT convolution in fp16, for inputs with layout ``[B, L, H]``.
-
-    Wrapper around :func:`causal_fftconv1d_fp16_bhl` that handles the BLH <-> BHL reshape.
-    """
-    x = rearrange(x, "b l h -> b h l")
-    kernel = rearrange(kernel, "b l h -> b h l")
-    y = causal_fftconv1d_fp16_bhl(x, kernel, shortcut)
-    return rearrange(y, "b h l -> b l h")
-
-
-###############################################################################
-# 2D
-###############################################################################
-
-
-def fftconv2d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution in fp16 with power-of-2 padding and ortho normalization.
-
-    Drop-in replacement for ``fftconv2d_fp32_bhl``.  Casts *x* and *kernel* to fp16
-    internally; shortcut is never cast.  Returns in the **original** dtype of *x*.
-
-    Args:
-        x: Input tensor ``[B, H, X_in, Y_in]`` (any dtype).
-        kernel: Kernel tensor ``[1, H, K_x, K_y]`` or ``[B, H, K_x, K_y]`` (same dtype as *x*).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, X_in, Y_in]`` in the original dtype of *x*.
-    """
-    x_fp16 = x.to(torch.float16)
-    k_fp16 = kernel.to(torch.float16)
-
-    _B, _hidden_dim, X_in, Y_in = x.shape
-    _, _, K_x, K_y = kernel.shape
-
-    fft_shape = (
-        _next_power_of_2(min(X_in + (K_x + 1) // 2, 2 * X_in)),
-        _next_power_of_2(min(Y_in + (K_y + 1) // 2, 2 * Y_in)),
-    )
-
-    # sqrt(N) correction: ortho computes conv/sqrt(N), we need conv
-    sqrt_N = math.sqrt(fft_shape[0] * fft_shape[1])
-
-    fft_x = torch.fft.rfft2(x_fp16, s=fft_shape, dim=(2, 3), norm="ortho")
-    fft_k = torch.fft.rfft2(k_fp16, s=fft_shape, dim=(2, 3), norm="ortho")
-
-    fft_x.mul_(fft_k)
-
-    crop_start_x = K_x // 2
-    crop_start_y = K_y // 2
-
-    y = torch.fft.irfft2(fft_x, s=fft_shape, dim=(2, 3), norm="ortho")[
-        ..., crop_start_x : crop_start_x + X_in, crop_start_y : crop_start_y + Y_in
-    ]
-
-    # Upcast before scaling to avoid fp16 overflow (sqrt_N can be large)
-    y = y.to(x.dtype)
-    y = y * sqrt_N
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        y = y + rearrange(shortcut, "h -> 1 h 1 1") * x
-
-    return y
-
-
-def fftconv2d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution in fp16, for inputs with layout ``[B, X, Y, H]``.
-
-    Wrapper around :func:`fftconv2d_fp16_bhl` that handles the BLH <-> BHL reshape.
-    """
-    x = rearrange(x, "b x y h -> b h x y")
-    kernel = rearrange(kernel, "b x y h -> b h x y")
-    y = fftconv2d_fp16_bhl(x, kernel, shortcut)
-    return rearrange(y, "b h x y -> b x y h")
-
-
-###############################################################################
-# 3D
-###############################################################################
-
-
-def fftconv3d_fp16_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """3D FFT convolution in fp16 with power-of-2 padding and ortho normalization.
-
-    Drop-in replacement for ``fftconv3d_fp32_bhl``.  Casts *x* and *kernel* to fp16
-    internally; shortcut is never cast.  Returns in the **original** dtype of *x*.
-
-    Args:
-        x: Input tensor ``[B, H, X_in, Y_in, Z_in]`` (any dtype).
-        kernel: Kernel tensor ``[1, H, K_x, K_y, K_z]`` or ``[B, H, K_x, K_y, K_z]`` (same dtype as *x*).
-        shortcut: Optional per-channel shortcut ``[H]`` (same dtype as *x*).
-
-    Returns:
-        Tensor ``[B, H, X_in, Y_in, Z_in]`` in the original dtype of *x*.
-    """
-    x_fp16 = x.to(torch.float16)
-    k_fp16 = kernel.to(torch.float16)
-
-    _B, _hidden_dim, X_in, Y_in, Z_in = x.shape
-    _, _, K_x, K_y, K_z = kernel.shape
-
-    fft_shape = (
-        _next_power_of_2(min(X_in + (K_x + 1) // 2, 2 * X_in)),
-        _next_power_of_2(min(Y_in + (K_y + 1) // 2, 2 * Y_in)),
-        _next_power_of_2(min(Z_in + (K_z + 1) // 2, 2 * Z_in)),
-    )
-
-    sqrt_N = math.sqrt(fft_shape[0] * fft_shape[1] * fft_shape[2])
-
-    fft_x = torch.fft.rfftn(x_fp16, s=fft_shape, dim=(2, 3, 4), norm="ortho")
-    fft_k = torch.fft.rfftn(k_fp16, s=fft_shape, dim=(2, 3, 4), norm="ortho")
-    fft_x.mul_(fft_k)
-
-    crop_start_x = K_x // 2
-    crop_start_y = K_y // 2
-    crop_start_z = K_z // 2
-
-    y = torch.fft.irfftn(fft_x, s=fft_shape, dim=(2, 3, 4), norm="ortho")[
-        :,
-        :,
-        crop_start_x : crop_start_x + X_in,
-        crop_start_y : crop_start_y + Y_in,
-        crop_start_z : crop_start_z + Z_in,
-    ]
-
-    # Upcast before scaling to avoid fp16 overflow (sqrt_N can be large)
-    y = y.to(x.dtype)
-    y = y * sqrt_N
-    if shortcut is not None:
-        assert shortcut.dtype == x.dtype, f"shortcut.dtype ({shortcut.dtype}) must match x.dtype ({x.dtype})"
-        y = y + rearrange(shortcut, "h -> 1 h 1 1 1") * x
-
-    return y
-
-
-def fftconv3d_fp16_bhl_w_reshape(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """3D FFT convolution in fp16, for inputs with layout ``[B, X, Y, Z, H]``.
-
-    Wrapper around :func:`fftconv3d_fp16_bhl` that handles the BLH <-> BHL reshape.
-    """
-    x = rearrange(x, "b x y z h -> b h x y z")
-    kernel = rearrange(kernel, "b x y z h -> b h x y z")
-    y = fftconv3d_fp16_bhl(x, kernel, shortcut)
-    return rearrange(y, "b h x y z -> b x y z h")
-
-
-###############################################################################
-# Chunked (memory-efficient) variants
-#
-# Same channel-chunking strategy as fftconv_chunked.py but calling the fp16
-# base functions.  This combines the ~36% fp16 memory savings with the ~26%
-# chunking savings for maximum memory reduction.
-###############################################################################
-
-_DEFAULT_CHUNK_SIZE = 128
-
-
-def fftconv1d_fp16_bhl_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """1D FFT convolution in fp16 with channel chunking (BHL layout).
-
-    Args:
-        x: Input tensor ``[B, H, L]``.
-        kernel: Kernel tensor ``[1|B, H, K]``.
-        shortcut: Optional per-channel shortcut ``[H]``.
-        chunk_size: Channels per chunk (None = default 128).
-
-    Returns:
-        Tensor ``[B, H, L]`` in the original dtype of *x*.
-    """
-    if chunk_size is None:
-        chunk_size = _DEFAULT_CHUNK_SIZE
-    H = x.shape[1]
-    if H <= chunk_size:
-        return fftconv1d_fp16_bhl(x, kernel, shortcut)
-    outputs = []
-    for start in range(0, H, chunk_size):
-        end = min(start + chunk_size, H)
-        outputs.append(
-            fftconv1d_fp16_bhl(
-                x[:, start:end],
-                kernel[:, start:end],
-                shortcut[start:end] if shortcut is not None else None,
-            )
-        )
-    return torch.cat(outputs, dim=1)
-
-
-def fftconv1d_fp16_bhl_w_reshape_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """1D FFT convolution in fp16 with chunking, for ``[B, L, H]`` inputs."""
-    x = rearrange(x, "b l h -> b h l")
-    kernel = rearrange(kernel, "b l h -> b h l")
-    y = fftconv1d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size)
-    return rearrange(y, "b h l -> b l h")
-
-
-def causal_fftconv1d_fp16_bhl_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """Causal 1D FFT convolution in fp16 with channel chunking (BHL layout).
-
-    Args:
-        x: Input tensor ``[B, H, L]``.
-        kernel: Kernel tensor ``[1|B, H, K]``.
-        shortcut: Optional per-channel shortcut ``[H]``.
-        chunk_size: Channels per chunk (None = default 128).
-
-    Returns:
-        Tensor ``[B, H, L]`` in the original dtype of *x*.
-    """
-    if chunk_size is None:
-        chunk_size = _DEFAULT_CHUNK_SIZE
-    H = x.shape[1]
-    if H <= chunk_size:
-        return causal_fftconv1d_fp16_bhl(x, kernel, shortcut)
-    outputs = []
-    for start in range(0, H, chunk_size):
-        end = min(start + chunk_size, H)
-        outputs.append(
-            causal_fftconv1d_fp16_bhl(
-                x[:, start:end],
-                kernel[:, start:end],
-                shortcut[start:end] if shortcut is not None else None,
-            )
-        )
-    return torch.cat(outputs, dim=1)
-
-
-def causal_fftconv1d_fp16_bhl_w_reshape_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """Causal 1D FFT convolution in fp16 with chunking, for ``[B, L, H]`` inputs."""
-    x = rearrange(x, "b l h -> b h l")
-    kernel = rearrange(kernel, "b l h -> b h l")
-    y = causal_fftconv1d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size)
-    return rearrange(y, "b h l -> b l h")
-
-
-def fftconv2d_fp16_bhl_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution in fp16 with channel chunking (BHL layout).
-
-    Args:
-        x: Input tensor ``[B, H, X_in, Y_in]``.
-        kernel: Kernel tensor ``[1|B, H, K_x, K_y]``.
-        shortcut: Optional per-channel shortcut ``[H]``.
-        chunk_size: Channels per chunk (None = default 128).
-
-    Returns:
-        Tensor ``[B, H, X_in, Y_in]`` in the original dtype of *x*.
-    """
-    if chunk_size is None:
-        chunk_size = _DEFAULT_CHUNK_SIZE
-    H = x.shape[1]
-    if H <= chunk_size:
-        return fftconv2d_fp16_bhl(x, kernel, shortcut)
-    outputs = []
-    for start in range(0, H, chunk_size):
-        end = min(start + chunk_size, H)
-        outputs.append(
-            fftconv2d_fp16_bhl(
-                x[:, start:end],
-                kernel[:, start:end],
-                shortcut[start:end] if shortcut is not None else None,
-            )
-        )
-    return torch.cat(outputs, dim=1)
-
-
-def fftconv2d_fp16_bhl_w_reshape_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution in fp16 with chunking, for ``[B, X, Y, H]`` inputs."""
-    x = rearrange(x, "b x y h -> b h x y")
-    kernel = rearrange(kernel, "b x y h -> b h x y")
-    y = fftconv2d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size)
-    return rearrange(y, "b h x y -> b x y h")
-
-
-def fftconv3d_fp16_bhl_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """3D FFT convolution in fp16 with channel chunking (BHL layout).
-
-    Args:
-        x: Input tensor ``[B, H, X_in, Y_in, Z_in]``.
-        kernel: Kernel tensor ``[1|B, H, K_x, K_y, K_z]``.
-        shortcut: Optional per-channel shortcut ``[H]``.
-        chunk_size: Channels per chunk (None = default 128).
-
-    Returns:
-        Tensor ``[B, H, X_in, Y_in, Z_in]`` in the original dtype of *x*.
-    """
-    if chunk_size is None:
-        chunk_size = _DEFAULT_CHUNK_SIZE
-    H = x.shape[1]
-    if H <= chunk_size:
-        return fftconv3d_fp16_bhl(x, kernel, shortcut)
-    outputs = []
-    for start in range(0, H, chunk_size):
-        end = min(start + chunk_size, H)
-        outputs.append(
-            fftconv3d_fp16_bhl(
-                x[:, start:end],
-                kernel[:, start:end],
-                shortcut[start:end] if shortcut is not None else None,
-            )
-        )
-    return torch.cat(outputs, dim=1)
-
-
-def fftconv3d_fp16_bhl_w_reshape_chunked(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-    chunk_size: int | None = None,
-) -> torch.Tensor:
-    """3D FFT convolution in fp16 with chunking, for ``[B, X, Y, Z, H]`` inputs."""
-    x = rearrange(x, "b x y z h -> b h x y z")
-    kernel = rearrange(kernel, "b x y z h -> b h x y z")
-    y = fftconv3d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size)
-    return rearrange(y, "b h x y z -> b x y z h")
diff --git a/nvsubquadratic/ops/fftconv_multihead.py b/nvsubquadratic/ops/fftconv_multihead.py
deleted file mode 100644
index 829628b0..00000000
--- a/nvsubquadratic/ops/fftconv_multihead.py
+++ /dev/null
@@ -1,294 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-r"""Multi-head 2D FFT convolution operators with dense within-head channel mixing.
-
-Motivation
-----------
-The standard FFT convolutions in :mod:`nvsubquadratic.ops.fftconv` are
-*depthwise*: each output channel is the convolution of a single input
-channel with its own kernel, i.e. the kernel has shape ``[H, K_x, K_y]`` and
-the frequency-domain product is element-wise across channels. This is fast
-but lets no information flow between channels — channel mixing has to be
-done by a separate pointwise layer (typically a 1x1 conv or MLP).
-
-The **multi-head** variant in this module bundles depthwise spatial mixing
-and dense channel mixing into a single op, in the spirit of multi-head
-attention:
-
-- Channels are split into ``num_heads`` groups of ``head_dim`` each.
-- Within a head, the convolution is *dense* across channels: every output
-  channel sees every input channel through its own learned 2D kernel.
-- Across heads, channels remain isolated (no cross-head mixing).
-
-Frequency-domain operation
---------------------------
-Concretely, if :math:`\hat{x} \in \mathbb{C}^{B \times n \times d \times F_x \times F_y}`
-is the rFFT of the input (with :math:`d = \text{head\_dim}`,
-:math:`n = \text{num\_heads}`), and
-:math:`\hat{K} \in \mathbb{C}^{n \times d_o \times d_i \times F_x \times F_y}`
-is the rFFT of the kernel, then per spatial frequency bin :math:`(f_x, f_y)`
-each head applies a dense :math:`d_o \times d_i` matrix to the input vector:
-
-.. math::
-    \hat{y}_{b, n, o, f_x, f_y}
-        = \sum_{i} \hat{K}_{n, o, i, f_x, f_y} \, \hat{x}_{b, n, i, f_x, f_y}
-
-Implemented as a single :func:`torch.einsum` over the frequency-domain
-tensors. The inverse rFFT then materialises the spatial output.
-
-Low-rank variant
-----------------
-For large ``head_dim``, the dense kernel ``[d_o, d_i, K_x, K_y]`` has
-:math:`d^2 K_x K_y` parameters and a matching memory/compute cost in
-frequency domain. The low-rank functions factor the kernel as
-:math:`K = U V` with rank ``r < d``:
-
-.. math::
-    \hat{y} = \hat{U} \,(\hat{V} \hat{x})
-
-This drops the per-frequency-bin cost from :math:`O(d^2)` to :math:`O(2 d r)`
-and the parameter count from :math:`d^2 K_x K_y` to :math:`2 d r K_x K_y`,
-without materialising the full :math:`d \times d` kernel spectrum.
-
-Shape conventions
------------------
-- Input: ``x: [B, num_heads, head_dim, H, W]``
-- Dense kernel: ``[num_heads, head_dim_out, head_dim_in, K_x, K_y]``
-- Low-rank factors: ``U: [n, d, r, K_x, K_y]``, ``V: [n, r, d, K_x, K_y]``
-- Output: ``[B, num_heads, head_dim, H, W]``
-
-Linear vs. circular
--------------------
-Functions without the ``_circular`` suffix produce same-aligned *linear*
-convolutions (zero-padded FFTs with crop). The ``_circular_*`` variants
-compute the periodic convolution at the same spatial size (no padding,
-no crop) and expect ``K_x == H``, ``K_y == W``.
-"""
-
-__all__ = [
-    "fftconv2d_multihead_bhl",
-    "fftconv2d_multihead_circular_bhl",
-    "fftconv2d_multihead_lowrank_bhl",
-    "fftconv2d_multihead_lowrank_circular_bhl",
-]
-
-import torch
-
-
-def fftconv2d_multihead_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution with dense within-head channel mixing.
-
-    Applies a dense [head_dim x head_dim] convolution within each head,
-    similar to how attention mixes across positions but here mixing channels.
-
-    Args:
-        x: Input tensor [B, num_heads, head_dim, H, W], float32
-        kernel: Kernel tensor [num_heads, head_dim_out, head_dim_in, K_x, K_y], float32
-        shortcut: Optional per-channel scale [num_heads * head_dim], float32
-
-    Returns:
-        Output tensor [B, num_heads, head_dim, H, W]
-    """
-    assert x.dtype == torch.float32, f"x must be float32. Current dtype: {x.dtype}"
-    assert kernel.dtype == torch.float32, f"kernel must be float32. Current dtype: {kernel.dtype}"
-    if shortcut is not None:
-        assert shortcut.dtype == torch.float32, f"shortcut must be float32. Current dtype: {shortcut.dtype}"
-
-    _B, num_heads, head_dim, H, W = x.shape
-    N, head_dim_out, head_dim_in, K_x, K_y = kernel.shape
-
-    assert N == num_heads, f"Kernel num_heads ({N}) must match input ({num_heads})"
-    assert head_dim_out == head_dim_in == head_dim, (
-        f"Kernel head_dims ({head_dim_out}, {head_dim_in}) must match input head_dim ({head_dim})"
-    )
-    assert K_x <= H * 2, f"Kernel size K_x ({K_x}) must be <= 2 * H ({2 * H})"
-    assert K_y <= W * 2, f"Kernel size K_y ({K_y}) must be <= 2 * W ({2 * W})"
-
-    # Determine FFT size for linear convolution
-    fft_h = min(H + (K_x + 1) // 2, 2 * H)
-    fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-    # FFT of input: [B, num_heads, head_dim, fft_h, fft_w//2+1]
-    x_fft = torch.fft.rfft2(x, s=(fft_h, fft_w))
-
-    # FFT of kernel: [num_heads, head_dim_out, head_dim_in, fft_h, fft_w//2+1]
-    k_fft = torch.fft.rfft2(kernel, s=(fft_h, fft_w))
-
-    # Dense conv within heads using einsum
-    # x_fft: [B, num_heads, head_dim_in, fft_h, fft_w//2+1]
-    # k_fft: [num_heads, head_dim_out, head_dim_in, fft_h, fft_w//2+1]
-    # out_fft: [B, num_heads, head_dim_out, fft_h, fft_w//2+1]
-    out_fft = torch.einsum("bnihw,noihw->bnohw", x_fft, k_fft)
-
-    # Crop after inverse FFT
-    crop_h = K_x // 2
-    crop_w = K_y // 2
-
-    # Inverse FFT and crop
-    out_full = torch.fft.irfft2(out_fft, s=(fft_h, fft_w))
-    out = out_full[..., crop_h : crop_h + H, crop_w : crop_w + W]
-
-    # Add shortcut if provided
-    if shortcut is not None:
-        # shortcut: [num_heads * head_dim] -> [1, num_heads, head_dim, 1, 1]
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-    return out
-
-
-def fftconv2d_multihead_lowrank_bhl(
-    x: torch.Tensor,
-    kernel_u: torch.Tensor,
-    kernel_v: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D FFT convolution with low-rank within-head channel mixing.
-
-    Instead of a full [head_dim x head_dim] kernel, uses a rank-r factorization:
-        K = U @ V^T
-    where U: [num_heads, head_dim, rank, K_x, K_y] and
-          V: [num_heads, rank, head_dim, K_x, K_y].
-
-    The convolution becomes two smaller matmuls in the frequency domain:
-        z = V @ x  (sum over head_dim_in, output rank)
-        y = U @ z  (sum over rank, output head_dim_out)
-
-    Args:
-        x: Input tensor [B, num_heads, head_dim, H, W], float32
-        kernel_u: U factor [num_heads, head_dim, rank, K_x, K_y], float32
-        kernel_v: V factor [num_heads, rank, head_dim, K_x, K_y], float32
-        shortcut: Optional per-channel scale [num_heads * head_dim], float32
-
-    Returns:
-        Output tensor [B, num_heads, head_dim, H, W]
-    """
-    _B, num_heads, head_dim, H, W = x.shape
-    _, _, _rank, K_x, K_y = kernel_u.shape
-
-    # Determine FFT size for linear convolution
-    fft_h = min(H + (K_x + 1) // 2, 2 * H)
-    fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-    # FFT of input and kernel factors
-    x_fft = torch.fft.rfft2(x, s=(fft_h, fft_w))
-    u_fft = torch.fft.rfft2(kernel_u, s=(fft_h, fft_w))
-    v_fft = torch.fft.rfft2(kernel_v, s=(fft_h, fft_w))
-
-    # Two-step low-rank conv (equivalent to K_fft @ x_fft where K_fft = U_fft @ V_fft,
-    # but avoids materializing the full [head_dim x head_dim] K_fft).
-    # Step 1: z = V @ x — contract over head_dim_in
-    # x_fft: [B, n, head_dim, fft_h, fft_w'], v_fft: [n, rank, head_dim, fft_h, fft_w']
-    z_fft = torch.einsum("bnihw,nrihw->bnrhw", x_fft, v_fft)
-
-    # Step 2: y = U @ z — contract over rank
-    # z_fft: [B, n, rank, fft_h, fft_w'], u_fft: [n, head_dim, rank, fft_h, fft_w']
-    out_fft = torch.einsum("bnrhw,norhw->bnohw", z_fft, u_fft)
-
-    # Crop after inverse FFT
-    crop_h = K_x // 2
-    crop_w = K_y // 2
-    out_full = torch.fft.irfft2(out_fft, s=(fft_h, fft_w))
-    out = out_full[..., crop_h : crop_h + H, crop_w : crop_w + W]
-
-    if shortcut is not None:
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-    return out
-
-
-def fftconv2d_multihead_circular_bhl(
-    x: torch.Tensor,
-    kernel: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D circular FFT convolution with dense within-head channel mixing.
-
-    Same as fftconv2d_multihead_bhl but uses circular (periodic) convolution.
-    Kernel size should equal input size for circular conv.
-
-    Args:
-        x: Input tensor [B, num_heads, head_dim, H, W], float32
-        kernel: Kernel tensor [num_heads, head_dim_out, head_dim_in, H, W], float32
-        shortcut: Optional per-channel scale [num_heads * head_dim], float32
-
-    Returns:
-        Output tensor [B, num_heads, head_dim, H, W]
-    """
-    assert x.dtype == torch.float32, f"x must be float32. Current dtype: {x.dtype}"
-    assert kernel.dtype == torch.float32, f"kernel must be float32. Current dtype: {kernel.dtype}"
-
-    _B, num_heads, head_dim, H, W = x.shape
-
-    # FFT of input and kernel (circular: no padding needed)
-    x_fft = torch.fft.rfft2(x)
-    k_fft = torch.fft.rfft2(kernel, s=(H, W))
-
-    # Dense conv within heads
-    out_fft = torch.einsum("bnihw,noihw->bnohw", x_fft, k_fft)
-
-    # Inverse FFT (no cropping for circular)
-    out = torch.fft.irfft2(out_fft, s=(H, W))
-
-    # Add shortcut if provided
-    if shortcut is not None:
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-    return out
-
-
-def fftconv2d_multihead_lowrank_circular_bhl(
-    x: torch.Tensor,
-    kernel_u: torch.Tensor,
-    kernel_v: torch.Tensor,
-    shortcut: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """2D circular FFT convolution with low-rank within-head channel mixing.
-
-    Circular (periodic) variant of fftconv2d_multihead_lowrank_bhl.
-
-    Args:
-        x: Input tensor [B, num_heads, head_dim, H, W], float32
-        kernel_u: U factor [num_heads, head_dim, rank, H, W], float32
-        kernel_v: V factor [num_heads, rank, head_dim, H, W], float32
-        shortcut: Optional per-channel scale [num_heads * head_dim], float32
-
-    Returns:
-        Output tensor [B, num_heads, head_dim, H, W]
-    """
-    _B, num_heads, head_dim, H, W = x.shape
-
-    x_fft = torch.fft.rfft2(x)
-    u_fft = torch.fft.rfft2(kernel_u, s=(H, W))
-    v_fft = torch.fft.rfft2(kernel_v, s=(H, W))
-
-    # Two-step low-rank conv (avoids materializing full [head_dim x head_dim] K_fft)
-    z_fft = torch.einsum("bnihw,nrihw->bnrhw", x_fft, v_fft)
-    out_fft = torch.einsum("bnrhw,norhw->bnohw", z_fft, u_fft)
-
-    out = torch.fft.irfft2(out_fft, s=(H, W))
-
-    if shortcut is not None:
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out = out + x * shortcut_reshaped
-
-    return out
diff --git a/nvsubquadratic/ops/mixed_fftconv.py b/nvsubquadratic/ops/mixed_fftconv.py
index fd7d221f..176b13a6 100644
--- a/nvsubquadratic/ops/mixed_fftconv.py
+++ b/nvsubquadratic/ops/mixed_fftconv.py
@@ -65,8 +65,8 @@
 *internal* normalised form used here; YAML configs should always use the
 string list form.
 
-See ``docs/ops/MIXED_BC_PLAN.md`` for the per-axis algorithm, dataset
-motivation table, and deferred work (fp16, multi-head, per-face BCs).
+See ``docs/ops/mixed_boundary_conditions.md`` for the per-axis algorithm,
+dataset motivation table, and limitations.
 
 Algorithm overview
 ------------------
diff --git a/pyproject.toml b/pyproject.toml
index 87657631..befa3cb6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ description = "A unified PyTorch-native library for subquadratic alternatives to
 requires-python = ">=3.11"
 readme = "README.md"
 license = "Apache-2.0"
-license-files = ["LICENSE/License.txt", "LICENSE/third_party.txt"]
+license-files = ["LICENSE", "THIRD_PARTY_NOTICES.txt"]
 authors = [
     { name = "NVIDIA Corporation" },
 ]
@@ -48,8 +48,6 @@ dependencies = [
     "numpy>=1.21.0",
     "pyarrow>=14.0.0,<20.0.0",     # Constrained to be compatible with cudf/pylibcudf in NGC containers
     "datasets>=2.14.0",
-    "diffusers>=0.25.0",
-    "clean-fid>=0.1.35",
     "megatron-core",
     "pytorch-lightning>=2.0.0",
     "omegaconf>=2.3.0",
diff --git a/scripts/data/generate_jit_fid_stats.py b/scripts/data/generate_jit_fid_stats.py
deleted file mode 100644
index 8832ece9..00000000
--- a/scripts/data/generate_jit_fid_stats.py
+++ /dev/null
@@ -1,175 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Generate cached FID reference statistics for the JiT ImageNet diffusion pipeline.
-
-Builds an ImageNet reference subset, feeds it through the
-``torch-fidelity`` Inception feature extractor, and saves the
-(mean, cov) pair as an ``.npz`` so subsequent FID evaluations can skip
-the per-run feature extraction.
-
-Targets: any CUDA GPU (H100 / A100 / Ampere+); ImageNet-1k on disk and
-``torch-fidelity`` installed.
-
-Usage:
-    PYTHONPATH=. conda run -n nv-subq python \\
-        scripts/data/generate_jit_fid_stats.py --output-dir /path/to/fid_stats
-
-Output: ``.npz`` files under ``--output-dir`` (one per requested
-resolution / split).
-"""
-
-import argparse
-import os
-from pathlib import Path
-
-import numpy as np
-import torch
-from torch.utils.data import Dataset, Subset
-from torch_fidelity.metric_fid import fid_featuresdict_to_statistics
-from torch_fidelity.utils import create_feature_extractor, extract_featuresdict_from_input_id_cached
-
-from experiments.datamodules._deprecated.ref_imagenet import ImageNetDataModule
-
-
-def _default_cache_dir() -> str:
-    return (
-        os.environ.get("IMAGENET_PATH")
-        or os.environ.get("IMAGENET_CACHE")
-        or os.environ.get("IMAGENET_CACHE_DIR")
-        or str(Path.cwd() / "imagenet")
-    )
-
-
-class _FIDReadyDataset(Dataset):
-    """Wrap ImageNetDataModule output and emit uint8 RGB tensors for torch-fidelity."""
-
-    def __init__(self, base_dataset: Dataset, datamodule: ImageNetDataModule) -> None:
-        self.base_dataset = base_dataset
-        self.datamodule = datamodule
-
-    def __len__(self) -> int:
-        return len(self.base_dataset)
-
-    def __getitem__(self, index: int) -> torch.Tensor:
-        image, _ = self.base_dataset[index]  # CHW, normalized for diffusion ([-1, 1])
-        image = self.datamodule.unnormalize(image)  # CHW, [0, 1]
-        image = torch.clamp(image * 255.0, 0.0, 255.0).to(torch.uint8)
-        return image
-
-
-def _build_reference_dataset(args: argparse.Namespace) -> Dataset:
-    datamodule = ImageNetDataModule(
-        data_dir=args.cache_dir,
-        batch_size=1,
-        num_workers=0,
-        pin_memory=False,
-        seed=args.seed,
-        image_size=args.resize_image_size,
-        final_image_size=args.image_size,
-        center_crop=True,
-        drop_labels=False,
-        hf_dataset_name="imagenet-1k",
-        hf_dataset_config=None,
-        hf_auth_token=os.environ.get("HF_TOKEN"),
-        task="generation",
-    )
-
-    if args.split == "train":
-        datamodule.setup(stage="fit")
-        if datamodule.train_dataset is None:
-            raise RuntimeError("ImageNetDataModule did not initialize the training dataset.")
-        base_dataset = datamodule.train_dataset
-    else:
-        datamodule.setup(stage="validate")
-        if datamodule.val_dataset is None:
-            raise RuntimeError("ImageNetDataModule did not initialize the validation dataset.")
-        base_dataset = datamodule.val_dataset
-
-    dataset: Dataset = _FIDReadyDataset(base_dataset, datamodule)
-    if args.max_samples is not None:
-        dataset = Subset(dataset, range(min(len(dataset), args.max_samples)))
-    return dataset
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description=(
-            "Generate FID reference statistics using the exact ImageNetDataModule generation preprocessing pipeline."
-        )
-    )
-    parser.add_argument("--cache_dir", type=str, default=_default_cache_dir())
-    parser.add_argument("--split", type=str, default="validation", choices=["train", "validation"])
-    parser.add_argument(
-        "--image_size",
-        type=int,
-        default=64,
-        help="Final generated image size (ImageNetDataModule.final_image_size).",
-    )
-    parser.add_argument(
-        "--resize_image_size",
-        type=int,
-        default=256,
-        help="Pre-crop resize size used by generation datamodule (ImageNetDataModule.image_size).",
-    )
-    parser.add_argument("--max_samples", type=int, default=None, help="Optional max number of samples to process.")
-    parser.add_argument("--batch_size", type=int, default=64, help="Feature extraction batch size.")
-    parser.add_argument("--seed", type=int, default=42, help="Torch random seed (affects train split flips).")
-    parser.add_argument(
-        "--output",
-        type=str,
-        default="examples/imagenet_diffusion/fid_stats/jit_in64_stats.npz",
-        help="Output .npz path for mu/sigma statistics.",
-    )
-    args = parser.parse_args()
-
-    torch.manual_seed(args.seed)
-    np.random.seed(args.seed)
-
-    print(
-        f"Generating FID statistics for split={args.split}, "
-        f"image_size={args.image_size}, resize_image_size={args.resize_image_size}"
-    )
-    print(f"Using ImageNet cache dir: {args.cache_dir}")
-
-    dataset = _build_reference_dataset(args)
-    print(f"Processing {len(dataset)} images...")
-
-    use_cuda = torch.cuda.is_available()
-    feature_extractor = create_feature_extractor(
-        "inception-v3-compat",
-        ["2048"],
-        cuda=use_cuda,
-        verbose=True,
-    )
-    featuresdict = extract_featuresdict_from_input_id_cached(
-        1,
-        feature_extractor,
-        input1=dataset,
-        cuda=use_cuda,
-        batch_size=args.batch_size,
-        verbose=True,
-    )
-
-    stats = fid_featuresdict_to_statistics(featuresdict, "2048")
-
-    output_path = Path(args.output).expanduser()
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    np.savez_compressed(output_path, mu=stats["mu"], sigma=stats["sigma"])
-    print(f"Saved stats to {output_path}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/evaluation/eval_cfid.py b/scripts/evaluation/eval_cfid.py
deleted file mode 100644
index dd309170..00000000
--- a/scripts/evaluation/eval_cfid.py
+++ /dev/null
@@ -1,280 +0,0 @@
-#!/usr/bin/env python
-
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Generate samples from a trained diffusion model and compute CleanFID."""
-
-from __future__ import annotations
-
-import argparse
-from pathlib import Path
-from typing import Any
-
-import torch
-from torchvision.utils import save_image
-
-from experiments.utils.checkpointing import download_checkpoint, load_checkpoint_state_dict
-from experiments.utils.cli import (
-    apply_config_overrides,
-    load_config_from_file,
-    verify_no_interpolator_overwrites,
-)
-from nvsubquadratic.lazy_config import instantiate
-from nvsubquadratic.metrics.cleanfid import compute_folder_fid
-
-
-# Set high precision for matrix multiplication (tensor cores)
-torch.set_float32_matmul_precision("high")
-
-
-def _parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Offline CleanFID evaluation helper.")
-    parser.add_argument("--config", required=True, help="Path to the experiment config.")
-    parser.add_argument("--checkpoint", type=str, default=None, help="Path to a .ckpt file.")
-    parser.add_argument(
-        "--wandb-run-path",
-        type=str,
-        default=None,
-        help="Optional W&B run path (entity/project/run_id) to download a checkpoint from.",
-    )
-    parser.add_argument(
-        "--wandb-alias",
-        type=str,
-        default="best",
-        help="Checkpoint alias to download when --wandb-run-path is provided.",
-    )
-    parser.add_argument(
-        "--num-samples",
-        type=int,
-        default=50_000,
-        help="Total number of samples to generate before computing CleanFID.",
-    )
-    parser.add_argument(
-        "--sample-batch-size",
-        type=int,
-        default=250,
-        help="Number of samples to draw per diffusion pass.",
-    )
-    parser.add_argument(
-        "--num-inference-steps",
-        type=int,
-        default=None,
-        help="Override the scheduler inference steps used while sampling.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=Path,
-        required=True,
-        help="Directory where generated PNG samples will be written.",
-    )
-    parser.add_argument(
-        "--fid-dataset-name",
-        type=str,
-        default="imagenet2012",
-        help="CleanFID reference dataset name.",
-    )
-    parser.add_argument(
-        "--fid-dataset-res",
-        type=int,
-        default=64,
-        help="Resolution of the CleanFID reference statistics.",
-    )
-    parser.add_argument(
-        "--fid-dataset-split",
-        type=str,
-        default="train",
-        help="Dataset split to use for CleanFID reference statistics.",
-    )
-    parser.add_argument(
-        "--use-ema",
-        action="store_true",
-        help="Force EMA weights for sampling if the checkpoint contains them.",
-    )
-    parser.add_argument(
-        "overrides",
-        nargs="*",
-        help="Optional config overrides in key=value format.",
-    )
-    return parser.parse_args()
-
-
-def _extract_example_shape(datamodule: Any) -> tuple[int, int, int]:
-    if hasattr(datamodule, "num_workers"):
-        setattr(datamodule, "num_workers", 0)
-    datamodule.prepare_data()
-    datamodule.setup(stage="fit")
-    loader = datamodule.train_dataloader()
-    batch = next(iter(loader))
-    if isinstance(batch, dict):
-        example = batch["input"]
-    elif isinstance(batch, (list, tuple)):
-        example = batch[0]
-    else:
-        raise ValueError("Unsupported batch structure while deriving example shape.")
-    if not torch.is_tensor(example):
-        raise TypeError("Datamodule inputs must be tensors.")
-    shape = tuple(example.shape[1:])
-    if len(shape) != 3:
-        raise ValueError(f"Expected image-like inputs with 3 dims; got {shape}.")
-    return shape  # type: ignore[return-value]
-
-
-def _resolve_checkpoint_path(args: argparse.Namespace) -> str:
-    if bool(args.checkpoint) == bool(args.wandb_run_path):
-        raise ValueError("Provide exactly one of --checkpoint or --wandb-run-path.")
-    if args.checkpoint:
-        return args.checkpoint
-    return download_checkpoint(run_path=args.wandb_run_path, alias=args.wandb_alias)
-
-
-def _set_example_shape(model: torch.nn.Module, shape: tuple[int, int, int]) -> None:
-    if not hasattr(model, "example_input_shape"):
-        raise AttributeError("Model does not expose example_input_shape.")
-    model.example_input_shape = torch.Size(shape)  # type: ignore[attr-defined]
-
-
-def _save_sample_batch(samples: torch.Tensor, output_dir: Path, start_idx: int) -> None:
-    samples = samples.detach().cpu()
-    samples = torch.clamp((samples + 1.0) / 2.0, 0.0, 1.0)
-    samples = samples.permute(0, 3, 1, 2).contiguous()  # B, C, H, W
-    if samples.shape[1] == 1:
-        samples = samples.repeat(1, 3, 1, 1)
-    for offset, tensor in enumerate(samples):
-        save_image(tensor, output_dir / f"{start_idx + offset:06d}.png")
-
-
-def main() -> None:
-    args = _parse_args()
-    ckpt_path = _resolve_checkpoint_path(args)
-
-    config = load_config_from_file(args.config)
-    verify_no_interpolator_overwrites(config, args.overrides)
-    config = apply_config_overrides(config, args.overrides)
-    if getattr(config, "diffusion", None) is None:
-        raise ValueError("Selected config does not define diffusion settings.")
-
-    # Override batch size to match the requested sample batch size for efficiency
-    if hasattr(config.dataset, "batch_size"):
-        config.dataset.batch_size = args.sample_batch_size
-
-    datamodule = instantiate(config.dataset)
-    example_shape = _extract_example_shape(datamodule)
-    if hasattr(datamodule, "teardown"):
-        datamodule.teardown(stage="fit")
-
-    network = instantiate(
-        config.net,
-        in_channels=getattr(datamodule, "input_channels", None),
-        out_channels=getattr(datamodule, "output_channels", None),
-    )
-    model = instantiate(config.lightning_wrapper_class, network=network, cfg=config)
-    _set_example_shape(model, example_shape)
-
-    # Optional compilation
-    if getattr(config, "compile", False):
-        print("[fid] Compiling model with torch.compile...")
-        model.network = torch.compile(model.network)
-
-    state_dict = load_checkpoint_state_dict(ckpt_path)
-    load_msg = model.load_state_dict(state_dict, strict=False)
-    if load_msg.missing_keys:
-        print(f"[load] Missing keys: {len(load_msg.missing_keys)}")
-    if load_msg.unexpected_keys:
-        print(f"[load] Unexpected keys: {len(load_msg.unexpected_keys)}")
-    if args.use_ema and getattr(model, "ema_enabled", False) and getattr(model, "_ema_model", None) is not None:
-        model._ema_has_been_updated = True  # type: ignore[attr-defined]
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    model.eval()
-
-    output_dir = args.output_dir.expanduser().resolve()
-    if output_dir.exists() and any(output_dir.iterdir()):
-        raise RuntimeError(f"Output directory must be empty: {output_dir}")
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    total = 0
-    inference_steps = args.num_inference_steps or getattr(config.diffusion, "num_inference_steps", None)
-    if inference_steps is None:
-        raise ValueError("Unable to determine num_inference_steps for sampling.")
-
-    print(f"[fid] Generating {args.num_samples} samples into {output_dir} ...")
-
-    # Determine precision for autocast
-    precision_str = getattr(config.train, "precision", "32-true")
-    if precision_str in ["bf16-mixed", "bf16"]:
-        autocast_dtype = torch.bfloat16
-    elif precision_str in ["16-mixed", "fp16"]:
-        autocast_dtype = torch.float16
-    else:
-        autocast_dtype = torch.float32
-
-    # Use the training dataloader to ensure we sample labels from the correct distribution
-    datamodule.prepare_data()
-    datamodule.setup(stage="fit")
-    loader = datamodule.train_dataloader()
-    iterator = iter(loader)
-
-    with torch.inference_mode(), torch.autocast("cuda", dtype=autocast_dtype, enabled=autocast_dtype != torch.float32):
-        while total < args.num_samples:
-            try:
-                batch = next(iterator)
-            except StopIteration:
-                iterator = iter(loader)
-                batch = next(iterator)
-
-            # Handle different batch structures
-            if isinstance(batch, dict):
-                # We consume as many labels as we need to fill the batch or finish the quota
-                current_labels = batch.get("label")
-                if current_labels is not None:
-                    current_labels = current_labels.to(device)
-            elif isinstance(batch, (list, tuple)) and len(batch) > 1:
-                # Assumption: (data, label) tuple
-                current_labels = batch[1].to(device)
-            else:
-                current_labels = None
-
-            # Determine how many samples to generate in this pass
-            remaining = args.num_samples - total
-
-            # If we have labels, use their batch size (capped by remaining)
-            # If no labels, use sample-batch-size (capped by remaining)
-            if current_labels is not None:
-                current = min(current_labels.shape[0], remaining)
-                # Slice labels to match the number of samples we are generating
-                batch_labels = current_labels[:current]
-            else:
-                current = min(args.sample_batch_size, remaining)
-                batch_labels = None
-
-            samples = model.sample(num_samples=current, num_inference_steps=inference_steps, labels=batch_labels)
-            _save_sample_batch(samples, output_dir, total)
-            total += current
-            print(f"\r[sampling] {total}/{args.num_samples} images complete", end="")
-    print("\n[sampling] Done.")
-
-    score = compute_folder_fid(
-        output_dir,
-        dataset_name=args.fid_dataset_name,
-        dataset_resolution=args.fid_dataset_res,
-        dataset_split=args.fid_dataset_split,
-    )
-    print(f"[fid] CleanFID score: {score:.4f}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/slurm/diffusion/ccnn_jit_diff_baseline_snellius.sh b/scripts/slurm/diffusion/ccnn_jit_diff_baseline_snellius.sh
deleted file mode 100644
index b3c0a034..00000000
--- a/scripts/slurm/diffusion/ccnn_jit_diff_baseline_snellius.sh
+++ /dev/null
@@ -1,100 +0,0 @@
-#!/bin/bash
-#SBATCH -t 48:00:00
-#SBATCH --gres=gpu:4
-#SBATCH --partition=gpu_h100
-#SBATCH --chdir="/gpfs/home1/dknigge/nvsq_update_branch"
-#SBATCH --output=./runs/slurm/in1k_ccnn_jit_baseline_%j.out
-#SBATCH --error=./runs/slurm/in1k_ccnn_jit_baseline_%j.err
-#SBATCH --job-name='in1k-ccnn_jit_baseline'
-#SBATCH --dependency=singleton
-
-source /home/dknigge/.bashrc
-source /gpfs/home1/dknigge/nvsq_update_branch/.venv/bin/activate
-
-# HF dataset env vars (mirrors extract_imagenet_to_tar.py defaults)
-HF_DATASET="${IMAGENET_HF_DATASET:-imagenet-1k}"
-HF_CACHE="${IMAGENET_PATH:-$HOME/project_dir/huggingface/imagenet}"
-
-if [ -d /scratch-local ]; then
-  LOCAL_SCRATCH="/scratch-local/${USER}"
-else
-  LOCAL_SCRATCH="${TMPDIR:-/tmp}/${USER}"
-fi
-LOCAL_HF_CACHE="${LOCAL_SCRATCH}/hf_cache"
-SENTINEL="${LOCAL_HF_CACHE}/.hf_cache_complete"
-mkdir -p "${LOCAL_HF_CACHE}"
-
-# Capture the start time before staging so walltime budget is accurate
-JOB_START_TIMESTAMP=$(date +%s)
-
-echo "Staging HF cache to node-local storage: ${LOCAL_HF_CACHE}"
-if [ -f "${SENTINEL}" ]; then
-  echo "Sentinel found — local HF cache is complete; skipping staging."
-else
-  if [ ! -d "${HF_CACHE}" ]; then
-    echo "ERROR: HF cache not found at ${HF_CACHE}."
-    exit 1
-  fi
-  echo "Copying HF cache to local NVMe..."
-  rsync -a --info=progress2 "${HF_CACHE}/" "${LOCAL_HF_CACHE}/"
-  touch "${SENTINEL}"
-  echo "Staging complete."
-fi
-
-export IMAGENET_HF_DATASET="${HF_DATASET}"
-export IMAGENET_PATH="${LOCAL_HF_CACHE}"
-echo "Start time captured: ${JOB_START_TIMESTAMP}"
-WORKDIR="/gpfs/home1/dknigge/nvsq_update_branch"
-TIME_LIMIT_HOURS=48
-CONFIG_FILE="examples/imagenet_diffusion/ccnn_jit_baseline.py"
-CONFIG_OVERRIDES=(
-  "wandb.job_group=imagenet_diffusion_ccnn_jit_baseline"
-  "wandb.entity=implicit-long-convs"
-  # Note: Batch size is handled in config now (256).
-  # If we need to override batch size per GPU here we can, but let's stick to config file default.
-)
-EXPERIMENT_NAME="imagenet_diffusion_ccnn_jit_baseline"
-
-RUNS_DIR="${WORKDIR}/runs"
-mkdir -p "${RUNS_DIR}/slurm"
-RUN_NAME_HASH=$( (echo -n "${CONFIG_FILE} "; printf "%s " "${CONFIG_OVERRIDES[@]}") | md5sum | awk '{print $1}' | cut -c1-8)
-RUN_NAME="run_${RUN_NAME_HASH}"
-EXPERIMENT_DIR="${RUNS_DIR}/${EXPERIMENT_NAME}/${RUN_NAME}"
-
-mkdir -p "${EXPERIMENT_DIR}/checkpoints"
-
-# Ensure a stable W&B run ID so autoresume works across jobs
-if [ -f "${EXPERIMENT_DIR}/run.id" ]; then
-    RUN_ID=$(<"${EXPERIMENT_DIR}/run.id")
-    echo "Resuming with existing W&B run ID: ${RUN_ID}"
-else
-    # Simple random ID generation
-    RUN_ID=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 8 | head -n 1)
-    echo "${RUN_ID}" > "${EXPERIMENT_DIR}/run.id"
-    echo "Generated new W&B run ID: ${RUN_ID}"
-fi
-
-# Start fresh so optimizer/scheduler changes always take effect cleanly
-AUTORESUME_OVERRIDES=("autoresume.enabled=False" "experiment_dir=${EXPERIMENT_DIR}")
-echo "Starting fresh with autoresume disabled for this launch"
-
-export PYTHONPATH="${WORKDIR}:${PYTHONPATH:-}"
-
-# Redirect W&B cache to node-local scratch to avoid filling home quota
-export WANDB_CACHE_DIR="${LOCAL_SCRATCH}/wandb_cache"
-export WANDB_DATA_DIR="${LOCAL_SCRATCH}/wandb_data"
-mkdir -p "${WANDB_CACHE_DIR}" "${WANDB_DATA_DIR}"
-
-echo "Running on node: $SLURM_NODELIST"
-
-# NOTE: run.py does NOT support --experiment_dir flag in argparse.
-# We must pass it as a config override: experiment_dir=...
-PYTHON_CMD=(python experiments/run.py --config "${CONFIG_FILE}")
-PYTHON_CMD+=("experiment_dir=${EXPERIMENT_DIR}")
-PYTHON_CMD+=("${CONFIG_OVERRIDES[@]}")
-PYTHON_CMD+=("${AUTORESUME_OVERRIDES[@]}")
-# Pass walltime info
-PYTHON_CMD+=("train.run_start_time=${JOB_START_TIMESTAMP}" "train.run_time_limit_hours=${TIME_LIMIT_HOURS}")
-
-printf "Command: %q " "${PYTHON_CMD[@]}"; echo
-"${PYTHON_CMD[@]}"
diff --git a/scripts/slurm/diffusion/jit_diff_baseline_snellius.sh b/scripts/slurm/diffusion/jit_diff_baseline_snellius.sh
deleted file mode 100644
index e1e11e72..00000000
--- a/scripts/slurm/diffusion/jit_diff_baseline_snellius.sh
+++ /dev/null
@@ -1,96 +0,0 @@
-#!/bin/bash
-#SBATCH -t 48:00:00
-#SBATCH --gres=gpu:4
-#SBATCH --partition=gpu_h100
-#SBATCH --chdir="/gpfs/home1/dknigge/nvsq_update_branch"
-#SBATCH --output=./runs/slurm/in1k_jit_baseline_%j.out
-#SBATCH --error=./runs/slurm/in1k_jit_baseline_%j.err
-#SBATCH --job-name='in1k-jit_baseline'
-#SBATCH --dependency=singleton
-
-source /home/dknigge/.bashrc
-source /gpfs/home1/dknigge/nvsq_update_branch/.venv/bin/activate
-
-# HF dataset env vars (mirrors extract_imagenet_to_tar.py defaults)
-HF_DATASET="${IMAGENET_HF_DATASET:-imagenet-1k}"
-HF_CACHE="${IMAGENET_PATH:-$HOME/project_dir/huggingface/imagenet}"
-
-LOCAL_SCRATCH="/scratch-local/${USER}"
-LOCAL_HF_CACHE="${LOCAL_SCRATCH}/hf_cache"
-SENTINEL="${LOCAL_HF_CACHE}/.hf_cache_complete"
-mkdir -p "${LOCAL_HF_CACHE}"
-
-# Capture the start time before staging so walltime budget is accurate
-JOB_START_TIMESTAMP=$(date +%s)
-
-echo "Staging HF cache to node-local storage: ${LOCAL_HF_CACHE}"
-if [ -f "${SENTINEL}" ]; then
-  echo "Sentinel found — local HF cache is complete; skipping staging."
-else
-  if [ ! -d "${HF_CACHE}" ]; then
-    echo "ERROR: HF cache not found at ${HF_CACHE}."
-    exit 1
-  fi
-  echo "Copying HF cache to local NVMe..."
-  rsync -a --info=progress2 "${HF_CACHE}/" "${LOCAL_HF_CACHE}/"
-  touch "${SENTINEL}"
-  echo "Staging complete."
-fi
-
-export IMAGENET_HF_DATASET="${HF_DATASET}"
-export IMAGENET_PATH="${LOCAL_HF_CACHE}"
-echo "Start time captured: ${JOB_START_TIMESTAMP}"
-WORKDIR="/gpfs/home1/dknigge/nvsq_update_branch"
-TIME_LIMIT_HOURS=48
-CONFIG_FILE="examples/imagenet_diffusion/jit_baseline.py"
-CONFIG_OVERRIDES=(
-  "wandb.job_group=in1k-jit_baseline"
-  "wandb.entity=implicit-long-convs"
-  # Note: Batch size is handled in config now (256).
-  # If we need to override batch size per GPU here we can, but let's stick to config file default.
-)
-EXPERIMENT_NAME="imagenet_diffusion_jit_baseline"
-
-RUNS_DIR="${WORKDIR}/runs"
-mkdir -p "${RUNS_DIR}/slurm"
-RUN_NAME_HASH=$( (echo -n "${CONFIG_FILE} "; printf "%s " "${CONFIG_OVERRIDES[@]}") | md5sum | awk '{print $1}' | cut -c1-8)
-RUN_NAME="run_${RUN_NAME_HASH}"
-EXPERIMENT_DIR="${RUNS_DIR}/${EXPERIMENT_NAME}/${RUN_NAME}"
-
-mkdir -p "${EXPERIMENT_DIR}/checkpoints"
-
-# Ensure a stable W&B run ID so autoresume works across jobs
-if [ -f "${EXPERIMENT_DIR}/run.id" ]; then
-    RUN_ID=$(<"${EXPERIMENT_DIR}/run.id")
-    echo "Resuming with existing W&B run ID: ${RUN_ID}"
-else
-    # Simple random ID generation
-    RUN_ID=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 8 | head -n 1)
-    echo "${RUN_ID}" > "${EXPERIMENT_DIR}/run.id"
-    echo "Generated new W&B run ID: ${RUN_ID}"
-fi
-
-# Start fresh so optimizer/scheduler changes always take effect cleanly
-AUTORESUME_OVERRIDES=("autoresume.enabled=False" "experiment_dir=${EXPERIMENT_DIR}")
-echo "Starting fresh with autoresume disabled for this launch"
-
-export PYTHONPATH="${WORKDIR}:${PYTHONPATH:-}"
-
-# Redirect W&B cache to node-local scratch to avoid filling home quota
-export WANDB_CACHE_DIR="${LOCAL_SCRATCH}/wandb_cache"
-export WANDB_DATA_DIR="${LOCAL_SCRATCH}/wandb_data"
-mkdir -p "${WANDB_CACHE_DIR}" "${WANDB_DATA_DIR}"
-
-echo "Running on node: $SLURM_NODELIST"
-
-# NOTE: run.py does NOT support --experiment_dir flag in argparse.
-# We must pass it as a config override: experiment_dir=...
-PYTHON_CMD=(python experiments/run.py --config "${CONFIG_FILE}")
-PYTHON_CMD+=("experiment_dir=${EXPERIMENT_DIR}")
-PYTHON_CMD+=("${CONFIG_OVERRIDES[@]}")
-PYTHON_CMD+=("${AUTORESUME_OVERRIDES[@]}")
-# Pass walltime info
-PYTHON_CMD+=("train.run_start_time=${JOB_START_TIMESTAMP}" "train.run_time_limit_hours=${TIME_LIMIT_HOURS}")
-
-printf "Command: %q " "${PYTHON_CMD[@]}"; echo
-"${PYTHON_CMD[@]}"
diff --git a/tests/README.md b/tests/README.md
index e145d263..96e5e313 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -23,10 +23,6 @@ tests/
 │   ├── test_torch_compile.py            # torch.compile compatibility
 │   ├── test_vit5_components.py          # ViT-5 components (RMSNorm, LayerScale, etc.)
 │   └── torchrun_sequence_mixer_cp_test.py  # Multi-GPU sequence mixer test
-├── networks/                            # Tests for nvsubquadratic/networks/
-│   ├── test_diffusion_wrapper.py        # DiffusionWrapper tests
-│   ├── test_diffusion_fid.py            # Diffusion FID evaluation
-│   └── test_hf_diffusers_wrapper.py     # HuggingFace diffusers wrapper
 ├── parallel/                            # Tests for nvsubquadratic/parallel/
 │   └── test_a2a_comms.py                # Zigzag-splitting helpers for AllToAllSingle
 ├── test_basic.py                        # Basic import/sanity checks
diff --git a/tests/modules/test_ckconv_nd_mixed_bc.py b/tests/modules/test_ckconv_nd_mixed_bc.py
index 0d8aa52d..1c5bdce4 100644
--- a/tests/modules/test_ckconv_nd_mixed_bc.py
+++ b/tests/modules/test_ckconv_nd_mixed_bc.py
@@ -282,18 +282,6 @@ def test_subq_ops_with_per_axis_padding_raises(self):
                 fft_backend="subq_ops",
             )
 
-    def test_fp16_with_per_axis_padding_raises(self):
-        with pytest.raises(NotImplementedError, match=r"use_fp16_fft is not supported"):
-            CKConvND(
-                data_dim=2,
-                hidden_dim=HIDDEN_DIM,
-                kernel_cfg=_make_kernel_cfg(data_dim=2),
-                mask_cfg=LazyConfig(torch.nn.Identity)(),
-                grid_type=None,
-                fft_padding=["circular", "zero"],
-                use_fp16_fft=True,
-            )
-
     def test_existing_circular_with_double_grid_still_rejected(self):
         # Legacy single-mode constraint untouched.
         with pytest.raises(AssertionError, match=r"requires grid_type='single'"):
diff --git a/tests/modules/test_ckconv_nd_subq.py b/tests/modules/test_ckconv_nd_subq.py
index 87ab8def..aaca6a0a 100644
--- a/tests/modules/test_ckconv_nd_subq.py
+++ b/tests/modules/test_ckconv_nd_subq.py
@@ -268,20 +268,6 @@ def test_rejects_causal_with_2d(self):
                 fft_backend="subq_ops",
             )
 
-    def test_rejects_fp16_fft(self):
-        """use_fp16_fft=True is rejected."""
-        with pytest.raises(AssertionError, match="does not support fp16 FFT"):
-            CKConvND(
-                data_dim=2,
-                hidden_dim=32,
-                kernel_cfg=LazyConfig(torch.nn.Identity)(),
-                mask_cfg=LazyConfig(torch.nn.Identity)(),
-                grid_type="double",
-                fft_padding="zero",
-                use_fp16_fft=True,
-                fft_backend="subq_ops",
-            )
-
     def test_rejects_invalid_backend(self):
         """Invalid fft_backend string is rejected."""
         with pytest.raises(AssertionError, match="Invalid fft_backend"):
diff --git a/tests/modules/test_torch_compile.py b/tests/modules/test_torch_compile.py
index 35d4c5f4..26fb0202 100644
--- a/tests/modules/test_torch_compile.py
+++ b/tests/modules/test_torch_compile.py
@@ -23,9 +23,21 @@
 import torch.nn as nn
 
 import nvsubquadratic.ops.fftconv as _fftconv
-from examples.mnist_classification.ccnn_4_160_hyena_rope_qknorm import get_config
+from experiments.datamodules.mnist import MNISTDataModule
+from experiments.default_cfg import ExperimentConfig, SchedulerConfig, TrainConfig, WandbConfig
+from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper
 from experiments.utils.cli import apply_config_overrides
-from nvsubquadratic.lazy_config import instantiate
+from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig, instantiate
+from nvsubquadratic.modules.ckconv_nd import CKConvND
+from nvsubquadratic.modules.hyena_nd import Hyena
+from nvsubquadratic.modules.kernels_nd import SIRENKernelND
+from nvsubquadratic.modules.masks_nd import GaussianModulationND
+from nvsubquadratic.modules.mlp import MLP
+from nvsubquadratic.modules.residual_block import ResidualBlock
+from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer
+from nvsubquadratic.networks.classification_resnet import ClassificationResNet
+from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init
+from nvsubquadratic.utils.qk_norm import L2Norm
 
 
 _ALLCLOSE_RTOL = 1e-4
@@ -34,6 +46,110 @@
 _GRAD_ATOL = 5e-4
 
 
+def _build_mnist_hyena_config() -> ExperimentConfig:
+    """Build a small MNIST-shaped Hyena classification config for the compile tests.
+
+    Relocated verbatim from the former ``examples/mnist_classification`` recipe so
+    this test is self-contained and does not depend on an example config. The
+    model (4 blocks, hidden dim 160, 2D circular Hyena) is intentionally tiny —
+    it is a ``torch.compile`` compatibility vehicle, not a trained model.
+    """
+    config = ExperimentConfig()
+
+    config.dataset = LazyConfig(MNISTDataModule)(
+        data_dir=".data/mnist",
+        data_type="image",
+        batch_size=128,
+        num_workers=0,
+        pin_memory=False,
+        use_deterministic_worker_init=True,
+        seed=config.seed,
+        task="classification",
+    )
+
+    config.net = LazyConfig(ClassificationResNet)(
+        in_channels=1,
+        out_channels=10,
+        num_blocks=4,
+        hidden_dim=160,
+        data_dim=2,
+        in_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.in_channels}", out_features="${net.hidden_dim}"),
+        out_proj_cfg=LazyConfig(torch.nn.Linear)(in_features="${net.hidden_dim}", out_features="${net.out_channels}"),
+        norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"),
+        block_cfg=LazyConfig(ResidualBlock)(
+            sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)(
+                hidden_dim="${net.hidden_dim}",
+                mixer_cfg=LazyConfig(Hyena)(
+                    global_conv_cfg=LazyConfig(CKConvND)(
+                        data_dim="${net.data_dim}",
+                        hidden_dim="${net.hidden_dim}",
+                        kernel_cfg=LazyConfig(SIRENKernelND)(
+                            data_dim="${net.data_dim}",
+                            out_dim="${net.hidden_dim}",
+                            mlp_hidden_dim=32,
+                            num_layers=3,
+                            embedding_dim=32,
+                            omega_0=100.0,
+                            L_cache=32,
+                            use_bias=True,
+                            hidden_omega_0=1.0,
+                        ),
+                        mask_cfg=LazyConfig(GaussianModulationND)(
+                            data_dim="${net.data_dim}",
+                            num_channels="${net.hidden_dim}",
+                            min_attenuation_at_step=0.1,
+                            max_attenuation_at_limit=0.95,
+                            init_extent=1.0,
+                            parametrization="direct",
+                        ),
+                        grid_type="single",
+                        fft_padding="circular",
+                    ),
+                    short_conv_cfg=LazyConfig(torch.nn.Conv2d)(
+                        in_channels="3 * ${net.hidden_dim}",
+                        out_channels="3 * ${net.hidden_dim}",
+                        kernel_size=3,
+                        groups="3 * ${net.hidden_dim}",
+                        padding=1,
+                        bias=False,
+                    ),
+                    gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(),
+                    pixelhyena_norm_cfg=LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels="${net.hidden_dim}"),
+                    qk_norm_cfg=LazyConfig(L2Norm)(),
+                ),
+                init_method_in=small_init,
+                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
+            ),
+            sequence_mixer_norm_cfg="${net.norm_cfg}",
+            condition_mixer_cfg=LazyConfig(torch.nn.Identity)(),
+            condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(),
+            mlp_cfg=LazyConfig(MLP)(
+                dim="${net.hidden_dim}",
+                activation="glu",
+                expansion_factor=1.0,
+                dropout_cfg=LazyConfig(torch.nn.Dropout)(p="${net.block_cfg.dropout_cfg.p}"),
+                init_method_in=small_init,
+                init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers="${net.num_blocks}"),
+            ),
+            mlp_norm_cfg="${net.norm_cfg}",
+            dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.1),
+        ),
+        dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=0.0),
+    )
+
+    config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)()
+    config.optimizer = LazyConfig(torch.optim.AdamW)(params=PLACEHOLDER, lr=0.001, weight_decay=0.01)
+    config.train = TrainConfig(batch_size="${dataset.batch_size}", iterations=100_000, grad_clip=10.0)
+    config.scheduler = SchedulerConfig(
+        name="cosine",
+        warmup_iterations_percentage=0.05,
+        total_iterations="${train.iterations}",
+    )
+    config.wandb = WandbConfig(job_group="mnist_classification_compile_test", project="nvsubquadratic")
+
+    return config
+
+
 @pytest.fixture(autouse=True)
 def _enable_compile_compatible_fft():
     """Enable real-arithmetic complex multiply so torch.compile/Inductor works.
@@ -49,8 +165,8 @@ def _enable_compile_compatible_fft():
 
 @pytest.fixture
 def mnist_hyena_config():
-    """Load and resolve MNIST Hyena configuration."""
-    config = get_config()
+    """Build and resolve the self-contained compile-test Hyena configuration."""
+    config = _build_mnist_hyena_config()
     return apply_config_overrides(config, [])
 
 
diff --git a/tests/networks/test_diffusion_fid.py b/tests/networks/test_diffusion_fid.py
deleted file mode 100644
index 6c2d6279..00000000
--- a/tests/networks/test_diffusion_fid.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from experiments.default_cfg import DiffusionConfig, DiffusionExperimentConfig
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-from nvsubquadratic.metrics.cleanfid import compute_folder_fid
-
-
-class _IdentityBackbone(torch.nn.Module):
-    hidden_dim = 4
-
-    def forward(self, input_and_condition):
-        return {"logits": input_and_condition["input"]}
-
-
-def _make_cfg() -> DiffusionExperimentConfig:
-    cfg = DiffusionExperimentConfig()
-    cfg.diffusion = DiffusionConfig()
-    cfg.diffusion.num_train_timesteps = 4
-    cfg.diffusion.num_inference_steps = 2
-    cfg.diffusion.log_samples = False
-    cfg.diffusion.num_samples = 2
-    cfg.diffusion.use_classifier_free_guidance = False
-    cfg.diffusion.num_classes = None
-    cfg.diffusion.fid_num_inference_steps = 1
-    cfg.optimizer = SimpleNamespace(weight_decay=0.0)
-    return cfg
-
-
-def test_diffusion_wrapper_validation_step():
-    cfg = _make_cfg()
-    module = DiffusionWrapper(network=_IdentityBackbone(), cfg=cfg)
-    module.log = lambda *args, **kwargs: None
-    module.trainer = SimpleNamespace(sanity_checking=False)
-
-    batch = {"input": torch.zeros(2, 4, 4, 1)}
-    loss = module.validation_step(batch, batch_idx=0)
-
-    assert isinstance(loss, torch.Tensor)
-
-
-def test_compute_folder_fid_missing_directory(tmp_path):
-    with pytest.raises(FileNotFoundError):
-        compute_folder_fid(tmp_path / "missing", dataset_name="imagenet2012", dataset_resolution=64)
diff --git a/tests/networks/test_diffusion_wrapper.py b/tests/networks/test_diffusion_wrapper.py
deleted file mode 100644
index 59d821f3..00000000
--- a/tests/networks/test_diffusion_wrapper.py
+++ /dev/null
@@ -1,123 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Unit tests for the DiffusionWrapper with and without classifier-free guidance."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import pytest
-import torch
-
-from experiments.default_cfg import DiffusionConfig, DiffusionExperimentConfig
-from experiments.lightning_wrappers.diffusion_wrapper import DiffusionWrapper
-
-
-class RecordingDenoiser(torch.nn.Module):
-    """Minimal denoiser that records every conditioning vector it receives."""
-
-    def __init__(self, hidden_dim: int, channels: int) -> None:
-        super().__init__()
-        self.hidden_dim = hidden_dim
-        self.channels = channels
-        self.calls: list[torch.Tensor] = []
-        # Register a parameter so Lightning can infer the device even though the model is almost empty.
-        self.offset = torch.nn.Parameter(torch.zeros(1))
-
-    def forward(self, data: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        condition = data["condition"]
-        # Store a detached copy so tests can inspect the unconditional vs. conditional branches.
-        self.calls.append(condition.detach().cpu())
-        logits = torch.tanh(data["input"] + self.offset)
-        return {"logits": logits}
-
-
-def _make_wrapper(
-    *,
-    num_classes: int | None,
-    use_cfg: bool,
-    guidance_scale: float = 1.0,
-    condition_dropout_prob: float = 0.0,
-    num_inference_steps: int = 2,
-    diffusion_overrides: dict[str, Any] | None = None,
-) -> DiffusionWrapper:
-    """Helper to instantiate a DiffusionWrapper with a lightweight recording network."""
-
-    diff_cfg = DiffusionConfig()
-    diff_cfg.num_train_timesteps = 8
-    diff_cfg.num_inference_steps = num_inference_steps
-    diff_cfg.num_samples = 2
-    diff_cfg.log_samples = False
-    diff_cfg.ema_enabled = False
-    diff_cfg.use_classifier_free_guidance = use_cfg
-    diff_cfg.guidance_scale = guidance_scale
-    diff_cfg.condition_dropout_prob = condition_dropout_prob
-    diff_cfg.num_classes = num_classes
-    if diffusion_overrides:
-        for key, value in diffusion_overrides.items():
-            setattr(diff_cfg, key, value)
-
-    exp_cfg = DiffusionExperimentConfig(diffusion=diff_cfg)
-    exp_cfg.train.iterations = 10
-    exp_cfg.train.batch_size = 1
-    exp_cfg.train.grad_clip = 0.0
-
-    network = RecordingDenoiser(hidden_dim=8, channels=3)
-    wrapper = DiffusionWrapper(network=network, cfg=exp_cfg)
-    wrapper.to(torch.device("cpu"))
-    return wrapper
-
-
-def test_class_conditioning_requires_labels() -> None:
-    """Ensure the wrapper refuses to run class-conditioned diffusion without labels."""
-    wrapper = _make_wrapper(num_classes=10, use_cfg=True, guidance_scale=3.0)
-    images = torch.randn(2, 4, 4, 3)
-    batch = {"input": images}
-    with pytest.raises(RuntimeError, match="Class conditioning requires datamodule"):
-        wrapper._shared_step(batch)  # type: ignore[arg-type]
-
-
-def test_cfg_sampling_invokes_both_branches() -> None:
-    """Classifier-free guidance should evaluate unconditional and conditional paths per step."""
-    torch.manual_seed(0)
-    wrapper = _make_wrapper(num_classes=4, use_cfg=True, guidance_scale=2.5, num_inference_steps=2)
-    wrapper.example_input_shape = torch.Size((4, 4, 3))
-
-    samples = wrapper.sample(num_samples=1, labels=torch.tensor([1]))
-    assert samples.shape == (1, 4, 4, 3)
-    # Heun step (i=0): forward(t=0.0) outside cfg_interval → 1 call,
-    #   forward(t=0.5) inside cfg_interval → 2 calls (uncond + cond).
-    # Final Euler step: forward(t=0.5) inside cfg_interval → 2 calls.
-    # Total: 1 + 2 + 2 = 5.
-    assert len(wrapper.network.calls) == 5  # type: ignore[attr-defined]
-    # First CFG pair starts at index 1 (index 0 is the non-CFG forward at t=0.0).
-    unconditional_condition = wrapper.network.calls[1]  # type: ignore[attr-defined]
-    conditional_condition = wrapper.network.calls[2]  # type: ignore[attr-defined]
-    assert not torch.allclose(unconditional_condition, conditional_condition)
-
-
-def test_sampling_without_guidance_uses_single_branch() -> None:
-    """When guidance is disabled the denoiser should run exactly once per timestep."""
-    wrapper = _make_wrapper(num_classes=4, use_cfg=False, num_inference_steps=3)
-    wrapper.example_input_shape = torch.Size((4, 4, 3))
-    wrapper.network.calls.clear()  # type: ignore[attr-defined]
-
-    samples = wrapper.sample(num_samples=1, labels=torch.tensor([2]))
-    assert samples.shape == (1, 4, 4, 3)
-    # Heun steps (i=0, i=1): 2 forward calls each = 4.
-    # Final Euler step: 1 forward call.
-    # Total: 4 + 1 = 5.
-    assert len(wrapper.network.calls) == 5  # type: ignore[attr-defined]
diff --git a/tests/networks/test_hf_diffusers_wrapper.py b/tests/networks/test_hf_diffusers_wrapper.py
deleted file mode 100644
index 5ca650ee..00000000
--- a/tests/networks/test_hf_diffusers_wrapper.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import sys
-import types
-from pathlib import Path
-
-import torch
-
-from nvsubquadratic.networks.huggingface_diffusers import (
-    DiffusersDiTWrapper,
-    DiffusersUVitWrapper,
-    HuggingFaceDiTConfig,
-    HuggingFaceUVitConfig,
-)
-
-
-PROJECT_ROOT = Path(__file__).resolve().parents[1]
-if str(PROJECT_ROOT) not in sys.path:
-    sys.path.insert(0, str(PROJECT_ROOT))
-
-
-class _DummyWrapper:
-    """Minimal stub mimicking the Lightning diffusion wrapper interface."""
-
-    def __init__(self):
-        def _condition(_, timesteps: torch.LongTensor) -> torch.Tensor:
-            # Mirror the Lightning behaviour by returning an embedding tensor; we only care about
-            # forwarding the timesteps to registered callbacks.
-            return torch.zeros((timesteps.shape[0], 4), dtype=torch.float32, device=timesteps.device)
-
-        self._condition_from_timesteps = types.MethodType(_condition, self)
-
-
-def _tiny_config() -> HuggingFaceDiTConfig:
-    return HuggingFaceDiTConfig(
-        sample_size=8,
-        patch_size=2,
-        in_channels=1,
-        out_channels=1,
-        num_layers=1,
-        num_attention_heads=2,
-        attention_head_dim=8,
-        num_embeds_ada_norm=32,
-        activation_fn="gelu-approximate",
-        norm_type="ada_norm_zero",
-    )
-
-
-def test_hf_wrapper_tracks_timesteps():
-    model = DiffusersDiTWrapper(_tiny_config(), in_channels=1, out_channels=1)
-    dummy_wrapper = _DummyWrapper()
-
-    model.hf_register_diffusion_wrapper(dummy_wrapper)
-
-    timesteps = torch.tensor([3, 7], dtype=torch.long)
-    _ = dummy_wrapper._condition_from_timesteps(timesteps)
-
-    assert model._latest_timesteps is not None
-    assert torch.equal(model._latest_timesteps, timesteps)
-
-
-def test_hf_wrapper_forward_matches_input_shape():
-    model = DiffusersDiTWrapper(_tiny_config(), in_channels=1, out_channels=1)
-    dummy_wrapper = _DummyWrapper()
-    model.hf_register_diffusion_wrapper(dummy_wrapper)
-
-    batch_size = 2
-    inputs = torch.randn(
-        batch_size, model.hf_config.sample_size, model.hf_config.sample_size, model.hf_config.in_channels
-    )
-    dummy_condition = torch.zeros(batch_size, model.hidden_dim)
-
-    timesteps = torch.tensor([5, 11], dtype=torch.long)
-    model._latest_timesteps = timesteps
-
-    outputs = model({"input": inputs, "condition": dummy_condition})
-
-    assert "logits" in outputs
-    assert outputs["logits"].shape == inputs.shape
-
-
-def test_hf_wrapper_accepts_channel_overrides():
-    cfg = _tiny_config()
-    cfg.in_channels = 2
-    model = DiffusersDiTWrapper(cfg, in_channels=3, out_channels=5)
-    assert model.hf_config.in_channels == 3
-    assert model.hf_config.out_channels == 5
-
-
-def test_extra_repr_handles_dict_config():
-    cfg = _tiny_config()
-    repr_str = repr(DiffusersDiTWrapper(cfg, in_channels=1, out_channels=1))
-    assert "DiffusersDiTWrapper" in repr_str
-
-
-def test_specialised_config_variants():
-    dit_cfg = HuggingFaceDiTConfig()
-    uvit_cfg = HuggingFaceUVitConfig()
-    assert dit_cfg.num_layers > 0
-    assert uvit_cfg.hidden_size > 0
-
-
-def test_uvit_wrapper_builds_with_defaults():
-    cfg = HuggingFaceUVitConfig(
-        sample_size=4,
-        in_channels=1,
-        out_channels=1,
-        hidden_size=64,
-        cond_embed_dim=16,
-        encoder_hidden_size=16,
-        block_out_channels=64,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        intermediate_size=128,
-        layer_norm_eps=1e-5,
-        micro_cond_encode_dim=8,
-        micro_cond_embed_dim=16,
-        codebook_size=32,
-        vocab_size=33,
-    )
-
-    model = DiffusersUVitWrapper(cfg)
-    dummy_wrapper = _DummyWrapper()
-    model.hf_register_diffusion_wrapper(dummy_wrapper)
-
-    timesteps = torch.tensor([1, 2], dtype=torch.long)
-    dummy_wrapper._condition_from_timesteps(timesteps)
-
-    inputs = torch.randn(2, cfg.sample_size, cfg.sample_size, cfg.in_channels)
-    cond = torch.zeros(2, cfg.cond_embed_dim)
-    outputs = model({"input": inputs, "condition": cond})
-
-    assert "logits" in outputs
-    assert outputs["logits"].shape[0] == inputs.shape[0]
diff --git a/tests/ops/test_circular_fftconv_fp16.py b/tests/ops/test_circular_fftconv_fp16.py
deleted file mode 100644
index 65790ae9..00000000
--- a/tests/ops/test_circular_fftconv_fp16.py
+++ /dev/null
@@ -1,458 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for FP16 circular FFT convolution operators (1D, 2D, 3D).
-
-These tests verify that fp16 circular FFT convolutions:
-1. Produce outputs matching the fp32 circular reference within expected tolerance
-2. Preserve correct output shapes for all dimensions
-3. Handle both BHL and BLH (w_reshape) layouts
-4. Work with and without the per-channel shortcut
-5. Reject mismatched shortcut dtypes
-6. Reject non-power-of-2 spatial dimensions
-7. Reject unsupported kernel sizes (1D: only K=L or K=L-1)
-8. Preserve caller dtype
-9. Produce correct gradients matching the fp32 backward pass
-
-Kernel sizes are restricted to K=L (full-size) and K=L-1 (one shorter)
-per spatial axis, matching what the continuous kernel networks produce
-in production.  See ``nvsubquadratic/ops/FP16_FFTCONV_DERIVATION.md``
-for the mathematical background.
-
-All tests require CUDA (cuFFT only supports fp16 FFT).
-
-See tests/README.md for test suites, markers, and SLURM usage.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from nvsubquadratic.ops.circular_fftconv import (
-    circular_fftconv1d_fp32_bhl as circular_fftconv1d_f32,
-)
-from nvsubquadratic.ops.circular_fftconv import (
-    circular_fftconv2d_fp32_bhl as circular_fftconv2d_f32,
-)
-from nvsubquadratic.ops.circular_fftconv import (
-    circular_fftconv3d_fp32_bhl as circular_fftconv3d_f32,
-)
-from nvsubquadratic.ops.circular_fftconv_fp16 import (
-    circular_fftconv1d_fp16_bhl,
-    circular_fftconv1d_fp16_bhl_w_reshape,
-    circular_fftconv2d_fp16_bhl,
-    circular_fftconv2d_fp16_bhl_w_reshape,
-    circular_fftconv3d_fp16_bhl,
-    circular_fftconv3d_fp16_bhl_w_reshape,
-)
-
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fp16 FFT (cuFFT)")
-
-# fp16 has ~0.1% relative error vs f32 due to reduced precision and ortho normalization.
-# K==L (full-size kernel) cases in higher dimensions accumulate more error.
-RTOL_FP16 = 0.15
-ATOL_FP16 = 0.4
-# Backward accumulates more error through the chain rule, especially for large
-# 3D K=L kernels (16^3 spatial elements).  Observed max atol ~0.31 for 3D K=L=16.
-RTOL_FP16_BWD = 0.20
-ATOL_FP16_BWD = 0.40
-
-
-@pytest.fixture
-def device() -> str:
-    return "cuda"
-
-
-###############################################################################
-# 1D
-###############################################################################
-
-
-class TestCircularFP16Conv1D:
-    """Tests for 1D fp16 circular FFT convolution.
-
-    The dual-centering implementation only supports K=L (full-size kernel)
-    and K=L-1 (one element shorter), which are the only cases produced by
-    the continuous kernel networks in production.
-    """
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [
-            (2, 32, 64, 64),  # K=L
-            (2, 32, 64, 63),  # K=L-1
-            (2, 16, 128, 128),  # K=L
-            (2, 16, 128, 127),  # K=L-1
-            (1, 64, 256, 256),  # K=L
-            (4, 16, 128, 127),  # K=L-1
-        ],
-    )
-    def test_fp16_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """FP16 output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, L, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, K, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv1d_f32(x, kernel, shortcut)
-        y_fp16 = circular_fftconv1d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 1D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 32, 64, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 32, 63, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv1d_f32(x, kernel, None)
-        y_fp16 = circular_fftconv1d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly and matches direct BHL call."""
-        torch.manual_seed(42)
-        x_blh = torch.randn(2, 64, 32, device=device, dtype=torch.float32)  # [B, L, H]
-        k_blh = torch.randn(1, 63, 32, device=device, dtype=torch.float32)  # [1, K, H], K=L-1
-        shortcut = torch.randn(32, device=device, dtype=torch.float32)
-
-        y = circular_fftconv1d_fp16_bhl_w_reshape(x_blh, k_blh, shortcut)
-
-        assert y.shape == (2, 64, 32)  # [B, L, H]
-
-    def test_returns_in_caller_dtype(self, device: str) -> None:
-        """Output dtype matches x's dtype."""
-        torch.manual_seed(42)
-
-        x_f32 = torch.randn(2, 32, 64, device=device, dtype=torch.float32)
-        kernel_f32 = torch.randn(1, 32, 64, device=device, dtype=torch.float32)
-        y1 = circular_fftconv1d_fp16_bhl(x_f32, kernel_f32, None)
-        assert y1.dtype == torch.float32
-
-        x_bf16 = x_f32.to(torch.bfloat16)
-        kernel_bf16 = kernel_f32.to(torch.bfloat16)
-        y2 = circular_fftconv1d_fp16_bhl(x_bf16, kernel_bf16, None)
-        assert y2.dtype == torch.bfloat16
-
-    def test_rejects_mismatched_shortcut_dtype(self, device: str) -> None:
-        """Mismatched shortcut dtype raises AssertionError."""
-        x = torch.randn(2, 32, 64, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 32, 64, device=device, dtype=torch.float16)
-        shortcut_bad = torch.randn(32, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match=r"shortcut\.dtype"):
-            circular_fftconv1d_fp16_bhl(x, kernel, shortcut_bad)
-
-    def test_rejects_non_power_of_2(self, device: str) -> None:
-        """Non-power-of-2 L raises AssertionError."""
-        x = torch.randn(2, 16, 100, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 16, 100, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="power of 2"):
-            circular_fftconv1d_fp16_bhl(x, kernel, None)
-
-    def test_rejects_arbitrary_kernel_size(self, device: str) -> None:
-        """K not in {L, L-1} raises AssertionError."""
-        x = torch.randn(2, 16, 128, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 16, 7, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="FP16 centering requires"):
-            circular_fftconv1d_fp16_bhl(x, kernel, None)
-
-    def test_does_not_mutate_inputs(self, device: str) -> None:
-        """FP16 1D must not modify the caller's x or kernel tensors."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 16, 64, device=device, dtype=torch.float16)
-        k = torch.randn(1, 16, 64, device=device, dtype=torch.float16)
-        x_orig = x.clone()
-        k_orig = k.clone()
-
-        circular_fftconv1d_fp16_bhl(x, k, None)
-
-        torch.testing.assert_close(x, x_orig)
-        torch.testing.assert_close(k, k_orig)
-
-    def test_does_not_mutate_inputs_fp32(self, device: str) -> None:
-        """FP16 1D must not modify the caller's fp32 x or kernel tensors."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 16, 64, device=device, dtype=torch.float32)
-        k = torch.randn(1, 16, 64, device=device, dtype=torch.float32)
-        x_orig = x.clone()
-        k_orig = k.clone()
-
-        circular_fftconv1d_fp16_bhl(x, k, None)
-
-        torch.testing.assert_close(x, x_orig)
-        torch.testing.assert_close(k, k_orig)
-
-    def test_batched_kernel(self, device: str) -> None:
-        """FP16 1D supports batched kernels [B, H, K]."""
-        torch.manual_seed(42)
-        B, H, L, K = 2, 32, 64, 64
-        x = torch.randn(B, H, L, device=device, dtype=torch.float32)
-        kernel = torch.randn(B, H, K, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv1d_f32(x, kernel, None)
-        y_fp16 = circular_fftconv1d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [(2, 16, 64, 63), (1, 32, 128, 128)],
-    )
-    def test_backward_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """FP16 circular 1D gradients match fp32 circular reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, L, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, K, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = circular_fftconv1d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = circular_fftconv1d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# 2D
-###############################################################################
-
-
-class TestCircularFP16Conv2D:
-    """Tests for 2D fp16 circular FFT convolution.
-
-    Uses K=L (full-size) and K=L-1 (one shorter) per spatial axis,
-    matching the kernel sizes produced by continuous kernel networks.
-    """
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Kx,Ky",
-        [
-            (2, 16, 32, 32, 31, 31),  # K=L-1 both axes
-            (2, 32, 64, 64, 64, 64),  # K=L both axes
-            (1, 8, 16, 16, 16, 16),  # K=L both axes
-            (4, 8, 32, 32, 32, 32),  # K=L both axes
-        ],
-    )
-    def test_fp16_vs_f32(self, device: str, B: int, H: int, X: int, Y: int, Kx: int, Ky: int) -> None:
-        """FP16 2D output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, X, Y, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, Kx, Ky, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv2d_f32(x, kernel, shortcut)
-        y_fp16 = circular_fftconv2d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 2D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 16, 32, 32, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 16, 31, 31, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv2d_f32(x, kernel, None)
-        y_fp16 = circular_fftconv2d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly."""
-        torch.manual_seed(42)
-        x_blh = torch.randn(2, 32, 32, 16, device=device, dtype=torch.float32)  # [B, X, Y, H]
-        k_blh = torch.randn(1, 31, 31, 16, device=device, dtype=torch.float32)  # K=L-1
-        shortcut = torch.randn(16, device=device, dtype=torch.float32)
-
-        y = circular_fftconv2d_fp16_bhl_w_reshape(x_blh, k_blh, shortcut)
-        assert y.shape == (2, 32, 32, 16)
-
-    def test_does_not_mutate_inputs(self, device: str) -> None:
-        """FP16 2D must not modify the caller's x or kernel tensors."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 8, 16, 16, device=device, dtype=torch.float16)
-        k = torch.randn(1, 8, 16, 16, device=device, dtype=torch.float16)
-        x_orig = x.clone()
-        k_orig = k.clone()
-
-        circular_fftconv2d_fp16_bhl(x, k, None)
-
-        torch.testing.assert_close(x, x_orig)
-        torch.testing.assert_close(k, k_orig)
-
-    def test_rejects_arbitrary_kernel_size(self, device: str) -> None:
-        """K_d not in {N_d, N_d-1} raises AssertionError."""
-        x = torch.randn(2, 8, 32, 32, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 8, 7, 7, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="FP16 centering requires"):
-            circular_fftconv2d_fp16_bhl(x, kernel, None)
-
-    def test_rejects_non_power_of_2(self, device: str) -> None:
-        """Non-power-of-2 spatial dims raise AssertionError."""
-        x = torch.randn(2, 8, 14, 14, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 8, 14, 14, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="power"):
-            circular_fftconv2d_fp16_bhl(x, kernel, None)
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Kx,Ky",
-        [(2, 8, 16, 16, 15, 15), (1, 16, 32, 32, 32, 32)],
-    )
-    def test_backward_vs_f32(self, device: str, B: int, H: int, X: int, Y: int, Kx: int, Ky: int) -> None:
-        """FP16 circular 2D gradients match fp32 circular reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, X, Y, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, Kx, Ky, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = circular_fftconv2d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = circular_fftconv2d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# 3D
-###############################################################################
-
-
-class TestCircularFP16Conv3D:
-    """Tests for 3D fp16 circular FFT convolution.
-
-    Uses K=L (full-size) and K=L-1 (one shorter) per spatial axis,
-    matching the kernel sizes produced by continuous kernel networks.
-    """
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Z,Kx,Ky,Kz",
-        [
-            (2, 4, 16, 16, 16, 15, 15, 15),  # K=L-1 all axes
-            (1, 8, 8, 8, 8, 8, 8, 8),  # K=L all axes
-            (2, 4, 16, 16, 16, 16, 16, 16),  # K=L all axes
-        ],
-    )
-    def test_fp16_vs_f32(self, device: str, B: int, H: int, X: int, Y: int, Z: int, Kx: int, Ky: int, Kz: int) -> None:
-        """FP16 3D output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, X, Y, Z, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, Kx, Ky, Kz, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv3d_f32(x, kernel, shortcut)
-        y_fp16 = circular_fftconv3d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 3D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 4, 16, 16, 16, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 4, 15, 15, 15, device=device, dtype=torch.float32)
-
-        y_f32 = circular_fftconv3d_f32(x, kernel, None)
-        y_fp16 = circular_fftconv3d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly."""
-        torch.manual_seed(42)
-        x_blh = torch.randn(2, 16, 16, 16, 4, device=device, dtype=torch.float32)
-        k_blh = torch.randn(1, 15, 15, 15, 4, device=device, dtype=torch.float32)  # K=L-1
-        shortcut = torch.randn(4, device=device, dtype=torch.float32)
-
-        y = circular_fftconv3d_fp16_bhl_w_reshape(x_blh, k_blh, shortcut)
-        assert y.shape == (2, 16, 16, 16, 4)
-
-    def test_does_not_mutate_inputs(self, device: str) -> None:
-        """FP16 3D must not modify the caller's x or kernel tensors."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 4, 8, 8, 8, device=device, dtype=torch.float16)
-        k = torch.randn(1, 4, 8, 8, 8, device=device, dtype=torch.float16)
-        x_orig = x.clone()
-        k_orig = k.clone()
-
-        circular_fftconv3d_fp16_bhl(x, k, None)
-
-        torch.testing.assert_close(x, x_orig)
-        torch.testing.assert_close(k, k_orig)
-
-    def test_rejects_arbitrary_kernel_size(self, device: str) -> None:
-        """K_d not in {N_d, N_d-1} raises AssertionError."""
-        x = torch.randn(2, 4, 16, 16, 16, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 4, 5, 5, 5, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="FP16 centering requires"):
-            circular_fftconv3d_fp16_bhl(x, kernel, None)
-
-    def test_rejects_non_power_of_2(self, device: str) -> None:
-        """Non-power-of-2 spatial dims raise AssertionError."""
-        x = torch.randn(2, 4, 12, 12, 12, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 4, 12, 12, 12, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match="power"):
-            circular_fftconv3d_fp16_bhl(x, kernel, None)
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Z,Kx,Ky,Kz",
-        [(1, 4, 8, 8, 8, 7, 7, 7), (2, 4, 16, 16, 16, 16, 16, 16)],
-    )
-    def test_backward_vs_f32(
-        self, device: str, B: int, H: int, X: int, Y: int, Z: int, Kx: int, Ky: int, Kz: int
-    ) -> None:
-        """FP16 circular 3D gradients match fp32 circular reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, X, Y, Z, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, Kx, Ky, Kz, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = circular_fftconv3d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = circular_fftconv3d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
diff --git a/tests/ops/test_fftconv_fp16.py b/tests/ops/test_fftconv_fp16.py
deleted file mode 100644
index 20dd40bd..00000000
--- a/tests/ops/test_fftconv_fp16.py
+++ /dev/null
@@ -1,739 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for FP16 FFT convolution operators (standard and chunked).
-
-These tests verify that fp16 FFT convolutions:
-1. Produce outputs matching f32 reference within expected numerical tolerance
-2. Preserve correct output shapes for all dimensions (1D, 2D, 3D)
-3. Handle both BHL and BLH (w_reshape) layouts
-4. Work with and without the per-channel shortcut
-5. Support channel chunking with fp16 precision
-6. Produce consistent results across chunked and non-chunked fp16 paths
-7. Support causal 1D convolutions in fp16
-
-All tests require CUDA (cuFFT only supports fp16 FFT).
-
-See tests/README.md for test suites, markers, and SLURM usage.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from nvsubquadratic.ops.fftconv import (
-    causal_fftconv1d_fp32_bhl as causal_fftconv1d_f32,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv1d_fp32_bhl as fftconv1d_f32,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv1d_fp32_bhl_w_reshape as fftconv1d_f32_w_reshape,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv2d_fp32_bhl as fftconv2d_f32,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv2d_fp32_bhl_w_reshape as fftconv2d_f32_w_reshape,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv3d_fp32_bhl as fftconv3d_f32,
-)
-from nvsubquadratic.ops.fftconv import (
-    fftconv3d_fp32_bhl_w_reshape as fftconv3d_f32_w_reshape,
-)
-from nvsubquadratic.ops.fftconv_fp16 import (
-    causal_fftconv1d_fp16_bhl,
-    causal_fftconv1d_fp16_bhl_chunked,
-    fftconv1d_fp16_bhl,
-    fftconv1d_fp16_bhl_chunked,
-    fftconv1d_fp16_bhl_w_reshape,
-    fftconv2d_fp16_bhl,
-    fftconv2d_fp16_bhl_chunked,
-    fftconv2d_fp16_bhl_w_reshape,
-    fftconv3d_fp16_bhl,
-    fftconv3d_fp16_bhl_chunked,
-    fftconv3d_fp16_bhl_w_reshape,
-)
-
-
-# cuFFT fp16 is CUDA-only
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fp16 FFT (cuFFT)")
-
-# Tolerances: fp16 has ~0.1% relative error vs f32 due to reduced precision
-# and ortho normalization. These are empirically validated thresholds.
-RTOL_FP16 = 0.1
-ATOL_FP16 = 0.1
-# Backward tolerances — fp16 backward is noisier than forward.
-RTOL_FP16_BWD = 0.15
-ATOL_FP16_BWD = 0.15
-# Chunked fp16 vs non-chunked fp16 should be exact (same computation)
-ATOL_CHUNKED = 1e-5
-
-
-@pytest.fixture
-def device() -> str:
-    """Return CUDA device."""
-    return "cuda"
-
-
-###############################################################################
-# 1D non-causal
-###############################################################################
-
-
-class TestFP16FFTConv1D:
-    """Tests for 1D fp16 FFT convolution."""
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [
-            (2, 32, 64, 7),
-            (2, 64, 256, 32),
-            (1, 128, 512, 64),
-            (4, 16, 128, 15),
-        ],
-    )
-    def test_fp16_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """FP16 output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, L, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, K, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv1d_f32(x, kernel, shortcut)
-        y_fp16 = fftconv1d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 1D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 32, 64, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 32, 7, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv1d_f32(x, kernel, None)
-        y_fp16 = fftconv1d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly and matches f32 w_reshape."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 64, 32, device=device, dtype=torch.float32)  # [B, L, H]
-        kernel = torch.randn(1, 7, 32, device=device, dtype=torch.float32)  # [1, K, H]
-        shortcut = torch.randn(32, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv1d_f32_w_reshape(x, kernel, shortcut)
-        y_fp16 = fftconv1d_fp16_bhl_w_reshape(x, kernel, shortcut)
-
-        assert y_fp16.shape == (2, 64, 32)  # [B, L, H]
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_accepts_any_input_dtype(self, device: str) -> None:
-        """FP16 functions accept any input dtype and return in the caller's dtype."""
-        x = torch.randn(2, 32, 64, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 32, 7, device=device, dtype=torch.float32)
-
-        y = fftconv1d_fp16_bhl(x, kernel, None)
-        assert y.dtype == torch.float32
-        assert y.shape == (2, 32, 64)
-
-    def test_returns_in_caller_dtype(self, device: str) -> None:
-        """Output dtype matches x's dtype when all inputs share the same dtype."""
-        torch.manual_seed(42)
-
-        # fp16 inputs → fp16 output
-        x_fp16 = torch.randn(2, 32, 64, device=device, dtype=torch.float16)
-        kernel_fp16 = torch.randn(1, 32, 7, device=device, dtype=torch.float16)
-        shortcut_fp16 = torch.randn(32, device=device, dtype=torch.float16)
-
-        y1 = fftconv1d_fp16_bhl(x_fp16, kernel_fp16, shortcut_fp16)
-        assert y1.dtype == torch.float16
-
-        # f32 inputs → f32 output
-        x_f32 = torch.randn(2, 32, 64, device=device, dtype=torch.float32)
-        kernel_f32 = torch.randn(1, 32, 7, device=device, dtype=torch.float32)
-        shortcut_f32 = torch.randn(32, device=device, dtype=torch.float32)
-        y2 = fftconv1d_fp16_bhl(x_f32, kernel_f32, shortcut_f32)
-        assert y2.dtype == torch.float32
-
-    def test_rejects_mismatched_shortcut_dtype(self, device: str) -> None:
-        """Mismatched shortcut dtype raises AssertionError."""
-        x = torch.randn(2, 32, 64, device=device, dtype=torch.float16)
-        kernel_fp16 = torch.randn(1, 32, 7, device=device, dtype=torch.float16)
-        shortcut_f32 = torch.randn(32, device=device, dtype=torch.float32)
-
-        with pytest.raises(AssertionError, match=r"shortcut\.dtype"):
-            fftconv1d_fp16_bhl(x, kernel_fp16, shortcut_f32)
-
-    @pytest.mark.parametrize("chunk_size", [16, 32, 64])
-    def test_chunked_matches_standard(self, device: str, chunk_size: int) -> None:
-        """Chunked fp16 produces same result as non-chunked fp16."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 128, 256, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 128, 32, device=device, dtype=torch.float16)
-        shortcut = torch.randn(128, device=device, dtype=torch.float16)
-
-        y_std = fftconv1d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = fftconv1d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=chunk_size)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    def test_chunked_h_not_divisible(self, device: str) -> None:
-        """Chunked fp16 handles H not divisible by chunk_size."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 100, 128, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 100, 16, device=device, dtype=torch.float16)
-        shortcut = torch.randn(100, device=device, dtype=torch.float16)
-
-        y_std = fftconv1d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = fftconv1d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=64)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    def test_batched_kernel(self, device: str) -> None:
-        """FP16 1D supports batched kernels [B, H, K]."""
-        torch.manual_seed(42)
-        B, H, L, K = 2, 32, 64, 7
-        x = torch.randn(B, H, L, device=device, dtype=torch.float32)
-        kernel = torch.randn(B, H, K, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv1d_f32(x, kernel, None)
-        y_fp16 = fftconv1d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [(2, 16, 64, 15), (1, 32, 128, 7)],
-    )
-    def test_backward_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """FP16 gradients w.r.t. x and kernel match fp32 reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, L, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, K, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = fftconv1d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = fftconv1d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# 1D causal
-###############################################################################
-
-
-class TestFP16CausalFFTConv1D:
-    """Tests for causal 1D fp16 FFT convolution."""
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [
-            (2, 32, 64, 7),
-            (2, 64, 256, 32),
-            (1, 128, 512, 64),
-        ],
-    )
-    def test_causal_fp16_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """Causal fp16 matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, L, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, K, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = causal_fftconv1d_f32(x, kernel, shortcut)
-        y_fp16 = causal_fftconv1d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_causal_is_causal(self, device: str) -> None:
-        """Causal fp16 output at position i depends only on input[0..i].
-
-        Perturb input at position L//2 and verify only outputs at >= L//2 change.
-        """
-        torch.manual_seed(42)
-        B, H, L, K = 1, 16, 64, 15
-        x = torch.randn(B, H, L, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, H, K, device=device, dtype=torch.float16)
-
-        y_orig = causal_fftconv1d_fp16_bhl(x, kernel, None)
-
-        x_perturbed = x.clone()
-        x_perturbed[:, :, L // 2] += 1.0
-        y_pert = causal_fftconv1d_fp16_bhl(x_perturbed, kernel, None)
-
-        # Positions before the perturbation should be unchanged
-        diff_before = (y_orig[..., : L // 2] - y_pert[..., : L // 2]).abs().max()
-        diff_after = (y_orig[..., L // 2 :] - y_pert[..., L // 2 :]).abs().max()
-
-        assert diff_before < 0.05, f"Causal violation: diff before perturbation = {diff_before}"
-        assert diff_after > 0.01, "Perturbation had no effect after position L//2"
-
-    @pytest.mark.parametrize("chunk_size", [16, 32])
-    def test_causal_chunked(self, device: str, chunk_size: int) -> None:
-        """Chunked causal fp16 matches non-chunked causal fp16."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 64, 128, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 64, 16, device=device, dtype=torch.float16)
-        shortcut = torch.randn(64, device=device, dtype=torch.float16)
-
-        y_std = causal_fftconv1d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = causal_fftconv1d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=chunk_size)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    @pytest.mark.parametrize(
-        "B,H,L,K",
-        [(2, 16, 64, 15), (1, 32, 128, 7)],
-    )
-    def test_backward_vs_f32(self, device: str, B: int, H: int, L: int, K: int) -> None:
-        """Causal FP16 gradients match fp32 reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, L, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, K, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = causal_fftconv1d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = causal_fftconv1d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# 2D
-###############################################################################
-
-
-class TestFP16FFTConv2D:
-    """Tests for 2D fp16 FFT convolution."""
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Kx,Ky",
-        [
-            (2, 32, 14, 14, 7, 7),
-            (2, 64, 14, 14, 27, 27),
-            (1, 128, 28, 28, 13, 13),
-            (4, 16, 32, 32, 5, 5),
-        ],
-    )
-    def test_fp16_vs_f32(self, device: str, B: int, H: int, X: int, Y: int, Kx: int, Ky: int) -> None:
-        """FP16 2D output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, X, Y, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, Kx, Ky, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv2d_f32(x, kernel, shortcut)
-        y_fp16 = fftconv2d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 2D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 32, 14, 14, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 32, 7, 7, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv2d_f32(x, kernel, None)
-        y_fp16 = fftconv2d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly and matches f32 w_reshape."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 14, 14, 32, device=device, dtype=torch.float32)  # [B, X, Y, H]
-        kernel = torch.randn(1, 7, 7, 32, device=device, dtype=torch.float32)
-        shortcut = torch.randn(32, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv2d_f32_w_reshape(x, kernel, shortcut)
-        y_fp16 = fftconv2d_fp16_bhl_w_reshape(x, kernel, shortcut)
-
-        assert y_fp16.shape == (2, 14, 14, 32)  # [B, X, Y, H]
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    @pytest.mark.parametrize("chunk_size", [16, 32, 64])
-    def test_chunked_matches_standard(self, device: str, chunk_size: int) -> None:
-        """Chunked fp16 2D produces same result as non-chunked fp16."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 128, 14, 14, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 128, 7, 7, device=device, dtype=torch.float16)
-        shortcut = torch.randn(128, device=device, dtype=torch.float16)
-
-        y_std = fftconv2d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = fftconv2d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=chunk_size)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    def test_chunked_h_not_divisible(self, device: str) -> None:
-        """Chunked fp16 2D handles H not divisible by chunk_size."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 100, 14, 14, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 100, 7, 7, device=device, dtype=torch.float16)
-        shortcut = torch.randn(100, device=device, dtype=torch.float16)
-
-        y_std = fftconv2d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = fftconv2d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=64)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Kx,Ky",
-        [(2, 8, 16, 16, 5, 5), (1, 16, 32, 32, 7, 7)],
-    )
-    def test_backward_vs_f32(self, device: str, B: int, H: int, X: int, Y: int, Kx: int, Ky: int) -> None:
-        """FP16 2D gradients match fp32 reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, X, Y, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, Kx, Ky, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = fftconv2d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = fftconv2d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# 3D
-###############################################################################
-
-
-class TestFP16FFTConv3D:
-    """Tests for 3D fp16 FFT convolution."""
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Z,Kx,Ky,Kz",
-        [
-            (2, 16, 8, 8, 8, 5, 5, 5),
-            (1, 32, 8, 16, 16, 3, 7, 7),
-            (2, 8, 16, 16, 16, 3, 3, 3),
-        ],
-    )
-    def test_fp16_vs_f32(
-        self,
-        device: str,
-        B: int,
-        H: int,
-        X: int,
-        Y: int,
-        Z: int,
-        Kx: int,
-        Ky: int,
-        Kz: int,
-    ) -> None:
-        """FP16 3D output matches f32 reference within tolerance."""
-        torch.manual_seed(42)
-        x = torch.randn(B, H, X, Y, Z, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, H, Kx, Ky, Kz, device=device, dtype=torch.float32)
-        shortcut = torch.randn(H, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv3d_f32(x, kernel, shortcut)
-        y_fp16 = fftconv3d_fp16_bhl(x, kernel, shortcut)
-
-        assert y_fp16.shape == y_f32.shape
-        assert y_fp16.dtype == x.dtype
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_no_shortcut(self, device: str) -> None:
-        """FP16 3D works without shortcut."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 16, 8, 8, 8, device=device, dtype=torch.float32)
-        kernel = torch.randn(1, 16, 3, 3, 3, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv3d_f32(x, kernel, None)
-        y_fp16 = fftconv3d_fp16_bhl(x, kernel, None)
-
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    def test_w_reshape_layout(self, device: str) -> None:
-        """BLH wrapper reshapes correctly and matches f32 w_reshape."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 8, 8, 8, 16, device=device, dtype=torch.float32)  # [B, X, Y, Z, H]
-        kernel = torch.randn(1, 3, 3, 3, 16, device=device, dtype=torch.float32)
-        shortcut = torch.randn(16, device=device, dtype=torch.float32)
-
-        y_f32 = fftconv3d_f32_w_reshape(x, kernel, shortcut)
-        y_fp16 = fftconv3d_fp16_bhl_w_reshape(x, kernel, shortcut)
-
-        assert y_fp16.shape == (2, 8, 8, 8, 16)  # [B, X, Y, Z, H]
-        torch.testing.assert_close(y_fp16, y_f32, rtol=RTOL_FP16, atol=ATOL_FP16)
-
-    @pytest.mark.parametrize("chunk_size", [8, 16])
-    def test_chunked_matches_standard(self, device: str, chunk_size: int) -> None:
-        """Chunked fp16 3D produces same result as non-chunked fp16."""
-        torch.manual_seed(42)
-        x = torch.randn(2, 32, 8, 8, 8, device=device, dtype=torch.float16)
-        kernel = torch.randn(1, 32, 3, 3, 3, device=device, dtype=torch.float16)
-        shortcut = torch.randn(32, device=device, dtype=torch.float16)
-
-        y_std = fftconv3d_fp16_bhl(x, kernel, shortcut)
-        y_chunk = fftconv3d_fp16_bhl_chunked(x, kernel, shortcut, chunk_size=chunk_size)
-
-        torch.testing.assert_close(y_chunk, y_std, atol=ATOL_CHUNKED, rtol=0)
-
-    @pytest.mark.parametrize(
-        "B,H,X,Y,Z,Kx,Ky,Kz",
-        [(1, 4, 8, 8, 8, 3, 3, 3), (2, 8, 16, 16, 16, 5, 5, 5)],
-    )
-    def test_backward_vs_f32(
-        self, device: str, B: int, H: int, X: int, Y: int, Z: int, Kx: int, Ky: int, Kz: int
-    ) -> None:
-        """FP16 3D gradients match fp32 reference gradients."""
-        torch.manual_seed(42)
-        x1 = torch.randn(B, H, X, Y, Z, device=device, dtype=torch.float32, requires_grad=True)
-        k1 = torch.randn(1, H, Kx, Ky, Kz, device=device, dtype=torch.float32, requires_grad=True)
-
-        y1 = fftconv3d_fp16_bhl(x1, k1)
-        grad_output = torch.randn_like(y1)
-        y1.backward(grad_output)
-
-        x2 = x1.detach().clone().requires_grad_(True)
-        k2 = k1.detach().clone().requires_grad_(True)
-        y2 = fftconv3d_f32(x2, k2)
-        y2.backward(grad_output)
-
-        torch.testing.assert_close(x1.grad, x2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-        torch.testing.assert_close(k1.grad, k2.grad, rtol=RTOL_FP16_BWD, atol=ATOL_FP16_BWD)
-
-
-###############################################################################
-# CKConvND integration
-###############################################################################
-
-
-class TestCKConvNDFP16Integration:
-    """Test that CKConvND correctly selects fp16 FFT functions."""
-
-    def test_fp16_flag_selects_fp16_table(self) -> None:
-        """CKConvND with use_fp16_fft=True uses the fp16 function table."""
-        from nvsubquadratic.modules.ckconv_nd import (
-            FFT_FUNCTIONS_FP16,
-            FFT_FUNCTIONS_FP16_CHUNKED,
-        )
-
-        for padding_key in ("zero", "causal"):
-            for dim in FFT_FUNCTIONS_FP16.get(padding_key, {}):
-                w_reshape_fn, bhl_fn = FFT_FUNCTIONS_FP16[padding_key][dim]
-                assert "fp16" in w_reshape_fn.__name__
-                assert "fp16" in bhl_fn.__name__
-                assert "chunked" not in w_reshape_fn.__name__
-                assert "chunked" not in bhl_fn.__name__
-
-        for padding_key in ("zero", "causal"):
-            for dim in FFT_FUNCTIONS_FP16_CHUNKED.get(padding_key, {}):
-                w_reshape_fn, bhl_fn = FFT_FUNCTIONS_FP16_CHUNKED[padding_key][dim]
-                assert "fp16" in w_reshape_fn.__name__
-                assert "fp16" in bhl_fn.__name__
-                assert "chunked" in w_reshape_fn.__name__
-                assert "chunked" in bhl_fn.__name__
-
-    def test_circular_padding_warns(self) -> None:
-        """CKConvND with use_fp16_fft=True and circular padding emits a warning."""
-        import warnings
-
-        from nvsubquadratic.lazy_config import LazyConfig
-        from nvsubquadratic.modules.ckconv_nd import CKConvND
-
-        with warnings.catch_warnings(record=True) as w:
-            warnings.simplefilter("always")
-            CKConvND(
-                data_dim=2,
-                hidden_dim=32,
-                kernel_cfg=LazyConfig(torch.nn.Identity)(),
-                mask_cfg=LazyConfig(torch.nn.Identity)(),
-                grid_type="single",
-                fft_padding="circular",
-                use_fp16_fft=True,
-            )
-            assert any("power-of-2" in str(warning.message) for warning in w)
-
-
-###############################################################################
-# Nightly: full-model fp16 FFT validation
-###############################################################################
-
-
-@pytest.mark.nightly
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.skipif(
-    "WANDB_API_KEY" not in __import__("os").environ,
-    reason="WANDB_API_KEY not set (run `source .env`)",
-)
-@pytest.mark.skipif(
-    not __import__("os").path.isdir("/shared/data/image_datasets/imagenet"),
-    reason="ImageNet not found",
-)
-class TestNightlyFP16Validation:
-    """Nightly test: validate GAP model with fp16 FFT on ImageNet.
-
-    Downloads the best GAP checkpoint from W&B and runs validation twice
-    (f32 FFT and fp16 FFT) to confirm accuracy parity.
-    """
-
-    def test_fp16_fft_accuracy_parity(self) -> None:
-        """FP16 FFT produces same validation accuracy as f32 on GAP Hyena model.
-
-        W&B run: tcji9tfx (v2 Hyena-GAP, ~81.5% top-1)
-        """
-        import re
-        import sys
-        from pathlib import Path
-
-        sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-
-        import pytorch_lightning as pl
-
-        from experiments.utils.checkpointing import (
-            StripCompiledPrefix,
-            download_checkpoint,
-            load_checkpoint_state_dict,
-        )
-        from nvsubquadratic.lazy_config import instantiate
-        from nvsubquadratic.modules.ckconv_nd import FFT_FUNCTIONS_FP16, CKConvND
-
-        WANDB_ENTITY = "implicit-long-convs"
-        WANDB_PROJECT = "nvsubquadratic"
-
-        _SIREN_RE = re.compile(r"(\.kernel\.kernel_network)\.(\d+)\.(weight|bias)$")
-
-        def _remap_siren(sd: dict) -> dict:
-            out = {}
-            for k, v in sd.items():
-                m = _SIREN_RE.search(k)
-                if m:
-                    idx = int(m.group(2)) // 2
-                    out[k[: m.start()] + f".kernel.hidden_linears.{idx}.{m.group(3)}"] = v
-                else:
-                    out[k] = v
-            return out
-
-        from examples.vit5_imagenet.v2.vit5_small_pretrain_hyena_gap_apex_gated_ema import (
-            get_config,
-        )
-
-        config = get_config()
-        config.train.do = False
-        config.debug = True
-        config.compile = False
-        # Disable local staging so the test works without NVMe space
-        if hasattr(config.dataset, "local_staging_dir"):
-            config.dataset.local_staging_dir = None
-
-        import nvsubquadratic.ops.fftconv as _fftconv
-
-        _fftconv.COMPILE_COMPATIBLE = True
-
-        pl.seed_everything(config.seed, workers=True)
-        torch.set_float32_matmul_precision("high")
-
-        datamodule = instantiate(config.dataset)
-        datamodule.prepare_data()
-        datamodule.setup()
-
-        network = instantiate(config.net)
-        model = instantiate(config.lightning_wrapper_class, network=network, cfg=config)
-
-        # The checkpoint (tcji9tfx) was trained before bias was removed from
-        # patch_embed and out_proj (weight-decay hygiene in #74).  Re-add the
-        # bias parameters so the checkpoint loads exactly and accuracy is valid.
-        net = model.network
-        pe = net.patch_embed
-        net.patch_embed = torch.nn.Conv2d(
-            pe.in_channels,
-            pe.out_channels,
-            kernel_size=pe.kernel_size,
-            stride=pe.stride,
-            padding=pe.padding,
-            bias=True,
-        )
-        op = net.out_proj
-        net.out_proj = torch.nn.Linear(op.in_features, op.out_features, bias=True)
-
-        run_path = f"{WANDB_ENTITY}/{WANDB_PROJECT}/tcji9tfx"
-        ckpt_path = download_checkpoint(run_path=run_path, alias="best")
-        state_dict = load_checkpoint_state_dict(ckpt_path)
-        state_dict = _remap_siren(state_dict)
-        strip = StripCompiledPrefix()
-        state_dict = strip(state_dict=state_dict, model=model)
-        model.load_state_dict(state_dict, strict=True)
-
-        trainer = pl.Trainer(
-            accelerator="gpu",
-            devices=1,
-            precision="bf16-mixed",
-            logger=False,
-            enable_checkpointing=False,
-            enable_progress_bar=True,
-        )
-
-        # Run validation with f32 FFT (baseline)
-        results_f32 = trainer.test(model, datamodule=datamodule)
-        acc_f32 = results_f32[0]["test/acc"]
-
-        # Switch all CKConvND modules to fp16 FFT
-        count = 0
-        for module in model.modules():
-            if isinstance(module, CKConvND):
-                module.use_fp16_fft = True
-                eff_pad = "causal" if module.is_causal else module.fft_padding
-                module.fftconv_fn, module.fftconv_fn_bhl_input = FFT_FUNCTIONS_FP16[eff_pad][module.data_dim]
-                count += 1
-
-        assert count > 0, "No CKConvND modules found — model structure may have changed"
-
-        # Run validation with fp16 FFT
-        results_fp16 = trainer.test(model, datamodule=datamodule)
-        acc_fp16 = results_fp16[0]["test/acc"]
-
-        # Accuracy should be within 0.5% — the original validation showed identical results
-        assert abs(acc_f32 - acc_fp16) < 0.005, (
-            f"FP16 FFT accuracy regression: f32={acc_f32:.4f}, fp16={acc_fp16:.4f}, diff={abs(acc_f32 - acc_fp16):.4f}"
-        )
-
-        # Both should exceed the known baseline (~81.5% top-1)
-        assert acc_f32 >= 0.81, f"f32 baseline regression: {acc_f32:.4f} < 0.81"
-        assert acc_fp16 >= 0.81, f"fp16 regression: {acc_fp16:.4f} < 0.81"
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
diff --git a/tests/test_circular_fftconv_fp16.py b/tests/test_circular_fftconv_fp16.py
deleted file mode 100644
index eabbb9e4..00000000
--- a/tests/test_circular_fftconv_fp16.py
+++ /dev/null
@@ -1,306 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for the dual-centered FP16 circular FFT convolution (1D, 2D, 3D).
-
-Validates the implementations in ``circular_fftconv_fp16.py`` against the
-FP32 reference implementations from ``circular_fftconv.py``.  Each test
-verifies:
-
-1. **No NaNs** — the dual-centering technique prevents overflow.
-2. **Low relative error** — mean relative error vs. FP32 stays below 5%
-   across a range of input means/stds (including pathological cases like
-   ``mean=50, std=5`` that overflow in naive FP16).
-3. **API contract** — output dtype/shape match input, shortcut is applied,
-   ``_w_reshape`` wrappers give equivalent results.
-
-Requires a CUDA GPU.  Run via::
-
-    conda run -n nv-subq pytest tests/test_circular_fftconv_fp16.py -v
-
-Or inside the container for GPU tests::
-
-    srun --gres=gpu:1 ... pytest tests/test_circular_fftconv_fp16.py -v
-"""
-
-import pytest
-import torch
-
-from nvsubquadratic.ops.circular_fftconv import (
-    circular_fftconv1d_fp32_bhl,
-    circular_fftconv2d_fp32_bhl,
-    circular_fftconv3d_fp32_bhl,
-)
-from nvsubquadratic.ops.circular_fftconv_fp16 import (
-    circular_fftconv1d_fp16_bhl,
-    circular_fftconv1d_fp16_bhl_w_reshape,
-    circular_fftconv2d_fp16_bhl,
-    circular_fftconv2d_fp16_bhl_w_reshape,
-    circular_fftconv3d_fp16_bhl,
-    circular_fftconv3d_fp16_bhl_w_reshape,
-)
-
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-DEVICE = "cuda"
-
-# (mean, std) pairs chosen to exercise the overflow paths that the dual-
-# centering technique was designed to fix.  (0,1) is the benign baseline;
-# (50,5) causes DC-bin overflow in naive FP16 even with ortho normalization.
-MEAN_STD_COMBOS = [(0.0, 1.0), (1.0, 1.0), (10.0, 1.0), (50.0, 5.0)]
-
-# 5% relative error tolerance — generous; typical is <0.1% for most shapes.
-REL_ERR_TOL = 0.05
-
-
-def _rel_err(y_fp16: torch.Tensor, y_ref: torch.Tensor) -> float:
-    """Mean relative error, clamped to avoid div-by-zero."""
-    return (y_fp16 - y_ref).abs().div(y_ref.abs().clamp(min=1e-6)).mean().item()
-
-
-# ─── 1D ──────────────────────────────────────────────────────────────────────
-
-
-class TestCircularFftconv1dFp16:
-    """Tests for 1D fp16 circular FFT convolution."""
-
-    @pytest.mark.parametrize("mean,std", MEAN_STD_COMBOS)
-    @pytest.mark.parametrize("L,K", [(256, 255), (256, 256), (1024, 1023), (4096, 4095)])
-    def test_no_nan_and_close_to_fp32(self, mean, std, L, K):
-        B, H = 2, 32
-        torch.manual_seed(42)
-        x = (torch.randn(B, H, L, device=DEVICE) * std + mean).float()
-        k = (torch.randn(1, H, K, device=DEVICE) * std + mean).float()
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_ref = circular_fftconv1d_fp32_bhl(x, k, sc)
-        y_fp16 = circular_fftconv1d_fp16_bhl(x, k, sc)
-
-        assert not y_fp16.isnan().any(), f"NaN detected (mean={mean}, std={std}, L={L}, K={K})"
-        assert _rel_err(y_fp16, y_ref) < REL_ERR_TOL, (
-            f"Relative error too large (mean={mean}, std={std}, L={L}, K={K})"
-        )
-
-    def test_output_dtype_matches_input(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 16, 256, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 16, 255, device=DEVICE, dtype=torch.float32)
-        y = circular_fftconv1d_fp16_bhl(x, k)
-        assert y.dtype == x.dtype
-
-    def test_output_shape(self):
-        torch.manual_seed(0)
-        x = torch.randn(4, 32, 512, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 32, 511, device=DEVICE, dtype=torch.float32)
-        y = circular_fftconv1d_fp16_bhl(x, k)
-        assert y.shape == x.shape
-
-    def test_w_reshape_matches_bhl(self):
-        """The _w_reshape wrapper should give the same result as manual reshape + bhl."""
-        torch.manual_seed(0)
-        B, H, L, K = 2, 16, 256, 255
-        x_bhl = torch.randn(B, H, L, device=DEVICE, dtype=torch.float32)
-        k_bhl = torch.randn(1, H, K, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_bhl = circular_fftconv1d_fp16_bhl(x_bhl, k_bhl, sc)
-
-        x_blh = x_bhl.permute(0, 2, 1).contiguous()
-        k_blh = k_bhl.permute(0, 2, 1).contiguous()
-        y_blh = circular_fftconv1d_fp16_bhl_w_reshape(x_blh, k_blh, sc)
-        y_from_blh = y_blh.permute(0, 2, 1)
-
-        torch.testing.assert_close(y_bhl, y_from_blh, atol=1e-4, rtol=1e-3)
-
-    def test_shortcut_applied(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 16, 256, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 16, 255, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(16, device=DEVICE, dtype=torch.float32)
-
-        y_no_sc = circular_fftconv1d_fp16_bhl(x, k)
-        y_sc = circular_fftconv1d_fp16_bhl(x, k, sc)
-        assert not torch.allclose(y_no_sc, y_sc), "Shortcut should change the output"
-
-
-# ─── 2D ──────────────────────────────────────────────────────────────────────
-
-
-class TestCircularFftconv2dFp16:
-    """Tests for 2D fp16 circular FFT convolution."""
-
-    @pytest.mark.parametrize("mean,std", MEAN_STD_COMBOS)
-    @pytest.mark.parametrize(
-        "X,Y,Kx,Ky",
-        [
-            (32, 32, 31, 31),
-            (32, 32, 32, 32),
-            (64, 64, 63, 63),
-            (128, 128, 127, 127),
-        ],
-    )
-    def test_no_nan_and_close_to_fp32(self, mean, std, X, Y, Kx, Ky):
-        B, H = 2, 8
-        torch.manual_seed(42)
-        x = (torch.randn(B, H, X, Y, device=DEVICE) * std + mean).float()
-        k = (torch.randn(1, H, Kx, Ky, device=DEVICE) * std + mean).float()
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_ref = circular_fftconv2d_fp32_bhl(x, k, sc)
-        y_fp16 = circular_fftconv2d_fp16_bhl(x, k, sc)
-
-        assert not y_fp16.isnan().any(), f"NaN detected (mean={mean}, X={X}, K={Kx})"
-        assert _rel_err(y_fp16, y_ref) < REL_ERR_TOL, f"Relative error too large (mean={mean}, X={X}, K={Kx})"
-
-    def test_output_dtype_and_shape(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 8, 64, 64, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 8, 63, 63, device=DEVICE, dtype=torch.float32)
-        y = circular_fftconv2d_fp16_bhl(x, k)
-        assert y.dtype == x.dtype
-        assert y.shape == x.shape
-
-    def test_w_reshape_matches_bhl(self):
-        torch.manual_seed(0)
-        B, H, X, Y, Kx, Ky = 2, 8, 32, 32, 31, 31
-        x_bhl = torch.randn(B, H, X, Y, device=DEVICE, dtype=torch.float32)
-        k_bhl = torch.randn(1, H, Kx, Ky, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_bhl = circular_fftconv2d_fp16_bhl(x_bhl, k_bhl, sc)
-
-        x_blh = x_bhl.permute(0, 2, 3, 1).contiguous()
-        k_blh = k_bhl.permute(0, 2, 3, 1).contiguous()
-        y_blh = circular_fftconv2d_fp16_bhl_w_reshape(x_blh, k_blh, sc)
-        y_from_blh = y_blh.permute(0, 3, 1, 2)
-
-        torch.testing.assert_close(y_bhl, y_from_blh, atol=1e-4, rtol=1e-3)
-
-    def test_shortcut_applied(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 8, 32, 32, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 8, 31, 31, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(8, device=DEVICE, dtype=torch.float32)
-
-        y_no_sc = circular_fftconv2d_fp16_bhl(x, k)
-        y_sc = circular_fftconv2d_fp16_bhl(x, k, sc)
-        assert not torch.allclose(y_no_sc, y_sc), "Shortcut should change the output"
-
-
-# ─── 3D ──────────────────────────────────────────────────────────────────────
-
-
-class TestCircularFftconv3dFp16:
-    """Tests for 3D fp16 circular FFT convolution."""
-
-    @pytest.mark.parametrize("mean,std", MEAN_STD_COMBOS)
-    @pytest.mark.parametrize(
-        "X,Y,Z,Kx,Ky,Kz",
-        [
-            (16, 16, 16, 15, 15, 15),
-            (16, 16, 16, 16, 16, 16),
-            (32, 32, 32, 31, 31, 31),
-        ],
-    )
-    def test_no_nan_and_close_to_fp32(self, mean, std, X, Y, Z, Kx, Ky, Kz):
-        B, H = 2, 4
-        torch.manual_seed(42)
-        x = (torch.randn(B, H, X, Y, Z, device=DEVICE) * std + mean).float()
-        k = (torch.randn(1, H, Kx, Ky, Kz, device=DEVICE) * std + mean).float()
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_ref = circular_fftconv3d_fp32_bhl(x, k, sc)
-        y_fp16 = circular_fftconv3d_fp16_bhl(x, k, sc)
-
-        assert not y_fp16.isnan().any(), f"NaN detected (mean={mean}, X={X}, K={Kx})"
-        assert _rel_err(y_fp16, y_ref) < REL_ERR_TOL, f"Relative error too large (mean={mean}, X={X}, K={Kx})"
-
-    def test_output_dtype_and_shape(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 4, 16, 16, 16, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 4, 15, 15, 15, device=DEVICE, dtype=torch.float32)
-        y = circular_fftconv3d_fp16_bhl(x, k)
-        assert y.dtype == x.dtype
-        assert y.shape == x.shape
-
-    def test_w_reshape_matches_bhl(self):
-        torch.manual_seed(0)
-        B, H = 2, 4
-        X, Y, Z, Kx, Ky, Kz = 16, 16, 16, 15, 15, 15
-        x_bhl = torch.randn(B, H, X, Y, Z, device=DEVICE, dtype=torch.float32)
-        k_bhl = torch.randn(1, H, Kx, Ky, Kz, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(H, device=DEVICE, dtype=torch.float32)
-
-        y_bhl = circular_fftconv3d_fp16_bhl(x_bhl, k_bhl, sc)
-
-        x_blh = x_bhl.permute(0, 2, 3, 4, 1).contiguous()
-        k_blh = k_bhl.permute(0, 2, 3, 4, 1).contiguous()
-        y_blh = circular_fftconv3d_fp16_bhl_w_reshape(x_blh, k_blh, sc)
-        y_from_blh = y_blh.permute(0, 4, 1, 2, 3)
-
-        torch.testing.assert_close(y_bhl, y_from_blh, atol=1e-3, rtol=1e-2)
-
-    def test_shortcut_applied(self):
-        torch.manual_seed(0)
-        x = torch.randn(2, 4, 16, 16, 16, device=DEVICE, dtype=torch.float32)
-        k = torch.randn(1, 4, 15, 15, 15, device=DEVICE, dtype=torch.float32)
-        sc = torch.randn(4, device=DEVICE, dtype=torch.float32)
-
-        y_no_sc = circular_fftconv3d_fp16_bhl(x, k)
-        y_sc = circular_fftconv3d_fp16_bhl(x, k, sc)
-        assert not torch.allclose(y_no_sc, y_sc), "Shortcut should change the output"
-
-
-# ─── Geo cache ────────────────────────────────────────────────────────────────
-
-
-class TestCenteringCorrectionCache:
-    """Tests for the geometry correction caches."""
-
-    def test_2d_cache_returns_none_when_no_correction_needed(self):
-        from nvsubquadratic.ops.circular_fftconv_fp16 import _centering_geo_cache_2d
-
-        assert _centering_geo_cache_2d.get(32, 32, 32, 32, torch.device(DEVICE)) is None
-
-    def test_3d_cache_returns_none_when_no_correction_needed(self):
-        from nvsubquadratic.ops.circular_fftconv_fp16 import _centering_geo_cache_3d
-
-        assert _centering_geo_cache_3d.get(16, 16, 16, 16, 16, 16, torch.device(DEVICE)) is None
-
-    def test_2d_cache_returns_tensor_when_correction_needed(self):
-        from nvsubquadratic.ops.circular_fftconv_fp16 import _centering_geo_cache_2d
-
-        geo = _centering_geo_cache_2d.get(31, 31, 32, 32, torch.device(DEVICE))
-        assert geo is not None
-        assert geo.shape == (32, 17)  # [X, Y//2+1]
-        assert geo.dtype == torch.complex64
-
-    def test_3d_cache_returns_tensor_when_correction_needed(self):
-        from nvsubquadratic.ops.circular_fftconv_fp16 import _centering_geo_cache_3d
-
-        geo = _centering_geo_cache_3d.get(15, 15, 15, 16, 16, 16, torch.device(DEVICE))
-        assert geo is not None
-        assert geo.shape == (16, 16, 9)  # [X, Y, Z//2+1]
-        assert geo.dtype == torch.complex64
-
-    def test_2d_cache_hit(self):
-        """Second call with same args should return the same tensor object (cache hit)."""
-        from nvsubquadratic.ops.circular_fftconv_fp16 import _centering_geo_cache_2d
-
-        dev = torch.device(DEVICE)
-        geo1 = _centering_geo_cache_2d.get(63, 63, 64, 64, dev)
-        geo2 = _centering_geo_cache_2d.get(63, 63, 64, 64, dev)
-        assert geo1 is geo2
diff --git a/tests/test_multihead_lowrank.py b/tests/test_multihead_lowrank.py
deleted file mode 100644
index d6d876a8..00000000
--- a/tests/test_multihead_lowrank.py
+++ /dev/null
@@ -1,314 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for low-rank multi-head CKConv and FFT convolution.
-
-Tests verify:
-1. Low-rank FFT conv ops produce correct shapes
-2. Low-rank FFT conv matches full-rank when rank == head_dim (exact equivalence)
-3. CKConvMultiheadND with kernel_rank produces correct shapes and gradients
-4. CKConvMultiheadND with kernel_rank + FiLM (batched kernels)
-5. CKConvMultiheadND without kernel_rank (full-rank) still works
-6. Low-rank reduces parameter count vs full-rank
-
-Run:
-    PYTHONPATH=. python -m pytest tests/test_multihead_lowrank.py -v
-"""
-
-import pytest
-import torch
-
-from nvsubquadratic.lazy_config import LazyConfig
-from nvsubquadratic.modules.ckconv_multihead_nd import CKConvMultiheadND
-from nvsubquadratic.modules.film import KernelFiLMGenerator
-from nvsubquadratic.modules.kernels_nd import SIRENKernelND
-from nvsubquadratic.ops.fftconv_multihead import (
-    fftconv2d_multihead_bhl,
-    fftconv2d_multihead_lowrank_bhl,
-    fftconv2d_multihead_lowrank_circular_bhl,
-)
-
-
-# ─── Fixtures ───────────────────────────────────────────────────────────────────
-
-
-@pytest.fixture
-def multihead_dims():
-    """Standard dimensions used across tests."""
-    return {"B": 2, "num_heads": 6, "head_dim": 64, "H": 14, "W": 14, "hidden_dim": 384}
-
-
-@pytest.fixture
-def film_cfg():
-    """FiLM generator config for conditioned kernels."""
-    return LazyConfig(KernelFiLMGenerator)(
-        cond_dim=384,
-        kernel_hidden_dim=32,
-        num_film_layers=2,
-        film_hidden_dim=64,
-    )
-
-
-def _make_ckconv(*, hidden_dim, num_heads, kernel_rank=None, film_cfg=None):
-    """Helper to build a CKConvMultiheadND with given rank and optional FiLM."""
-    head_dim = hidden_dim // num_heads
-    if kernel_rank is not None:
-        out_dim = num_heads * 2 * kernel_rank * head_dim
-    else:
-        out_dim = num_heads * head_dim * head_dim
-
-    return CKConvMultiheadND(
-        data_dim=2,
-        hidden_dim=hidden_dim,
-        num_heads=num_heads,
-        kernel_cfg=LazyConfig(SIRENKernelND)(
-            data_dim=2,
-            out_dim=out_dim,
-            mlp_hidden_dim=32,
-            num_layers=3,
-            embedding_dim=32,
-            omega_0=10.0,
-            L_cache=15,
-            use_bias=True,
-            hidden_omega_0=1.0,
-            film_cfg=film_cfg,
-        ),
-        mask_cfg=LazyConfig(torch.nn.Identity)(),
-        grid_type="double",
-        fft_padding="zero",
-        kernel_rank=kernel_rank,
-    )
-
-
-# ─── 1. Low-rank FFT conv op tests ──────────────────────────────────────────────
-
-
-class TestLowRankFFTConv:
-    def test_lowrank_output_shape(self, multihead_dims):
-        """Low-rank FFT conv produces the correct output shape."""
-        d = multihead_dims
-        rank = 8
-        x = torch.randn(d["B"], d["num_heads"], d["head_dim"], d["H"], d["W"])
-        kernel_u = torch.randn(d["num_heads"], d["head_dim"], rank, 28, 28)
-        kernel_v = torch.randn(d["num_heads"], rank, d["head_dim"], 28, 28)
-        shortcut = torch.randn(d["hidden_dim"])
-
-        out = fftconv2d_multihead_lowrank_bhl(x, kernel_u, kernel_v, shortcut)
-        assert out.shape == (d["B"], d["num_heads"], d["head_dim"], d["H"], d["W"])
-
-    def test_lowrank_circular_output_shape(self, multihead_dims):
-        """Circular low-rank FFT conv produces the correct output shape."""
-        d = multihead_dims
-        rank = 4
-        x = torch.randn(d["B"], d["num_heads"], d["head_dim"], d["H"], d["W"])
-        kernel_u = torch.randn(d["num_heads"], d["head_dim"], rank, d["H"], d["W"])
-        kernel_v = torch.randn(d["num_heads"], rank, d["head_dim"], d["H"], d["W"])
-        shortcut = torch.randn(d["hidden_dim"])
-
-        out = fftconv2d_multihead_lowrank_circular_bhl(x, kernel_u, kernel_v, shortcut)
-        assert out.shape == (d["B"], d["num_heads"], d["head_dim"], d["H"], d["W"])
-
-    def test_lowrank_matches_manual_freq_domain_computation(self):
-        """Verify low-rank FFT conv matches manual K_fft = U_fft @ V_fft computation.
-
-        In frequency domain, the low-rank conv computes:
-            z_fft[n,r,f] = sum_i V_fft[n,r,i,f] * x_fft[n,i,f]
-            y_fft[n,o,f] = sum_r U_fft[n,o,r,f] * z_fft[n,r,f]
-
-        This is equivalent to full-rank with K_fft[n,o,i,f] = sum_r U_fft[n,o,r,f] * V_fft[n,r,i,f].
-        We verify by constructing K_fft explicitly and comparing.
-        """
-        B, num_heads, head_dim, H, W = 1, 2, 8, 7, 7
-        K_x, K_y = 14, 14
-        rank = 4
-
-        x = torch.randn(B, num_heads, head_dim, H, W)
-        kernel_u = torch.randn(num_heads, head_dim, rank, K_x, K_y) * 0.01
-        kernel_v = torch.randn(num_heads, rank, head_dim, K_x, K_y) * 0.01
-        shortcut = torch.randn(num_heads * head_dim)
-
-        # Compute via low-rank function
-        out_lr = fftconv2d_multihead_lowrank_bhl(x, kernel_u, kernel_v, shortcut)
-
-        # Compute manually: K_fft = U_fft @ V_fft, then full-rank conv with K_fft
-        fft_h = min(H + (K_x + 1) // 2, 2 * H)
-        fft_w = min(W + (K_y + 1) // 2, 2 * W)
-
-        x_fft = torch.fft.rfft2(x, s=(fft_h, fft_w))
-        u_fft = torch.fft.rfft2(kernel_u, s=(fft_h, fft_w))
-        v_fft = torch.fft.rfft2(kernel_v, s=(fft_h, fft_w))
-
-        # K_fft[n,o,i,f1,f2] = sum_r U_fft[n,o,r,f1,f2] * V_fft[n,r,i,f1,f2]
-        k_fft = torch.einsum("norhw,nrihw->noihw", u_fft, v_fft)
-
-        # Full-rank conv with synthesized K_fft
-        out_fft = torch.einsum("bnihw,noihw->bnohw", x_fft, k_fft)
-        crop_h = K_x // 2
-        crop_w = K_y // 2
-        out_manual_full = torch.fft.irfft2(out_fft, s=(fft_h, fft_w))
-        out_manual = out_manual_full[..., crop_h : crop_h + H, crop_w : crop_w + W]
-        # Add shortcut
-        shortcut_reshaped = shortcut.view(1, num_heads, head_dim, 1, 1)
-        out_manual = out_manual + x * shortcut_reshaped
-
-        torch.testing.assert_close(out_lr, out_manual, atol=1e-4, rtol=1e-4)
-
-    def test_lowrank_matches_fullrank_when_rank_equals_head_dim(self):
-        """Low-rank with rank == head_dim exactly matches full-rank conv.
-
-        The two-step low-rank conv computes y = U @ (V @ x) which by
-        associativity equals y = (U @ V) @ x = K @ x. When rank == head_dim,
-        setting U = spatial delta (identity in freq domain) and V = K gives
-        K_eff = I @ K = K, so the output matches full-rank exactly.
-        """
-        B, num_heads, head_dim, H, W = 2, 2, 8, 7, 7
-        K_x, K_y = 14, 14
-        rank = head_dim  # full rank
-
-        x = torch.randn(B, num_heads, head_dim, H, W)
-        # Full-rank kernel: [num_heads, head_dim, head_dim, K_x, K_y]
-        kernel_full = torch.randn(num_heads, head_dim, head_dim, K_x, K_y) * 0.01
-        shortcut = torch.randn(num_heads * head_dim)
-
-        # U = identity at each spatial position: [num_heads, head_dim, rank, K_x, K_y]
-        # V = kernel: [num_heads, rank, head_dim, K_x, K_y]
-        # But U and V are in spatial domain and K_fft = FFT(U) @ FFT(V) ≠ FFT(U @ V)
-        # So we need to set U to a spatial delta (identity only at center, zero elsewhere)
-        # and V = kernel, so that FFT(U) = I (constant) and K_fft = I @ FFT(V) = FFT(V) = FFT(K).
-        kernel_u = torch.zeros(num_heads, head_dim, rank, K_x, K_y)
-        # Place identity at spatial center (the DC/origin point for the kernel)
-        for n in range(num_heads):
-            for d in range(head_dim):
-                kernel_u[n, d, d, 0, 0] = 1.0  # delta at origin
-        kernel_v = kernel_full.clone()  # V has shape [num_heads, head_dim_out=rank, head_dim_in, K_x, K_y]
-
-        out_full = fftconv2d_multihead_bhl(x, kernel_full, shortcut)
-        out_lr = fftconv2d_multihead_lowrank_bhl(x, kernel_u, kernel_v, shortcut)
-
-        torch.testing.assert_close(out_lr, out_full, atol=1e-4, rtol=1e-4)
-
-    def test_lowrank_gradient_flow(self, multihead_dims):
-        """Gradients flow through low-rank FFT conv."""
-        d = multihead_dims
-        rank = 8
-        x = torch.randn(d["B"], d["num_heads"], d["head_dim"], d["H"], d["W"], requires_grad=True)
-        kernel_u = torch.randn(d["num_heads"], d["head_dim"], rank, 28, 28, requires_grad=True)
-        kernel_v = torch.randn(d["num_heads"], rank, d["head_dim"], 28, 28, requires_grad=True)
-
-        out = fftconv2d_multihead_lowrank_bhl(x, kernel_u, kernel_v)
-        out.sum().backward()
-
-        assert x.grad is not None
-        assert kernel_u.grad is not None
-        assert kernel_v.grad is not None
-
-
-# ─── 2. CKConvMultiheadND module tests ──────────────────────────────────────────
-
-
-class TestCKConvMultiheadNDLowRank:
-    def test_lowrank_output_shape_blh(self, multihead_dims):
-        """Low-rank CKConvMultiheadND produces correct shape with BLH input."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8)
-        x = torch.randn(d["B"], d["H"], d["W"], d["hidden_dim"])
-        out = model(x)
-        assert out.shape == x.shape
-
-    def test_lowrank_output_shape_bhl(self, multihead_dims):
-        """Low-rank CKConvMultiheadND produces correct shape with BHL input."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8)
-        x = torch.randn(d["B"], d["hidden_dim"], d["H"], d["W"])
-        out = model(x, is_bhl_input=True)
-        assert out.shape == x.shape
-
-    def test_lowrank_gradient_flow(self, multihead_dims):
-        """All parameters receive gradients in the low-rank path."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8)
-        x = torch.randn(d["B"], d["H"], d["W"], d["hidden_dim"])
-        out = model(x)
-        out.sum().backward()
-
-        for name, p in model.named_parameters():
-            if p.requires_grad:
-                assert p.grad is not None, f"No gradient for {name}"
-
-    def test_lowrank_with_film(self, multihead_dims, film_cfg):
-        """Low-rank works with FiLM conditioning (batched kernels)."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8, film_cfg=film_cfg)
-        x = torch.randn(d["B"], d["H"], d["W"], d["hidden_dim"])
-        cond = torch.randn(d["B"], d["hidden_dim"])
-        out = model(x, conditioning=cond)
-        assert out.shape == x.shape
-
-        out.sum().backward()
-        for name, p in model.named_parameters():
-            if p.requires_grad:
-                assert p.grad is not None, f"No gradient for {name}"
-
-    def test_fullrank_still_works(self, multihead_dims, film_cfg):
-        """Full-rank path (kernel_rank=None) is not broken."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=None, film_cfg=film_cfg)
-        x = torch.randn(d["B"], d["H"], d["W"], d["hidden_dim"])
-        cond = torch.randn(d["B"], d["hidden_dim"])
-        out = model(x, conditioning=cond)
-        assert out.shape == x.shape
-
-        out.sum().backward()
-        for name, p in model.named_parameters():
-            if p.requires_grad:
-                assert p.grad is not None, f"No gradient for {name}"
-
-    def test_lowrank_reduces_params(self, multihead_dims):
-        """Low-rank model has fewer parameters than full-rank."""
-        d = multihead_dims
-        model_lr = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8)
-        model_full = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=None)
-
-        params_lr = sum(p.numel() for p in model_lr.parameters())
-        params_full = sum(p.numel() for p in model_full.parameters())
-
-        assert params_lr < params_full, (
-            f"Low-rank ({params_lr}) should have fewer params than full-rank ({params_full})"
-        )
-        # With rank=8 and head_dim=64, SIREN output is 4x smaller
-        # so the SIREN out_linear should be ~4x smaller
-        ratio = params_full / params_lr
-        assert ratio > 2.0, f"Expected significant param reduction, got only {ratio:.1f}x"
-
-    def test_extra_repr_includes_rank(self, multihead_dims):
-        """extra_repr shows kernel_rank when set."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=8)
-        assert "kernel_rank=8" in model.extra_repr()
-
-    def test_extra_repr_no_rank_when_fullrank(self, multihead_dims):
-        """extra_repr does not show kernel_rank when using full-rank."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=None)
-        assert "kernel_rank" not in model.extra_repr()
-
-    @pytest.mark.parametrize("rank", [1, 4, 8, 16])
-    def test_various_ranks(self, multihead_dims, rank):
-        """Different rank values all produce correct output."""
-        d = multihead_dims
-        model = _make_ckconv(hidden_dim=d["hidden_dim"], num_heads=d["num_heads"], kernel_rank=rank)
-        x = torch.randn(d["B"], d["H"], d["W"], d["hidden_dim"])
-        out = model(x)
-        assert out.shape == x.shape