diff --git a/ci/pipeline.yml b/ci/pipeline.yml
index 285c34da..2d0a300b 100644
--- a/ci/pipeline.yml
+++ b/ci/pipeline.yml
@@ -35,5 +35,5 @@ test_job:
     SLURM_JOB_NUM_NODES: 1
     SLURM_PARTITION: normal
     SLURM_NTASKS: 1
-    SLURM_TIMELIMIT: '00:40:00'
+    SLURM_TIMELIMIT: '02:30:00'
     GIT_STRATEGY: fetch
diff --git a/setup.py b/setup.py
index 4af4b874..c9ccba30 100644
--- a/setup.py
+++ b/setup.py
@@ -139,7 +139,7 @@ def run(self):
             "sphericart": [
                 "sphericart/lib/*",
                 "sphericart/include/*",
-            ]
+            ],
         },
         extras_require=extras_require,
     )
diff --git a/sphericart-jax/CMakeLists.txt b/sphericart-jax/CMakeLists.txt
index 8842acfd..9a00a818 100644
--- a/sphericart-jax/CMakeLists.txt
+++ b/sphericart-jax/CMakeLists.txt
@@ -60,7 +60,7 @@ if(CMAKE_CUDA_COMPILER AND SPHERICART_ENABLE_CUDA)
     
     enable_language(CUDA)
     include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
-    set(CUDA_SOURCES ${CMAKE_CURRENT_LIST_DIR}/src/sphericart_jax_cuda.cu ${CMAKE_CURRENT_LIST_DIR}/src/sphericart_jax_cuda.cpp)
+    set(CUDA_SOURCES ${CMAKE_CURRENT_LIST_DIR}/src/sphericart_jax_cuda.cpp)
     pybind11_add_module(sphericart_jax_cuda ${CUDA_SOURCES})
     set_target_properties(sphericart_jax_cuda PROPERTIES CUDA_ARCHITECTURES native)
     install(TARGETS sphericart_jax_cuda DESTINATION sphericart_jax_cuda)
diff --git a/sphericart-jax/include/sphericart/sphericart_jax_cuda.hpp b/sphericart-jax/include/sphericart/sphericart_jax_cuda.hpp
deleted file mode 100644
index 7bb0b109..00000000
--- a/sphericart-jax/include/sphericart/sphericart_jax_cuda.hpp
+++ /dev/null
@@ -1,47 +0,0 @@
-// This file is needed as a workaround for pybind11 not accepting cuda files.
-// Note that all the templated functions are split into separate functions so
-// that they can be compiled in the `.cu` file.
-
-#ifndef _SPHERICART_JAX_CUDA_HPP_
-#define _SPHERICART_JAX_CUDA_HPP_
-
-#include <cuda_runtime.h>
-#include <cuda_runtime_api.h>
-
-struct SphDescriptor {
-    std::int64_t n_samples;
-    std::int64_t lmax;
-};
-
-namespace sphericart_jax {
-
-namespace cuda {
-
-void cuda_spherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_spherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_dspherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_dspherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_ddspherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_ddspherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_solid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_solid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_dsolid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_dsolid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_ddsolid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-void cuda_ddsolid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len);
-
-} // namespace cuda
-} // namespace sphericart_jax
-
-#endif
diff --git a/sphericart-jax/setup.py b/sphericart-jax/setup.py
index 9100b520..11fa1bb6 100644
--- a/sphericart-jax/setup.py
+++ b/sphericart-jax/setup.py
@@ -93,6 +93,6 @@ def run(self):
             "sphericart-jax": [
                 "sphericart/jax/lib/*",
                 "sphericart/jax/include/*",
-            ]
+            ],
         },
     )
diff --git a/sphericart-jax/src/sphericart_jax_cpu.cpp b/sphericart-jax/src/sphericart_jax_cpu.cpp
index 5571ff3e..b5335805 100644
--- a/sphericart-jax/src/sphericart_jax_cpu.cpp
+++ b/sphericart-jax/src/sphericart_jax_cpu.cpp
@@ -90,7 +90,6 @@ void cpu_sph_with_hessians(void* out_tuple, const void** in) {
 }
 
 // Registration of the custom calls with pybind11
-
 pybind11::dict Registrations() {
     pybind11::dict dict;
     dict["cpu_spherical_f32"] = EncapsulateFunction(cpu_sph<sphericart::SphericalHarmonics, float>);
diff --git a/sphericart-jax/src/sphericart_jax_cuda.cpp b/sphericart-jax/src/sphericart_jax_cuda.cpp
index 61824f53..23ba7f76 100644
--- a/sphericart-jax/src/sphericart_jax_cuda.cpp
+++ b/sphericart-jax/src/sphericart_jax_cuda.cpp
@@ -2,38 +2,116 @@
 // devices. It is exposed as a standard pybind11 module defining "capsule"
 // objects containing our methods. For simplicity, we export a separate capsule
 // for each supported dtype.
-// This file is separated from `sphericart_jax_cuda.cu` because pybind11 does
-// not accept cuda files.
 
 #include <cstdlib>
 #include <map>
 #include <mutex>
 #include <tuple>
 
+#include "sphericart_cuda.hpp"
 #include "sphericart/pybind11_kernel_helpers.hpp"
-#include "sphericart/sphericart_jax_cuda.hpp"
 
-using namespace sphericart_jax;
-using namespace sphericart_jax::cuda;
+struct SphDescriptor {
+    std::int64_t n_samples;
+    std::int64_t lmax;
+};
 
-namespace {
+namespace sphericart_jax {
+namespace cuda {
 
-// Registration of the custom calls with pybind11
+template <template <typename> class C, typename T>
+using CacheMapCUDA = std::map<size_t, std::unique_ptr<C<T>>>;
+
+template <template <typename> class C, typename T>
+std::unique_ptr<C<T>>& _get_or_create_sph_cuda(size_t l_max) {
+    // Static map to cache instances based on parameters
+    static CacheMapCUDA<C, T> sph_cache;
+    static std::mutex cache_mutex;
+
+    // Check if instance exists in cache, if not create and store it
+    std::lock_guard<std::mutex> lock(cache_mutex);
+    auto it = sph_cache.find(l_max);
+    if (it == sph_cache.end()) {
+        it = sph_cache.insert({l_max, std::make_unique<C<T>>(l_max)}).first;
+    }
+    return it->second;
+}
+
+template <template <typename> class C, typename T>
+inline void cuda_sph(void* stream, void** in, const char* opaque, std::size_t opaque_len) {
+    // Parse the inputs
+    const T* xyz = reinterpret_cast<const T*>(in[0]);
+    T* sph = reinterpret_cast<T*>(in[1]);
+
+    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
+    const std::int64_t n_samples = d.n_samples;
+    const std::int64_t lmax = d.lmax;
+
+    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
+
+    calculator->compute(xyz, n_samples, sph, stream);
+}
 
+template <template <typename> class C, typename T>
+inline void cuda_sph_with_gradients(
+    void* stream, void** in, const char* opaque, std::size_t opaque_len
+) {
+    // Parse the inputs
+    const T* xyz = reinterpret_cast<const T*>(in[0]);
+    T* sph = reinterpret_cast<T*>(in[1]);
+    T* dsph = reinterpret_cast<T*>(in[2]);
+
+    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
+    const std::int64_t n_samples = d.n_samples;
+    const std::int64_t lmax = d.lmax;
+
+    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
+    calculator->compute_with_gradients(xyz, n_samples, sph, dsph, stream);
+}
+
+template <template <typename> class C, typename T>
+inline void cuda_sph_with_hessians(
+    void* stream, void** in, const char* opaque, std::size_t opaque_len
+) {
+    // Parse the inputs
+    const T* xyz = reinterpret_cast<const T*>(in[0]);
+    T* sph = reinterpret_cast<T*>(in[1]);
+    T* dsph = reinterpret_cast<T*>(in[2]);
+    T* ddsph = reinterpret_cast<T*>(in[3]);
+
+    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
+    const std::int64_t n_samples = d.n_samples;
+    const std::int64_t lmax = d.lmax;
+
+    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
+    calculator->compute_with_hessians(xyz, n_samples, sph, dsph, ddsph, stream);
+}
+
+// Registration of the custom calls with pybind11
 pybind11::dict Registrations() {
     pybind11::dict dict;
-    dict["cuda_spherical_f32"] = EncapsulateFunction(cuda_spherical_f32);
-    dict["cuda_spherical_f64"] = EncapsulateFunction(cuda_spherical_f64);
-    dict["cuda_dspherical_f32"] = EncapsulateFunction(cuda_dspherical_f32);
-    dict["cuda_dspherical_f64"] = EncapsulateFunction(cuda_dspherical_f64);
-    dict["cuda_ddspherical_f32"] = EncapsulateFunction(cuda_ddspherical_f32);
-    dict["cuda_ddspherical_f64"] = EncapsulateFunction(cuda_ddspherical_f64);
-    dict["cuda_solid_f32"] = EncapsulateFunction(cuda_solid_f32);
-    dict["cuda_solid_f64"] = EncapsulateFunction(cuda_solid_f64);
-    dict["cuda_dsolid_f32"] = EncapsulateFunction(cuda_dsolid_f32);
-    dict["cuda_dsolid_f64"] = EncapsulateFunction(cuda_dsolid_f64);
-    dict["cuda_ddsolid_f32"] = EncapsulateFunction(cuda_ddsolid_f32);
-    dict["cuda_ddsolid_f64"] = EncapsulateFunction(cuda_ddsolid_f64);
+    dict["cuda_spherical_f32"] =
+        EncapsulateFunction(cuda_sph<sphericart::cuda::SphericalHarmonics, float>);
+    dict["cuda_spherical_f64"] =
+        EncapsulateFunction(cuda_sph<sphericart::cuda::SphericalHarmonics, double>);
+    dict["cuda_dspherical_f32"] =
+        EncapsulateFunction(cuda_sph_with_gradients<sphericart::cuda::SphericalHarmonics, float>);
+    dict["cuda_dspherical_f64"] =
+        EncapsulateFunction(cuda_sph_with_gradients<sphericart::cuda::SphericalHarmonics, double>);
+    dict["cuda_ddspherical_f32"] =
+        EncapsulateFunction(cuda_sph_with_hessians<sphericart::cuda::SphericalHarmonics, float>);
+    dict["cuda_ddspherical_f64"] =
+        EncapsulateFunction(cuda_sph_with_hessians<sphericart::cuda::SphericalHarmonics, double>);
+    dict["cuda_solid_f32"] = EncapsulateFunction(cuda_sph<sphericart::cuda::SolidHarmonics, float>);
+    dict["cuda_solid_f64"] = EncapsulateFunction(cuda_sph<sphericart::cuda::SolidHarmonics, double>);
+    dict["cuda_dsolid_f32"] =
+        EncapsulateFunction(cuda_sph_with_gradients<sphericart::cuda::SolidHarmonics, float>);
+    dict["cuda_dsolid_f64"] =
+        EncapsulateFunction(cuda_sph_with_gradients<sphericart::cuda::SolidHarmonics, double>);
+    dict["cuda_ddsolid_f32"] =
+        EncapsulateFunction(cuda_sph_with_hessians<sphericart::cuda::SolidHarmonics, float>);
+    dict["cuda_ddsolid_f64"] =
+        EncapsulateFunction(cuda_sph_with_hessians<sphericart::cuda::SolidHarmonics, double>);
     return dict;
 }
 
@@ -44,4 +122,5 @@ PYBIND11_MODULE(sphericart_jax_cuda, m) {
     });
 }
 
-} // namespace
+} // namespace cuda
+} // namespace sphericart_jax
diff --git a/sphericart-jax/src/sphericart_jax_cuda.cu b/sphericart-jax/src/sphericart_jax_cuda.cu
deleted file mode 100644
index d77ceb41..00000000
--- a/sphericart-jax/src/sphericart_jax_cuda.cu
+++ /dev/null
@@ -1,148 +0,0 @@
-// This file defines the Python interface to the XLA custom calls on CUDA
-// devices. It is exposed as a standard pybind11 module defining "capsule"
-// objects containing our methods. For simplicity, we export a separate capsule
-// for each supported dtype.
-
-#include <cstdlib>
-#include <map>
-#include <mutex>
-#include <tuple>
-
-#include <cuda_runtime.h>
-#include <cuda_runtime_api.h>
-
-#include "sphericart_cuda.hpp"
-#include "sphericart/pybind11_kernel_helpers.hpp"
-#include "sphericart/sphericart_jax_cuda.hpp"
-
-namespace sphericart_jax {
-namespace cuda {
-
-template <template <typename> class C, typename T>
-using CacheMapCUDA = std::map<size_t, std::unique_ptr<C<T>>>;
-
-template <template <typename> class C, typename T>
-std::unique_ptr<C<T>>& _get_or_create_sph_cuda(size_t l_max) {
-    // Static map to cache instances based on parameters
-    static CacheMapCUDA<C, T> sph_cache;
-    static std::mutex cache_mutex;
-
-    // Check if instance exists in cache, if not create and store it
-    std::lock_guard<std::mutex> lock(cache_mutex);
-    auto it = sph_cache.find(l_max);
-    if (it == sph_cache.end()) {
-        it = sph_cache.insert({l_max, std::make_unique<C<T>>(l_max)}).first;
-    }
-    return it->second;
-}
-
-template <template <typename> class C, typename T>
-inline void cuda_sph(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    // Parse the inputs
-    const T* xyz = reinterpret_cast<const T*>(in[0]);
-    T* sph = reinterpret_cast<T*>(in[1]);
-
-    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
-    const std::int64_t n_samples = d.n_samples;
-    const std::int64_t lmax = d.lmax;
-
-    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
-
-    calculator->compute(xyz, n_samples, sph, reinterpret_cast<void*>(stream));
-}
-
-template <template <typename> class C, typename T>
-inline void cuda_sph_with_gradients(
-    cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len
-) {
-    // Parse the inputs
-    const T* xyz = reinterpret_cast<const T*>(in[0]);
-    T* sph = reinterpret_cast<T*>(in[1]);
-    T* dsph = reinterpret_cast<T*>(in[2]);
-
-    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
-    const std::int64_t n_samples = d.n_samples;
-    const std::int64_t lmax = d.lmax;
-
-    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
-    calculator->compute_with_gradients(xyz, n_samples, sph, dsph, reinterpret_cast<void*>(stream));
-}
-
-template <template <typename> class C, typename T>
-inline void cuda_sph_with_hessians(
-    cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len
-) {
-    // Parse the inputs
-    const T* xyz = reinterpret_cast<const T*>(in[0]);
-    T* sph = reinterpret_cast<T*>(in[1]);
-    T* dsph = reinterpret_cast<T*>(in[2]);
-    T* ddsph = reinterpret_cast<T*>(in[3]);
-
-    const SphDescriptor& d = *UnpackDescriptor<SphDescriptor>(opaque, opaque_len);
-    const std::int64_t n_samples = d.n_samples;
-    const std::int64_t lmax = d.lmax;
-
-    auto& calculator = _get_or_create_sph_cuda<C, T>(lmax);
-    calculator->compute_with_hessians(
-        xyz, n_samples, sph, dsph, ddsph, reinterpret_cast<void*>(stream)
-    );
-}
-
-void cuda_spherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph<sphericart::cuda::SphericalHarmonics, float>(stream, in, opaque, opaque_len);
-}
-
-void cuda_spherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph<sphericart::cuda::SphericalHarmonics, double>(stream, in, opaque, opaque_len);
-}
-
-void cuda_dspherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_gradients<sphericart::cuda::SphericalHarmonics, float>(
-        stream, in, opaque, opaque_len
-    );
-}
-
-void cuda_dspherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_gradients<sphericart::cuda::SphericalHarmonics, double>(
-        stream, in, opaque, opaque_len
-    );
-}
-
-void cuda_ddspherical_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_hessians<sphericart::cuda::SphericalHarmonics, float>(
-        stream, in, opaque, opaque_len
-    );
-}
-
-void cuda_ddspherical_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_hessians<sphericart::cuda::SphericalHarmonics, double>(
-        stream, in, opaque, opaque_len
-    );
-}
-
-void cuda_solid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph<sphericart::cuda::SolidHarmonics, float>(stream, in, opaque, opaque_len);
-}
-
-void cuda_solid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph<sphericart::cuda::SolidHarmonics, double>(stream, in, opaque, opaque_len);
-}
-
-void cuda_dsolid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_gradients<sphericart::cuda::SolidHarmonics, float>(stream, in, opaque, opaque_len);
-}
-
-void cuda_dsolid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_gradients<sphericart::cuda::SolidHarmonics, double>(stream, in, opaque, opaque_len);
-}
-
-void cuda_ddsolid_f32(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_hessians<sphericart::cuda::SolidHarmonics, float>(stream, in, opaque, opaque_len);
-}
-
-void cuda_ddsolid_f64(cudaStream_t stream, void** in, const char* opaque, std::size_t opaque_len) {
-    cuda_sph_with_hessians<sphericart::cuda::SolidHarmonics, double>(stream, in, opaque, opaque_len);
-}
-
-} // namespace cuda
-} // namespace sphericart_jax
diff --git a/sphericart-torch/python/tests/test_autograd.py b/sphericart-torch/python/tests/test_autograd.py
index 0723378e..855bc54f 100644
--- a/sphericart-torch/python/tests/test_autograd.py
+++ b/sphericart-torch/python/tests/test_autograd.py
@@ -18,8 +18,8 @@ def test_autograd_cartesian(xyz):
         l_max=4, backward_second_derivatives=True
     )
 
-    def compute(xyz):
-        sph = calculator.compute(xyz=xyz)
+    def compute(xyz_in):
+        sph = calculator.compute(xyz=xyz_in)
         assert torch.linalg.norm(sph) != 0.0
         return sph
 
@@ -27,9 +27,9 @@ def compute(xyz):
     assert torch.autograd.gradgradcheck(compute, xyz, fast_mode=True)
 
     if torch.cuda.is_available():
-        xyz = xyz.to(device="cuda")
-        assert torch.autograd.gradcheck(compute, xyz, fast_mode=True)
-        assert torch.autograd.gradgradcheck(compute, xyz, fast_mode=True)
+        xyz_cuda = xyz.detach().cuda().requires_grad_(True)
+        assert torch.autograd.gradcheck(compute, xyz_cuda, fast_mode=True)
+        assert torch.autograd.gradgradcheck(compute, xyz_cuda, fast_mode=True)
 
 
 def test_autograd_normalized(xyz):
@@ -37,8 +37,8 @@ def test_autograd_normalized(xyz):
         l_max=4, backward_second_derivatives=True
     )
 
-    def compute(xyz):
-        sph = calculator.compute(xyz=xyz)
+    def compute(xyz_in):
+        sph = calculator.compute(xyz=xyz_in)
         assert torch.linalg.norm(sph) != 0.0
         return sph
 
@@ -46,9 +46,9 @@ def compute(xyz):
     assert torch.autograd.gradgradcheck(compute, xyz, fast_mode=True)
 
     if torch.cuda.is_available():
-        xyz = xyz.to(device="cuda")
-        assert torch.autograd.gradcheck(compute, xyz, fast_mode=True)
-        assert torch.autograd.gradgradcheck(compute, xyz, fast_mode=True)
+        xyz_cuda = xyz.detach().cuda().requires_grad_(True)
+        assert torch.autograd.gradcheck(compute, xyz_cuda, fast_mode=True)
+        assert torch.autograd.gradgradcheck(compute, xyz_cuda, fast_mode=True)
 
 
 def test_autograd_hessian(xyz):
diff --git a/sphericart-torch/python/tests/test_cuda.py b/sphericart-torch/python/tests/test_cuda.py
index 706d0b27..f58ecbab 100644
--- a/sphericart-torch/python/tests/test_cuda.py
+++ b/sphericart-torch/python/tests/test_cuda.py
@@ -20,10 +20,6 @@ def test_cpu_vs_cuda(xyz):
         sph, grad_sph = calculator.compute_with_gradients(xyz)
         sph_cuda, grad_sph_cuda = calculator.compute_with_gradients(xyz_cuda)
 
-        print(xyz[0])
-        print((sph - sph_cuda.to("cpu")).flatten().max())
-        print((sph - sph_cuda.to("cpu"))[0])
-        print((sph - sph_cuda.to("cpu"))[1])
         assert torch.allclose(sph, sph_cuda.to("cpu"))
         assert torch.allclose(grad_sph, grad_sph_cuda.to("cpu"))
 
diff --git a/sphericart-torch/setup.py b/sphericart-torch/setup.py
index cd479054..b23abeff 100644
--- a/sphericart-torch/setup.py
+++ b/sphericart-torch/setup.py
@@ -89,6 +89,6 @@ def run(self):
             "sphericart-torch": [
                 "sphericart/torch/lib/*",
                 "sphericart/torch/include/*",
-            ]
+            ],
         },
     )
