Merge quantem-cuda kernel dispatch (upstream PR #243) onto the benchmark stack

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]

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",

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 = {}

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

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):