diff --git a/src/quantem/core/config.py b/src/quantem/core/config.py
index c2eda10b..7213f3b5 100644
--- a/src/quantem/core/config.py
+++ b/src/quantem/core/config.py
@@ -34,6 +34,15 @@
     if "cuda" in str(e):
         NUM_DEVICES = 0
     _defaults["has_cupy"] = False
+try:
+    import quantem.cuda  # type: ignore  # noqa: F401
+
+    _defaults["has_quantem_cuda"] = True
+except ModuleNotFoundError:
+    _defaults["has_quantem_cuda"] = False
+except Exception:
+    # installed but unloadable (e.g. libcudart missing at runtime)
+    _defaults["has_quantem_cuda"] = False
 
 
 defaults: list[Mapping] = [_defaults]
diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py
index bc46821f..21e85deb 100644
--- a/src/quantem/core/ml/models/kplanes.py
+++ b/src/quantem/core/ml/models/kplanes.py
@@ -10,6 +10,8 @@
 import torch.nn.functional as F
 from torch import nn
 
+from quantem.core import config
+
 from .model_base import PPLR, TensorDecompositionModel
 from .so3params import SO3ParamQuat, SO3ParamR9SVD
 
@@ -144,12 +146,15 @@ def interpolate_ms_features(
     pts: torch.Tensor,
     ms_grids: nn.ParameterList,
 ) -> torch.Tensor:
-    mat_mode = [[0, 1], [0, 2], [1, 2]]
+    # Plane axis layout: XY=(0,1), XZ=(0,2), YZ=(1,2). Stacking views avoids the
+    # per-call index-tensor allocation (a host-to-device copy) that list-based
+    # advanced indexing does three times per forward.
+    x, y, z = pts.unbind(-1)
     coord_plane = torch.stack(
         [
-            pts[:, mat_mode[0]],
-            pts[:, mat_mode[1]],
-            pts[:, mat_mode[2]],
+            torch.stack((x, y), dim=-1),
+            torch.stack((x, z), dim=-1),
+            torch.stack((y, z), dim=-1),
         ]
     ).view(3, -1, 1, 2)
 
@@ -243,6 +248,12 @@ def __init__(
             nn.init.zeros_(out.bias)
             layers.append(out)
             self.sigma_net = nn.Sequential(*layers)
+        else:
+            # Linear head fallback, matching KPlanesTILTED._build_sigma_net and
+            # CPTilted: forward/get_params reference sigma_net unconditionally.
+            self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True)
+            nn.init.normal_(self.sigma_net.weight, std=0.01)
+            nn.init.zeros_(self.sigma_net.bias)
 
     def get_densities(self, coords: torch.Tensor):
         """Computes and returns densities"""