diff --git a/sphericart/CMakeLists.txt b/sphericart/CMakeLists.txt
index e7b6f8c0..a5092bbf 100644
--- a/sphericart/CMakeLists.txt
+++ b/sphericart/CMakeLists.txt
@@ -2,6 +2,50 @@ cmake_minimum_required(VERSION 3.10)
 
 project(sphericart LANGUAGES C CXX)
 
+#[[
+ This function wraps the input file with the appropriate string to allow us to 
+ static-initialize string variables, for example:
+     static const char* CUDA_CODE =
+#include "generated/wrapped_sphericart_impl.cu"
+        ;
+]]
+function(make_includeable INPUT_FILE OUTPUT_FILE)
+    if(NOT EXISTS ${INPUT_FILE})
+        message(FATAL_ERROR "Error: The input file '${INPUT_FILE}' does not exist.")
+    endif()
+    file(READ ${INPUT_FILE} content)
+    # Format the content to be included as a raw string in C++
+    set(content "R\"======(\n${content}\n)======\"")
+    # Write the formatted content to the output file
+    file(WRITE ${OUTPUT_FILE} "${content}")
+endfunction()
+
+#[[
+ This function wraps prepends a list of headers to an input file. Used as a preprocessing
+ step to prepare the sphericart_impl.cu file for NVRTC JIT compilation.
+]]
+function(prepend_headers_to_source INPUT_FILE OUTPUT_FILE HEADERS)
+    # Read the current content of the source file
+    if(NOT EXISTS ${INPUT_FILE})
+        message(FATAL_ERROR "Error: The input file '${INPUT_FILE}' does not exist.")
+    endif()
+    file(READ ${INPUT_FILE} CONTENTS)
+    # Initialize a variable to hold the accumulated headers
+    set(HEADER_CONTENTS "")
+    # Loop through each header and append its contents to HEADER_CONTENTS
+    foreach(HEADER IN LISTS HEADERS)
+        if(NOT EXISTS ${INPUT_FILE})
+            message(FATAL_ERROR "Error: The input file '${INPUT_FILE}' does not exist.")
+        endif()
+        file(READ ${HEADER} SINGLE_HEADER_CONTENT)
+        set(HEADER_CONTENTS "${HEADER_CONTENTS}\n${SINGLE_HEADER_CONTENT}")
+    endforeach()
+    # Prepend the accumulated headers to the original source file content
+    set(NEW_CONTENTS "${HEADER_CONTENTS}\n${CONTENTS}")
+    # Write the new contents back to the source file
+    file(WRITE ${OUTPUT_FILE} "${NEW_CONTENTS}")
+endfunction()
+
 file(READ ${PROJECT_SOURCE_DIR}/VERSION SPHERICART_VERSION)
 string(STRIP ${SPHERICART_VERSION} SPHERICART_VERSION)
 string(REGEX REPLACE "^([0-9]+)\\..*" "\\1" SPHERICART_VERSION_MAJOR "${SPHERICART_VERSION}")
@@ -43,6 +87,7 @@ include(CheckLanguage)
 check_language(CUDA)
 if(CMAKE_CUDA_COMPILER)
     enable_language(CUDA)
+    find_package(CUDAToolkit REQUIRED)
     set(CUDA_USE_STATIC_CUDA_RUNTIME OFF CACHE BOOL "" FORCE)
 else()
     message(STATUS "Could not find a CUDA compiler")
@@ -57,8 +102,10 @@ list(APPEND COMMON_SOURCES "include/cuda_base.hpp")
 list(APPEND COMMON_SOURCES "include/sphericart_cuda.hpp")
 
 if (CMAKE_CUDA_COMPILER AND SPHERICART_ENABLE_CUDA)
-    list(APPEND COMMON_SOURCES "src/cuda_base.cu")
-    list(APPEND COMMON_SOURCES "src/sphericart_cuda.cu")
+    list(APPEND COMMON_SOURCES "include/dynamic_cuda.hpp")
+    list(APPEND COMMON_SOURCES "include/cuda_cache.hpp")
+    list(APPEND COMMON_SOURCES "src/cuda_base.cpp")
+    list(APPEND COMMON_SOURCES "src/sphericart_cuda.cpp")
 else()
     list(APPEND COMMON_SOURCES "src/cuda_stub.cpp")
     list(APPEND COMMON_SOURCES "src/sphericart_cuda_stub.cpp")
@@ -67,10 +114,33 @@ endif()
 add_library(sphericart ${COMMON_SOURCES})
 
 if (CMAKE_CUDA_COMPILER AND SPHERICART_ENABLE_CUDA)
