Dispatch interpolate_ms_features_tilted to the fused quantem-cuda kernel

src/quantem/core/ml/models/kplanes.py

@@ -13,6 +13,11 @@
 from .model_base import PPLR, TensorDecompositionModel
 from .so3params import SO3ParamQuat, SO3ParamR9SVD
 
+try:  # optional fused CUDA kernels (pip install quantem-cuda)
+    from quantem.cuda.core.ml import kplanes_tilted_fuse as _kplanes_tilted_fuse_cuda
+except ImportError:
+    _kplanes_tilted_fuse_cuda = None
+
 """
 K-planes utility functions
 """
@@ -326,6 +331,23 @@ def interpolate_ms_features_tilted(
     T = rotation_matrices.shape[0]
     B = pts.shape[0]
 
+    # Fused CUDA path: one kernel per level instead of einsum + grid_sample +
+    # prod. Falls through to torch for CPU/non-fp32/odd shapes.
+    if (
+        _kplanes_tilted_fuse_cuda is not None
+        and pts.is_cuda
+        and pts.dtype == torch.float32
+        and pts.ndim == 2
+        and pts.shape[-1] == 3
+        and rotation_matrices.dtype == torch.float32
+        and rotation_matrices.ndim == 3
+        and rotation_matrices.shape[-2:] == (3, 3)
+        and all(g.dtype == torch.float32 and g.ndim == 4 and g.shape[0] == 3 * T for g in ms_grids)
+    ):
+        return torch.cat(
+            [_kplanes_tilted_fuse_cuda(pts, rotation_matrices, g) for g in ms_grids], dim=-1
+        )
+
     # (T, B, 3)  — rotate all points by all rotations at once
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 

tests/ml/test_kplanes.py

@@ -18,3 +18,106 @@ 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))
+
+
+def _tilted_fused_available():
+    import torch
+
+    from quantem.core.ml.models import kplanes as kplanes_mod
+
+    return torch.cuda.is_available() and kplanes_mod._kplanes_tilted_fuse_cuda is not None
+
+
+class TestTiltedFusedDispatchParity:
+    """interpolate_ms_features_tilted must give the same result through the
+    fused quantem-cuda kernel and the torch fallback path."""
+
+    def _inputs(self, device="cuda", dtype=None, T=3):
+        import torch
+        from torch import nn
+
+        dtype = dtype or torch.float32
+        gen = torch.Generator(device=device).manual_seed(0)
+        pts = (
+            torch.empty(513, 3, device=device, dtype=torch.float32).uniform_(
+                -1.1, 1.1, generator=gen
+            )
+        ).to(dtype)
+        rotations = (
+            torch.empty(T, 3, 3, device=device, dtype=torch.float32).uniform_(
+                -1.0, 1.0, generator=gen
+            )
+        ).to(dtype)
+        grids = nn.ParameterList(
+            nn.Parameter(
+                torch.empty(3 * T, C, R, R, device=device, dtype=torch.float32)
+                .uniform_(0.1, 0.5, generator=gen)
+                .to(dtype)
+            )
+            for C, R in ((4, 16), (4, 32))
+        )
+        return pts, grids, rotations
+
+    @pytest.mark.skipif(
+        not _tilted_fused_available(), reason="requires a CUDA device and quantem-cuda"
+    )
+    def test_forward_matches_torch_path(self, monkeypatch):
+        import torch
+
+        from quantem.core.ml.models import kplanes as kplanes_mod
+        from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+        pts, grids, rotations = self._inputs()
+        out_fused = interpolate_ms_features_tilted(pts, grids, rotations)
+        monkeypatch.setattr(kplanes_mod, "_kplanes_tilted_fuse_cuda", None)
+        out_torch = interpolate_ms_features_tilted(pts, grids, rotations)
+        assert out_fused.shape == out_torch.shape
+        torch.testing.assert_close(out_fused, out_torch, rtol=1e-4, atol=5e-6)
+
+    @pytest.mark.skipif(
+        not _tilted_fused_available(), reason="requires a CUDA device and quantem-cuda"
+    )
+    def test_gradients_match_torch_path(self, monkeypatch):
+        import torch
+
+        from quantem.core.ml.models import kplanes as kplanes_mod
+        from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+        pts, grids, rotations = self._inputs()
+        pts.requires_grad_(True)
+        rotations.requires_grad_(True)
+        T = rotations.shape[0]
+        upstream = torch.randn(513, T * sum(g.shape[1] for g in grids), device="cuda")
+
+        def run():
+            for t in (pts, rotations, *grids):
+                t.grad = None
+            (interpolate_ms_features_tilted(pts, grids, rotations) * upstream).sum().backward()
+            return [pts.grad.clone(), rotations.grad.clone()] + [g.grad.clone() for g in grids]
+
+        g_fused = run()
+        monkeypatch.setattr(kplanes_mod, "_kplanes_tilted_fuse_cuda", None)
+        g_torch = run()
+        for gf, gt in zip(g_fused, g_torch):
+            torch.testing.assert_close(gf, gt, rtol=1e-3, atol=1e-5)
+
+    @pytest.mark.skipif(
+        not _tilted_fused_available(), reason="requires a CUDA device and quantem-cuda"
+    )
+    def test_non_fp32_falls_back(self):
+        import torch
+
+        from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+        pts, grids, rotations = self._inputs(dtype=torch.float64)
+        out = interpolate_ms_features_tilted(pts, grids, rotations)
+        assert out.dtype == torch.float64
+
+    def test_cpu_path_unaffected(self):
+        import torch
+
+        from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+        pts, grids, rotations = self._inputs(device="cpu")
+        out = interpolate_ms_features_tilted(pts, grids, rotations)
+        assert out.shape == (513, 3 * 8) and out.dtype == torch.float32