@@ -328,10 +339,31 @@ def interpolate_ms_features_tilted(
     block before concatenation. Used for coarse-to-fine band-limiting: ramp the finer
     scales on over training so the model commits to view-consistent low-frequency structure
     first (a classical limited-angle prior). None => all scales fully on (default behaviour).
+
+    When the optional ``quantem-cuda`` package is installed, the tensors live
+    on a CUDA device, and the ``use_cuda_kernels`` config option is true
+    (default), each scale dispatches to the fused CUDA kernel (rotate +
+    grid_sample + Hadamard product in one launch, analytic backward) —
+    identical semantics to the torch path below, including scale_gates.
     """
     T = rotation_matrices.shape[0]
     B = pts.shape[0]
 
+    if (
+        pts.is_cuda
+        and pts.dtype == torch.float32
+        and rotation_matrices.dtype == torch.float32
+        and all(g.dtype == torch.float32 for g in ms_grids)
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        from quantem.cuda.core.ml import kplanes_tilted_fuse
+
+        feats = [kplanes_tilted_fuse(pts, rotation_matrices, g) for g in ms_grids]
+        if scale_gates is not None:
+            feats = [f * scale_gates[si] for si, f in enumerate(feats)]
+        return torch.cat(feats, dim=-1)
+
     # (T, B, 3)  — rotate all points by all rotations at once
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 
@@ -651,33 +683,29 @@ def interpolate_ms_features_cp_tilted(
     # Rotate all points by all rotations: (T, B, 3)
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 
+    # For each transform t, we need three 1D samples: at x_t, y_t, z_t.
+    # Lay them out as (3T, B) coords, matching each line grid's first dim.
+    # Axis order per transform: x, y, z.
+    coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B)
+
+    # grid_sample wants 4D input for 2D sampling: sample the (3T, C, 1, L) lines
+    # with 2D coords whose y is fixed at 0. The grid only depends on the points,
+    # so it is built once here rather than per scale inside the loop.
+    grid = torch.stack(
+        [
+            coords_1d,  # x
+            torch.zeros_like(coords_1d),  # y
+        ],
+        dim=-1,
+    ).unsqueeze(1)  # (3T, 1, B, 2)
+
     per_scale_features = []
     for line_coef in ms_grids:
         # line_coef: (3T, C, L)  — three 1D feature lines per transform (x, y, z)
-        C, _ = line_coef.shape[1], line_coef.shape[2]
-
-        # For each transform t, we need three 1D samples: at x_t, y_t, z_t.
-        # Lay them out as (3T, B) coords, matching line_coef's first dim.
-        # Axis order per transform: x, y, z.
-        coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B)
-        # coords_1d: (3T, B), each row is samples along one axis for one transform
-
-        # grid_sample wants 4D input for 2D sampling, or we can use 1D via a
-        # (3T, C, 1, L) reshape and pass 2D coords with y fixed at 0.
-        # Simpler: use F.grid_sample with a 4D trick, or just do manual linear interp.
-        # Here's the grid_sample way:
-        line_coef_4d = line_coef.unsqueeze(2)  # (3T, C, 1, L)
-        # grid: need (3T, Hout=1, Wout=B, 2), with x = coord, y = 0
-        grid = torch.stack(
-            [
-                coords_1d,  # x
-                torch.zeros_like(coords_1d),  # y
-            ],
-            dim=-1,
-        ).unsqueeze(1)  # (3T, 1, B, 2)
+        C = line_coef.shape[1]
 
         sampled = F.grid_sample(
-            line_coef_4d,
+            line_coef.unsqueeze(2),  # (3T, C, 1, L)
             grid,
             align_corners=True,
             mode="bilinear",
diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py
index af125391..05c66ddb 100644
--- a/src/quantem/core/ml/optimizer_mixin.py
+++ b/src/quantem/core/ml/optimizer_mixin.py
@@ -662,7 +662,9 @@ def set_optimizer(self, opt_params: OptimizerParamsType | dict | None = None) ->
         for key, tensors in groups.items():
             for p in tensors:
                 p.requires_grad_(True)
-            param_groups.append({"params": tensors, **specs[key].params()})
+            # "name" records the group key so reconnect_optimizer_to_parameters
+            # can re-join hyperparameters to groups by key, not position.
+            param_groups.append({"params": tensors, **specs[key].params(), "name": key})
         self._optimizer = self._build_optimizer(spec_list[0], param_groups)
 
     def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
@@ -672,11 +674,15 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
         so each group's ``lr`` etc. overrides the optimizer-level default. ``NoneOptimizer`` must
         have been filtered out by the caller.
         """
+        # Fused Adam/AdamW runs the whole step in one CUDA kernel (~2x faster than the
+        # default foreach path on large grids, same update rule); only valid when every
+        # parameter lives on a CUDA device.
+        fused = all(p.is_cuda for group in param_groups for p in group["params"])
         match opt_params:
             case OptimizerParams.Adam():
-                return torch.optim.Adam(param_groups)
+                return torch.optim.Adam(param_groups, fused=fused)
             case OptimizerParams.AdamW():
-                return torch.optim.AdamW(param_groups)
+                return torch.optim.AdamW(param_groups, fused=fused)
             case OptimizerParams.SGD():
                 return torch.optim.SGD(param_groups)
             case OptimizerParams.NoneOptimizer():
@@ -804,14 +810,26 @@ def reconnect_optimizer_to_parameters(self) -> None:
         old_hyperparams = [
             {k: v for k, v in pg.items() if k != "params"} for pg in self._optimizer.param_groups
         ]
+        old_by_name = {hp["name"]: hp for hp in old_hyperparams if "name" in hp}
 
         self._optimizer.param_groups.clear()
-        for tensors in new_groups.values():
-            self._optimizer.add_param_group({"params": tensors})
-
-        # Restore per-group hyperparameters by index
-        for new_pg, old_pg in zip(self._optimizer.param_groups, old_hyperparams):
-            new_pg.update(old_pg)
+        for key, tensors in new_groups.items():
+            self._optimizer.add_param_group({"params": tensors, "name": key})
+
+        if old_by_name:
+            # Re-join hyperparameters to groups by key: group order/membership
+            # from get_optimization_parameters() is not contractual, and index
+            # alignment silently attaches the wrong lr when it changes. Groups
+            # with no old counterpart keep the optimizer defaults.
+            for new_pg in self._optimizer.param_groups:
+                hp = old_by_name.get(new_pg["name"])
+                if hp is not None:
+                    new_pg.update(hp)
+        else:
+            # Optimizers restored from checkpoints predating group names:
+            # index alignment is the only association available.
+            for new_pg, old_pg in zip(self._optimizer.param_groups, old_hyperparams):
+                new_pg.update(old_pg)
 
         # Remap state for tensors that survived
         new_state = {}
diff --git a/src/quantem/core/quantem.yaml b/src/quantem/core/quantem.yaml
index 80f34587..0e4913a4 100644
--- a/src/quantem/core/quantem.yaml
+++ b/src/quantem/core/quantem.yaml
@@ -14,6 +14,10 @@ dtype_complex: complex64
 # global verbosity, not currently used
 
 verbose: 1
+
+# Dispatch to the fused quantem-cuda kernels when installed (pip install quantem[cuda])
+# and the tensors are on a CUDA device. Set false to force the pure-torch paths.
+use_cuda_kernels: true
 cupy:
   # The size of the fft cache in MB used by cupy
   # https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index d311d8dd..17c4e72c 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -396,6 +396,30 @@ def get_tv_loss(
         return loss
 
     def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> torch.Tensor:
+        # Identical math via the fused quantem-cuda L1 kernel when available.
+        # Weighted size-1 axes fall through to torch so degenerate inputs
+        # behave exactly as before.
+        if (
+            array.ndim == 3
+            and array.is_cuda
+            and array.dtype == torch.float32
+            and config.get("has_quantem_cuda")
+            and config.get("use_cuda_kernels", default=True)
+            and not (weight[0] > 0 and array.shape[0] == 1)
+            and not (weight[1] > 0 and (array.shape[1] == 1 or array.shape[2] == 1))
+        ):
+            from quantem.cuda.core import tv_loss_l1_3d
+
+            s, h, w_ = array.shape
+            wts = array.new_tensor([weight[0], weight[1], weight[1]])
+            active = int((wts > 0).sum())
+            if active == 0:
+                return self._get_zero_loss_tensor()
+            counts = array.new_tensor(
+                [(s - 1) * h * w_, s * (h - 1) * w_, s * h * (w_ - 1)]
+            ).clamp(min=1)
+            return (wts * tv_loss_l1_3d(array) / counts).sum() / active
+
         loss = self._get_zero_loss_tensor()
         calc_dim = 0
         for dim in range(array.ndim):
diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py
index 2066c5b4..daec3891 100644
--- a/src/quantem/tomography/object_models.py
+++ b/src/quantem/tomography/object_models.py
@@ -1,3 +1,4 @@
+import weakref
 from abc import abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
@@ -19,6 +20,7 @@
 from quantem.core.utils.rng import RNGMixin
 from quantem.tomography.dataset_models import TomographyINRPretrainDataset
 from quantem.tomography.tomography_context import ReconstructionContext
+from quantem.tomography.utils import tv_loss_vol_sq
 
 
 class ObjConstraintParams:
@@ -455,11 +457,9 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         # TV over the three trailing spatial dims, leaving any leading channel/batch axes
         # intact. Works for a 3-D volume, obj_view's [1, D, H, W], and a multimodal
         # [C, D, H, W] (channels = elemental compositions), matching the INR / tensor-decomp
-        # convention where the object carries a leading channel dimension.
-        tv_d = torch.pow(ctx.obj[..., 1:, :, :] - ctx.obj[..., :-1, :, :], 2).sum()
-        tv_h = torch.pow(ctx.obj[..., :, 1:, :] - ctx.obj[..., :, :-1, :], 2).sum()
-        tv_w = torch.pow(ctx.obj[..., :, :, 1:] - ctx.obj[..., :, :, :-1], 2).sum()
-        tv_loss = tv_d + tv_h + tv_w
+        # convention where the object carries a leading channel dimension. tv_loss_vol_sq
+        # dispatches to the fused quantem-cuda kernel when available.
+        tv_loss = tv_loss_vol_sq(ctx.obj)
 
         return tv_loss * self.constraints.tv_vol / ctx.obj.numel()
 
@@ -473,6 +473,14 @@ def to(self, device: str | torch.device):
         return self
 
 
+# torch.compile artifacts keyed by the live model object. Kept outside the
+# instances so AutoSerialize never sees them and reset()/rebuild_model()
+# (which swap the model object) naturally invalidate the cache.
+_compiled_forward_cache: "weakref.WeakKeyDictionary[nn.Module, Callable]" = (
+    weakref.WeakKeyDictionary()
+)
+
+
 class ObjectINR(ObjectConstraints, DDPMixin):
     DEFAULT_CONSTRAINTS = ObjConstraintParams.ObjINRConstraints()
 
@@ -482,6 +490,7 @@ def __init__(
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
         model: nn.Module | None = None,
+        compile_model: bool = False,
         _token: object | None = None,
     ):
         super().__init__(
@@ -492,6 +501,7 @@ def __init__(
         )
         self._pretrain_losses = []
         self._pretrain_lrs = []
+        self._compile_model = bool(compile_model)
         self.constraints: ObjConstraintParams.ObjINRConstraints = self.DEFAULT_CONSTRAINTS.copy()
         # Register the network submodule (important: real nn.Module attribute)
         if model is not None:
@@ -505,18 +515,36 @@ def from_model(
         shape: tuple[int, int, int],
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
+        compile_model: bool = False,
     ):
         obj_model = cls(
             shape=shape,
             device=device,
             rng=rng,
             model=model,  # ✅ build/register in __init__
+            compile_model=compile_model,
         )
 
         obj_model.setup_distributed(device=device)
         obj_model.to(device)
         return obj_model
 
+    def _model_call(self, coords: torch.Tensor) -> torch.Tensor:
+        """Invoke the model, through torch.compile when compile_model was set.
+
+        Compiles the bound __call__ (a plain function, so nothing extra is
+        registered on the module tree or picked up by AutoSerialize) and
+        caches per model object.
+        """
+        model = self.model
+        if not getattr(self, "_compile_model", False):
+            return model(coords)
+        fn = _compiled_forward_cache.get(model)
+        if fn is None:
+            fn = torch.compile(model.__call__)
+            _compiled_forward_cache[model] = fn
+        return fn(coords)
+
     # --- Properties ---
 
     @property
@@ -593,6 +621,10 @@ def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor:
 
         return pred
 
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """Hook for the training loop: returns None (INR TV uses autograd, no tap merging)."""
+        return None
+
     # --- Define get_tv_loss ---
 
     def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
@@ -704,7 +736,10 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
         """forward pass for the INR model"""
         assert coords is not None, "ObjectINR.forward requires coords"
 
-        all_densities = self.model(coords)
+        # TV's autograd.grad recompute stays on the eager model (double
+        # backward through compiled graphs is not reliable); only this main
+        # forward goes through the compiled path.
+        all_densities = self._model_call(coords)
 
         if all_densities.dim() > 1:
             all_densities = all_densities.squeeze(-1)
@@ -726,6 +761,40 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
 
         return all_densities
 
+    def forward_with_tv_taps(
+        self, coords: torch.Tensor, tap_coords: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Single model call covering the main batch and the volume-TV tap points.
+
+        The out-of-bounds mask and hard constraints apply to the main chunk
+        only; tap densities are returned raw (border-clamped), matching the
+        fallback path in ``get_volume_tv_loss``.
+        """
+        merged = self.model(torch.cat([coords, tap_coords], dim=0))
+        if isinstance(merged, tuple):
+            merged = merged[0]
+        main, taps = merged[: coords.shape[0]], merged[coords.shape[0] :]
+
+        if main.dim() > 1:
+            main = main.squeeze(-1)
+        valid_mask = (
+            (coords[:, 0] >= -1)
+            & (coords[:, 0] <= 1)
+            & (coords[:, 1] >= -1)
+            & (coords[:, 1] <= 1)
+            & (coords[:, 2] >= -1)
+            & (coords[:, 2] <= 1)
+        ).float()
+        if main.dim() > 1:
+            valid_mask = valid_mask.unsqueeze(-1)
+        main = main * valid_mask
+        main = self.apply_hard_constraints(main)
+
+        if taps.dim() == 1:
+            taps = taps.unsqueeze(-1)
+        return main, taps
+
     # Pretrain Loop
 
     def pretrain(
@@ -963,6 +1032,30 @@ def from_model(
         obj_model.to(device)
         return obj_model
 
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """
+        Sample the finite-difference tap coordinates for the volume TV loss.
+
+        Returns a (4*n, 3) tensor [base; base+h*ex; base+h*ey; base+h*ez] for n
+        sampled base points, or None when tv_vol == 0.  The training loop
+        concatenates this to all_coords so the TV taps are evaluated in the
+        same model call as the main forward pass.
+        """
+        if self.constraints.tv_vol == 0:
+            return None
+        model = _unwrap(self.model)
+        h = 2.0 / min(model.resolution)
+        num_tv_samples = min(10_000, coords.shape[0])
+        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+        tv_coords = coords[tv_indices]  # (n, 3)
+        ex = torch.zeros(3, device=tv_coords.device)
+        ex[0] = h
+        ey = torch.zeros(3, device=tv_coords.device)
+        ey[1] = h
+        ez = torch.zeros(3, device=tv_coords.device)
+        ez[2] = h
+        return torch.cat([tv_coords, tv_coords + ex, tv_coords + ey, tv_coords + ez], dim=0)
+
     # --- Constraints ---
 
     def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
@@ -1001,7 +1094,9 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         if self.constraints.tv_plane > 0:
             tv_loss = tv_loss + self._get_plane_tv_loss()
         if self.constraints.tv_vol > 0:
-            tv_loss = tv_loss + self.get_volume_tv_loss(ctx.coords)
+            tv_loss = tv_loss + self.get_volume_tv_loss(
+                ctx.coords, precomputed_tap_densities=ctx.tv_tap_densities
+            )
         return tv_loss
 
     def _get_plane_tv_loss(self) -> torch.Tensor:
@@ -1029,33 +1124,48 @@ def _get_plane_tv_loss(self) -> torch.Tensor:
 
         return self.constraints.tv_plane * torch.stack(per_level).sum()
 
-    def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor:
+    def get_volume_tv_loss(
+        self,
+        coords: torch.Tensor,
+        precomputed_tap_densities: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Isotropic volume TV via finite differences. Same form as the autograd
         version (L1 of gradient L2-norm) but avoids double-backward, so it
         works for KPlanesTILTED, CPTilted, and anything else.
-        """
-        num_tv_samples = min(10_000, coords.shape[0])
-        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
-        tv_coords = coords[tv_indices]  # (N, 3)
 
+        When *precomputed_tap_densities* is provided (a (4N, C) tensor from the
+        merged single-pass forward in the training loop), the model call is
+        skipped entirely and the supplied values are used directly.  When absent
+        the existing batched 4N-point fallback path runs unchanged.
+        """
         model = _unwrap(self.model)
         h = 2.0 / min(model.resolution)
 
-        # Evaluate the base points and the three axis-shifted copies in a single
-        # batched forward (4N points) instead of 4 sequential model calls.
-        offsets = h * torch.eye(3, device=tv_coords.device, dtype=tv_coords.dtype)  # (3, 3)
-        all_coords = torch.cat(
-            [tv_coords, (tv_coords.unsqueeze(0) + offsets.unsqueeze(1)).reshape(-1, 3)]
-        )  # (4N, 3)
+        if precomputed_tap_densities is not None:
+            # (4N, C) tap densities from the merged single-pass forward in the
+            # training loop; layout [base; +ex; +ey; +ez] per sample_tv_tap_coords.
+            all_pred = precomputed_tap_densities
+        else:
+            num_tv_samples = min(10_000, coords.shape[0])
+            tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+            tv_coords = coords[tv_indices]  # (N, 3)
+
+            # Evaluate the base points and the three axis-shifted copies in a single
+            # batched forward (4N points) instead of 4 sequential model calls.
+            offsets = h * torch.eye(3, device=tv_coords.device, dtype=tv_coords.dtype)  # (3, 3)
+            all_coords = torch.cat(
+                [tv_coords, (tv_coords.unsqueeze(0) + offsets.unsqueeze(1)).reshape(-1, 3)]
+            )  # (4N, 3)
+
+            all_pred = model(all_coords)
+            if isinstance(all_pred, tuple):
+                all_pred = all_pred[0]
 
-        all_pred = model(all_coords)
-        if isinstance(all_pred, tuple):
-            all_pred = all_pred[0]
         if all_pred.dim() == 1:
             all_pred = all_pred.unsqueeze(-1)  # (4N, 1)
 
-        n = tv_coords.shape[0]
+        n = all_pred.shape[0] // 4
         pred = all_pred[:n]  # (N, C)
         shifted_pred = all_pred[n:].view(3, n, -1)  # (3, N, C)
 
diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py
index c2a9e5f2..73763ce4 100644
--- a/src/quantem/tomography/tomography.py
+++ b/src/quantem/tomography/tomography.py
@@ -276,7 +276,14 @@ def reconstruct(
                 ):
                     all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
-                    all_densities = self.obj_model.forward(all_coords)
+                    tap_coords = self.obj_model.sample_tv_tap_coords(all_coords)
+                    if tap_coords is not None:
+                        all_densities, tv_tap_raw = self.obj_model.forward_with_tv_taps(
+                            all_coords, tap_coords
+                        )
+                    else:
+                        all_densities = self.obj_model.forward(all_coords)
+                        tv_tap_raw = None
 
                     integrated_densities = self.dset.integrate_rays(
                         all_densities,
@@ -294,6 +301,7 @@ def reconstruct(
                         pred=pred,
                         all_densities=all_densities,
                         target=target,
+                        tv_tap_densities=tv_tap_raw,
                     )
                 )
 
@@ -343,14 +351,13 @@ def reconstruct(
                     )
                     prev_R = R_now.clone()
 
+            # One stacked all_reduce and one host sync instead of three of each.
+            losses = torch.stack([total_loss, consistency_loss, epoch_soft_constraint_loss])
             if self.world_size > 1:
-                dist.all_reduce(total_loss, dist.ReduceOp.AVG)
-                dist.all_reduce(consistency_loss, dist.ReduceOp.AVG)
-                dist.all_reduce(epoch_soft_constraint_loss, dist.ReduceOp.AVG)
-
-            total_loss = total_loss.item() / len(self.dataloader)
-            consistency_loss = consistency_loss.item() / len(self.dataloader)
-            epoch_soft_constraint_loss = epoch_soft_constraint_loss.item() / len(self.dataloader)
+                dist.all_reduce(losses, dist.ReduceOp.AVG)
+            total_loss, consistency_loss, epoch_soft_constraint_loss = (
+                losses / len(self.dataloader)
+            ).tolist()
 
             self.step_schedulers(loss=total_loss)
             # TODO: Maybe reorganize the losses so that the order makes sense lol.
@@ -384,15 +391,9 @@ def reconstruct(
                         loss_func=loss_func,
                     )
 
-            metrics = torch.tensor(
-                [total_loss, consistency_loss, epoch_soft_constraint_loss], device=self.device
-            )
-
-            if self.world_size > 1:
-                dist.all_reduce(metrics, dist.ReduceOp.AVG)
-
-            total_loss, consistency_loss, epoch_soft_constraint_loss = metrics.tolist()
-
+            # The three losses were already rank-averaged (and batch-normalized) right
+            # after the batch loop; re-reducing identical values here was a redundant
+            # all_reduce plus an extra host sync per epoch.
             pbar.set_description(
                 f"Reconstruction | Loss: {total_loss:.5e}, Consistency Loss: {consistency_loss:.5e}, Soft Constraint Loss: {epoch_soft_constraint_loss:.5e}"
             )
diff --git a/src/quantem/tomography/tomography_context.py b/src/quantem/tomography/tomography_context.py
index c550c997..b62078e2 100644
--- a/src/quantem/tomography/tomography_context.py
+++ b/src/quantem/tomography/tomography_context.py
@@ -1,9 +1,10 @@
 from dataclasses import dataclass
 from typing import Optional
-from quantem.core.ml.constraints import BaseContext
 
 import torch
 
+from quantem.core.ml.constraints import BaseContext
+
 
 @dataclass
 class ReconstructionContext(BaseContext):
@@ -30,3 +31,4 @@ class ReconstructionContext(BaseContext):
     all_densities: Optional[torch.Tensor] = None
     obj: Optional[torch.Tensor] = None
     target: Optional[torch.Tensor] = None
+    tv_tap_densities: Optional[torch.Tensor] = None
diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py
index b1105d3c..553a27b4 100644
--- a/src/quantem/tomography/utils.py
+++ b/src/quantem/tomography/utils.py
@@ -1,6 +1,8 @@
 import torch
 import torch.nn.functional as F
 
+from quantem.core import config
+
 # --- Projection Operator Utils ---
 
 
@@ -72,6 +74,42 @@ def differentiable_rotx_vectorized(mags, theta, mode="bilinear"):
     return rotated.permute(0, 2, 3, 1)  # back to (B, Z, Y, X)
 
 
+def tv_loss_vol_sq(obj: torch.Tensor) -> torch.Tensor:
+    """Squared-anisotropic volume TV: sum of squared forward differences.
+
+    Computes ``Σ (Δd)² + Σ (Δh)² + Σ (Δw)²`` over the three trailing
+    spatial dims, leaving any leading channel/batch axes intact (they are
+    included in the sum). This is the unnormalized ``tv_vol`` regularizer;
+    callers apply their own ``weight / numel`` scaling.
+
+    When the optional ``quantem-cuda`` package is installed
+    (``pip install quantem[cuda]``), the tensor is on a CUDA device, and the
+    ``use_cuda_kernels`` config option is true (default), this dispatches to
+    the fused CUDA forward/backward kernel — identical math, one kernel
+    launch instead of several large intermediates.
+
+    Args:
+        obj: Tensor of shape ``[..., D, H, W]`` (ndim >= 3).
+
+    Returns:
+        0-dim tensor on the same device as ``obj``; differentiable.
+    """
+    if (
+        obj.is_cuda
+        and obj.dtype == torch.float32
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        from quantem.cuda.core import tv_loss_sq_3d
+
+        return tv_loss_sq_3d(obj)
+
+    tv_d = torch.pow(obj[..., 1:, :, :] - obj[..., :-1, :, :], 2).sum()
+    tv_h = torch.pow(obj[..., :, 1:, :] - obj[..., :, :-1, :], 2).sum()
+    tv_w = torch.pow(obj[..., :, :, 1:] - obj[..., :, :, :-1], 2).sum()
+    return tv_d + tv_h + tv_w
+
+
 def tv_loss_1d(x: torch.Tensor, reduction: str = "mean") -> torch.Tensor:
     """
     1D Total Variation Loss.
diff --git a/tests/diffractive_imaging/test_tv_cuda_dispatch.py b/tests/diffractive_imaging/test_tv_cuda_dispatch.py
new file mode 100644
index 00000000..e320ceb3
--- /dev/null
+++ b/tests/diffractive_imaging/test_tv_cuda_dispatch.py
@@ -0,0 +1,103 @@
+"""Tests for ``ObjectConstraints._calc_tv_loss``'s optional quantem-cuda dispatch.
+
+The constraint must produce the same loss and gradients on every path: pure
+torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent), and
+the fused L1 kernel when quantem-cuda is installed. The dispatch itself is
+asserted by monkeypatching the kernel entry point, so these tests are
+meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.diffractive_imaging.object_models import ObjectConstraints
+
+requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+class _Host:
+    """Minimal stand-in providing what _calc_tv_loss needs from self."""
+
+    def __init__(self, device):
+        self.device = device
+
+    _get_zero_loss_tensor = ObjectConstraints._get_zero_loss_tensor
+    _calc_tv_loss = ObjectConstraints._calc_tv_loss
+
+
+def _phase(device, shape=(8, 24, 20), seed=0, requires_grad=False):
+    gen = torch.Generator().manual_seed(seed)
+    arr = torch.rand(shape, generator=gen, dtype=torch.float32)
+    return arr.to(device).requires_grad_(requires_grad)
+
+
+WEIGHTS = [(5.0, 0.1), (5.0, 0.0), (0.0, 0.3), (0.0, 0.0)]
+
+
+@requires_gpu
+@pytest.mark.parametrize("weight", WEIGHTS)
+def test_gpu_matches_cpu_reference(weight):
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu"), weight)
+    actual = _Host("cuda")._calc_tv_loss(_phase("cuda"), weight)
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
+
+
+@requires_gpu
+@pytest.mark.parametrize("weight", WEIGHTS[:3])
+def test_gpu_grads_match_cpu_reference(weight):
+    a_cpu = _phase("cpu", requires_grad=True)
+    _Host("cpu")._calc_tv_loss(a_cpu, weight).backward()
+    a_gpu = _phase("cuda", requires_grad=True)
+    _Host("cuda")._calc_tv_loss(a_gpu, weight).backward()
+    torch.testing.assert_close(a_gpu.grad.cpu(), a_cpu.grad, rtol=1e-4, atol=1e-7)
+
+
+@requires_gpu
+@pytest.mark.parametrize("shape", [(1, 24, 20), (24, 20)])
+def test_degenerate_and_2d_match_cpu(shape):
+    # num_slices == 1 (weight[0] zeroed upstream, as get_tv_loss does) and
+    # plain 2-D arrays keep exact parity whichever path runs.
+    weight = (0.0, 0.3)
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu", shape=shape), weight)
+    actual = _Host("cuda")._calc_tv_loss(_phase("cuda", shape=shape), weight)
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel(monkeypatch):
+    import quantem.cuda.core
+
+    calls = []
+    real = quantem.cuda.core.tv_loss_l1_3d
+
+    def spy(volume):
+        calls.append(volume.shape)
+        return real(volume)
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_l1_3d", spy)
+    _Host("cuda")._calc_tv_loss(_phase("cuda"), (5.0, 0.1))
+    assert len(calls) == 1
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core
+
+    def boom(volume):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_l1_3d", boom)
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu"), (5.0, 0.1))
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        actual = _Host("cuda")._calc_tv_loss(_phase("cuda"), (5.0, 0.1))
+    finally:
+        config.set({"use_cuda_kernels": True})
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
diff --git a/tests/ml/test_kplanes.py b/tests/ml/test_kplanes.py
index bdd865bc..f6938519 100644
--- a/tests/ml/test_kplanes.py
+++ b/tests/ml/test_kplanes.py
@@ -18,3 +18,24 @@ def test_anisotropic_raises(self):
             KPlanes(M_features=2, resolution=(16, 16, 8))
         with pytest.raises(ValueError, match="isotropic"):
             KPlanesTILTED(M_features=2, T=2, resolution=(16, 8, 16))
+
+
+class TestDefaultHeadConstruction:
+    """Regression: KPlanes(use_hybrid_mlp=False) -- the constructor default --
+    built no sigma_net at all, so forward / get_params / ObjectTensorDecomp
+    .from_model crashed with AttributeError. KPlanesTILTED and CPTilted both
+    fall back to a linear head; plain KPlanes must do the same."""
+
+    def test_default_get_params(self):
+        model = KPlanes(M_features=2, resolution=(8, 8, 8))
+        params = model.get_params()
+        assert set(params) == set(model.param_keys)
+        assert all(len(v) > 0 for v in params.values())
+
+    def test_default_forward(self):
+        import torch
+
+        model = KPlanes(M_features=2, resolution=(8, 8, 8))
+        out = model(torch.rand(5, 3) * 2 - 1)
+        assert out.shape == (5, 1)
+        assert torch.isfinite(out).all()
diff --git a/tests/ml/test_kplanes_cuda_dispatch.py b/tests/ml/test_kplanes_cuda_dispatch.py
new file mode 100644
index 00000000..da61f6db
--- /dev/null
+++ b/tests/ml/test_kplanes_cuda_dispatch.py
@@ -0,0 +1,106 @@
+"""Tests for ``interpolate_ms_features_tilted``'s optional quantem-cuda dispatch.
+
+The function must produce the same features and gradients on every path:
+pure torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent),
+and the fused CUDA kernel when quantem-cuda is installed. The dispatch
+itself is asserted by monkeypatching the kernel entry point, so these tests
+are meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+from torch import nn
+
+from quantem.core import config
+from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+def _inputs(device, requires_grad=False, seed=0, B=64, T=3, C=5, scales=(9, 17)):
+    gen = torch.Generator().manual_seed(seed)
+    pts = torch.rand(B, 3, generator=gen) * 2.2 - 1.1  # some points outside [-1, 1]
+    rotations = torch.rand(T, 3, 3, generator=gen) * 2 - 1
+    grids = nn.ParameterList(
+        nn.Parameter(torch.rand(3 * T, C, s, s, generator=gen) * 0.4 + 0.1) for s in scales
+    )
+    pts = pts.to(device).requires_grad_(requires_grad)
+    rotations = rotations.to(device).requires_grad_(requires_grad)
+    grids = grids.to(device)
+    if not requires_grad:
+        for g in grids:
+            g.requires_grad_(False)
+    return pts, rotations, grids
+
+
+@requires_gpu
+def test_gpu_matches_cpu_reference():
+    pts_c, rot_c, grids_c = _inputs("cpu")
+    expected = interpolate_ms_features_tilted(pts_c, grids_c, rot_c)
+    pts_g, rot_g, grids_g = _inputs("cuda")
+    actual = interpolate_ms_features_tilted(pts_g, grids_g, rot_g).cpu()
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+
+
+@requires_gpu
+def test_gpu_grads_match_cpu_reference():
+    pts_c, rot_c, grids_c = _inputs("cpu", requires_grad=True)
+    interpolate_ms_features_tilted(pts_c, grids_c, rot_c).square().sum().backward()
+    pts_g, rot_g, grids_g = _inputs("cuda", requires_grad=True)
+    interpolate_ms_features_tilted(pts_g, grids_g, rot_g).square().sum().backward()
+    torch.testing.assert_close(pts_g.grad.cpu(), pts_c.grad, rtol=1e-3, atol=1e-5)
+    torch.testing.assert_close(rot_g.grad.cpu(), rot_c.grad, rtol=1e-3, atol=1e-5)
+    for gg, gc in zip(grids_g, grids_c):
+        torch.testing.assert_close(gg.grad.cpu(), gc.grad, rtol=1e-3, atol=1e-5)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel_per_scale(monkeypatch):
+    import quantem.cuda.core.ml
+
+    calls = []
+    real = quantem.cuda.core.ml.kplanes_tilted_fuse
+
+    def spy(pts, rotations, plane):
+        calls.append(plane.shape)
+        return real(pts, rotations, plane)
+
+    monkeypatch.setattr(quantem.cuda.core.ml, "kplanes_tilted_fuse", spy)
+    pts, rot, grids = _inputs("cuda")
+    interpolate_ms_features_tilted(pts, grids, rot)
+    assert len(calls) == len(grids)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core.ml
+
+    def boom(pts, rotations, plane):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core.ml, "kplanes_tilted_fuse", boom)
+    pts, rot, grids = _inputs("cuda")
+    pts_c, rot_c, grids_c = _inputs("cpu")
+    expected = interpolate_ms_features_tilted(pts_c, grids_c, rot_c)
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        actual = interpolate_ms_features_tilted(pts, grids, rot).cpu()
+    finally:
+        config.set({"use_cuda_kernels": True})
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+
+
+@requires_gpu
+def test_non_fp32_takes_torch_path():
+    pts, rot, grids = _inputs("cuda")
+    pts64 = pts.double()
+    rot64 = rot.double()
+    grids64 = nn.ParameterList(nn.Parameter(g.double()) for g in grids)
+    out = interpolate_ms_features_tilted(pts64, grids64, rot64)
+    assert out.dtype == torch.float64
diff --git a/tests/ml/test_optimizermixin.py b/tests/ml/test_optimizermixin.py
index 7dc6f4e7..8cbc1bc9 100644
--- a/tests/ml/test_optimizermixin.py
+++ b/tests/ml/test_optimizermixin.py
@@ -499,3 +499,65 @@ def test_explicit_values_still_win(self):
         assert p.params(base_LR=1.0)["min_lr"] == 1e-6
         e = SchedulerParams.Exponential(gamma=0.8)
         assert e.params(base_LR=1.0, num_iter=10)["gamma"] == 0.8
+
+
+@requires_torch
+class TestReconnectHyperparamAlignment:
+    """Regression: reconnect_optimizer_to_parameters restored per-group
+    hyperparameters by zip-index against get_optimization_parameters(),
+    silently attaching the wrong lr to every group whenever the group dict's
+    order (or membership) changed between optimizer creation and reconnect."""
+
+    def _model(self):
+        model = _FakeModel({"a": [_param()], "b": [_param(2.0)]})
+        model.set_optimizer({"a": OptimizerParams.SGD(lr=1e-2), "b": OptimizerParams.SGD(lr=1e-3)})
+        return model
+
+    def _lrs_by_name(self, model):
+        return {pg["name"]: pg["lr"] for pg in model.optimizer.param_groups}
+
+    def test_groups_carry_names(self):
+        model = self._model()
+        assert self._lrs_by_name(model) == {"a": 1e-2, "b": 1e-3}
+
+    def test_reordered_groups_keep_their_lr(self):
+        model = self._model()
+        # Simulate get_optimization_parameters returning a different order at
+        # reconnect time (group membership/order is not contractual).
+        model._groups = {"b": model._groups["b"], "a": model._groups["a"]}
+        model.reconnect_optimizer_to_parameters()
+        assert self._lrs_by_name(model) == {"a": 1e-2, "b": 1e-3}
+
+    def test_added_group_gets_defaults_others_keep_lr(self):
+        model = self._model()
+        model._groups = dict(model._groups)
+        model._groups["c"] = [_param(3.0)]
+        model.reconnect_optimizer_to_parameters()
+        lrs = self._lrs_by_name(model)
+        assert lrs["a"] == 1e-2 and lrs["b"] == 1e-3
+        assert "c" in lrs  # present, with optimizer defaults
+
+    def test_state_survives_reconnect(self):
+        # Adam: SGD without momentum keeps no per-param state to preserve.
+        model = _FakeModel({"a": [_param()], "b": [_param(2.0)]})
+        model.set_optimizer(
+            {"a": OptimizerParams.Adam(lr=1e-2), "b": OptimizerParams.Adam(lr=1e-3)}
+        )
+        params = [pg["params"][0] for pg in model.optimizer.param_groups]
+        for p in params:
+            p.grad = torch.ones_like(p)
+        model.optimizer.step()
+        assert len(model.optimizer.state) > 0
+        before = {id(p): dict(model.optimizer.state[p]) for p in params}
+        model.reconnect_optimizer_to_parameters()
+        for p in params:
+            assert id(p) in before and p in model.optimizer.state
+
+    def test_legacy_groups_without_names_fall_back_to_index(self):
+        model = self._model()
+        # Optimizers restored from older checkpoints carry no group names.
+        for pg in model.optimizer.param_groups:
+            del pg["name"]
+        model.reconnect_optimizer_to_parameters()
+        lrs = [pg["lr"] for pg in model.optimizer.param_groups]
+        assert lrs == [1e-2, 1e-3]
diff --git a/tests/tomography/test_cuda_dispatch.py b/tests/tomography/test_cuda_dispatch.py
new file mode 100644
index 00000000..db6e9691
--- /dev/null
+++ b/tests/tomography/test_cuda_dispatch.py
@@ -0,0 +1,91 @@
+"""Tests for ``tv_loss_vol_sq`` and its optional quantem-cuda dispatch.
+
+The helper must produce the same value and gradients on every path: pure
+torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent), and
+the fused CUDA kernel when quantem-cuda is installed. The dispatch itself is
+asserted by monkeypatching the kernel entry point, so these tests are
+meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.tomography.utils import tv_loss_vol_sq
+
+from .conftest import requires_gpu
+
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+def tv_vol_sq_ref(obj: torch.Tensor) -> torch.Tensor:
+    tv_d = torch.pow(obj[..., 1:, :, :] - obj[..., :-1, :, :], 2).sum()
+    tv_h = torch.pow(obj[..., :, 1:, :] - obj[..., :, :-1, :], 2).sum()
+    tv_w = torch.pow(obj[..., :, :, 1:] - obj[..., :, :, :-1], 2).sum()
+    return tv_d + tv_h + tv_w
+
+
+@pytest.mark.parametrize("shape", [(7, 6, 5), (2, 7, 6, 5)])
+def test_cpu_matches_reference(shape):
+    obj = torch.rand(shape, generator=torch.Generator().manual_seed(0))
+    torch.testing.assert_close(tv_loss_vol_sq(obj), tv_vol_sq_ref(obj))
+
+
+def test_constant_volume_is_zero():
+    assert tv_loss_vol_sq(torch.ones(8, 8, 8)).item() == 0.0
+
+
+@requires_gpu
+@pytest.mark.parametrize("shape", [(7, 6, 5), (2, 7, 6, 5)])
+def test_gpu_matches_cpu_reference(shape):
+    obj = torch.rand(shape, generator=torch.Generator().manual_seed(1))
+    expected = tv_vol_sq_ref(obj)
+    actual = tv_loss_vol_sq(obj.cuda()).cpu()
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+
+
+@requires_gpu
+def test_gpu_grad_matches_cpu_reference():
+    obj = torch.rand(6, 7, 8, generator=torch.Generator().manual_seed(2))
+    v_cpu = obj.clone().requires_grad_(True)
+    tv_vol_sq_ref(v_cpu).backward()
+    v_gpu = obj.cuda().requires_grad_(True)
+    tv_loss_vol_sq(v_gpu).backward()
+    torch.testing.assert_close(v_gpu.grad.cpu(), v_cpu.grad, rtol=1e-4, atol=1e-6)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel(monkeypatch):
+    import quantem.cuda.core
+
+    calls = []
+    real = quantem.cuda.core.tv_loss_sq_3d
+
+    def spy(volume):
+        calls.append(volume.shape)
+        return real(volume)
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_sq_3d", spy)
+    tv_loss_vol_sq(torch.rand(4, 4, 4, device="cuda"))
+    assert len(calls) == 1
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core
+
+    def boom(volume):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_sq_3d", boom)
+    obj = torch.rand(4, 4, 4, device="cuda")
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        torch.testing.assert_close(tv_loss_vol_sq(obj), tv_vol_sq_ref(obj))
+    finally:
+        config.set({"use_cuda_kernels": True})
diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py
index 268fa2c6..66316f5b 100644
--- a/tests/tomography/test_dataset_models.py
+++ b/tests/tomography/test_dataset_models.py
@@ -176,10 +176,15 @@ def test_transform_batch_rays_identity_at_zero_pose(self):
         assert torch.allclose(out, rays, atol=1e-5)
 
     def test_integrate_rays_sums_with_step_size(self):
+        # integrate_rays became an instance method dispatching on ray_sampling
+        # (box_fixed_ds vs legacy); the default legacy path keeps the original
+        # step-size summation this test pins down.
         B, S = 3, 5
-        out = TomographyINRDataset.integrate_rays(
-            torch.ones(B, S), num_samples_per_ray=S, target_values_len=B
+        stack = _stack(nang=3, n=4)
+        d = TomographyINRDataset.from_data(
+            stack, np.linspace(-60, 60, 3, dtype="f4"), ray_sampling="legacy"
         )
+        out = d.integrate_rays(torch.ones(B * S), num_samples_per_ray=S, target_values_len=B)
         step = 2.0 / (S - 1)
         assert out.shape == (B,)
         assert torch.allclose(out, torch.full((B,), S * step))
diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py
index da01c7d8..c614aeb9 100644
--- a/tests/tomography/test_object_models.py
+++ b/tests/tomography/test_object_models.py
@@ -309,3 +309,28 @@ def test_normalize_optimizer_params_rejects_wrong_keys(self, torch_device):
             obj._normalize_optimizer_params(
                 {"grids": OptimizerParams.Adam(), "wrong": OptimizerParams.Adam()}
             )
+
+
+class TestObjectINRCompileOptIn:
+    """compile_model=True must change performance only, never results."""
+
+    def _make(self, compile_model, device):
+        from quantem.core.ml.inr import HSiren
+
+        torch.manual_seed(0)
+        model = HSiren(hidden_layers=1, hidden_features=16, alpha=1)
+        return ObjectINR.from_model(
+            model, shape=(8, 8, 8), device=device, compile_model=compile_model
+        )
+
+    def test_default_is_eager(self, torch_device):
+        obj = self._make(False, torch_device)
+        assert obj._compile_model is False
+
+    @pytest.mark.slow
+    def test_compiled_forward_matches_eager(self, torch_device):
+        torch.manual_seed(1)
+        coords = (torch.rand(64, 3) * 2 - 1).to(torch_device)
+        out_eager = self._make(False, torch_device).forward(coords)
+        out_compiled = self._make(True, torch_device).forward(coords)
+        torch.testing.assert_close(out_compiled, out_eager, rtol=1e-5, atol=1e-6)