-    set_target_properties(sphericart PROPERTIES CUDA_ARCHITECTURES all)
+    #include the build/generated folder for compilation
+    target_include_directories(sphericart PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
+    target_include_directories(sphericart PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
+    target_link_libraries(sphericart PRIVATE ${CMAKE_DL_LIBS})
+
+    set(SPHERICART_CUDA_HEADERS
+        "include/sphericart_impl.cuh"
+        "include/templates_core.hpp"
+        "include/macros.hpp"
+    )
+
+    # Prepend headers to the source file
+    prepend_headers_to_source(
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/sphericart_impl.cu"
+        "${CMAKE_CURRENT_BINARY_DIR}/generated/tmp.cu" 
+        "${SPHERICART_CUDA_HEADERS}"
+        )
+
+    # Make the source file includeable by converting it to a string literal
+    make_includeable(
+        "${CMAKE_CURRENT_BINARY_DIR}/generated/tmp.cu"
+        "${CMAKE_CURRENT_BINARY_DIR}/generated/wrapped_sphericart_impl.cu"
+    )
+
+    file(REMOVE "${CMAKE_CURRENT_BINARY_DIR}/generated/tmp.cuh")
 endif()
 
-
 set_target_properties(sphericart PROPERTIES
     VERSION ${SPHERICART_VERSION}
     SOVERSION ${SPHERICART_VERSION_MAJOR}.${SPHERICART_VERSION_MINOR}
@@ -191,4 +261,4 @@ install(FILES "${PROJECT_BINARY_DIR}/sphericart-config-version.cmake"
         DESTINATION ${LIB_INSTALL_DIR}/cmake/sphericart)
 
 install(DIRECTORY ${PROJECT_SOURCE_DIR}/include/ DESTINATION ${INCLUDE_INSTALL_DIR})
-install(DIRECTORY ${PROJECT_BINARY_DIR}/include/ DESTINATION ${INCLUDE_INSTALL_DIR})
+install(DIRECTORY ${PROJECT_BINARY_DIR}/include/ DESTINATION ${INCLUDE_INSTALL_DIR})
\ No newline at end of file
diff --git a/sphericart/include/cuda_base.hpp b/sphericart/include/cuda_base.hpp
index e5609a6b..1c938f7c 100644
--- a/sphericart/include/cuda_base.hpp
+++ b/sphericart/include/cuda_base.hpp
@@ -1,19 +1,6 @@
 #ifndef SPHERICART_CUDA_BASE_HPP
 #define SPHERICART_CUDA_BASE_HPP
 
-#include "sphericart.hpp"
-
-#define CUDA_CHECK(call)                                                                           \
-    do {                                                                                           \
-        cudaError_t cudaStatus = (call);                                                           \
-        if (cudaStatus != cudaSuccess) {                                                           \
-            std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << " - "                  \
-                      << cudaGetErrorString(cudaStatus) << std::endl;                              \
-            cudaDeviceReset();                                                                     \
-            exit(EXIT_FAILURE);                                                                    \
-        }                                                                                          \
-    } while (0)
-
 namespace sphericart {
 
 namespace cuda {
@@ -66,9 +53,9 @@ namespace cuda {
  */
 template <typename scalar_t>
 void spherical_harmonics_cuda_base(
-    const scalar_t* __restrict__ xyz,
+    const scalar_t* xyz,
     const int nedges,
-    const scalar_t* __restrict__ prefactors,
+    const scalar_t* prefactors,
     const int nprefactors,
     const int64_t l_max,
     const bool normalize,
@@ -76,19 +63,19 @@ void spherical_harmonics_cuda_base(
     const int64_t GRID_DIM_Y,
     const bool gradients,
     const bool hessian,
-    scalar_t* __restrict__ sph,
-    scalar_t* __restrict__ dsph,
-    scalar_t* __restrict__ ddsph,
+    scalar_t* sph,
+    scalar_t* dsph,
+    scalar_t* ddsph,
     void* cuda_stream = nullptr
 );
 
 template <typename scalar_t>
 void spherical_harmonics_backward_cuda_base(
-    const scalar_t* __restrict__ dsph,
-    const scalar_t* __restrict__ sph_grad,
+    const scalar_t* dsph,
+    const scalar_t* sph_grad,
     const int nedges,
     const int ntotal,
-    scalar_t* __restrict__ xyz_grad,
+    scalar_t* xyz_grad,
     void* cuda_stream = nullptr
 );
 
@@ -114,15 +101,6 @@ void spherical_harmonics_backward_cuda_base(
  * @param current_shared_mem_alloc
  *        the current size of the shared memory allocation.
  */
-int adjust_shared_memory(
-    size_t element_size,
-    int64_t l_max,
-    int64_t GRID_DIM_X,
-    int64_t GRID_DIM_Y,
-    bool requires_grad,
-    bool requires_hessian,
-    int64_t current_shared_mem_alloc
-);
 
 } // namespace cuda
 } // namespace sphericart
diff --git a/sphericart/include/cuda_cache.hpp b/sphericart/include/cuda_cache.hpp
new file mode 100644
index 00000000..563bf5cb
--- /dev/null
+++ b/sphericart/include/cuda_cache.hpp
@@ -0,0 +1,431 @@
+// Cache Manager Header
+#ifndef SPHERICART_CUDA_CACHE_HPP
+#define SPHERICART_CUDA_CACHE_HPP
+
+#include <fstream>
+#include <vector>
+#include <unordered_map>
+#include <string>
+#include <memory>
+#include <typeinfo>
+
+#include <nvrtc.h>
+#include <cuda.h>
+#include <cxxabi.h>
+
+#include "dynamic_cuda.hpp"
+
+// Helper function to demangle the type name if necessary
+std::string demangleTypeName(const std::string& name) {
+#if defined(__GNUC__) || defined(__clang__)
+    int status = 0;
+    std::unique_ptr<char, void (*)(void*)> demangled_name(
+        abi::__cxa_demangle(name.c_str(), 0, 0, &status), std::free
+    );
+    return (status == 0) ? demangled_name.get() : name;
+#else
+    throw std::runtime_error("demangling not supported using this toolchain.");
+#endif
+}
+
+// Base case: No template arguments, return function name without any type information
+std::string getKernelName(const std::string& fn_name) { return fn_name; }
+
+// Function to get type name of a single type
+template <typename T> std::string typeName() { return demangleTypeName(typeid(T).name()); }
+
+// Variadic template function to build type list
+template <typename T, typename... Ts> void buildTemplateTypes(std::string& base) {
+    base += typeName<T>(); // Add the first type
+    // If there are more types, add a comma and recursively call for the remaining types
+    if constexpr (sizeof...(Ts) > 0) {
+        base += ", ";
+        buildTemplateTypes<Ts...>(base); // Recursively call for the rest of the types
+    }
+}
+
+// Helper function to start building the types
+template <typename T, typename... Ts> std::string buildTemplateTypes() {
+    std::string result;
+    buildTemplateTypes<T, Ts...>(result); // Use recursive variadic template
+    return result;
+}
+
+/*
+Function to get the kernel name with the list of templated types if any:
+*/
+template <typename T, typename... Ts> std::string getKernelName(const std::string& fn_name) {
+    std::string type_list = buildTemplateTypes<T, Ts...>(); // Build type list
+    return fn_name + "<" + type_list + ">"; // Return function name with type list in angle brackets
+}
+
+// Function to load CUDA source code from a file
+std::string load_cuda_source(const std::string& filename) {
+    std::ifstream file(filename);
+    if (!file.is_open()) {
+        throw std::runtime_error("Failed to open file: " + filename);
+    }
+    std::ostringstream ss;
+    ss << file.rdbuf();
+    return ss.str();
+}
+
+/*
+Container class for the cached kernels. Provides functionality for launching compiled kernels as
+well as automatically resizing dynamic shared memory allocations, when needed. Kernels are compiled
+on first launch.
+*/
+class CachedKernel {
+
+  public:
+    CachedKernel(
+        std::string kernel_name,
+        std::string kernel_code,
+        std::string source_name,
+        std::vector<std::string> options
+    ) {
+        this->kernel_name = kernel_name;
+        this->kernel_code = kernel_code;
+        this->source_name = source_name;
+        this->options = options;
+    }
+
+    CachedKernel() = default;
+
+    // Copy constructor
+    CachedKernel(const CachedKernel&) = default;
+
+    // Copy assignment operator
+    CachedKernel& operator=(const CachedKernel&) = default;
+
+    inline void setFuncAttribute(CUfunction_attribute attribute, int value) const {
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuFuncSetAttribute(function, attribute, value));
+    }
+
+    int getFuncAttribute(CUfunction_attribute attribute) const {
+        int value;
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuFuncGetAttribute(&value, attribute, function));
+        return value;
+    }
+
+    /*
+    launches the kernel, and optionally synchronizes until control can be passed back to host.
+    */
+    void launch(
+        dim3 grid,
+        dim3 block,
+        size_t shared_mem_size,
+        void* cuda_stream,
+        std::vector<void*> args,
+        bool synchronize = true
+    ) {
+
+        if (!compiled) {
+            this->compileKernel(args);
+        }
+
+        CUcontext currentContext = nullptr;
+        // Get current context
+        CUresult result = CUDA_DRIVER_INSTANCE.cuCtxGetCurrent(&currentContext);
+
+        if (result != CUDA_SUCCESS || !currentContext) {
+            throw std::runtime_error("CachedKernel::launch error getting current context.");
+        }
+
+        if (currentContext != context) {
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxSetCurrent(context));
+        }
+
+        this->checkAndAdjustSharedMem(shared_mem_size);
+
+        cudaStream_t cstream = reinterpret_cast<cudaStream_t>(cuda_stream);
+
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuLaunchKernel(
+            function,
+            grid.x,
+            grid.y,
+            grid.z,
+            block.x,
+            block.y,
+            block.z,
+            shared_mem_size,
+            cstream,
+            args.data(),
+            0
+        ));
+
+        if (synchronize) {
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxSynchronize());
+        }
+
+        if (currentContext != context) {
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxSetCurrent(currentContext));
+        }
+    }
+
+  private:
+    /*
+    The default shared memory space on most recent NVIDIA cards is defaulted
+    49152 bytes. This method
+    attempts to adjust the shared memory to fit the requested configuration if
+    the kernel launch parameters exceeds the default 49152 bytes.
+    */
+    void checkAndAdjustSharedMem(int query_shared_mem_size) {
+        if (current_smem_size == 0) {
+            CUdevice cuDevice;
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxGetDevice(&cuDevice));
+
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuDeviceGetAttribute(
+                &max_smem_size_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, cuDevice
+            ));
+
+            int reserved_smem_per_block = 0;
+
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuDeviceGetAttribute(
+                &reserved_smem_per_block, CU_DEVICE_ATTRIBUTE_RESERVED_SHARED_MEMORY_PER_BLOCK, cuDevice
+            ));
+
+            int curr_max_smem_per_block = 0;
+
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuDeviceGetAttribute(
+                &curr_max_smem_per_block, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, cuDevice
+            ));
+
+            current_smem_size = (curr_max_smem_per_block - reserved_smem_per_block);
+        }
+
+        if (query_shared_mem_size > current_smem_size) {
+
+            if (query_shared_mem_size > max_smem_size_optin) {
+                throw std::runtime_error(
+                    "CachedKernel::launch requested more smem than available on card."
+                );
+            } else {
+                CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuFuncSetAttribute(
+                    function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, query_shared_mem_size
+                ));
+                current_smem_size = query_shared_mem_size;
+            }
+        }
+    }
+
+    /*
+        Compiles the kernel "kernel_name" located in source file "kernel_code", which additional
+        parameters "options" passed to NVRTC_INSTANCE. Will auto-detect the compute capability of
+       the available card. args for the launch need to be queried as we need to grab the CUcontext
+       in which these ptrs exist.
+        */
+    void compileKernel(std::vector<void*>& kernel_args) {
+
+        this->initCudaDriver();
+
+        CUcontext currentContext = nullptr;
+
+        for (size_t ptr_id = 0; ptr_id < kernel_args.size(); ptr_id++) {
+            unsigned int memtype = 0;
+            CUdeviceptr device_ptr = *reinterpret_cast<CUdeviceptr*>(kernel_args[ptr_id]);
+
+            CUresult res = CUDA_DRIVER_INSTANCE.cuPointerGetAttribute(
+                &memtype, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, device_ptr
+            );
+
+            if (res == CUDA_SUCCESS && memtype == CU_MEMORYTYPE_DEVICE) {
+                CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuPointerGetAttribute(
+                    &currentContext, CU_POINTER_ATTRIBUTE_CONTEXT, device_ptr
+                ));
+
+                if (currentContext) {
+                    break;
+                }
+            }
+        }
+
+        CUcontext query = nullptr;
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxGetCurrent(&query));
+
+        if (query != currentContext) {
+            CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxSetCurrent(currentContext));
+        }
+
+        CUdevice cuDevice;
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuCtxGetDevice(&cuDevice));
+
+        nvrtcProgram prog;
+
+        NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcCreateProgram(
+            &prog, this->kernel_code.c_str(), this->source_name.c_str(), 0, nullptr, nullptr
+        ));
+
+        NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcAddNameExpression(prog, this->kernel_name.c_str()));
+
+        std::vector<const char*> c_options;
+        c_options.reserve(this->options.size());
+        for (const auto& option : this->options) {
+            c_options.push_back(option.c_str());
+        }
+
+        int major = 0;
+        int minor = 0;
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuDeviceGetAttribute(
+            &major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, cuDevice
+        ));
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuDeviceGetAttribute(
+            &minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, cuDevice
+        ));
+        int arch = major * 10 + minor;
+        std::string smbuf = "--gpu-architecture=sm_" + std::to_string(arch);
+        c_options.push_back(smbuf.c_str());
+
+        nvrtcResult compileResult =
+            NVRTC_INSTANCE.nvrtcCompileProgram(prog, c_options.size(), c_options.data());
+        if (compileResult != NVRTC_SUCCESS) {
+            size_t logSize;
+            NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcGetProgramLogSize(prog, &logSize));
+            std::string log(logSize, '\0');
+            NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcGetProgramLog(prog, &log[0]));
+            throw std::runtime_error(
+                "KernelFactory::compileAndCacheKernel: Failed to compile CUDA program:\n" + log
+            );
+        }
+
+        // Get PTX code
+        size_t ptxSize;
+        NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcGetPTXSize(prog, &ptxSize));
+        std::vector<char> ptxCode(ptxSize);
+        NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcGetPTX(prog, ptxCode.data()));
+
+        CUmodule module;
+
+        CUresult cuResult =
+            CUDA_DRIVER_INSTANCE.cuModuleLoadDataEx(&module, ptxCode.data(), 0, 0, 0);
+
+        if (cuResult != CUDA_SUCCESS) {
+            throw std::runtime_error(
+                "KernelFactory::compileAndCacheKernel: Failed to load PTX code into CUDA "
+                "module "
+                "(error code: " +
+                std::to_string(cuResult) + ")"
+            );
+        }
+
+        const char* lowered_name;
+        NVRTC_SAFE_CALL(
+            NVRTC_INSTANCE.nvrtcGetLoweredName(prog, this->kernel_name.c_str(), &lowered_name)
+        );
+        CUfunction kernel;
+        CUDADRIVER_SAFE_CALL(CUDA_DRIVER_INSTANCE.cuModuleGetFunction(&kernel, module, lowered_name)
+        );
+
+        this->module = module;
+        this->function = kernel;
+        this->context = currentContext;
+        this->compiled = true;
+
+        NVRTC_SAFE_CALL(NVRTC_INSTANCE.nvrtcDestroyProgram(&prog));
+    }
+
+    void initCudaDriver() {
+
+        int deviceCount = 0;
+        // Check if CUDA has already been initialized
+        CUresult res = CUDA_DRIVER_INSTANCE.cuDeviceGetCount(&deviceCount);
+        if (res == CUDA_ERROR_NOT_INITIALIZED) {
+            // CUDA hasn't been initialized, so we initialize it now
+            res = CUDA_DRIVER_INSTANCE.cuInit(0);
+            if (res != CUDA_SUCCESS) {
+                throw std::runtime_error(
+                    "KernelFactory::initCudaDriver: Failed to initialize CUDA CUDA_DRIVER_INSTANCE."
+                );
+                return;
+            }
+        }
+    }
+
+    int current_smem_size = 0;
+    int max_smem_size_optin = 0;
+    CUmodule module = nullptr;
+    CUfunction function = nullptr;
+    CUcontext context = nullptr;
+    bool compiled = false;
+
+    std::string kernel_name;
+    std::string kernel_code;
+    std::string source_name;
+    std::vector<std::string> options;
+};
+
+/*
+Factory class to create and store compiled cuda kernels for caching as a simple name-based hashmap.
+Allows both compiling from a source file, or for compiling from a variable containing CUDA code.
+*/
+class KernelFactory {
+
+  public:
+    static KernelFactory& instance() {
+        static KernelFactory instance;
+        return instance;
+    }
+
+    void cacheKernel(
+        const std::string& kernel_name,
+        const std::string& source_path,
+        const std::string& source_name,
+        const std::vector<std::string>& options
+    ) {
+        kernel_cache[kernel_name] =
+            std::make_unique<CachedKernel>(kernel_name, source_path, source_name, options);
+    }
+
+    bool hasKernel(const std::string& kernel_name) const {
+        return kernel_cache.find(kernel_name) != kernel_cache.end();
+    }
+
+    CachedKernel* getKernel(const std::string& kernel_name) const {
+        auto it = kernel_cache.find(kernel_name);
+        if (it != kernel_cache.end()) {
+            return it->second.get();
+        }
+        throw std::runtime_error("Kernel not found in cache.");
+    }
+
+    /*
+    Tries to retrieve the kernel "kernel_name". If not found, instantiate it and save to cache.
+    */
+    CachedKernel* createFromSource(
+        const std::string& kernel_name,
+        const std::string& source_path,
+        const std::string& source_name,
+        const std::vector<std::string>& options
+    ) {
+        if (!this->hasKernel(kernel_name)) {
+            std::string kernel_code = load_cuda_source(source_path);
+            this->cacheKernel(kernel_name, kernel_code, source_name, options);
+        }
+        return this->getKernel(kernel_name);
+    }
+
+    /*
+    Tries to retrieve the kernel "kernel_name". If not found, instantiate it and save to cache.
+    */
+    CachedKernel* create(
+        const std::string& kernel_name,
+        const std::string& source_variable,
+        const std::string& source_name,
+        const std::vector<std::string>& options
+    ) {
+        if (!this->hasKernel(kernel_name)) {
+            this->cacheKernel(kernel_name, source_variable, source_name, options);
+        }
+
+        return this->getKernel(kernel_name);
+    }
+
+  private:
+    KernelFactory() {}
+    std::unordered_map<std::string, std::unique_ptr<CachedKernel>> kernel_cache;
+
+    KernelFactory(const KernelFactory&) = delete;
+    KernelFactory& operator=(const KernelFactory&) = delete;
+};
+
+#endif
\ No newline at end of file
diff --git a/sphericart/include/dynamic_cuda.hpp b/sphericart/include/dynamic_cuda.hpp
new file mode 100644
index 00000000..eda6b150
--- /dev/null
+++ b/sphericart/include/dynamic_cuda.hpp
@@ -0,0 +1,369 @@
+// DynamicCUDAHeader.hpp
+#ifndef DYNAMIC_CUDA_HEADER_HPP
+#define DYNAMIC_CUDA_HEADER_HPP
+
+#ifdef __linux__
+#include <dlfcn.h>
+#else
+#error "Platform not supported"
+#endif
+#include <stdexcept>
+#include <string>
+#include <functional>
+#include <unordered_map>
+#include <any>
+#include <sstream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <nvrtc.h>
+
+#define NVRTC_SAFE_CALL(x)                                                                         \
+    do {                                                                                           \
+        nvrtcResult result = x;                                                                    \
+        if (result != NVRTC_SUCCESS) {                                                             \
+            std::ostringstream errorMsg;                                                           \
+            errorMsg << "\nerror: " #x " failed with error "                                       \
+                     << NVRTC_INSTANCE.nvrtcGetErrorString(result) << '\n'                         \
+                     << "File: " << __FILE__ << '\n'                                               \
+                     << "Line: " << static_cast<int>(__LINE__) << '\n';                            \
+            throw std::runtime_error(errorMsg.str());                                              \
+        }                                                                                          \
+    } while (0)
+
+#define CUDADRIVER_SAFE_CALL(x)                                                                    \
+    do {                                                                                           \
+        CUresult result = x;                                                                       \
+        if (result != CUDA_SUCCESS) {                                                              \
+            const char* msg;                                                                       \
+            CUDA_DRIVER_INSTANCE.cuGetErrorName(result, &msg);                                     \
+            std::ostringstream errorMsg;                                                           \
+            errorMsg << "\nerror: " #x " failed with error " << (msg ? msg : "Unknown error")      \
+                     << '\n'                                                                       \
+                     << "File: " << __FILE__ << '\n'                                               \
+                     << "Line: " << static_cast<int>(__LINE__) << '\n';                            \
+            throw std::runtime_error(errorMsg.str());                                              \
+        }                                                                                          \
+    } while (0)
+
+#define CUDART_SAFE_CALL(call)                                                                     \
+    do {                                                                                           \
+        cudaError_t cudaStatus = (call);                                                           \
+        if (cudaStatus != cudaSuccess) {                                                           \
+            std::ostringstream errorMsg;                                                           \
+            const char* error = CUDART_INSTANCE.cudaGetErrorString(cudaStatus);                    \
+            errorMsg << "\nfailed with error " << (error ? error : "Unknown error") << '\n'        \
+                     << "File: " << __FILE__ << '\n'                                               \
+                     << "Line: " << static_cast<int>(__LINE__) << '\n';                            \
+            throw std::runtime_error(errorMsg.str());                                              \
+        }                                                                                          \
+    } while (0)
+
+// Define a template to dynamically load symbols
+template <typename FuncType> FuncType load(void* handle, const char* functionName) {
+    auto func = reinterpret_cast<FuncType>(dlsym(handle, functionName));
+    if (!func) {
+        throw std::runtime_error(std::string("Failed to load function: ") + functionName);
+    }
+    return func;
+}
+
+/*
+This class allows us to dynamically load the CUDA runtime and reference the functions contained
+within the libcudart.so library (see CUDA Runtime API:
+https://docs.nvidia.com/cuda/cuda-runtime-api/index.html).
+*/
+class CUDART {
+  public:
+    bool loaded() { return cudartHandle != nullptr; }
+
+    using cudaGetDeviceCount_t = cudaError_t (*)(int*);
+    using cudaGetDevice_t = cudaError_t (*)(int*);
+    using cudaSetDevice_t = cudaError_t (*)(int);
+    using cudaMalloc_t = cudaError_t (*)(void**, size_t);
+    using cudaMemcpy_t = cudaError_t (*)(void*, const void*, size_t, cudaMemcpyKind);
+    using cudaGetErrorName_t = const char* (*)(cudaError_t);
+    using cudaGetErrorString_t = const char* (*)(cudaError_t);
+    using cudaDeviceSynchronize_t = cudaError_t (*)(void);
+    using cudaPointerGetAttributes_t = cudaError_t (*)(cudaPointerAttributes*, const void*);
+    using cudaFree_t = cudaError_t (*)(void*);
+
+    cudaGetDeviceCount_t cudaGetDeviceCount;
+    cudaGetDevice_t cudaGetDevice;
+    cudaSetDevice_t cudaSetDevice;
+    cudaMalloc_t cudaMalloc;
+    cudaMemcpy_t cudaMemcpy;
+    cudaGetErrorName_t cudaGetErrorName;
+    cudaGetErrorString_t cudaGetErrorString;
+    cudaDeviceSynchronize_t cudaDeviceSynchronize;
+    cudaPointerGetAttributes_t cudaPointerGetAttributes;
+    cudaFree_t cudaFree;
+
+    CUDART() {
+#ifdef __linux__
+        cudartHandle = dlopen("libcudart.so", RTLD_NOW);
+#else
+#error "Platform not supported"
+#endif
+        if (cudartHandle) {
+            // load cudart function pointers using template
+            cudaGetDeviceCount = load<cudaGetDeviceCount_t>(cudartHandle, "cudaGetDeviceCount");
+            cudaGetDevice = load<cudaGetDevice_t>(cudartHandle, "cudaGetDevice");
+            cudaSetDevice = load<cudaSetDevice_t>(cudartHandle, "cudaSetDevice");
+            cudaMalloc = load<cudaMalloc_t>(cudartHandle, "cudaMalloc");
+            cudaMemcpy = load<cudaMemcpy_t>(cudartHandle, "cudaMemcpy");
+            cudaGetErrorName = load<cudaGetErrorName_t>(cudartHandle, "cudaGetErrorName");
+            cudaGetErrorString = load<cudaGetErrorString_t>(cudartHandle, "cudaGetErrorString");
+            cudaDeviceSynchronize =
+                load<cudaDeviceSynchronize_t>(cudartHandle, "cudaDeviceSynchronize");
+            cudaPointerGetAttributes =
+                load<cudaPointerGetAttributes_t>(cudartHandle, "cudaPointerGetAttributes");
+            cudaFree = load<cudaFree_t>(cudartHandle, "cudaFree");
+        }
+    }
+
+    ~CUDART() {
+#ifdef __linux__
+        if (cudartHandle) {
+            dlclose(cudartHandle);
+        }
+#else
+#error "Platform not supported"
+#endif
+    }
+
+    // Prevent copying
+    CUDART(const CUDART&) = delete;
+    CUDART& operator=(const CUDART&) = delete;
+
+    void* cudartHandle = nullptr;
+};
+
+/*
+This class allows us to dynamically load the CUDA Driver and reference the functions contained
+within the libcuda.so library (CUDA Driver API:
+https://docs.nvidia.com/cuda/cuda-driver-api/index.html).
+*/
+class CUDADriver {
+
+  public:
+    bool loaded() { return cudaHandle != nullptr; }
+
+    using cuInit_t = CUresult (*)(unsigned int);
+    using cuDeviceGetCount_t = CUresult (*)(int*);
+    using cuDevicePrimaryCtxRetain_t = CUresult (*)(CUcontext*, CUdevice);
+    using cuDevicePrimaryCtxRelease_t = CUresult (*)(CUdevice);
+    using cuCtxCreate_t = CUresult (*)(CUcontext*, unsigned int, CUdevice);
+    using cuCtxDestroy_t = CUresult (*)(CUcontext);
+    using cuCtxGetCurrent_t = CUresult (*)(CUcontext*);
+    using cuCtxSetCurrent_t = CUresult (*)(CUcontext);
+    using cuModuleLoadDataEx_t = CUresult (*)(CUmodule*, const void*, unsigned int, int*, int*);
+    using cuModuleGetFunction_t = CUresult (*)(CUfunction*, CUmodule, const char*);
+    using cuFuncSetAttribute_t = CUresult (*)(CUfunction, CUfunction_attribute, int);
+    using cuFuncGetAttribute_t = CUresult (*)(int*, CUfunction_attribute, CUfunction);
+    using cuCtxGetDevice_t = CUresult (*)(CUdevice*);
+    using cuDeviceGetAttribute_t = CUresult (*)(int*, CUdevice_attribute, CUdevice);
+    using cuDeviceGetName_t = CUresult (*)(char*, int, CUdevice);
+    using cuDeviceTotalMem_t = CUresult (*)(size_t*, CUdevice);
+    using cuLaunchKernel_t =
+        CUresult (*)(CUfunction, unsigned int, unsigned int, unsigned int, unsigned int, unsigned int, unsigned int, size_t, CUstream, void**, void*);
+    using cuStreamCreate_t = CUresult (*)(CUstream*, unsigned int);
+    using cuStreamDestroy_t = CUresult (*)(CUstream);
+    using cuCtxSynchronize_t = CUresult (*)(void);
+    using cuGetErrorName_t = CUresult (*)(CUresult, const char**);
+    using cuCtxPushCurrent_t = CUresult (*)(CUcontext);
+    using cuPointerGetAttribute_t = CUresult (*)(void*, CUpointer_attribute, CUdeviceptr);
+
+    cuInit_t cuInit;
+    cuDeviceGetCount_t cuDeviceGetCount;
+    cuCtxCreate_t cuCtxCreate;
+    cuCtxDestroy_t cuCtxDestroy;
+    cuDevicePrimaryCtxRetain_t cuDevicePrimaryCtxRetain;
+    cuDevicePrimaryCtxRelease_t cuDevicePrimaryCtxRelease;
+    cuCtxGetCurrent_t cuCtxGetCurrent;
+    cuCtxSetCurrent_t cuCtxSetCurrent;
+    cuModuleLoadDataEx_t cuModuleLoadDataEx;
+    cuModuleGetFunction_t cuModuleGetFunction;
+    cuFuncSetAttribute_t cuFuncSetAttribute;
+    cuFuncGetAttribute_t cuFuncGetAttribute;
+    cuCtxGetDevice_t cuCtxGetDevice;
+    cuDeviceGetAttribute_t cuDeviceGetAttribute;
+    cuDeviceGetName_t cuDeviceGetName;
+    cuDeviceTotalMem_t cuDeviceTotalMem;
+    cuLaunchKernel_t cuLaunchKernel;
+    cuStreamCreate_t cuStreamCreate;
+    cuStreamDestroy_t cuStreamDestroy;
+    cuGetErrorName_t cuGetErrorName;
+    cuCtxSynchronize_t cuCtxSynchronize;
+    cuCtxPushCurrent_t cuCtxPushCurrent;
+    cuPointerGetAttribute_t cuPointerGetAttribute;
+
+    CUDADriver() {
+#ifdef __linux__
+        cudaHandle = dlopen("libcuda.so", RTLD_NOW);
+#else
+#error "Platform not supported"
+#endif
+        if (cudaHandle) {
+            // Load CUDA driver function pointers using template
+            cuInit = load<cuInit_t>(cudaHandle, "cuInit");
+            cuDeviceGetCount = load<cuDeviceGetCount_t>(cudaHandle, "cuDeviceGetCount");
+            cuCtxCreate = load<cuCtxCreate_t>(cudaHandle, "cuCtxCreate");
+            cuCtxDestroy = load<cuCtxDestroy_t>(cudaHandle, "cuCtxDestroy");
+            cuDevicePrimaryCtxRetain =
+                load<cuDevicePrimaryCtxRetain_t>(cudaHandle, "cuDevicePrimaryCtxRetain");
+            cuDevicePrimaryCtxRelease =
+                load<cuDevicePrimaryCtxRelease_t>(cudaHandle, "cuDevicePrimaryCtxRelease");
+            cuCtxGetCurrent = load<cuCtxGetCurrent_t>(cudaHandle, "cuCtxGetCurrent");
+            cuCtxSetCurrent = load<cuCtxSetCurrent_t>(cudaHandle, "cuCtxSetCurrent");
+            cuModuleLoadDataEx = load<cuModuleLoadDataEx_t>(cudaHandle, "cuModuleLoadDataEx");
+            cuModuleGetFunction = load<cuModuleGetFunction_t>(cudaHandle, "cuModuleGetFunction");
+            cuFuncSetAttribute = load<cuFuncSetAttribute_t>(cudaHandle, "cuFuncSetAttribute");
+            cuFuncGetAttribute = load<cuFuncGetAttribute_t>(cudaHandle, "cuFuncGetAttribute");
+            cuCtxGetDevice = load<cuCtxGetDevice_t>(cudaHandle, "cuCtxGetDevice");
+            cuDeviceGetAttribute = load<cuDeviceGetAttribute_t>(cudaHandle, "cuDeviceGetAttribute");
+            cuDeviceGetName = load<cuDeviceGetName_t>(cudaHandle, "cuDeviceGetName");
+            cuDeviceTotalMem = load<cuDeviceTotalMem_t>(cudaHandle, "cuDeviceTotalMem");
+            cuLaunchKernel = load<cuLaunchKernel_t>(cudaHandle, "cuLaunchKernel");
+            cuStreamCreate = load<cuStreamCreate_t>(cudaHandle, "cuStreamCreate");
+            cuStreamDestroy = load<cuStreamDestroy_t>(cudaHandle, "cuStreamDestroy");
+            cuCtxSynchronize = load<cuCtxSynchronize_t>(cudaHandle, "cuCtxSynchronize");
+            cuGetErrorName = load<cuGetErrorName_t>(cudaHandle, "cuGetErrorName");
+            cuCtxPushCurrent = load<cuCtxPushCurrent_t>(cudaHandle, "cuCtxPushCurrent");
+            cuPointerGetAttribute =
+                load<cuPointerGetAttribute_t>(cudaHandle, "cuPointerGetAttribute");
+        }
+    }
+
+    ~CUDADriver() {
+#ifdef __linux__
+        if (cudaHandle) {
+            dlclose(cudaHandle);
+        }
+#else
+#error "Platform not supported"
+#endif
+    }
+
+    // Prevent copying
+    CUDADriver(const CUDADriver&) = delete;
+    CUDADriver& operator=(const CUDADriver&) = delete;
+
+    void* cudaHandle = nullptr;
+};
+
+/*
+This class allows us to dynamically load NVRTC and reference the functions contained within the
+libnvrtc.so library (see NVRTC API: https://docs.nvidia.com/cuda/nvrtc/index.html).
+*/
+class NVRTC {
+
+  public:
+    bool loaded() { return nvrtcHandle != nullptr; }
+
+    using nvrtcCreateProgram_t =
+        nvrtcResult (*)(nvrtcProgram*, const char*, const char*, int, const char*[], const char*[]);
+    using nvrtcCompileProgram_t = nvrtcResult (*)(nvrtcProgram, int, const char*[]);
+    using nvrtcGetPTX_t = nvrtcResult (*)(nvrtcProgram, char*);
+    using nvrtcGetPTXSize_t = nvrtcResult (*)(nvrtcProgram, size_t*);
+    using nvrtcGetProgramLog_t = nvrtcResult (*)(nvrtcProgram, char*);
+    using nvrtcGetProgramLogSize_t = nvrtcResult (*)(nvrtcProgram, size_t*);
+    using nvrtcAddNameExpression_t = nvrtcResult (*)(nvrtcProgram, const char* const);
+    using nvrtcGetLoweredName_t = nvrtcResult (*)(nvrtcProgram, const char*, const char**);
+    using nvrtcDestroyProgram_t = nvrtcResult (*)(nvrtcProgram*);
+    using nvrtcGetErrorString_t = const char* (*)(nvrtcResult);
+
+    nvrtcCreateProgram_t nvrtcCreateProgram;
+    nvrtcCompileProgram_t nvrtcCompileProgram;
+    nvrtcGetPTX_t nvrtcGetPTX;
+    nvrtcGetPTXSize_t nvrtcGetPTXSize;
+    nvrtcGetProgramLog_t nvrtcGetProgramLog;
+    nvrtcGetProgramLogSize_t nvrtcGetProgramLogSize;
+    nvrtcGetLoweredName_t nvrtcGetLoweredName;
+    nvrtcAddNameExpression_t nvrtcAddNameExpression;
+    nvrtcDestroyProgram_t nvrtcDestroyProgram;
+    nvrtcGetErrorString_t nvrtcGetErrorString;
+
+    NVRTC() {
+#ifdef __linux__
+        nvrtcHandle = dlopen("libnvrtc.so", RTLD_NOW);
+#else
+#error "Platform not supported"
+#endif
+
+        if (nvrtcHandle) {
+            // Load NVRTC function pointers using template
+            nvrtcCreateProgram = load<nvrtcCreateProgram_t>(nvrtcHandle, "nvrtcCreateProgram");
+            nvrtcCompileProgram = load<nvrtcCompileProgram_t>(nvrtcHandle, "nvrtcCompileProgram");
+            nvrtcGetPTX = load<nvrtcGetPTX_t>(nvrtcHandle, "nvrtcGetPTX");
+            nvrtcGetPTXSize = load<nvrtcGetPTXSize_t>(nvrtcHandle, "nvrtcGetPTXSize");
+            nvrtcGetProgramLog = load<nvrtcGetProgramLog_t>(nvrtcHandle, "nvrtcGetProgramLog");
+            nvrtcGetProgramLogSize =
+                load<nvrtcGetProgramLogSize_t>(nvrtcHandle, "nvrtcGetProgramLogSize");
+            nvrtcGetLoweredName = load<nvrtcGetLoweredName_t>(nvrtcHandle, "nvrtcGetLoweredName");
+            nvrtcAddNameExpression =
+                load<nvrtcAddNameExpression_t>(nvrtcHandle, "nvrtcAddNameExpression");
+            nvrtcDestroyProgram = load<nvrtcDestroyProgram_t>(nvrtcHandle, "nvrtcDestroyProgram");
+            nvrtcGetErrorString = load<nvrtcGetErrorString_t>(nvrtcHandle, "nvrtcGetErrorString");
+        }
+    }
+
+    ~NVRTC() {
+#ifdef __linux__
+        if (nvrtcHandle) {
+            dlclose(nvrtcHandle);
+        }
+#else
+#error "Platform not supported"
+#endif
+    }
+
+    // Prevent copying
+    NVRTC(const NVRTC&) = delete;
+    NVRTC& operator=(const NVRTC&) = delete;
+
+    void* nvrtcHandle = nullptr;
+};
+
+/*
+This implements the Schawrz counter idiom to ensure propper constructor/destructor ordering - we
+don't want these classes destroyed before jax or pytorch classes/objects are.
+*/
+
+// Static memory buffer for each class
+static std::aligned_storage<sizeof(CUDART), alignof(CUDART)>::type cudartBuffer;
+static std::aligned_storage<sizeof(CUDADriver), alignof(CUDADriver)>::type cudaDriverBuffer;
+static std::aligned_storage<sizeof(NVRTC), alignof(NVRTC)>::type nvrtcBuffer;
+
+/*
+global references that should be used by any dependent code. Use inline instead of extern as we want
+to define the implementation in a single-header.
+*/
+inline CUDART& CUDART_INSTANCE = reinterpret_cast<CUDART&>(cudartBuffer);
+inline CUDADriver& CUDA_DRIVER_INSTANCE = reinterpret_cast<CUDADriver&>(cudaDriverBuffer);
+inline NVRTC& NVRTC_INSTANCE = reinterpret_cast<NVRTC&>(nvrtcBuffer);
+
+static int nifty_counter = 0;
+
+// Counter class for initializing and destroying static objects
+static struct CUDAInitializer {
+
+    CUDAInitializer() {
+        if (nifty_counter++ == 0) {
+            new (&CUDART_INSTANCE) CUDART();
+            new (&CUDA_DRIVER_INSTANCE) CUDADriver();
+            new (&NVRTC_INSTANCE) NVRTC();
+        }
+    }
+
+    ~CUDAInitializer() {
+        if (--nifty_counter == 0) {
+            CUDART_INSTANCE.~CUDART();
+            CUDA_DRIVER_INSTANCE.~CUDADriver();
+            NVRTC_INSTANCE.~NVRTC();
+        }
+    }
+
+} cudaInitializer;
+
+#endif // DYNAMIC_CUDA_HEADER_HPP
\ No newline at end of file
diff --git a/sphericart/include/macros.hpp b/sphericart/include/macros.hpp
index fdb24179..f5408b0b 100644
--- a/sphericart/include/macros.hpp
+++ b/sphericart/include/macros.hpp
@@ -34,8 +34,19 @@
 // we need this monstruosity to make sure that literals are not treated as
 // double precision, as that slows down greatly floating-point execution.
 // basically this static casts the constant to the type of the target
-// pointer
-#define cast(sph, lit) static_cast<std::remove_pointer_t<decltype(sph)>>(lit)
+// pointer. We do not use the remove_pointer struct provided by the standard library here:
+// https://en.cppreference.com/w/cpp/types/remove_pointer, as this fails to compile with NVRTC
+template <typename T> struct remove_pointer {
+    using type = T;
+};
+
+template <typename T> struct remove_pointer<T*> {
+    using type = T;
+};
+
+template <typename T> using remove_pointer_t = typename remove_pointer<T>::type;
+
+#define cast(sph, lit) static_cast<remove_pointer_t<decltype(sph)>>(lit)
 
 #define COMPUTE_SPH_L0(sph_i, SPH_IDX) (sph_i)[SPH_IDX(0)] = cast(sph_i, 0.282094791773878);
 
diff --git a/sphericart/include/sphericart_cuda.hpp b/sphericart/include/sphericart_cuda.hpp
index 04e1bd8b..b50c50a5 100644
--- a/sphericart/include/sphericart_cuda.hpp
+++ b/sphericart/include/sphericart_cuda.hpp
@@ -146,10 +146,10 @@ template <typename T> class SphericalHarmonics {
   private:
     size_t l_max; // maximum l value computed by this class
     size_t nprefactors;
-    bool normalized;               // should we normalize the input vectors?
-    T* prefactors_cpu = nullptr;   // host prefactors buffer
-    T** prefactors_cuda = nullptr; // storage space for prefactors
-    int device_count = 0;          // number of visible GPU devices
+    bool normalized;              // should we normalize the input vectors?
+    T* prefactors_cpu = nullptr;  // host prefactors buffer
+    T* prefactors_cuda = nullptr; // storage space for prefactors
+    int device_count = 0;         // number of visible GPU devices
     int64_t CUDA_GRID_DIM_X_ = 8;
     int64_t CUDA_GRID_DIM_Y_ = 8;
     bool cached_compute_with_gradients = false;
diff --git a/sphericart/include/sphericart_impl.cuh b/sphericart/include/sphericart_impl.cuh
new file mode 100644
index 00000000..293d3b70
--- /dev/null
+++ b/sphericart/include/sphericart_impl.cuh
@@ -0,0 +1,101 @@
+#ifndef SPHERICART_IMPL_CUH
+#define SPHERICART_IMPL_CUH
+
+#define CUDA_DEVICE_PREFIX __device__
+
+/*
+    Clears the shared memory buffers for the spherical harmonics and gradients
+   if required.
+*/
+template <typename scalar_t>
+__device__ inline void clear_buffers(
+    int nelements,
+    scalar_t* sph,
+    scalar_t* dsph_x,
+    scalar_t* dsph_y,
+    scalar_t* dsph_z,
+
+    scalar_t* dsph_dxdx,
+    scalar_t* dsph_dxdy,
+    scalar_t* dsph_dxdz,
+
+    scalar_t* dsph_dydx,
+    scalar_t* dsph_dydy,
+    scalar_t* dsph_dydz,
+
+    scalar_t* dsph_dzdx,
+    scalar_t* dsph_dzdy,
+    scalar_t* dsph_dzdz,
+    bool requires_grad,
+    bool requires_hessian
+);
+
+/*
+    Writes out the shared memory buffers to global memory, as well as applying
+   normalisation if necessary.
+*/
+template <typename scalar_t>
+__device__ inline void write_buffers(
+    int edge_idx,
+    int nedges,
+    scalar_t x,
+    scalar_t y,
+    scalar_t z,
+    scalar_t ir,
+    int n_elements,
+    int offset,
+    scalar_t* buffer_sph,
+
+    scalar_t* buffer_dsph_x,
+    scalar_t* buffer_dsph_y,
+    scalar_t* buffer_dsph_z,
+
+    scalar_t* buffer_dsph_dxdx,
+    scalar_t* buffer_dsph_dxdy,
+    scalar_t* buffer_dsph_dxdz,
+
+    scalar_t* buffer_dsph_dydx,
+    scalar_t* buffer_dsph_dydy,
+    scalar_t* buffer_dsph_dydz,
+
+    scalar_t* buffer_dsph_dzdx,
+    scalar_t* buffer_dsph_dzdy,
+    scalar_t* buffer_dsph_dzdz,
+    scalar_t* sph,
+    scalar_t* dsph,
+    scalar_t* ddsph,
+    int n_total,
+    bool requires_grad,
+    bool requires_hessian,
+    bool normalize
+);
+
+/*
+    CUDA kernel for computing Cartesian spherical harmonics and their
+   derivatives.
+*/
+template <typename scalar_t>
+__global__ void spherical_harmonics_kernel(
+    scalar_t* xyz,
+    int nedges,
+    scalar_t* prefactors,
+    int nprefactors,
+    int lmax,
+    int ntotal,
+    bool requires_grad,
+    bool requires_hessian,
+    bool normalize,
+    scalar_t* sph,
+    scalar_t* dsph,
+    scalar_t* ddsph
+);
+
+/*
+    CUDA kernel to computes the backwards pass for autograd.
+*/
+template <typename scalar_t>
+__global__ void backward_kernel(
+    scalar_t* dsph, scalar_t* sph_grad, int nedges, int n_total, scalar_t* xyz_grad
+);
+
+#endif
\ No newline at end of file
diff --git a/sphericart/include/templates.hpp b/sphericart/include/templates.hpp
index e4d0bdc6..9a74c6a7 100644
--- a/sphericart/include/templates.hpp
+++ b/sphericart/include/templates.hpp
@@ -14,6 +14,8 @@
     different numbers of terms computed with hard-coded expressions.
 */
 
+#include "templates_core.hpp"
+
 #include <cmath>
 #include <vector>
 
@@ -325,415 +327,6 @@ void hardcoded_sph(
     }
 }
 
-int inline dummy_idx(int i) { return i; }
-
-/** Computes the sph and their derivatives for a given Cartesian point and a
- * given l. The template implementation supports different floating poitn types
- * T, determines whether to compute derivatives (DO_DERIVATIVES), assumes that
- * l is greater than HARDCODED_LMAX. GET_INDEX is a function that might allow
- * to map differently the indices in the spherical harmonics (used in the CUDA
- * implementation).
- */
-template <
-    typename T,
-    bool DO_DERIVATIVES,
-    bool DO_SECOND_DERIVATIVES,
-    int HARDCODED_LMAX,
-    int (*GET_INDEX)(int) = dummy_idx>
-CUDA_DEVICE_PREFIX static inline void generic_sph_l_channel(
-    int l,
-    [[maybe_unused]] T x, // these might be unused for low LMAX. not worth a
-                          // full separate implementation
-    [[maybe_unused]] T y,
-    [[maybe_unused]] T z,
-    [[maybe_unused]] T rxy,
-    const T* pk,
-    const T* qlmk,
-    T* c,
-    T* s,
-    T* twomz,
-    T* sph_i,
-    [[maybe_unused]] T* dx_sph_i,
-    [[maybe_unused]] T* dy_sph_i,
-    [[maybe_unused]] T* dz_sph_i,
-    [[maybe_unused]] T* dxdx_sph_i,
-    [[maybe_unused]] T* dxdy_sph_i,
-    [[maybe_unused]] T* dxdz_sph_i,
-    [[maybe_unused]] T* dydx_sph_i,
-    [[maybe_unused]] T* dydy_sph_i,
-    [[maybe_unused]] T* dydz_sph_i,
-    [[maybe_unused]] T* dzdx_sph_i,
-    [[maybe_unused]] T* dzdy_sph_i,
-    [[maybe_unused]] T* dzdz_sph_i
-) {
-    /*
-    This is the main low-level code to compute sph and dsph for an arbitrary l.
-    The code is a bit hard to follow because of (too?) many micro-optimizations.
-    Sine and cosine terms are precomputed. Here the Qlm modifield Legendre
-    polynomials are evaluated, and combined with the other terms and the
-    prefactors. The core iteration is an iteration down to lower values of m,
-
-    Qlm = A z Ql(m+1) + B rxy^2 Ql(m+2)
-
-    1. the special cases with l=+-m and +-1 are done separately, also because
-    they initialize the recursive expression
-    2. we assume that some lower l are done with hard-coding, and HARDCODED_LMAX
-    is passed as a template parameter. this is used to split the loops over m in
-    a part with fixed size, known at compile time, and one with variable length.
-    3. we do the recursion using stack variables and never store the full Qlm
-    array
-    4. we compute separately Qlm and Q(l-1) - the latter needed for derivatives
-    rather than reuse the calculation from another l channel. It appears that
-    the simplification in memory access makes this beneficial, with the added
-    advantage that each l channel can be computed independently
-
-    Template parameters:
-    typename T: float type (e.g. single/double precision)
-    bool DO_DERIVATIVES: should we evaluate the derivatives?
-    bool DO_SECOND_DERIVATIVES: should we evaluate the second derivatives?
-    bool NORMALIZED: should we normalize the input positions?
-
-    Actual parameters:
-    l: which l we are computing
-    x,y,z: the Cartesian coordinates of the point
-    rxy: sqrt(x^2+y^2), precomputed because it's used for all l
-    pk, qlmk: prefactors used in the calculation of Ylm and Qlm, respectively
-    c,s: the c_k and s_k cos-like and sin-like terms combined with the Qlm to
-    compute Ylm twomz: 2*m*z, these are also computed once and reused for all l
-    sph_i, d[x,y,z]sph_i: storage locations of the output arrays for Ylm and
-    dYlm/d[x,y,z] d[x,y,z]d[x,y,z]sph_i: storage locations of the output arrays
-    for the second derivatives
-
-    */
-    static_assert(
-        !(DO_SECOND_DERIVATIVES && !DO_DERIVATIVES),
-        "Cannot calculate second derivatives without first derivatives"
-    );
-
-    // working space for the recursive evaluation of Qlm and Q(l-1)m.
-    // qlm_[0,1,2] correspond to the current Qlm, Ql(m+1) and Ql(m+2), and the
-    // ql1m_[0,1,2] hold the same but for l-1
-    // ql2m_[0,1,2] hold the same but for l-2
-    [[maybe_unused]] T qlm_2, qlm_1, qlm_0;
-    [[maybe_unused]] T ql1m_2, ql1m_1, ql1m_0;
-    [[maybe_unused]] T ql2m_2, ql2m_1, ql2m_0;
-
-    [[maybe_unused]] T x2, y2,
-        xy; // for second derivatives. we could get them from parent but not
-            // worth the additional complexity
-    if constexpr (DO_SECOND_DERIVATIVES) {
-        x2 = x * x;
-        y2 = y * y;
-        xy = x * y;
-    }
-
-    // m=+-l
-    qlm_2 = qlmk[l]; // fetches the pre-computed Qll
-    auto pq = qlm_2 * pk[l];
-    sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l)];
-    sph_i[GET_INDEX(+l)] = pq * c[GET_INDEX(l)];
-
-    if constexpr (DO_DERIVATIVES) {
-        pq *= l;
-        dx_sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l - 1)];
-        dy_sph_i[GET_INDEX(-l)] = dx_sph_i[GET_INDEX(l)] = pq * c[GET_INDEX(l - 1)];
-        dy_sph_i[GET_INDEX(l)] = -dx_sph_i[GET_INDEX(-l)];
-        dz_sph_i[GET_INDEX(-l)] = 0;
-        dz_sph_i[GET_INDEX(l)] = 0;
-        ql1m_2 = 0;
-
-        if constexpr (DO_SECOND_DERIVATIVES) {
-            pq *= (l - 1);
-            dxdx_sph_i[GET_INDEX(l)] = pq * c[GET_INDEX(l - 2)];
-            dxdx_sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l - 2)];
-            dxdy_sph_i[GET_INDEX(l)] = dydx_sph_i[GET_INDEX(l)] = dydy_sph_i[GET_INDEX(-l)] =
-                -dxdx_sph_i[GET_INDEX(-l)];
-            dxdy_sph_i[GET_INDEX(-l)] = dydx_sph_i[GET_INDEX(-l)] = dxdx_sph_i[GET_INDEX(l)];
-            dxdz_sph_i[GET_INDEX(l)] = dzdx_sph_i[GET_INDEX(l)] = 0.0;
-            dxdz_sph_i[GET_INDEX(-l)] = dzdx_sph_i[GET_INDEX(-l)] = 0.0;
-            dydy_sph_i[GET_INDEX(l)] = -dxdx_sph_i[GET_INDEX(l)];
-            dydz_sph_i[GET_INDEX(l)] = dzdy_sph_i[GET_INDEX(l)] = 0.0;
-            dydz_sph_i[GET_INDEX(-l)] = dzdy_sph_i[GET_INDEX(-l)] = 0.0;
-            dzdz_sph_i[GET_INDEX(l)] = dzdz_sph_i[GET_INDEX(-l)] = 0.0;
-            ql2m_2 = 0.0;
-        }
-    }
-
-    // m = +-(l-1)
-    qlm_1 = -z * qlm_2;
-    pq = qlm_1 * pk[l - 1];
-    sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 1)];
-    sph_i[GET_INDEX(+l - 1)] = pq * c[GET_INDEX(l - 1)];
-
-    if constexpr (DO_DERIVATIVES) {
-        pq *= (l - 1);
-        dx_sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 2)];
-        dy_sph_i[GET_INDEX(-l + 1)] = dx_sph_i[GET_INDEX(l - 1)] = pq * c[GET_INDEX(l - 2)];
-        dy_sph_i[GET_INDEX(l - 1)] = -dx_sph_i[GET_INDEX(-l + 1)];
-
-        // uses Q(l-1)(l-1) to initialize the Qlm  recursion
-        ql1m_1 = qlmk[-1];
-        auto pdq = pk[l - 1] * (l + l - 1) * ql1m_1;
-        dz_sph_i[GET_INDEX(-l + 1)] = pdq * s[GET_INDEX(l - 1)];
-        dz_sph_i[GET_INDEX(l - 1)] = pdq * c[GET_INDEX(l - 1)];
-
-        if constexpr (DO_SECOND_DERIVATIVES) {
-            pq *= (l - 2);
-            if (l == 2) { // this is a special case for second derivatives
-                dxdx_sph_i[GET_INDEX(l - 1)] = 0;
-                dxdx_sph_i[GET_INDEX(-l + 1)] = 0;
-            } else {
-                dxdx_sph_i[GET_INDEX(l - 1)] = pq * c[GET_INDEX(l - 3)];
-                dxdx_sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 3)];
-            }
-            dxdy_sph_i[GET_INDEX(l - 1)] = dydx_sph_i[GET_INDEX(l - 1)] =
-                dydy_sph_i[GET_INDEX(-l + 1)] = -dxdx_sph_i[GET_INDEX(-l + 1)];
-            dxdy_sph_i[GET_INDEX(-l + 1)] = dydx_sph_i[GET_INDEX(-l + 1)] =
-                dxdx_sph_i[GET_INDEX(l - 1)];
-            auto temp = -pk[l - 1] * (l - 1) * qlm_2; // this is p[l-1]*q[l-1][l-1]*(2*l-1)*(l-1) =
-                                                      // p[l-1]*(l-1)*Q_ll
-            dxdz_sph_i[GET_INDEX(l - 1)] = dzdx_sph_i[GET_INDEX(l - 1)] = temp * c[GET_INDEX(l - 2)];
-            dxdz_sph_i[GET_INDEX(-l + 1)] = dzdx_sph_i[GET_INDEX(-l + 1)] =
-                temp * s[GET_INDEX(l - 2)];
-            dydy_sph_i[GET_INDEX(l - 1)] = -dxdx_sph_i[GET_INDEX(l - 1)];
-            dydz_sph_i[GET_INDEX(l - 1)] = dzdy_sph_i[GET_INDEX(l - 1)] =
-                -dxdz_sph_i[GET_INDEX(-l + 1)];
-            dydz_sph_i[GET_INDEX(-l + 1)] = dzdy_sph_i[GET_INDEX(-l + 1)] =
-                dxdz_sph_i[GET_INDEX(l - 1)];
-            dzdz_sph_i[GET_INDEX(l - 1)] = dzdz_sph_i[GET_INDEX(-l + 1)] = 0.0;
-            ql2m_1 = 0.0;
-        }
-    }
-
-    // and now do the other m's, decrementally
-    for (auto m = l - 2; m > HARDCODED_LMAX - 1; --m) {
-        qlm_0 = qlmk[m] * (twomz[GET_INDEX(m)] * qlm_1 + rxy * qlm_2);
-        qlm_2 = qlm_1;
-        qlm_1 = qlm_0; // shift
-        pq = qlm_0 * pk[m];
-        sph_i[GET_INDEX(-m)] = pq * s[GET_INDEX(m)];
-        sph_i[GET_INDEX(+m)] = pq * c[GET_INDEX(m)];
-
-        if constexpr (DO_DERIVATIVES) {
-            ql1m_0 = qlmk[m - l] * (twomz[GET_INDEX(m)] * ql1m_1 + rxy * ql1m_2);
-            ql1m_2 = ql1m_1;
-            ql1m_1 = ql1m_0; // shift
-
-            pq *= m;
-            auto pqs = pq * s[GET_INDEX(m - 1)], pqc = pq * c[GET_INDEX(m - 1)];
-            auto pdq = pk[m] * ql1m_2;
-            auto pdqx = pdq * x;
-            dx_sph_i[GET_INDEX(-m)] = (pdqx * s[GET_INDEX(m)] + pqs);
-            dx_sph_i[GET_INDEX(+m)] = (pdqx * c[GET_INDEX(m)] + pqc);
-            auto pdqy = pdq * y;
-            dy_sph_i[GET_INDEX(-m)] = (pdqy * s[GET_INDEX(m)] + pqc);
-            dy_sph_i[GET_INDEX(m)] = (pdqy * c[GET_INDEX(m)] - pqs);
-            pdq = pk[m] * (l + m) * ql1m_1;
-            dz_sph_i[GET_INDEX(-m)] = pdq * s[GET_INDEX(m)];
-            dz_sph_i[GET_INDEX(m)] = pdq * c[GET_INDEX(m)];
-
-            if constexpr (DO_SECOND_DERIVATIVES) {
-                if (m == l - 2) {
-                    // In this case, the recursion still needs to be initialized
-                    // using Q(l-2)(l-2)
-                    ql2m_0 = qlmk[-l - 1];
-                } else {
-                    // Recursion
-                    ql2m_0 = qlmk[m - 2 * l + 1] * (twomz[GET_INDEX(m)] * ql2m_1 + rxy * ql2m_2);
-                }
-
-                pq /= m;
-                auto pql1m_1 = pk[m] * ql1m_2; // Note the index discrepancy: ql1m_1
-                                               // was already shifted above to ql1m_2
-                auto pql2m_2 = pk[m] * ql2m_2;
-                auto pql2m_0 = pk[m] * ql2m_0;
-                auto pql1m_0 = pk[m] * ql1m_1; // Note the index discrepancy: ql1m_0
-                                               // was already shifted above to ql1m_1
-                auto pql2m_1 = pk[m] * ql2m_1;
-
-                // Diagonal hessian terms
-                T mmpqc2 = 0.0;
-                T mmpqs2 = 0.0;
-                if (m != 1) {
-                    mmpqc2 = m * (m - 1) * pq * c[GET_INDEX(m - 2)];
-                    mmpqs2 = m * (m - 1) * pq * s[GET_INDEX(m - 2)];
-                }
-
-                dxdx_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
-                                           x2 * pql2m_2 * c[GET_INDEX(m)] +
-                                           2 * m * x * pql1m_1 * c[GET_INDEX(m - 1)] + mmpqc2;
-                dxdx_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
-                                            x2 * pql2m_2 * s[GET_INDEX(m)] +
-                                            2 * m * x * pql1m_1 * s[GET_INDEX(m - 1)] + mmpqs2;
-                dydy_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
-                                           y2 * pql2m_2 * c[GET_INDEX(m)] -
-                                           2 * m * y * pql1m_1 * s[GET_INDEX(m - 1)] - mmpqc2;
-                dydy_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
-                                            y2 * pql2m_2 * s[GET_INDEX(m)] +
-                                            2 * m * y * pql1m_1 * c[GET_INDEX(m - 1)] - mmpqs2;
-                dzdz_sph_i[GET_INDEX(m)] = (l + m) * (l + m - 1) * pql2m_0 * c[GET_INDEX(m)];
-                dzdz_sph_i[GET_INDEX(-m)] = (l + m) * (l + m - 1) * pql2m_0 * s[GET_INDEX(m)];
-
-                // Off-diagonal terms. Note that these are symmetric
-                dxdy_sph_i[GET_INDEX(m)] = dydx_sph_i[GET_INDEX(m)] =
-                    xy * pql2m_2 * c[GET_INDEX(m)] + y * pql1m_1 * m * c[GET_INDEX(m - 1)] -
-                    x * pql1m_1 * m * s[GET_INDEX(m - 1)] - mmpqs2;
-                dxdy_sph_i[GET_INDEX(-m)] = dydx_sph_i[GET_INDEX(-m)] =
-                    xy * pql2m_2 * s[GET_INDEX(m)] + y * pql1m_1 * m * s[GET_INDEX(m - 1)] +
-                    x * pql1m_1 * m * c[GET_INDEX(m - 1)] + mmpqc2;
-                dxdz_sph_i[GET_INDEX(m)] = dzdx_sph_i[GET_INDEX(m)] =
-                    x * (l + m) * pql2m_1 * c[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
-                dxdz_sph_i[GET_INDEX(-m)] = dzdx_sph_i[GET_INDEX(-m)] =
-                    x * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
-                dydz_sph_i[GET_INDEX(m)] = dzdy_sph_i[GET_INDEX(m)] =
-                    y * (l + m) * pql2m_1 * c[GET_INDEX(m)] -
-                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
-                dydz_sph_i[GET_INDEX(-m)] = dzdy_sph_i[GET_INDEX(-m)] =
-                    y * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
-
-                ql2m_2 = ql2m_1;
-                ql2m_1 = ql2m_0; // shift at the end because we need all three
-                                 // at the same time
-            }
-        }
-    }
-    for (auto m = HARDCODED_LMAX - 1; m > 0; --m) {
-        qlm_0 = qlmk[m] * (twomz[GET_INDEX(m)] * qlm_1 + rxy * qlm_2);
-        qlm_2 = qlm_1;
-        qlm_1 = qlm_0; // shift
-
-        pq = qlm_0 * pk[m];
-        sph_i[GET_INDEX(-m)] = pq * s[GET_INDEX(m)];
-        sph_i[GET_INDEX(+m)] = pq * c[GET_INDEX(m)];
-
-        if constexpr (DO_DERIVATIVES) {
-            ql1m_0 = qlmk[m - l] * (twomz[GET_INDEX(m)] * ql1m_1 + rxy * ql1m_2);
-            ql1m_2 = ql1m_1;
-            ql1m_1 = ql1m_0; // shift
-
-            pq *= m;
-            auto pqs = pq * s[GET_INDEX(m - 1)], pqc = pq * c[GET_INDEX(m - 1)];
-            auto pdq = pk[m] * ql1m_2;
-            auto pdqx = pdq * x;
-            dx_sph_i[GET_INDEX(-m)] = (pdqx * s[GET_INDEX(m)] + pqs);
-            dx_sph_i[GET_INDEX(+m)] = (pdqx * c[GET_INDEX(m)] + pqc);
-            auto pdqy = pdq * y;
-            dy_sph_i[GET_INDEX(-m)] = (pdqy * s[GET_INDEX(m)] + pqc);
-            dy_sph_i[GET_INDEX(m)] = (pdqy * c[GET_INDEX(m)] - pqs);
-            pdq = pk[m] * (l + m) * ql1m_1;
-            dz_sph_i[GET_INDEX(-m)] = pdq * s[GET_INDEX(m)];
-            dz_sph_i[GET_INDEX(m)] = pdq * c[GET_INDEX(m)];
-
-            if constexpr (DO_SECOND_DERIVATIVES) {
-                if (m == l - 2) {
-                    // In this case, the recursion still needs to be initialized
-                    // using Q(l-2)(l-2)
-                    ql2m_0 = qlmk[-l - 1];
-                } else {
-                    // Recursion
-                    ql2m_0 = qlmk[m - 2 * l + 1] * (twomz[GET_INDEX(m)] * ql2m_1 + rxy * ql2m_2);
-                }
-
-                pq /= m;
-                auto pql1m_1 = pk[m] * ql1m_2; // Note the index discrepancy: ql1m_1
-                                               // was already shifted above to ql1m_2
-                auto pql2m_2 = pk[m] * ql2m_2;
-                auto pql2m_0 = pk[m] * ql2m_0;
-                auto pql1m_0 = pk[m] * ql1m_1; // Note the index discrepancy: ql1m_0
-                                               // was already shifted above to ql1m_1
-                auto pql2m_1 = pk[m] * ql2m_1;
-
-                // Diagonal hessian terms
-                T mmpqc2 = 0.0;
-                T mmpqs2 = 0.0;
-
-                if (m != 1) {
-                    mmpqc2 = m * (m - 1) * pq * c[GET_INDEX(m - 2)];
-                    mmpqs2 = m * (m - 1) * pq * s[GET_INDEX(m - 2)];
-                }
-
-                dxdx_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
-                                           x2 * pql2m_2 * c[GET_INDEX(m)] +
-                                           2 * m * x * pql1m_1 * c[GET_INDEX(m - 1)] + mmpqc2;
-                dxdx_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
-                                            x2 * pql2m_2 * s[GET_INDEX(m)] +
-                                            2 * m * x * pql1m_1 * s[GET_INDEX(m - 1)] + mmpqs2;
-                dydy_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
-                                           y2 * pql2m_2 * c[GET_INDEX(m)] -
-                                           2 * m * y * pql1m_1 * s[GET_INDEX(m - 1)] - mmpqc2;
-                dydy_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
-                                            y2 * pql2m_2 * s[GET_INDEX(m)] +
-                                            2 * m * y * pql1m_1 * c[GET_INDEX(m - 1)] - mmpqs2;
-                dzdz_sph_i[GET_INDEX(m)] = (l + m) * (l + m - 1) * pql2m_0 * c[GET_INDEX(m)];
-                dzdz_sph_i[GET_INDEX(-m)] = (l + m) * (l + m - 1) * pql2m_0 * s[GET_INDEX(m)];
-
-                // Off-diagonal terms. Note that these are symmetric
-                dxdy_sph_i[GET_INDEX(m)] = dydx_sph_i[GET_INDEX(m)] =
-                    xy * pql2m_2 * c[GET_INDEX(m)] + y * pql1m_1 * m * c[GET_INDEX(m - 1)] -
-                    x * pql1m_1 * m * s[GET_INDEX(m - 1)] - mmpqs2;
-                dxdy_sph_i[GET_INDEX(-m)] = dydx_sph_i[GET_INDEX(-m)] =
-                    xy * pql2m_2 * s[GET_INDEX(m)] + y * pql1m_1 * m * s[GET_INDEX(m - 1)] +
-                    x * pql1m_1 * m * c[GET_INDEX(m - 1)] + mmpqc2;
-                dxdz_sph_i[GET_INDEX(m)] = dzdx_sph_i[GET_INDEX(m)] =
-                    x * (l + m) * pql2m_1 * c[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
-                dxdz_sph_i[GET_INDEX(-m)] = dzdx_sph_i[GET_INDEX(-m)] =
-                    x * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
-                dydz_sph_i[GET_INDEX(m)] = dzdy_sph_i[GET_INDEX(m)] =
-                    y * (l + m) * pql2m_1 * c[GET_INDEX(m)] -
-                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
-                dydz_sph_i[GET_INDEX(-m)] = dzdy_sph_i[GET_INDEX(-m)] =
-                    y * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
-                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
-
-                ql2m_2 = ql2m_1;
-                ql2m_1 = ql2m_0; // shift at the end because we need all three
-                                 // at the same time
-            }
-        }
-    }
-
-    // m=0 is also a special case
-    qlm_0 = qlmk[0] * (twomz[GET_INDEX(0)] * qlm_1 + rxy * qlm_2);
-    sph_i[GET_INDEX(0)] = qlm_0 * pk[0];
-
-    if constexpr (DO_DERIVATIVES) {
-        ql1m_0 = qlmk[-l] * (twomz[GET_INDEX(0)] * ql1m_1 + rxy * ql1m_2);
-        // derivatives
-        dx_sph_i[GET_INDEX(0)] = pk[0] * x * ql1m_1;
-        dy_sph_i[GET_INDEX(0)] = pk[0] * y * ql1m_1;
-        dz_sph_i[GET_INDEX(0)] = pk[0] * l * ql1m_0;
-
-        if constexpr (DO_SECOND_DERIVATIVES) {
-            if (l == 2) {
-                // special case: recursion is not initialized yet
-                ql2m_0 = qlmk[-2 * l + 1];
-            } else {
-                ql2m_0 = qlmk[-2 * l + 1] * (twomz[GET_INDEX(0)] * ql2m_1 + rxy * ql2m_2);
-            }
-
-            auto pql1m_1 = pk[0] * ql1m_1;
-            auto pql2m_2 = pk[0] * ql2m_2;
-            auto pql2m_0 = pk[0] * ql2m_0;
-            auto pql2m_1 = pk[0] * ql2m_1;
-
-            // diagonal
-            dxdx_sph_i[GET_INDEX(0)] = pql1m_1 + x2 * pql2m_2;
-            dydy_sph_i[GET_INDEX(0)] = pql1m_1 + y2 * pql2m_2;
-            dzdz_sph_i[GET_INDEX(0)] = (l) * (l - 1) * pql2m_0;
-
-            // off-diagonal (symmetric)
-            dxdy_sph_i[GET_INDEX(0)] = dydx_sph_i[GET_INDEX(0)] = xy * pql2m_2;
-            dxdz_sph_i[GET_INDEX(0)] = dzdx_sph_i[GET_INDEX(0)] = x * l * pql2m_1;
-            dydz_sph_i[GET_INDEX(0)] = dzdy_sph_i[GET_INDEX(0)] = y * l * pql2m_1;
-        }
-    }
-}
-
 template <typename T, bool DO_DERIVATIVES, bool DO_SECOND_DERIVATIVES, bool NORMALIZED, int HARDCODED_LMAX>
 static inline void generic_sph_sample(
     const T* xyz_i,
diff --git a/sphericart/include/templates_core.hpp b/sphericart/include/templates_core.hpp
new file mode 100644
index 00000000..2512989c
--- /dev/null
+++ b/sphericart/include/templates_core.hpp
@@ -0,0 +1,429 @@
+#ifndef SPHERICART_TEMPLATES_CORE_HPP
+#define SPHERICART_TEMPLATES_CORE_HPP
+
+#ifndef CUDA_DEVICE_PREFIX
+#define CUDA_DEVICE_PREFIX
+#endif
+
+/*
+    Template implementation of Cartesian Ylm calculators.
+
+    The template functions use compile-time `if constexpr()` constructs to
+    implement calculators for spherical harmonics that can handle different
+    type of calls, e.g. with or without derivative calculations, and with
+    different numbers of terms computed with hard-coded expressions.
+*/
+
+// a SPH_IDX that does nothing
+#define DUMMY_SPH_IDX
+
+int inline CUDA_DEVICE_PREFIX dummy_idx(int i) { return i; }
+
+/** Computes the sph and their derivatives for a given Cartesian point and a
+ * given l. The template implementation supports different floating poitn types
+ * T, determines whether to compute derivatives (DO_DERIVATIVES), assumes that
+ * l is greater than HARDCODED_LMAX. GET_INDEX is a function that might allow
+ * to map differently the indices in the spherical harmonics (used in the CUDA
+ * implementation).
+ */
+template <
+    typename T,
+    bool DO_DERIVATIVES,
+    bool DO_SECOND_DERIVATIVES,
+    int HARDCODED_LMAX,
+    int (*GET_INDEX)(int) = dummy_idx>
+CUDA_DEVICE_PREFIX static inline void generic_sph_l_channel(
+    int l,
+    [[maybe_unused]] T x, // these might be unused for low LMAX. not worth a
+                          // full separate implementation
+    [[maybe_unused]] T y,
+    [[maybe_unused]] T z,
+    [[maybe_unused]] T rxy,
+    const T* pk,
+    const T* qlmk,
+    T* c,
+    T* s,
+    T* twomz,
+    T* sph_i,
+    [[maybe_unused]] T* dx_sph_i,
+    [[maybe_unused]] T* dy_sph_i,
+    [[maybe_unused]] T* dz_sph_i,
+    [[maybe_unused]] T* dxdx_sph_i,
+    [[maybe_unused]] T* dxdy_sph_i,
+    [[maybe_unused]] T* dxdz_sph_i,
+    [[maybe_unused]] T* dydx_sph_i,
+    [[maybe_unused]] T* dydy_sph_i,
+    [[maybe_unused]] T* dydz_sph_i,
+    [[maybe_unused]] T* dzdx_sph_i,
+    [[maybe_unused]] T* dzdy_sph_i,
+    [[maybe_unused]] T* dzdz_sph_i
+) {
+    /*
+    This is the main low-level code to compute sph and dsph for an arbitrary l.
+    The code is a bit hard to follow because of (too?) many micro-optimizations.
+    Sine and cosine terms are precomputed. Here the Qlm modifield Legendre
+    polynomials are evaluated, and combined with the other terms and the
+    prefactors. The core iteration is an iteration down to lower values of m,
+
+    Qlm = A z Ql(m+1) + B rxy^2 Ql(m+2)
+
+    1. the special cases with l=+-m and +-1 are done separately, also because
+    they initialize the recursive expression
+    2. we assume that some lower l are done with hard-coding, and HARDCODED_LMAX
+    is passed as a template parameter. this is used to split the loops over m in
+    a part with fixed size, known at compile time, and one with variable length.
+    3. we do the recursion using stack variables and never store the full Qlm
+    array
+    4. we compute separately Qlm and Q(l-1) - the latter needed for derivatives
+    rather than reuse the calculation from another l channel. It appears that
+    the simplification in memory access makes this beneficial, with the added
+    advantage that each l channel can be computed independently
+
+    Template parameters:
+    typename T: float type (e.g. single/double precision)
+    bool DO_DERIVATIVES: should we evaluate the derivatives?
+    bool DO_SECOND_DERIVATIVES: should we evaluate the second derivatives?
+    bool NORMALIZED: should we normalize the input positions?
+
+    Actual parameters:
+    l: which l we are computing
+    x,y,z: the Cartesian coordinates of the point
+    rxy: sqrt(x^2+y^2), precomputed because it's used for all l
+    pk, qlmk: prefactors used in the calculation of Ylm and Qlm, respectively
+    c,s: the c_k and s_k cos-like and sin-like terms combined with the Qlm to
+    compute Ylm twomz: 2*m*z, these are also computed once and reused for all l
+    sph_i, d[x,y,z]sph_i: storage locations of the output arrays for Ylm and
+    dYlm/d[x,y,z] d[x,y,z]d[x,y,z]sph_i: storage locations of the output arrays
+    for the second derivatives
+
+    */
+    static_assert(
+        !(DO_SECOND_DERIVATIVES && !DO_DERIVATIVES),
+        "Cannot calculate second derivatives without first derivatives"
+    );
+
+    // working space for the recursive evaluation of Qlm and Q(l-1)m.
+    // qlm_[0,1,2] correspond to the current Qlm, Ql(m+1) and Ql(m+2), and the
+    // ql1m_[0,1,2] hold the same but for l-1
+    // ql2m_[0,1,2] hold the same but for l-2
+    [[maybe_unused]] T qlm_2, qlm_1, qlm_0;
+    [[maybe_unused]] T ql1m_2, ql1m_1, ql1m_0;
+    [[maybe_unused]] T ql2m_2, ql2m_1, ql2m_0;
+
+    [[maybe_unused]] T x2, y2,
+        xy; // for second derivatives. we could get them from parent but not
+            // worth the additional complexity
+    if constexpr (DO_SECOND_DERIVATIVES) {
+        x2 = x * x;
+        y2 = y * y;
+        xy = x * y;
+    }
+
+    // m=+-l
+    qlm_2 = qlmk[l]; // fetches the pre-computed Qll
+    auto pq = qlm_2 * pk[l];
+    sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l)];
+    sph_i[GET_INDEX(+l)] = pq * c[GET_INDEX(l)];
+
+    if constexpr (DO_DERIVATIVES) {
+        pq *= l;
+        dx_sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l - 1)];
+        dy_sph_i[GET_INDEX(-l)] = dx_sph_i[GET_INDEX(l)] = pq * c[GET_INDEX(l - 1)];
+        dy_sph_i[GET_INDEX(l)] = -dx_sph_i[GET_INDEX(-l)];
+        dz_sph_i[GET_INDEX(-l)] = 0;
+        dz_sph_i[GET_INDEX(l)] = 0;
+        ql1m_2 = 0;
+
+        if constexpr (DO_SECOND_DERIVATIVES) {
+            pq *= (l - 1);
+            dxdx_sph_i[GET_INDEX(l)] = pq * c[GET_INDEX(l - 2)];
+            dxdx_sph_i[GET_INDEX(-l)] = pq * s[GET_INDEX(l - 2)];
+            dxdy_sph_i[GET_INDEX(l)] = dydx_sph_i[GET_INDEX(l)] = dydy_sph_i[GET_INDEX(-l)] =
+                -dxdx_sph_i[GET_INDEX(-l)];
+            dxdy_sph_i[GET_INDEX(-l)] = dydx_sph_i[GET_INDEX(-l)] = dxdx_sph_i[GET_INDEX(l)];
+            dxdz_sph_i[GET_INDEX(l)] = dzdx_sph_i[GET_INDEX(l)] = 0.0;
+            dxdz_sph_i[GET_INDEX(-l)] = dzdx_sph_i[GET_INDEX(-l)] = 0.0;
+            dydy_sph_i[GET_INDEX(l)] = -dxdx_sph_i[GET_INDEX(l)];
+            dydz_sph_i[GET_INDEX(l)] = dzdy_sph_i[GET_INDEX(l)] = 0.0;
+            dydz_sph_i[GET_INDEX(-l)] = dzdy_sph_i[GET_INDEX(-l)] = 0.0;
+            dzdz_sph_i[GET_INDEX(l)] = dzdz_sph_i[GET_INDEX(-l)] = 0.0;
+            ql2m_2 = 0.0;
+        }
+    }
+
+    // m = +-(l-1)
+    qlm_1 = -z * qlm_2;
+    pq = qlm_1 * pk[l - 1];
+    sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 1)];
+    sph_i[GET_INDEX(+l - 1)] = pq * c[GET_INDEX(l - 1)];
+
+    if constexpr (DO_DERIVATIVES) {
+        pq *= (l - 1);
+        dx_sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 2)];
+        dy_sph_i[GET_INDEX(-l + 1)] = dx_sph_i[GET_INDEX(l - 1)] = pq * c[GET_INDEX(l - 2)];
+        dy_sph_i[GET_INDEX(l - 1)] = -dx_sph_i[GET_INDEX(-l + 1)];
+
+        // uses Q(l-1)(l-1) to initialize the Qlm  recursion
+        ql1m_1 = qlmk[-1];
+        auto pdq = pk[l - 1] * (l + l - 1) * ql1m_1;
+        dz_sph_i[GET_INDEX(-l + 1)] = pdq * s[GET_INDEX(l - 1)];
+        dz_sph_i[GET_INDEX(l - 1)] = pdq * c[GET_INDEX(l - 1)];
+
+        if constexpr (DO_SECOND_DERIVATIVES) {
+            pq *= (l - 2);
+            if (l == 2) { // this is a special case for second derivatives
+                dxdx_sph_i[GET_INDEX(l - 1)] = 0;
+                dxdx_sph_i[GET_INDEX(-l + 1)] = 0;
+            } else {
+                dxdx_sph_i[GET_INDEX(l - 1)] = pq * c[GET_INDEX(l - 3)];
+                dxdx_sph_i[GET_INDEX(-l + 1)] = pq * s[GET_INDEX(l - 3)];
+            }
+            dxdy_sph_i[GET_INDEX(l - 1)] = dydx_sph_i[GET_INDEX(l - 1)] =
+                dydy_sph_i[GET_INDEX(-l + 1)] = -dxdx_sph_i[GET_INDEX(-l + 1)];
+            dxdy_sph_i[GET_INDEX(-l + 1)] = dydx_sph_i[GET_INDEX(-l + 1)] =
+                dxdx_sph_i[GET_INDEX(l - 1)];
+            auto temp = -pk[l - 1] * (l - 1) * qlm_2; // this is p[l-1]*q[l-1][l-1]*(2*l-1)*(l-1) =
+                                                      // p[l-1]*(l-1)*Q_ll
+            dxdz_sph_i[GET_INDEX(l - 1)] = dzdx_sph_i[GET_INDEX(l - 1)] = temp * c[GET_INDEX(l - 2)];
+            dxdz_sph_i[GET_INDEX(-l + 1)] = dzdx_sph_i[GET_INDEX(-l + 1)] =
+                temp * s[GET_INDEX(l - 2)];
+            dydy_sph_i[GET_INDEX(l - 1)] = -dxdx_sph_i[GET_INDEX(l - 1)];
+            dydz_sph_i[GET_INDEX(l - 1)] = dzdy_sph_i[GET_INDEX(l - 1)] =
+                -dxdz_sph_i[GET_INDEX(-l + 1)];
+            dydz_sph_i[GET_INDEX(-l + 1)] = dzdy_sph_i[GET_INDEX(-l + 1)] =
+                dxdz_sph_i[GET_INDEX(l - 1)];
+            dzdz_sph_i[GET_INDEX(l - 1)] = dzdz_sph_i[GET_INDEX(-l + 1)] = 0.0;
+            ql2m_1 = 0.0;
+        }
+    }
+
+    // and now do the other m's, decrementally
+    for (auto m = l - 2; m > HARDCODED_LMAX - 1; --m) {
+        qlm_0 = qlmk[m] * (twomz[GET_INDEX(m)] * qlm_1 + rxy * qlm_2);
+        qlm_2 = qlm_1;
+        qlm_1 = qlm_0; // shift
+        pq = qlm_0 * pk[m];
+        sph_i[GET_INDEX(-m)] = pq * s[GET_INDEX(m)];
+        sph_i[GET_INDEX(+m)] = pq * c[GET_INDEX(m)];
+
+        if constexpr (DO_DERIVATIVES) {
+            ql1m_0 = qlmk[m - l] * (twomz[GET_INDEX(m)] * ql1m_1 + rxy * ql1m_2);
+            ql1m_2 = ql1m_1;
+            ql1m_1 = ql1m_0; // shift
+
+            pq *= m;
+            auto pqs = pq * s[GET_INDEX(m - 1)], pqc = pq * c[GET_INDEX(m - 1)];
+            auto pdq = pk[m] * ql1m_2;
+            auto pdqx = pdq * x;
+            dx_sph_i[GET_INDEX(-m)] = (pdqx * s[GET_INDEX(m)] + pqs);
+            dx_sph_i[GET_INDEX(+m)] = (pdqx * c[GET_INDEX(m)] + pqc);
+            auto pdqy = pdq * y;
+            dy_sph_i[GET_INDEX(-m)] = (pdqy * s[GET_INDEX(m)] + pqc);
+            dy_sph_i[GET_INDEX(m)] = (pdqy * c[GET_INDEX(m)] - pqs);
+            pdq = pk[m] * (l + m) * ql1m_1;
+            dz_sph_i[GET_INDEX(-m)] = pdq * s[GET_INDEX(m)];
+            dz_sph_i[GET_INDEX(m)] = pdq * c[GET_INDEX(m)];
+
+            if constexpr (DO_SECOND_DERIVATIVES) {
+                if (m == l - 2) {
+                    // In this case, the recursion still needs to be initialized
+                    // using Q(l-2)(l-2)
+                    ql2m_0 = qlmk[-l - 1];
+                } else {
+                    // Recursion
+                    ql2m_0 = qlmk[m - 2 * l + 1] * (twomz[GET_INDEX(m)] * ql2m_1 + rxy * ql2m_2);
+                }
+
+                pq /= m;
+                auto pql1m_1 = pk[m] * ql1m_2; // Note the index discrepancy: ql1m_1
+                                               // was already shifted above to ql1m_2
+                auto pql2m_2 = pk[m] * ql2m_2;
+                auto pql2m_0 = pk[m] * ql2m_0;
+                auto pql1m_0 = pk[m] * ql1m_1; // Note the index discrepancy: ql1m_0
+                                               // was already shifted above to ql1m_1
+                auto pql2m_1 = pk[m] * ql2m_1;
+
+                // Diagonal hessian terms
+                T mmpqc2 = 0.0;
+                T mmpqs2 = 0.0;
+                if (m != 1) {
+                    mmpqc2 = m * (m - 1) * pq * c[GET_INDEX(m - 2)];
+                    mmpqs2 = m * (m - 1) * pq * s[GET_INDEX(m - 2)];
+                }
+
+                dxdx_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
+                                           x2 * pql2m_2 * c[GET_INDEX(m)] +
+                                           2 * m * x * pql1m_1 * c[GET_INDEX(m - 1)] + mmpqc2;
+                dxdx_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
+                                            x2 * pql2m_2 * s[GET_INDEX(m)] +
+                                            2 * m * x * pql1m_1 * s[GET_INDEX(m - 1)] + mmpqs2;
+                dydy_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
+                                           y2 * pql2m_2 * c[GET_INDEX(m)] -
+                                           2 * m * y * pql1m_1 * s[GET_INDEX(m - 1)] - mmpqc2;
+                dydy_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
+                                            y2 * pql2m_2 * s[GET_INDEX(m)] +
+                                            2 * m * y * pql1m_1 * c[GET_INDEX(m - 1)] - mmpqs2;
+                dzdz_sph_i[GET_INDEX(m)] = (l + m) * (l + m - 1) * pql2m_0 * c[GET_INDEX(m)];
+                dzdz_sph_i[GET_INDEX(-m)] = (l + m) * (l + m - 1) * pql2m_0 * s[GET_INDEX(m)];
+
+                // Off-diagonal terms. Note that these are symmetric
+                dxdy_sph_i[GET_INDEX(m)] = dydx_sph_i[GET_INDEX(m)] =
+                    xy * pql2m_2 * c[GET_INDEX(m)] + y * pql1m_1 * m * c[GET_INDEX(m - 1)] -
+                    x * pql1m_1 * m * s[GET_INDEX(m - 1)] - mmpqs2;
+                dxdy_sph_i[GET_INDEX(-m)] = dydx_sph_i[GET_INDEX(-m)] =
+                    xy * pql2m_2 * s[GET_INDEX(m)] + y * pql1m_1 * m * s[GET_INDEX(m - 1)] +
+                    x * pql1m_1 * m * c[GET_INDEX(m - 1)] + mmpqc2;
+                dxdz_sph_i[GET_INDEX(m)] = dzdx_sph_i[GET_INDEX(m)] =
+                    x * (l + m) * pql2m_1 * c[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
+                dxdz_sph_i[GET_INDEX(-m)] = dzdx_sph_i[GET_INDEX(-m)] =
+                    x * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
+                dydz_sph_i[GET_INDEX(m)] = dzdy_sph_i[GET_INDEX(m)] =
+                    y * (l + m) * pql2m_1 * c[GET_INDEX(m)] -
+                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
+                dydz_sph_i[GET_INDEX(-m)] = dzdy_sph_i[GET_INDEX(-m)] =
+                    y * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
+
+                ql2m_2 = ql2m_1;
+                ql2m_1 = ql2m_0; // shift at the end because we need all three
+                                 // at the same time
+            }
+        }
+    }
+    for (auto m = HARDCODED_LMAX - 1; m > 0; --m) {
+        qlm_0 = qlmk[m] * (twomz[GET_INDEX(m)] * qlm_1 + rxy * qlm_2);
+        qlm_2 = qlm_1;
+        qlm_1 = qlm_0; // shift
+
+        pq = qlm_0 * pk[m];
+        sph_i[GET_INDEX(-m)] = pq * s[GET_INDEX(m)];
+        sph_i[GET_INDEX(+m)] = pq * c[GET_INDEX(m)];
+
+        if constexpr (DO_DERIVATIVES) {
+            ql1m_0 = qlmk[m - l] * (twomz[GET_INDEX(m)] * ql1m_1 + rxy * ql1m_2);
+            ql1m_2 = ql1m_1;
+            ql1m_1 = ql1m_0; // shift
+
+            pq *= m;
+            auto pqs = pq * s[GET_INDEX(m - 1)], pqc = pq * c[GET_INDEX(m - 1)];
+            auto pdq = pk[m] * ql1m_2;
+            auto pdqx = pdq * x;
+            dx_sph_i[GET_INDEX(-m)] = (pdqx * s[GET_INDEX(m)] + pqs);
+            dx_sph_i[GET_INDEX(+m)] = (pdqx * c[GET_INDEX(m)] + pqc);
+            auto pdqy = pdq * y;
+            dy_sph_i[GET_INDEX(-m)] = (pdqy * s[GET_INDEX(m)] + pqc);
+            dy_sph_i[GET_INDEX(m)] = (pdqy * c[GET_INDEX(m)] - pqs);
+            pdq = pk[m] * (l + m) * ql1m_1;
+            dz_sph_i[GET_INDEX(-m)] = pdq * s[GET_INDEX(m)];
+            dz_sph_i[GET_INDEX(m)] = pdq * c[GET_INDEX(m)];
+
+            if constexpr (DO_SECOND_DERIVATIVES) {
+                if (m == l - 2) {
+                    // In this case, the recursion still needs to be initialized
+                    // using Q(l-2)(l-2)
+                    ql2m_0 = qlmk[-l - 1];
+                } else {
+                    // Recursion
+                    ql2m_0 = qlmk[m - 2 * l + 1] * (twomz[GET_INDEX(m)] * ql2m_1 + rxy * ql2m_2);
+                }
+
+                pq /= m;
+                auto pql1m_1 = pk[m] * ql1m_2; // Note the index discrepancy: ql1m_1
+                                               // was already shifted above to ql1m_2
+                auto pql2m_2 = pk[m] * ql2m_2;
+                auto pql2m_0 = pk[m] * ql2m_0;
+                auto pql1m_0 = pk[m] * ql1m_1; // Note the index discrepancy: ql1m_0
+                                               // was already shifted above to ql1m_1
+                auto pql2m_1 = pk[m] * ql2m_1;
+
+                // Diagonal hessian terms
+                T mmpqc2 = 0.0;
+                T mmpqs2 = 0.0;
+
+                if (m != 1) {
+                    mmpqc2 = m * (m - 1) * pq * c[GET_INDEX(m - 2)];
+                    mmpqs2 = m * (m - 1) * pq * s[GET_INDEX(m - 2)];
+                }
+
+                dxdx_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
+                                           x2 * pql2m_2 * c[GET_INDEX(m)] +
+                                           2 * m * x * pql1m_1 * c[GET_INDEX(m - 1)] + mmpqc2;
+                dxdx_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
+                                            x2 * pql2m_2 * s[GET_INDEX(m)] +
+                                            2 * m * x * pql1m_1 * s[GET_INDEX(m - 1)] + mmpqs2;
+                dydy_sph_i[GET_INDEX(m)] = pql1m_1 * c[GET_INDEX(m)] +
+                                           y2 * pql2m_2 * c[GET_INDEX(m)] -
+                                           2 * m * y * pql1m_1 * s[GET_INDEX(m - 1)] - mmpqc2;
+                dydy_sph_i[GET_INDEX(-m)] = pql1m_1 * s[GET_INDEX(m)] +
+                                            y2 * pql2m_2 * s[GET_INDEX(m)] +
+                                            2 * m * y * pql1m_1 * c[GET_INDEX(m - 1)] - mmpqs2;
+                dzdz_sph_i[GET_INDEX(m)] = (l + m) * (l + m - 1) * pql2m_0 * c[GET_INDEX(m)];
+                dzdz_sph_i[GET_INDEX(-m)] = (l + m) * (l + m - 1) * pql2m_0 * s[GET_INDEX(m)];
+
+                // Off-diagonal terms. Note that these are symmetric
+                dxdy_sph_i[GET_INDEX(m)] = dydx_sph_i[GET_INDEX(m)] =
+                    xy * pql2m_2 * c[GET_INDEX(m)] + y * pql1m_1 * m * c[GET_INDEX(m - 1)] -
+                    x * pql1m_1 * m * s[GET_INDEX(m - 1)] - mmpqs2;
+                dxdy_sph_i[GET_INDEX(-m)] = dydx_sph_i[GET_INDEX(-m)] =
+                    xy * pql2m_2 * s[GET_INDEX(m)] + y * pql1m_1 * m * s[GET_INDEX(m - 1)] +
+                    x * pql1m_1 * m * c[GET_INDEX(m - 1)] + mmpqc2;
+                dxdz_sph_i[GET_INDEX(m)] = dzdx_sph_i[GET_INDEX(m)] =
+                    x * (l + m) * pql2m_1 * c[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
+                dxdz_sph_i[GET_INDEX(-m)] = dzdx_sph_i[GET_INDEX(-m)] =
+                    x * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
+                dydz_sph_i[GET_INDEX(m)] = dzdy_sph_i[GET_INDEX(m)] =
+                    y * (l + m) * pql2m_1 * c[GET_INDEX(m)] -
+                    (l + m) * pql1m_0 * m * s[GET_INDEX(m - 1)];
+                dydz_sph_i[GET_INDEX(-m)] = dzdy_sph_i[GET_INDEX(-m)] =
+                    y * (l + m) * pql2m_1 * s[GET_INDEX(m)] +
+                    (l + m) * pql1m_0 * m * c[GET_INDEX(m - 1)];
+
+                ql2m_2 = ql2m_1;
+                ql2m_1 = ql2m_0; // shift at the end because we need all three
+                                 // at the same time
+            }
+        }
+    }
+
+    // m=0 is also a special case
+    qlm_0 = qlmk[0] * (twomz[GET_INDEX(0)] * qlm_1 + rxy * qlm_2);
+    sph_i[GET_INDEX(0)] = qlm_0 * pk[0];
+
+    if constexpr (DO_DERIVATIVES) {
+        ql1m_0 = qlmk[-l] * (twomz[GET_INDEX(0)] * ql1m_1 + rxy * ql1m_2);
+        // derivatives
+        dx_sph_i[GET_INDEX(0)] = pk[0] * x * ql1m_1;
+        dy_sph_i[GET_INDEX(0)] = pk[0] * y * ql1m_1;
+        dz_sph_i[GET_INDEX(0)] = pk[0] * l * ql1m_0;
+
+        if constexpr (DO_SECOND_DERIVATIVES) {
+            if (l == 2) {
+                // special case: recursion is not initialized yet
+                ql2m_0 = qlmk[-2 * l + 1];
+            } else {
+                ql2m_0 = qlmk[-2 * l + 1] * (twomz[GET_INDEX(0)] * ql2m_1 + rxy * ql2m_2);
+            }
+
+            auto pql1m_1 = pk[0] * ql1m_1;
+            auto pql2m_2 = pk[0] * ql2m_2;
+            auto pql2m_0 = pk[0] * ql2m_0;
+            auto pql2m_1 = pk[0] * ql2m_1;
+
+            // diagonal
+            dxdx_sph_i[GET_INDEX(0)] = pql1m_1 + x2 * pql2m_2;
+            dydy_sph_i[GET_INDEX(0)] = pql1m_1 + y2 * pql2m_2;
+            dzdz_sph_i[GET_INDEX(0)] = (l) * (l - 1) * pql2m_0;
+
+            // off-diagonal (symmetric)
+            dxdy_sph_i[GET_INDEX(0)] = dydx_sph_i[GET_INDEX(0)] = xy * pql2m_2;
+            dxdz_sph_i[GET_INDEX(0)] = dzdx_sph_i[GET_INDEX(0)] = x * l * pql2m_1;
+            dydz_sph_i[GET_INDEX(0)] = dzdy_sph_i[GET_INDEX(0)] = y * l * pql2m_1;
+        }
+    }
+}
+
+#endif
\ No newline at end of file
diff --git a/sphericart/src/cuda_base.cpp b/sphericart/src/cuda_base.cpp
new file mode 100644
index 00000000..49c4b2ce
--- /dev/null
+++ b/sphericart/src/cuda_base.cpp
@@ -0,0 +1,206 @@
+#define _SPHERICART_INTERNAL_IMPLEMENTATION
+
+#include <cmath>
+#include "cuda_cache.hpp"
+#include "cuda_base.hpp"
+
+#define HARDCODED_LMAX 1
+
+/*
+    Computes the total amount of shared memory space required by
+   spherical_harmonics_kernel.
+
+    For lmax <= HARCODED_LMAX, we need to store all (HARDCODED_LMAX + 1)**2
+   scalars in shared memory. For lmax > HARDCODED_LMAX, we only need to store
+   each spherical harmonics vector per sample in shared memory.
+*/
+static size_t total_buffer_size(
+    size_t l_max, size_t GRID_DIM_Y, size_t dtype_size, bool requires_grad, bool requires_hessian
+) {
+    int nl =
+        std::max(static_cast<size_t>((HARDCODED_LMAX + 1) * (HARDCODED_LMAX + 1)), 2 * l_max + 1);
+
+    size_t total_buff_size = 0;
+
+    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_c
+    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_s
+    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_twomz
+    total_buff_size += (l_max + 1) * (l_max + 2) * dtype_size; // buffer_prefactors
+    total_buff_size += GRID_DIM_Y * nl * dtype_size;           // buffer_sph_out
+
+    if (requires_grad) {
+        total_buff_size += 3 * GRID_DIM_Y * nl * dtype_size; // buffer_sph_derivs
+    }
+
+    if (requires_hessian) {
+        total_buff_size += 9 * GRID_DIM_Y * nl * dtype_size; // buffer_sph_hessian
+    }
+
+    return total_buff_size;
+}
+
+/*
+    Wrapper to compile and launch the CUDA kernel. outputs a vector containing the spherical
+   harmonics and their gradients if required to sph, dsph and ddsph pointers.
+
+    GRID_DIM_X is the number of threads to launch in the x dimension. Used to
+   parallelize over the sample dimension. GRID_DIM_Y is the number of threads to
+   launch in the y dimension. Used only to improve memory throughput on reads
+   and writes.
+
+    Total number of threads used is GRID_DIM_X * GRID_DIM_Y.
+
+    cuda_stream should be of type (void *), therefore if you want to pass in
+    a cudaStream_t, first do void * stream_ptr = reinterpret_cast<void *>
+   (stream);
+*/
+template <typename scalar_t>
+void sphericart::cuda::spherical_harmonics_cuda_base(
+    const scalar_t* xyz,
+    const int nedges,
+    const scalar_t* prefactors,
+    const int nprefactors,
+    const int64_t l_max,
+    const bool normalize,
+    const int64_t GRID_DIM_X,
+    const int64_t GRID_DIM_Y,
+    const bool gradients,
+    const bool hessian,
+    scalar_t* sph,
+    scalar_t* dsph,
+    scalar_t* ddsph,
+    void* cuda_stream
+) {
+
+    static const char* CUDA_CODE =
+#include "generated/wrapped_sphericart_impl.cu"
+        ;
+
+    int n_total = (l_max + 1) * (l_max + 1);
+    dim3 block_dim(GRID_DIM_X, GRID_DIM_Y);
+    auto find_num_blocks = [](int x, int bdim) { return (x + bdim - 1) / bdim; };
+    cudaStream_t cstream = reinterpret_cast<cudaStream_t>(cuda_stream);
+    dim3 grid_dim(find_num_blocks(nedges, GRID_DIM_Y));
+
+    size_t smem_size = total_buffer_size(l_max, GRID_DIM_Y, sizeof(scalar_t), gradients, hessian);
+
+    scalar_t* _xyz = const_cast<scalar_t*>(xyz);
+    scalar_t* _prefactors = const_cast<scalar_t*>(prefactors);
+    int _nprefactors = nprefactors;
+    int _lmax = l_max;
+    int _n_total = n_total;
+    int _nedges = nedges;
+    bool _gradients = gradients;
+    bool _hessian = hessian;
+    bool _normalize = normalize;
+
+    std::vector<void*> args = {
+        &_xyz,
+        &_nedges,
+        &_prefactors,
+        &_nprefactors,
+        &_lmax,
+        &_n_total,
+        &_gradients,
+        &_hessian,
+        &_normalize,
+        &sph,
+        &dsph,
+        &ddsph
+    };
+
+    std::string kernel_name = getKernelName<scalar_t>("spherical_harmonics_kernel");
+    auto& kernel_factory = KernelFactory::instance();
+
+    CachedKernel* kernel = kernel_factory.create(
+        kernel_name, std::string(CUDA_CODE), "wrapped_sphericart_impl.cu", {"--std=c++17"}
+    );
+
+    kernel->launch(grid_dim, block_dim, smem_size, cstream, args);
+}
+
+template void sphericart::cuda::spherical_harmonics_cuda_base<float>(
+    const float* xyz,
+    const int nedges,
+    const float* prefactors,
+    const int nprefactors,
+    const int64_t l_max,
+    const bool normalize,
+    const int64_t GRID_DIM_X,
+    const int64_t GRID_DIM_Y,
+    const bool gradients,
+    const bool hessian,
+    float* sph,
+    float* dsph,
+    float* ddsph,
+    void* cuda_stream
+);
+
+template void sphericart::cuda::spherical_harmonics_cuda_base<double>(
+    const double* xyz,
+    const int nedges,
+    const double* prefactors,
+    const int nprefactors,
+    const int64_t l_max,
+    const bool normalize,
+    const int64_t GRID_DIM_X,
+    const int64_t GRID_DIM_Y,
+    const bool gradients,
+    const bool hessian,
+    double* sph,
+    double* dsph,
+    double* ddsph,
+    void* cuda_stream
+);
+
+template <typename scalar_t>
+void sphericart::cuda::spherical_harmonics_backward_cuda_base(
+    const scalar_t* dsph,
+    const scalar_t* sph_grad,
+    const int nedges,
+    const int ntotal,
+    scalar_t* xyz_grad,
+    void* cuda_stream
+) {
+    static const char* CUDA_CODE =
+#include "generated/wrapped_sphericart_impl.cu"
+        ;
+
+    std::string kernel_name = getKernelName<scalar_t>("backward_kernel");
+
+    auto& kernel_factory = KernelFactory::instance();
+
+    dim3 block_dim(4, 32);
+    auto find_num_blocks = [](int x, int bdim) { return (x + bdim - 1) / bdim; };
+    dim3 grid_dim(find_num_blocks(nedges, 32), 3);
+    cudaStream_t cstream = reinterpret_cast<cudaStream_t>(cuda_stream);
+
+    int _n_total = ntotal;
+    int _nedges = nedges;
+
+    std::vector<void*> args = {&dsph, &sph_grad, &_nedges, &_n_total, &xyz_grad};
+
+    CachedKernel* kernel = kernel_factory.create(
+        kernel_name, std::string(CUDA_CODE), "wrapped_sphericart_impl.cu", {"--std=c++17"}
+    );
+
+    kernel->launch(grid_dim, block_dim, 0, cstream, args);
+}
+
+template void sphericart::cuda::spherical_harmonics_backward_cuda_base<float>(
+    const float* dsph,
+    const float* sph_grad,
+    const int nedges,
+    const int ntotal,
+    float* xyz_grad,
+    void* cuda_stream
+);
+
+template void sphericart::cuda::spherical_harmonics_backward_cuda_base<double>(
+    const double* dsph,
+    const double* sph_grad,
+    const int nedges,
+    const int ntotal,
+    double* xyz_grad,
+    void* cuda_stream
+);
diff --git a/sphericart/src/cuda_stub.cpp b/sphericart/src/cuda_stub.cpp
index 6bf0c025..525c7158 100644
--- a/sphericart/src/cuda_stub.cpp
+++ b/sphericart/src/cuda_stub.cpp
@@ -4,9 +4,9 @@
 
 template <typename scalar_t>
 void sphericart::cuda::spherical_harmonics_cuda_base(
-    const scalar_t* __restrict__ /*xyz*/,
+    const scalar_t* /*xyz*/,
     const int /*nedges*/,
-    const scalar_t* __restrict__ /*prefactors*/,
+    const scalar_t* /*prefactors*/,
     const int /*nprefactors*/,
     const int64_t /*l_max*/,
     const bool /*normalize*/,
@@ -14,9 +14,9 @@ void sphericart::cuda::spherical_harmonics_cuda_base(
     const int64_t /*GRID_DIM_Y*/,
     const bool /*gradients*/,
     const bool /*hessian*/,
-    scalar_t* __restrict__ /*sph*/,
-    scalar_t* __restrict__ /*dsph*/,
-    scalar_t* __restrict__ /*ddsph*/,
+    scalar_t* /*sph*/,
+    scalar_t* /*dsph*/,
+    scalar_t* /*ddsph*/,
     void* /*cuda_stream*/
 ) {
 
@@ -24,9 +24,9 @@ void sphericart::cuda::spherical_harmonics_cuda_base(
 }
 
 template void sphericart::cuda::spherical_harmonics_cuda_base<double>(
-    const double* __restrict__ xyz,
+    const double* xyz,
     const int nedges,
-    const double* __restrict__ prefactors,
+    const double* prefactors,
     const int nprefactors,
     const int64_t l_max,
     const bool normalize,
@@ -34,16 +34,16 @@ template void sphericart::cuda::spherical_harmonics_cuda_base<double>(
     const int64_t GRID_DIM_Y,
     const bool gradients,
     const bool hessian,
-    double* __restrict__ sph,
-    double* __restrict__ dsph,
-    double* __restrict__ ddsph,
+    double* sph,
+    double* dsph,
+    double* ddsph,
     void* cuda_stream
 );
 
 template void sphericart::cuda::spherical_harmonics_cuda_base<float>(
-    const float* __restrict__ xyz,
+    const float* xyz,
     const int nedges,
-    const float* __restrict__ prefactors,
+    const float* prefactors,
     const int nprefactors,
     const int64_t l_max,
     const bool normalize,
@@ -51,27 +51,27 @@ template void sphericart::cuda::spherical_harmonics_cuda_base<float>(
     const int64_t GRID_DIM_Y,
     const bool gradients,
     const bool hessian,
-    float* __restrict__ sph,
-    float* __restrict__ dsph,
-    float* __restrict__ ddsph,
+    float* sph,
+    float* dsph,
+    float* ddsph,
     void* cuda_stream
 );
 
 template <typename scalar_t>
 void sphericart::cuda::spherical_harmonics_backward_cuda_base(
-    const scalar_t* __restrict__ /*dsph*/,
-    const scalar_t* __restrict__ /*sph_grad*/,
+    const scalar_t* /*dsph*/,
+    const scalar_t* /*sph_grad*/,
     const int /*nedges*/,
     const int /*ntotal*/,
-    scalar_t* __restrict__ /*xyz_grad*/,
+    scalar_t* /*xyz_grad*/,
     void* /*cuda_stream*/
 ) {
     throw std::runtime_error("sphericart was not compiled with CUDA support");
 }
 
 template void sphericart::cuda::spherical_harmonics_backward_cuda_base<float>(
-    const float* __restrict__ dsph,
-    const float* __restrict__ sph_grad,
+    const float* dsph,
+    const float* sph_grad,
     const int nedges,
     const int ntotal,
     float* __restrict__ xyz_grad,
@@ -79,15 +79,10 @@ template void sphericart::cuda::spherical_harmonics_backward_cuda_base<float>(
 );
 
 template void sphericart::cuda::spherical_harmonics_backward_cuda_base<double>(
-    const double* __restrict__ dsph,
-    const double* __restrict__ sph_grad,
+    const double* dsph,
+    const double* sph_grad,
     const int nedges,
     const int ntotal,
     double* __restrict__ xyz_grad,
     void* cuda_stream
-);
-
-int sphericart::cuda::adjust_shared_memory(size_t, int64_t, int64_t, int64_t, bool, bool, int64_t) {
-    throw std::runtime_error("sphericart was not compiled with CUDA support");
-    return -1;
-}
+);
\ No newline at end of file
diff --git a/sphericart/src/sphericart_cuda.cpp b/sphericart/src/sphericart_cuda.cpp
new file mode 100644
index 00000000..d284c5a9
--- /dev/null
+++ b/sphericart/src/sphericart_cuda.cpp
@@ -0,0 +1,188 @@
+#include <iostream>
+#include <stdexcept>
+
+#define _SPHERICART_INTERNAL_IMPLEMENTATION
+
+#include "dynamic_cuda.hpp"
+#include "sphericart.hpp"
+#include "cuda_base.hpp"
+#include "sphericart_cuda.hpp"
+
+using namespace sphericart::cuda;
+
+/*
+    This code checks whether the cuda libraries we'd need to dynamically load are available on the host
+*/
+
+void checkCuda() {
+    if (!CUDA_DRIVER_INSTANCE.loaded()) {
+        throw std::runtime_error(
+            "Failed to load libcuda.so. Try appending the directory containing this library to "
+            "your $LD_LIBRARY_PATH environment variable."
+        );
+    }
+
+    if (!CUDART_INSTANCE.loaded()) {
+        throw std::runtime_error(
+            "Failed to load libcudart.so. Try appending the directory containing this library to "
+            "your $LD_LIBRARY_PATH environment variable."
+        );
+    }
+
+    if (!NVRTC_INSTANCE.loaded()) {
+        throw std::runtime_error(
+            "Failed to load libnvrtc.so. Try appending the directory containing this library to "
+            "your $LD_LIBRARY_PATH environment variable."
+        );
+    }
+}
+
+template <typename T> SphericalHarmonics<T>::SphericalHarmonics(size_t l_max) {
+    /*
+        This is the constructor of the SphericalHarmonics class. It initizlizes
+       buffer space, compute prefactors, and sets the function pointers that are
+       used for the actual calls
+    */
+
+    checkCuda();
+
+    this->l_max = (int)l_max;
+    this->nprefactors = (int)(l_max + 1) * (l_max + 2);
+    this->normalized = true; // SphericalHarmonics class
+    this->prefactors_cpu = new T[this->nprefactors];
+
+    CUDART_SAFE_CALL(CUDART_INSTANCE.cudaGetDeviceCount(&this->device_count));
+
+    // compute prefactors on host
+    compute_sph_prefactors<T>((int)l_max, this->prefactors_cpu);
+}
+
+template <typename T> SphericalHarmonics<T>::~SphericalHarmonics() {
+    // Destructor, frees the prefactors
+    if (this->prefactors_cpu != nullptr) {
+        delete[] (this->prefactors_cpu);
+        this->prefactors_cpu = nullptr;
+    }
+
+    if (this->prefactors_cuda) {
+        CUDART_SAFE_CALL(CUDART_INSTANCE.cudaFree(this->prefactors_cuda));
+    }
+}
+
+template <typename T>
+void SphericalHarmonics<T>::compute_internal(
+    const T* xyz,
+    const size_t n_samples,
+    bool compute_with_gradients,
+    bool compute_with_hessian,
+    T* sph,
+    T* dsph,
+    T* ddsph,
+    void* cuda_stream
+) {
+    if (n_samples == 0) {
+        // nothing to compute; we return here because some libraries (e.g. torch)
+        // seem to use nullptrs for tensors with 0 elements
+        return;
+    }
+
+    if (sph == nullptr) {
+        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
+                                 "sph ptr initialised, instead nullptr found. Initialise "
+                                 "sph with cudaMalloc.");
+    }
+
+    if (compute_with_gradients && dsph == nullptr) {
+        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
+                                 "dsph != nullptr since compute_with_gradients = true. "
+                                 "initialise dsph with cudaMalloc.");
+    }
+
+    if (compute_with_hessian && ddsph == nullptr) {
+        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
+                                 "ddsph != nullptr since compute_with_hessian = true. "
+                                 "initialise ddsph with cudaMalloc.");
+    }
+
+    cudaPointerAttributes attributes;
+
+    CUDART_SAFE_CALL(CUDART_INSTANCE.cudaPointerGetAttributes(&attributes, xyz));
+
+    int current_device;
+
+    CUDART_SAFE_CALL(CUDART_INSTANCE.cudaGetDevice(&current_device));
+
+    if (current_device != attributes.device) {
+        CUDART_SAFE_CALL(CUDART_INSTANCE.cudaSetDevice(attributes.device));
+    }
+
+    if (!this->prefactors_cuda) {
+        CUDART_SAFE_CALL(
+            CUDART_INSTANCE.cudaMalloc((void**)&this->prefactors_cuda, this->nprefactors * sizeof(T))
+        );
+
+        CUDART_SAFE_CALL(CUDART_INSTANCE.cudaMemcpy(
+            this->prefactors_cuda,
+            this->prefactors_cpu,
+            this->nprefactors * sizeof(T),
+            cudaMemcpyHostToDevice
+        ));
+    }
+
+    sphericart::cuda::spherical_harmonics_cuda_base<T>(
+        xyz,
+        n_samples,
+        this->prefactors_cuda,
+        this->nprefactors,
+        this->l_max,
+        this->normalized,
+        this->CUDA_GRID_DIM_X_,
+        this->CUDA_GRID_DIM_Y_,
+        compute_with_gradients,
+        compute_with_hessian,
+        sph,
+        dsph,
+        ddsph,
+        cuda_stream
+    );
+
+    if (current_device != attributes.device) {
+        CUDART_SAFE_CALL(CUDART_INSTANCE.cudaSetDevice(current_device));
+    }
+}
+template <typename T>
+void SphericalHarmonics<T>::compute(const T* xyz, const size_t n_samples, T* sph, void* cuda_stream) {
+    SphericalHarmonics<T>::compute_internal(
+        xyz, n_samples, false, false, sph, nullptr, nullptr, cuda_stream
+    );
+}
+
+template <typename T>
+void SphericalHarmonics<T>::compute_with_gradients(
+    const T* xyz, const size_t n_samples, T* sph, T* dsph, void* cuda_stream
+) {
+    SphericalHarmonics<T>::compute_internal(
+        xyz, n_samples, true, false, sph, dsph, nullptr, cuda_stream
+    );
+}
+
+template <typename T>
+void SphericalHarmonics<T>::compute_with_hessians(
+    const T* xyz, const size_t n_samples, T* sph, T* dsph, T* ddsph, void* cuda_stream
+) {
+    SphericalHarmonics<T>::compute_internal(
+        xyz, n_samples, true, true, sph, dsph, ddsph, cuda_stream
+    );
+}
+
+template <typename T>
+SolidHarmonics<T>::SolidHarmonics(size_t l_max) : SphericalHarmonics<T>(l_max) {
+    this->normalized = false; // SolidHarmonics class
+}
+
+// instantiates the SphericalHarmonics and SolidHarmonics classes
+// for basic floating point types
+template class sphericart::cuda::SphericalHarmonics<float>;
+template class sphericart::cuda::SphericalHarmonics<double>;
+template class sphericart::cuda::SolidHarmonics<float>;
+template class sphericart::cuda::SolidHarmonics<double>;
diff --git a/sphericart/src/sphericart_cuda.cu b/sphericart/src/sphericart_cuda.cu
deleted file mode 100644
index b6f4f124..00000000
--- a/sphericart/src/sphericart_cuda.cu
+++ /dev/null
@@ -1,248 +0,0 @@
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <iostream>
-#include <stdexcept>
-
-#define _SPHERICART_INTERNAL_IMPLEMENTATION
-#include "cuda_base.hpp"
-#include "sphericart_cuda.hpp"
-
-/*host macro that checks for errors in CUDA calls, and prints the file + line
- * and error string if one occurs
- */
-
-using namespace sphericart::cuda;
-
-#define CUDA_CHECK(call)                                                                           \
-    do {                                                                                           \
-        cudaError_t cudaStatus = (call);                                                           \
-        if (cudaStatus != cudaSuccess) {                                                           \
-            std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << " - "                  \
-                      << cudaGetErrorString(cudaStatus) << std::endl;                              \
-            cudaDeviceReset();                                                                     \
-            exit(EXIT_FAILURE);                                                                    \
-        }                                                                                          \
-    } while (0)
-
-template <typename T> SphericalHarmonics<T>::SphericalHarmonics(size_t l_max) {
-    /*
-        This is the constructor of the SphericalHarmonics class. It initizlizes
-       buffer space, compute prefactors, and sets the function pointers that are
-       used for the actual calls
-    */
-    this->l_max = (int)l_max;
-    this->nprefactors = (int)(l_max + 1) * (l_max + 2);
-    this->normalized = true; // SphericalHarmonics class
-    this->prefactors_cpu = new T[this->nprefactors];
-
-    CUDA_CHECK(cudaGetDeviceCount(&this->device_count));
-
-    // compute prefactors on host first
-    compute_sph_prefactors<T>((int)l_max, this->prefactors_cpu);
-
-    if (this->device_count) {
-        int current_device;
-
-        CUDA_CHECK(cudaGetDevice(&current_device));
-
-        // allocate prefactorts on every visible device and copy from host
-        this->prefactors_cuda = new T*[this->device_count];
-
-        for (int device = 0; device < this->device_count; device++) {
-            CUDA_CHECK(cudaSetDevice(device));
-            CUDA_CHECK(
-                cudaMalloc((void**)&this->prefactors_cuda[device], this->nprefactors * sizeof(T))
-            );
-            CUDA_CHECK(cudaMemcpy(
-                this->prefactors_cuda[device],
-                this->prefactors_cpu,
-                this->nprefactors * sizeof(T),
-                cudaMemcpyHostToDevice
-            ));
-        }
-
-        // initialise the currently available amount of shared memory on all visible devices
-        this->_current_shared_mem_allocation = adjust_shared_memory(
-            sizeof(T),
-            this->l_max,
-            this->CUDA_GRID_DIM_X_,
-            this->CUDA_GRID_DIM_Y_,
-            false,
-            false,
-            this->_current_shared_mem_allocation
-        );
-
-        // set the context back to the current device
-        CUDA_CHECK(cudaSetDevice(current_device));
-    }
-}
-
-template <typename T> SphericalHarmonics<T>::~SphericalHarmonics() {
-    // Destructor, frees the prefactors
-    if (this->prefactors_cpu != nullptr) {
-        delete[] (this->prefactors_cpu);
-        this->prefactors_cpu = nullptr;
-    }
-
-    if (this->device_count > 0) {
-
-        int current_device;
-
-        CUDA_CHECK(cudaGetDevice(&current_device));
-
-        for (int device = 0; device < this->device_count; device++) {
-            CUDA_CHECK(cudaSetDevice(device));
-            CUDA_CHECK(cudaDeviceSynchronize());
-            if (this->prefactors_cuda != nullptr && this->prefactors_cuda[device] != nullptr) {
-                CUDA_CHECK(cudaFree(this->prefactors_cuda[device]));
-                this->prefactors_cuda[device] = nullptr;
-            }
-        }
-        this->prefactors_cuda = nullptr;
-
-        CUDA_CHECK(cudaSetDevice(current_device));
-    }
-}
-
-template <typename T>
-void SphericalHarmonics<T>::compute_internal(
-    const T* xyz,
-    const size_t n_samples,
-    bool compute_with_gradients,
-    bool compute_with_hessian,
-    T* sph,
-    T* dsph,
-    T* ddsph,
-    void* cuda_stream
-) {
-    if (n_samples == 0) {
-        // nothing to compute; we return here because some libraries (e.g. torch)
-        // seem to use nullptrs for tensors with 0 elements
-        return;
-    }
-
-    if (sph == nullptr) {
-        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
-                                 "sph ptr initialised, instead nullptr found. Initialise "
-                                 "sph with cudaMalloc.");
-    }
-
-    if (compute_with_gradients && dsph == nullptr) {
-        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
-                                 "dsph != nullptr since compute_with_gradients = true. "
-                                 "initialise dsph with cudaMalloc.");
-    }
-
-    if (compute_with_hessian && ddsph == nullptr) {
-        throw std::runtime_error("sphericart::cuda::SphericalHarmonics::compute expected "
-                                 "ddsph != nullptr since compute_with_hessian = true. "
-                                 "initialise ddsph with cudaMalloc.");
-    }
-
-    if (this->cached_compute_with_gradients != compute_with_gradients ||
-        this->cached_compute_with_hessian != compute_with_hessian) {
-
-        this->_current_shared_mem_allocation = adjust_shared_memory(
-            sizeof(T),
-            this->l_max,
-            this->CUDA_GRID_DIM_X_,
-            this->CUDA_GRID_DIM_Y_,
-            compute_with_gradients,
-            compute_with_hessian,
-            this->_current_shared_mem_allocation
-        );
-
-        if (this->_current_shared_mem_allocation == -1) {
-
-            std::cerr << "Warning: Failed to update shared memory size, "
-                         "re-attempting with  GRID_DIM_Y = 4\n"
-                      << std::endl;
-
-            this->CUDA_GRID_DIM_Y_ = 4;
-            this->_current_shared_mem_allocation = adjust_shared_memory(
-                sizeof(T),
-                this->l_max,
-                this->CUDA_GRID_DIM_X_,
-                this->CUDA_GRID_DIM_Y_,
-                compute_with_gradients,
-                compute_with_hessian,
-                this->_current_shared_mem_allocation
-            );
-
-            if (this->_current_shared_mem_allocation == -1) {
-                throw std::runtime_error("Insufficient shared memory available to compute "
-                                         "spherical_harmonics with requested parameters.");
-            }
-        }
-
-        this->cached_compute_with_gradients = compute_with_gradients;
-        this->cached_compute_with_hessian = compute_with_hessian;
-    }
-
-    cudaPointerAttributes attributes;
-
-    CUDA_CHECK(cudaPointerGetAttributes(&attributes, xyz));
-
-    int current_device;
-
-    CUDA_CHECK(cudaGetDevice(&current_device));
-
-    if (current_device != attributes.device) {
-        CUDA_CHECK(cudaSetDevice(attributes.device));
-    }
-
-    sphericart::cuda::spherical_harmonics_cuda_base<T>(
-        xyz,
-        n_samples,
-        this->prefactors_cuda[attributes.device],
-        this->nprefactors,
-        this->l_max,
-        this->normalized,
-        this->CUDA_GRID_DIM_X_,
-        this->CUDA_GRID_DIM_Y_,
-        compute_with_gradients,
-        compute_with_hessian,
-        sph,
-        dsph,
-        ddsph,
-        cuda_stream
-    );
-
-    CUDA_CHECK(cudaSetDevice(current_device));
-}
-template <typename T>
-void SphericalHarmonics<T>::compute(const T* xyz, const size_t n_samples, T* sph, void* cuda_stream) {
-    SphericalHarmonics<T>::compute_internal(
-        xyz, n_samples, false, false, sph, nullptr, nullptr, cuda_stream
-    );
-}
-
-template <typename T>
-void SphericalHarmonics<T>::compute_with_gradients(
-    const T* xyz, const size_t n_samples, T* sph, T* dsph, void* cuda_stream
-) {
-    SphericalHarmonics<T>::compute_internal(
-        xyz, n_samples, true, false, sph, dsph, nullptr, cuda_stream
-    );
-}
-
-template <typename T>
-void SphericalHarmonics<T>::compute_with_hessians(
-    const T* xyz, const size_t n_samples, T* sph, T* dsph, T* ddsph, void* cuda_stream
-) {
-    SphericalHarmonics<T>::compute_internal(
-        xyz, n_samples, true, true, sph, dsph, ddsph, cuda_stream
-    );
-}
-
-template <typename T>
-SolidHarmonics<T>::SolidHarmonics(size_t l_max) : SphericalHarmonics<T>(l_max) {
-    this->normalized = false; // SolidHarmonics class
-}
-
-// instantiates the SphericalHarmonics and SolidHarmonics classes
-// for basic floating point types
-template class sphericart::cuda::SphericalHarmonics<float>;
-template class sphericart::cuda::SphericalHarmonics<double>;
-template class sphericart::cuda::SolidHarmonics<float>;
-template class sphericart::cuda::SolidHarmonics<double>;
diff --git a/sphericart/src/cuda_base.cu b/sphericart/src/sphericart_impl.cu
similarity index 66%
rename from sphericart/src/cuda_base.cu
rename to sphericart/src/sphericart_impl.cu
index 1e70c19d..e8ca608f 100644
--- a/sphericart/src/cuda_base.cu
+++ b/sphericart/src/sphericart_impl.cu
@@ -1,50 +1,13 @@
-#include <cmath>
-#include <iostream>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#define _SPHERICART_INTERNAL_IMPLEMENTATION
-#define CUDA_DEVICE_PREFIX __device__
-
-#include "cuda_base.hpp"
 
 #define HARDCODED_LMAX 1
 
 /* MASK used for warp reductions */
 #define FULL_MASK 0xffffffff
 
-#define CUDA_CHECK(call)                                                                           \
-    do {                                                                                           \
-        cudaError_t cudaStatus = (call);                                                           \
-        if (cudaStatus != cudaSuccess) {                                                           \
-            std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << " - "                  \
-                      << cudaGetErrorString(cudaStatus) << std::endl;                              \
-            cudaDeviceReset();                                                                     \
-            exit(EXIT_FAILURE);                                                                    \
-        }                                                                                          \
-    } while (0)
-
-#define CUDA_CHECK_KERNEL()                                                                        \
-    do {                                                                                           \
-        cudaDeviceSynchronize();                                                                   \
-        cudaError_t err = cudaGetLastError();                                                      \
-        if (err != cudaSuccess) {                                                                  \
-            fprintf(                                                                               \
-                stderr,                                                                            \
-                "CUDA error after kernel launch in %s at line %d: %s\n",                           \
-                __FILE__,                                                                          \
-                __LINE__,                                                                          \
-                cudaGetErrorString(err)                                                            \
-            );                                                                                     \
-            exit(EXIT_FAILURE);                                                                    \
-        }                                                                                          \
-    } while (0)
-
 /*
     Computes the index for buffer values which are shared across GRID_DIM_Y
 */
-__device__ int get_index(int i) { return i * blockDim.y + threadIdx.y; }
+__device__ inline int get_index(int i) { return i * blockDim.y + threadIdx.y; }
 
 /*
     Clears the shared memory buffers for the spherical harmonics and gradients
@@ -104,8 +67,8 @@ __device__ inline void clear_buffers(
 */
 template <typename scalar_t>
 __device__ inline void write_buffers(
-    size_t edge_idx,
-    size_t nedges,
+    int edge_idx,
+    int nedges,
     scalar_t x,
     scalar_t y,
     scalar_t z,
@@ -132,7 +95,7 @@ __device__ inline void write_buffers(
     scalar_t* sph,
     scalar_t* dsph,
     scalar_t* ddsph,
-    size_t n_total,
+    int n_total,
     bool requires_grad,
     bool requires_hessian,
     bool normalize
@@ -239,23 +202,23 @@ __device__ inline void write_buffers(
 */
 template <typename scalar_t>
 __global__ void spherical_harmonics_kernel(
-    const scalar_t* __restrict__ xyz,
+    scalar_t* xyz,
     int nedges,
-    const scalar_t* __restrict__ prefactors,
+    scalar_t* prefactors,
     int nprefactors,
     int lmax,
     int ntotal,
     bool requires_grad,
     bool requires_hessian,
     bool normalize,
-    scalar_t* __restrict__ sph,
-    scalar_t* __restrict__ dsph,
-    scalar_t* __restrict__ ddsph
+    scalar_t* sph,
+    scalar_t* dsph,
+    scalar_t* ddsph
 ) {
 
     extern __shared__ char buffer[];
 
-    size_t offset = 0;
+    int offset = 0;
 
     scalar_t* buffer_c = reinterpret_cast<scalar_t*>(buffer + offset);
     offset += blockDim.y * (lmax + 1) * sizeof(scalar_t);
@@ -317,7 +280,7 @@ __global__ void spherical_harmonics_kernel(
         offset += blockDim.y * nl * sizeof(scalar_t);
     }
 
-    size_t edge_idx = blockIdx.x * blockDim.y + threadIdx.y;
+    int edge_idx = blockIdx.x * blockDim.y + threadIdx.y;
 
     scalar_t x = 0.0;
     scalar_t y = 0.0;
@@ -647,205 +610,34 @@ __global__ void spherical_harmonics_kernel(
     }
 }
 
-/*
-    Computes the total amount of shared memory space required by
-   spherical_harmonics_kernel.
-
-    For lmax <= HARCODED_LMAX, we need to store all (HARDCODED_LMAX + 1)**2
-   scalars in shared memory. For lmax > HARDCODED_LMAX, we only need to store
-   each spherical harmonics vector per sample in shared memory.
-*/
-static size_t total_buffer_size(
-    size_t l_max,
-    size_t GRID_DIM_X,
-    size_t GRID_DIM_Y,
-    size_t dtype_size,
-    bool requires_grad,
-    bool requires_hessian
-) {
-    int nl = max(static_cast<size_t>((HARDCODED_LMAX + 1) * (HARDCODED_LMAX + 1)), 2 * l_max + 1);
-
-    size_t total_buff_size = 0;
-
-    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_c
-    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_s
-    total_buff_size += GRID_DIM_Y * (l_max + 1) * dtype_size;  // buffer_twomz
-    total_buff_size += (l_max + 1) * (l_max + 2) * dtype_size; // buffer_prefactors
-    total_buff_size += GRID_DIM_Y * nl * dtype_size;           // buffer_sph_out
-
-    if (requires_grad) {
-        total_buff_size += 3 * GRID_DIM_Y * nl * dtype_size; // buffer_sph_derivs
-    }
-
-    if (requires_hessian) {
-        total_buff_size += 9 * GRID_DIM_Y * nl * dtype_size; // buffer_sph_hessian
-    }
-
-    return total_buff_size;
-}
-
-/*
-    The default shared memory space on most recent NVIDIA cards is defaulted
-   49152 bytes, regarldess if there is more available per SM. This method
-   attempts to adjust the shared memory to fit the requested configuration if
-   the kernel launch parameters exceeds the default 49152 bytes.
-*/
-
-int sphericart::cuda::adjust_shared_memory(
-    size_t element_size,
-    int64_t l_max,
-    int64_t GRID_DIM_X,
-    int64_t GRID_DIM_Y,
+template __global__ void spherical_harmonics_kernel<float>(
+    float* xyz,
+    int nedges,
+    float* prefactors,
+    int nprefactors,
+    int lmax,
+    int ntotal,
     bool requires_grad,
     bool requires_hessian,
-    int64_t current_shared_mem_alloc
-) {
-
-    int device_count;
-
-    cudaGetDeviceCount(&device_count);
-
-    int current_device;
-    cudaGetDevice(&current_device);
-
-    cudaDeviceProp deviceProp;
-    cudaGetDeviceProperties(&deviceProp, current_device);
-
-    auto required_buff_size = total_buffer_size(
-        l_max, GRID_DIM_X, GRID_DIM_Y, element_size, requires_grad, requires_hessian
-    );
-
-    if (required_buff_size > current_shared_mem_alloc &&
-        required_buff_size > (deviceProp.sharedMemPerBlock - deviceProp.reservedSharedMemPerBlock)) {
-
-        if (required_buff_size > deviceProp.sharedMemPerBlockOptin) {
-            return -1; // failure - need to adjust parameters
-        }
-
-        // broadcast changes to all visible GPUs
-        for (int device = 0; device < device_count; device++) {
-
-            CUDA_CHECK(cudaSetDevice(device));
-
-            switch (element_size) {
-            case 8:
-                cudaFuncSetAttribute(
-                    spherical_harmonics_kernel<double>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    required_buff_size
-                );
-                break;
-            case 4:
-                cudaFuncSetAttribute(
-                    spherical_harmonics_kernel<float>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    required_buff_size
-                );
-                break;
-            }
-        }
-
-        CUDA_CHECK(cudaSetDevice(current_device));
-
-        return required_buff_size;
-
-    } else {
-        return (current_shared_mem_alloc >
-                (deviceProp.sharedMemPerBlock - deviceProp.reservedSharedMemPerBlock))
-                   ? current_shared_mem_alloc
-                   : (deviceProp.sharedMemPerBlock - deviceProp.reservedSharedMemPerBlock);
-    }
-}
-
-/*
-    Wrapper to launch the CUDA kernel. Returns a vector containing the spherical
-   harmonics and their gradients if required, otherwise returns the spherical
-   harmonics and an empty tensor.
-
-    GRID_DIM_X is the number of threads to launch in the x dimension. Used to
-   parallelize over the sample dimension. GRID_DIM_Y is the number of threads to
-   launch in the y dimension. Used only to improve memory throughput on reads
-   and writes.
-
-    Total number of threads used is GRID_DIM_X * GRID_DIM_Y.
-
-    cuda_stream should be of type (void *), therefore if you want to pass in
-    a cudaStream_t, first do void * stream_ptr = reinterpret_cast<void *>
-   (stream);
-*/
-
-template <typename scalar_t>
-void sphericart::cuda::spherical_harmonics_cuda_base(
-    const scalar_t* __restrict__ xyz,
-    const int nedges,
-    const scalar_t* __restrict__ prefactors,
-    const int nprefactors,
-    const int64_t l_max,
-    const bool normalize,
-    const int64_t GRID_DIM_X,
-    const int64_t GRID_DIM_Y,
-    const bool gradients,
-    const bool hessian,
-    scalar_t* __restrict__ sph,
-    scalar_t* __restrict__ dsph,
-    scalar_t* __restrict__ ddsph,
-    void* cuda_stream
-) {
-
-    int n_total = (l_max + 1) * (l_max + 1);
-
-    dim3 grid_dim(GRID_DIM_X, GRID_DIM_Y);
-
-    auto find_num_blocks = [](int x, int bdim) { return (x + bdim - 1) / bdim; };
-
-    cudaStream_t cstream = reinterpret_cast<cudaStream_t>(cuda_stream);
-
-    dim3 block_dim(find_num_blocks(nedges, GRID_DIM_Y));
-
-    size_t total_buff_size =
-        total_buffer_size(l_max, GRID_DIM_X, GRID_DIM_Y, sizeof(scalar_t), gradients, hessian);
-
-    spherical_harmonics_kernel<scalar_t><<<block_dim, grid_dim, total_buff_size, cstream>>>(
-        xyz, nedges, prefactors, nprefactors, l_max, n_total, gradients, hessian, normalize, sph, dsph, ddsph
-    );
-
-    CUDA_CHECK_KERNEL();
-
-    CUDA_CHECK(cudaStreamSynchronize(cstream));
-}
-
-template void sphericart::cuda::spherical_harmonics_cuda_base<float>(
-    const float* __restrict__ xyz,
-    const int nedges,
-    const float* __restrict__ prefactors,
-    const int nprefactors,
-    const int64_t l_max,
-    const bool normalize,
-    const int64_t GRID_DIM_X,
-    const int64_t GRID_DIM_Y,
-    const bool gradients,
-    const bool hessian,
-    float* __restrict__ sph,
-    float* __restrict__ dsph,
-    float* __restrict__ ddsph,
-    void* cuda_stream
+    bool normalize,
+    float* sph,
+    float* dsph,
+    float* ddsph
 );
 
-template void sphericart::cuda::spherical_harmonics_cuda_base<double>(
-    const double* __restrict__ xyz,
-    const int nedges,
-    const double* __restrict__ prefactors,
-    const int nprefactors,
-    const int64_t l_max,
-    const bool normalize,
-    const int64_t GRID_DIM_X,
-    const int64_t GRID_DIM_Y,
-    const bool gradients,
-    const bool hessian,
-    double* __restrict__ sph,
-    double* __restrict__ dsph,
-    double* __restrict__ ddsph,
-    void* cuda_stream
+template __global__ void spherical_harmonics_kernel<double>(
+    double* xyz,
+    int nedges,
+    double* prefactors,
+    int nprefactors,
+    int lmax,
+    int ntotal,
+    bool requires_grad,
+    bool requires_hessian,
+    bool normalize,
+    double* sph,
+    double* dsph,
+    double* ddsph
 );
 
 /*
@@ -853,14 +645,10 @@ template void sphericart::cuda::spherical_harmonics_cuda_base<double>(
 */
 template <typename scalar_t>
 __global__ void backward_kernel(
-    const scalar_t* __restrict__ dsph,
-    const scalar_t* __restrict__ sph_grad,
-    size_t nedges,
-    size_t n_total,
-    scalar_t* __restrict__ xyz_grad
+    scalar_t* dsph, scalar_t* sph_grad, int nedges, int n_total, scalar_t* xyz_grad
 ) {
 
-    size_t edge_idx = blockIdx.x * blockDim.y + threadIdx.y;
+    int edge_idx = blockIdx.x * blockDim.y + threadIdx.y;
 
     int spatial = blockIdx.y;
 
@@ -891,45 +679,10 @@ __global__ void backward_kernel(
     }
 }
 
-template <typename scalar_t>
-void sphericart::cuda::spherical_harmonics_backward_cuda_base(
-    const scalar_t* __restrict__ dsph,
-    const scalar_t* __restrict__ sph_grad,
-    const int nedges,
-    const int ntotal,
-    scalar_t* __restrict__ xyz_grad,
-    void* cuda_stream
-) {
-    dim3 grid_dim(4, 32);
-
-    auto find_num_blocks = [](int x, int bdim) { return (x + bdim - 1) / bdim; };
-
-    dim3 block_dim(find_num_blocks(nedges, 32), 3);
-
-    cudaStream_t cstream = reinterpret_cast<cudaStream_t>(cuda_stream);
-
-    backward_kernel<scalar_t>
-        <<<block_dim, grid_dim, 0, cstream>>>(dsph, sph_grad, nedges, ntotal, xyz_grad);
-
-    CUDA_CHECK_KERNEL();
-
-    cudaStreamSynchronize(cstream);
-}
-
-template void sphericart::cuda::spherical_harmonics_backward_cuda_base<float>(
-    const float* __restrict__ dsph,
-    const float* __restrict__ sph_grad,
-    const int nedges,
-    const int ntotal,
-    float* __restrict__ xyz_grad,
-    void* cuda_stream
+template __global__ void backward_kernel<float>(
+    float* dsph, float* sph_grad, int nedges, int n_total, float* xyz_grad
 );
 
-template void sphericart::cuda::spherical_harmonics_backward_cuda_base<double>(
-    const double* __restrict__ dsph,
-    const double* __restrict__ sph_grad,
-    const int nedges,
-    const int ntotal,
-    double* __restrict__ xyz_grad,
-    void* cuda_stream
-);
+template __global__ void backward_kernel<double>(
+    double* dsph, double* sph_grad, int nedges, int n_total, double* xyz_grad
+);
\ No newline at end of file
diff --git a/tox.ini b/tox.ini
index ef0d218a..aa83f3f9 100644
--- a/tox.ini
+++ b/tox.ini
@@ -83,7 +83,7 @@ commands =
 
     # Install this one manually. Listing it in the deps list above does not install jaxlib.
     # Note: jax[cuda12] is not available on Windows and MacOS.
-    bash -c "command -v nvcc && python -m pip install -U jax[cuda12] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U jax[cpu]"
+    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U jax[cpu]'
 
     pip install {[testenv]pip_install_flags} .
     pytest python
@@ -99,12 +99,15 @@ deps =
 
     numpy<2.0.0
     torch
-    jax[cpu]
     pytest
 
 passenv=
     PIP_EXTRA_INDEX_URL
 commands =
+    # Install this one manually. Listing it in the deps list above does not install jaxlib.
+    # Note: jax[cuda12] is not available on Windows and MacOS.
+    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U jax[cpu]'
+
     pip install {[testenv]pip_install_flags} .
     pip install {[testenv]pip_install_flags} ./sphericart-torch
     pip install {[testenv]pip_install_flags} ./sphericart-jax