Repository: SHI-Labs/NATTEN (language: python)

Issue:
Neighborhood Attention 3D
Sorry, could you share the  Neighborhood Attention 3D (CUDA)?  ice_rain123@foxmail.com


Discussion / hints:
Hello and thank you for your interest.

Unfortunately we don't have a stable release for the 3D module yet.
I'll mark this as an enhancement and close it when we actually release it.
Hi, Is there any update for the 3D version? Thx!
@alihassanijr Kindly let us know if any updates on 3D NA. Thanks in Advance.
Thank you for your interest @sheshap @btwbtm @wenyihong @icerain-alt .

We will release an initial version of 3D NA this week.

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1911170..fb7de5c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## [0.14.7]
+* Refactored kernels
+  * The backend is messy, particularly the CUDA backend. A step in the right direction is at least factoring out duplicated.
+  * Out of the 7 operations in NATTEN's backend, 6 have duplicates (really 3 underlying ops with different inputs.)
+  * See #26 for more details.
+* 3D NA
+
 ## [0.14.6] - 2023-03-21
 Just a really small update that syncs the changes to the private branch.
 It's mostly about the changed signature in both QK and AV (both 1D and 2D), where we now take in both kernel size and dilation.
diff --git a/README.md b/README.md
index 16f1363..d8a87a4 100644
--- a/README.md
+++ b/README.md
@@ -95,10 +95,10 @@ make test
 ## Catalog
 - [x] Neighborhood Attention 1D (CUDA)
 - [x] Neighborhood Attention 2D (CUDA)
-- [ ] Neighborhood Attention 3D (CUDA)
+- [x] Neighborhood Attention 3D (CUDA)
 - [x] Neighborhood Attention 1D (CPU)
 - [x] Neighborhood Attention 2D (CPU)
-- [ ] Neighborhood Attention 3D (CPU)
+- [x] Neighborhood Attention 3D (CPU)
 - [x] Dilation support
 - [x] Float16 support and utilization
 - [ ] BFloat16 support
@@ -106,15 +106,34 @@ make test
 - [ ] Windows builds
 
 ## Usage
-Simply import `NeighborhoodAttention1D` or `NeighborhoodAttention2D` from `natten`:
+Simply import `NeighborhoodAttention1D`, `NeighborhoodAttention2D`, or `NeighborhoodAttention3D` from `natten`:
 ```python
 from natten import NeighborhoodAttention1D
 from natten import NeighborhoodAttention2D
+from natten import NeighborhoodAttention3D
 
 na1d = NeighborhoodAttention1D(dim=128, kernel_size=7, dilation=2, num_heads=4)
 na2d = NeighborhoodAttention2D(dim=128, kernel_size=7, dilation=2, num_heads=4)
+na3d = NeighborhoodAttention3D(dim=128, kernel_size=7, dilation=2, num_heads=4)
 ```
 
+NA3D also supports different kernel size and dilation values for depth:
+```python
+na3d = NeighborhoodAttention3D(
+	dim=128,
+	kernel_size=7,
+	kernel_size_d=5,
+	dilation=2,
+	dilation_d=3,
+	num_heads=4)
+```
+
+Modules expect inputs of shape `[batch_size, *, dim]`:
+* NA1D: `[batch_size, sequence_length, dim]`
+* NA2D: `[batch_size, height, width, dim]`
+* NA3D: `[batch_size, depth, height, width, dim]`
+
+
 ### FLOPs
 We recommend counting flops through [fvcore](https://github.com/facebookresearch/fvcore).
 
diff --git a/assets/README_pypi.md b/assets/README_pypi.md
index 9f8122c..5a79028 100644
--- a/assets/README_pypi.md
+++ b/assets/README_pypi.md
@@ -95,10 +95,10 @@ make test
 ## Catalog
 - [x] Neighborhood Attention 1D (CUDA)
 - [x] Neighborhood Attention 2D (CUDA)
-- [ ] Neighborhood Attention 3D (CUDA)
+- [x] Neighborhood Attention 3D (CUDA)
 - [x] Neighborhood Attention 1D (CPU)
 - [x] Neighborhood Attention 2D (CPU)
-- [ ] Neighborhood Attention 3D (CPU)
+- [x] Neighborhood Attention 3D (CPU)
 - [x] Dilation support
 - [x] Float16 support and utilization
 - [ ] BFloat16 support
@@ -106,15 +106,34 @@ make test
 - [ ] Windows builds
 
 ## Usage
-Simply import `NeighborhoodAttention1D` or `NeighborhoodAttention2D` from `natten`:
+Simply import `NeighborhoodAttention1D`, `NeighborhoodAttention2D`, or `NeighborhoodAttention3D` from `natten`:
 ```python
 from natten import NeighborhoodAttention1D
 from natten import NeighborhoodAttention2D
+from natten import NeighborhoodAttention3D
 
 na1d = NeighborhoodAttention1D(dim=128, kernel_size=7, dilation=2, num_heads=4)
 na2d = NeighborhoodAttention2D(dim=128, kernel_size=7, dilation=2, num_heads=4)
+na3d = NeighborhoodAttention3D(dim=128, kernel_size=7, dilation=2, num_heads=4)
 ```
 
+NA3D also supports different kernel size and dilation values for depth:
+```python
+na3d = NeighborhoodAttention3D(
+	dim=128,
+	kernel_size=7,
+	kernel_size_d=5,
+	dilation=2,
+	dilation_d=3,
+	num_heads=4)
+```
+
+Modules expect inputs of shape `[batch_size, *, dim]`:
+* NA1D: `[batch_size, sequence_length, dim]`
+* NA2D: `[batch_size, height, width, dim]`
+* NA3D: `[batch_size, depth, height, width, dim]`
+
+
 ### FLOPs
 We recommend counting flops through [fvcore](https://github.com/facebookresearch/fvcore).
 
diff --git a/src/natten/__init__.py b/src/natten/__init__.py
index 8f086ea..ddef1d0 100644
--- a/src/natten/__init__.py
+++ b/src/natten/__init__.py
@@ -23,5 +23,6 @@
 
 from .natten1d import NeighborhoodAttention1D
 from .natten2d import NeighborhoodAttention2D
+from .natten3d import NeighborhoodAttention3D
 
 __version__ = "0.14.7.dev"
diff --git a/src/natten/csrc/cpu/inverse_neighborhood_3d.cpp b/src/natten/csrc/cpu/inverse_neighborhood_3d.cpp
new file mode 100644
index 0000000..39f3ab8
--- /dev/null
+++ b/src/natten/csrc/cpu/inverse_neighborhood_3d.cpp
@@ -0,0 +1,96 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Inverse-Neighborhood-Neighborhood CPU kernel for 3D data.
+           Applies inverse neighborhood attention weights to inverse neighborhood values.
+           Used to compute key and value grads.
+*/
+
+#include <torch/extension.h>
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+
+#include "natten_cpu_commons.h"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor6D = typename at::TensorAccessor<scalar_t, 6>;
+
+#define GRAIN_SIZE 0
+
+template <int KS, int DKS, int NS, int DNS, typename scalar_t>
+void inverse_neighborhood_3d(          // K-grad / V-grad
+    const Tensor6D<scalar_t> weights,  // d_attn / attn
+    const Tensor6D<scalar_t> values,   // query  / d_out
+    Tensor6D<scalar_t> output,         // d_key  / d_value
+    const int depth, 
+    const int height, 
+    const int width,
+    const int heads,
+    const int kernel_size_in,
+    const int kernel_size_d_in,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int batch_size) {
+    const int KERNEL_SIZE = (KS>1) ? KS : kernel_size_in;
+    const int KERNEL_SIZE_D = (DKS>1) ? DKS : kernel_size_d_in;
+    const int NEIGHBORHOOD_SIZE = (NS>0) ? NS : KERNEL_SIZE / 2;
+    const int NEIGHBORHOOD_SIZE_D = (DNS>0) ? DNS : KERNEL_SIZE_D / 2;
+    for (int b = 0; b < batch_size; b++) {
+        at::parallel_for(0, heads, GRAIN_SIZE, [&](int start, int end) {
+        for (int h = start; h < end; h++) {
+            for (int k = 0; k < depth; k++) {
+            const int nk = get_backward_window_start(k, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+            const int ek = get_backward_window_end(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+                for (int i = 0; i < height; i++) {
+                const int ni = get_backward_window_start(i, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int ei = get_backward_window_end(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                for (int j = 0; j < width; j++) {
+                    const int nj = get_backward_window_start(j, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                    const int ej = get_backward_window_end(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                    for (int d = 0; d < dim; d++) {
+                        const int weightsOffset = b * weights.stride(0) + h * weights.stride(1);
+                        const int outOffset = b * values.stride(0) + h * values.stride(1) + d;
+                        scalar_t output_update = scalar_t(0);
+                        for (int xk=nk; xk < ek; xk+=dilation_d){
+                        const int onk = get_window_start(xk, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+                        for (int xi=ni; xi < ei; xi+=dilation){
+                        const int oni = get_window_start(xi, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                        for (int xj=nj; xj < ej; xj+=dilation){
+                            const int onj = get_window_start(xj, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                            const int outIndex = outOffset + xk * values.stride(2) + xi * values.stride(3) + xj * values.stride(4);
+                            const int weightsIndex = weightsOffset + xk * weights.stride(2) + xi * weights.stride(3) + xj * weights.stride(4) + int((k-onk)/dilation_d)*(KERNEL_SIZE*KERNEL_SIZE) + int((i-oni)/dilation)*KERNEL_SIZE + int((j-onj)/dilation);
+                            output_update += values.data()[outIndex] * weights.data()[weightsIndex];
+                        }}}
+                        const int linearIndex = b*output.stride(0) + h*output.stride(1) + k*output.stride(2) + i*output.stride(3) + j*output.stride(4) + d;
+                        output.data()[linearIndex] = output_update;
+                    }
+                }}}
+        }});
+    }
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cpu/na3d.cpp b/src/natten/csrc/cpu/na3d.cpp
new file mode 100644
index 0000000..f55d6e4
--- /dev/null
+++ b/src/natten/csrc/cpu/na3d.cpp
@@ -0,0 +1,182 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood Attention 3D - CPU interface
+*/
+
+#include <torch/extension.h>
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+
+#include "natten_cpu_commons.h"
+#include "pointwise_neighborhood_3d.cpp"
+#include "neighborhood_neighborhood_3d.cpp"
+#include "inverse_neighborhood_3d.cpp"
+#include "rel_pos_bias_3d.cpp"
+
+namespace natten {
+
+torch::Tensor natten3dqkrpb_cpu_forward(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int batch_size = query.size(0);
+    int heads = query.size(1);
+    int depth = query.size(2);
+    int height = query.size(3);
+    int width = query.size(4);
+    int dim = query.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+
+    auto attn = torch::zeros(
+            {batch_size, heads, depth, height, width, kernel_size_d*kernel_size*kernel_size}, query.options());
+
+    AT_DISPATCH_FLOATING_TYPES(query.scalar_type(), "natten3dqkrpb_cpu_forward", ([&] {
+        if (rpb.has_value()) {
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   pointwise_neighborhood_3d_bias, 
+                    query.accessor<scalar_t, 6>(), key.accessor<scalar_t, 6>(), 
+                    rpb.value().accessor<scalar_t, 4>(), attn.accessor<scalar_t, 6>(), 
+                    depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+        } else {
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   pointwise_neighborhood_3d, 
+                    query.accessor<scalar_t, 6>(), key.accessor<scalar_t, 6>(), 
+                    attn.accessor<scalar_t, 6>(), 
+                    depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+        }
+    }));
+    return attn;
+}
+
+std::vector<torch::Tensor> natten3dqkrpb_cpu_backward(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int batch_size = query.size(0);
+    int heads = query.size(1);
+    int depth = query.size(2);
+    int height = query.size(3);
+    int width = query.size(4);
+    int dim = query.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    int RPB_MAX_D = kernel_size_d * 2 - 1;
+    int RPB_MAX = kernel_size * 2 - 1;
+   
+    auto d_query = torch::zeros_like(query);
+    auto d_key = torch::zeros_like(key);
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX_D, RPB_MAX, RPB_MAX}, d_attn.options());
+
+    AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten3dqkrpb_backward_cpu", ([&] {
+        const auto d_attn_a = d_attn.accessor<scalar_t, 6>();
+        const auto query_a = query.accessor<scalar_t, 6>();
+        const auto key_a = key.accessor<scalar_t, 6>();
+        auto d_query_a = d_query.accessor<scalar_t, 6>();
+        auto d_key_a = d_key.accessor<scalar_t, 6>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.accessor<scalar_t, 4>();
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   rel_pos_bias_gradient_3d, 
+                    d_rpb_a, d_attn_a, depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, batch_size);
+        }
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   neighborhood_neighborhood_3d, 
+                d_attn_a, key_a, d_query_a, depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   inverse_neighborhood_3d, 
+                d_attn_a, query_a, d_key_a, depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+    }));
+    return {d_query, d_key, d_rpb};
+}
+
+torch::Tensor natten3dav_cpu_forward(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    AT_ASSERTM(attn.device().is_cpu(), "attn must be a CPU tensor");
+    AT_ASSERTM(value.device().is_cpu(), "value must be a CPU tensor");
+    int batch_size = value.size(0);
+    int heads = value.size(1);
+    int depth = value.size(2);
+    int height = value.size(3);
+    int width = value.size(4);
+    int dim = value.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+
+    auto out = torch::zeros_like(value);
+
+    AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "natten3dav_forward_cpu", ([&] {
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   neighborhood_neighborhood_3d, 
+                attn.accessor<scalar_t, 6>(), value.accessor<scalar_t, 6>(), out.accessor<scalar_t, 6>(), 
+                depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+    }));
+    return out;
+}
+
+std::vector<torch::Tensor> natten3dav_cpu_backward(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    AT_ASSERTM(d_out.device().is_cpu(), "d_out must be a CPU tensor");
+    AT_ASSERTM(attn.device().is_cpu(), "attn must be a CPU tensor");
+    AT_ASSERTM(value.device().is_cpu(), "value must be a CPU tensor");
+    int batch_size = value.size(0);
+    int heads = value.size(1);
+    int depth = value.size(2);
+    int height = value.size(3);
+    int width = value.size(4);
+    int dim = value.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+
+    auto d_attn = torch::zeros_like(attn);
+    auto d_value = torch::zeros_like(value);
+
+    AT_DISPATCH_FLOATING_TYPES(d_attn.scalar_type(), "natten3dav_backward_cpu", ([&] {
+        const auto d_out_a = d_out.accessor<scalar_t, 6>();
+        const auto attn_a = attn.accessor<scalar_t, 6>();
+        const auto value_a = value.accessor<scalar_t, 6>();
+        auto d_attn_a = d_attn.accessor<scalar_t, 6>();
+        auto d_value_a = d_value.accessor<scalar_t, 6>();
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   pointwise_neighborhood_3d, 
+                d_out_a, value_a, d_attn_a, depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,   inverse_neighborhood_3d, 
+                attn_a, d_out_a, d_value_a, depth, height, width, heads, kernel_size, kernel_size_d, dilation, dilation_d, dim, batch_size);
+    }));
+    return {d_attn, d_value};
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cpu/natten_cpu_commons.h b/src/natten/csrc/cpu/natten_cpu_commons.h
index c211980..8dcccd8 100644
--- a/src/natten/csrc/cpu/natten_cpu_commons.h
+++ b/src/natten/csrc/cpu/natten_cpu_commons.h
@@ -72,6 +72,10 @@ inline int get_pb_start(const int index, const int length, const int KERNEL_SIZE
 
 #define CHECK_SEQUENCE(length, kernel_size, dilation) TORCH_CHECK(length >= kernel_size*dilation, "Input sequence length must be greater than or equal to kernel size x dilation.")
 #define CHECK_FEATMAP(height, width, kernel_size, dilation) TORCH_CHECK(height >= kernel_size*dilation && width >= kernel_size*dilation, "Input resolution must be greater than or equal to kernel size x dilation.")
+#define CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d) { \
+    CHECK_SEQUENCE(depth, kernel_size_d, dilation_d);                                             \
+    CHECK_FEATMAP(height, width, kernel_size, dilation);                                          \
+}                                                                                                 \
 
 
 #define _IN_LAUNCH_DNA_KNS(KS, NS, dilation, NAME, ...)                     \
@@ -158,4 +162,59 @@ inline int get_pb_start(const int index, const int length, const int KERNEL_SIZE
     }                                                                                                \
 })
 
+// 3D KERNEL LAUNCHER
+#define LAUNCH_NA_KDNDS_INN(kernel_size, KERNEL_SIZE_DPTH, NEIGH_SIZE_DPTH, NAME, ...)   \
+({                                                                                       \
+    switch (kernel_size) {                                                               \
+        case 3:                                                                          \
+            NAME<3, KERNEL_SIZE_DPTH, 1, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);        \
+            break;                                                                       \
+        case 5:                                                                          \
+            NAME<5, KERNEL_SIZE_DPTH, 2, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);        \
+            break;                                                                       \
+        case 7:                                                                          \
+            NAME<7, KERNEL_SIZE_DPTH, 3, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);        \
+            break;                                                                       \
+        case 9:                                                                          \
+            NAME<9, KERNEL_SIZE_DPTH, 4, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);        \
+            break;                                                                       \
+        case 11:                                                                         \
+            NAME<11, KERNEL_SIZE_DPTH, 5, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);       \
+            break;                                                                       \
+        case 13:                                                                         \
+            NAME<13, KERNEL_SIZE_DPTH, 6, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);       \
+            break;                                                                       \
+        default:                                                                         \
+            NAME<-1, KERNEL_SIZE_DPTH, -1, NEIGH_SIZE_DPTH, scalar_t>(__VA_ARGS__);      \
+            break;                                                                       \
+    }                                                                                    \
+})
+
+#define LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, NAME, ...)              \
+({                                                                          \
+    switch (kernel_size_d) {                                                \
+        case 3:                                                             \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 3, 1, NAME, __VA_ARGS__);      \
+            break;                                                          \
+        case 5:                                                             \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 5, 2, NAME, __VA_ARGS__);      \
+            break;                                                          \
+        case 7:                                                             \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 7, 3, NAME, __VA_ARGS__);      \
+            break;                                                          \
+        case 9:                                                             \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 9, 4, NAME, __VA_ARGS__);      \
+            break;                                                          \
+        case 11:                                                            \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 11, 5, NAME, __VA_ARGS__);     \
+            break;                                                          \
+        case 13:                                                            \
+            LAUNCH_NA_KDNDS_INN(kernel_size, 13, 6, NAME, __VA_ARGS__);     \
+            break;                                                          \
+        default:                                                            \
+            LAUNCH_NA_KDNDS_INN(kernel_size, -1, -1, NAME, __VA_ARGS__);    \
+            break;                                                          \
+    }                                                                       \
+})
+
 #endif
diff --git a/src/natten/csrc/cpu/neighborhood_neighborhood_3d.cpp b/src/natten/csrc/cpu/neighborhood_neighborhood_3d.cpp
new file mode 100644
index 0000000..ffb09b7
--- /dev/null
+++ b/src/natten/csrc/cpu/neighborhood_neighborhood_3d.cpp
@@ -0,0 +1,94 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood-Neighborhood CPU kernel for 3D data.
+           Applies neighborhood attention weights to neighborhood values.
+*/
+
+#include <torch/extension.h>
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+
+#include "natten_cpu_commons.h"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor6D = typename at::TensorAccessor<scalar_t, 6>;
+
+#define GRAIN_SIZE 0
+
+template <int KS, int DKS, int NS, int DNS, typename scalar_t>
+void neighborhood_neighborhood_3d(           // AV     / Q-grad
+    const Tensor6D<scalar_t> weights,        // attn   / d_attn
+    const Tensor6D<scalar_t> values,         // value  / key
+    Tensor6D<scalar_t> output,               // output / d_query
+    const int depth, 
+    const int height, 
+    const int width,
+    const int heads,
+    const int kernel_size_in,
+    const int kernel_size_d_in,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int batch_size) {
+    const int KERNEL_SIZE = (KS>1) ? KS : kernel_size_in;
+    const int KERNEL_SIZE_D = (DKS>1) ? DKS : kernel_size_d_in;
+    const int NEIGHBORHOOD_SIZE = (NS>0) ? NS : KERNEL_SIZE / 2;
+    const int NEIGHBORHOOD_SIZE_D = (DNS>0) ? DNS : KERNEL_SIZE_D / 2;
+    at::parallel_for(0, batch_size*heads*depth*height*width, GRAIN_SIZE, [&](int start, int end) {
+    for (int x = start; x < end; x++) {
+        int indtmp1 = x/width;
+        const int j = x - indtmp1 * width;
+        int indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+        const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        for (int d = 0; d < dim; d++) {
+            scalar_t updt = scalar_t(0);
+            int weightsOffset = b * weights.stride(0) + h * weights.stride(1) + k * weights.stride(2) + i * weights.stride(3) + j * weights.stride(4);
+            const int valuesOffset = b * values.stride(0) + h * values.stride(1) + d;
+            for (int xk=nk; xk < nk + KERNEL_SIZE_D * dilation_d; xk+=dilation_d){
+            for (int xi=ni; xi < ni + KERNEL_SIZE * dilation; xi+=dilation){
+            for (int xj=nj; xj < nj + KERNEL_SIZE * dilation; xj+=dilation){
+                const int valuesIndex = valuesOffset + xk * values.stride(2) + xi * values.stride(3) + xj * values.stride(4);
+                updt += weights.data()[weightsOffset] * values.data()[valuesIndex];
+                ++weightsOffset;
+            }}}
+            const int linearIndex = b*output.stride(0) + h*output.stride(1) + k*output.stride(2) + i*output.stride(3) + j*output.stride(4) + d*output.stride(5);
+            output.data()[linearIndex] = updt;
+        }
+    }});
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cpu/pointwise_neighborhood_3d.cpp b/src/natten/csrc/cpu/pointwise_neighborhood_3d.cpp
new file mode 100644
index 0000000..d7dbc5d
--- /dev/null
+++ b/src/natten/csrc/cpu/pointwise_neighborhood_3d.cpp
@@ -0,0 +1,152 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Pointwise-Neighborhood CPU kernel for 3D data.
+           Computes attention weights between query points and their corresponding
+           key neighborhood.
+           Extra kernel with fused bias (relative positional bias.)
+*/
+
+#include <torch/extension.h>
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+
+#include "natten_cpu_commons.h"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor4D = typename at::TensorAccessor<scalar_t, 4>;
+template<class scalar_t>
+using Tensor6D = typename at::TensorAccessor<scalar_t, 6>;
+
+#define GRAIN_SIZE 0
+
+template <int KS, int DKS, int NS, int DNS, typename scalar_t>
+void pointwise_neighborhood_3d(     // QK    / A-grad
+    const Tensor6D<scalar_t> query, // query / d_out
+    const Tensor6D<scalar_t> key,   // key   / value
+    Tensor6D<scalar_t> attn,        // attn  / d_attn
+    const int depth, 
+    const int height, 
+    const int width,
+    const int heads,
+    const int kernel_size_in,
+    const int kernel_size_d_in,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int batch_size) {
+    const int KERNEL_SIZE = (KS>1) ? KS : kernel_size_in;
+    const int KERNEL_SIZE_D = (DKS>1) ? DKS : kernel_size_d_in;
+    const int NEIGHBORHOOD_SIZE = (NS>0) ? NS : KERNEL_SIZE / 2;
+    const int NEIGHBORHOOD_SIZE_D = (DNS>0) ? DNS : KERNEL_SIZE_D / 2;
+    at::parallel_for(0, batch_size*heads*depth*height*width, GRAIN_SIZE, [&](int start, int end) {
+    for (int x = start; x < end; x++) {
+        int indtmp1 = x/width;
+        const int j = x - indtmp1 * width;
+        int indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+        const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        for (int kk = 0; kk < KERNEL_SIZE_D; kk++) {
+        for (int ki = 0; ki < KERNEL_SIZE; ki++) {
+        for (int kj = 0; kj < KERNEL_SIZE; kj++) {
+            scalar_t updt = scalar_t(0);
+            const int batchHeadOffset = b * query.stride(0) + h * query.stride(1);
+            const int queryOffset = batchHeadOffset + k * query.stride(2) + i * query.stride(3) + j * query.stride(4);
+            const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * key.stride(2) + (ki*dilation+ni) * key.stride(3) + (kj*dilation+nj) * key.stride(4);
+            for (int dimOffset=0; dimOffset < dim; ++dimOffset)
+                updt += query.data()[queryOffset+dimOffset] * key.data()[keyOffset+dimOffset];
+            const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + kk*(KERNEL_SIZE*KERNEL_SIZE)+ki*KERNEL_SIZE+kj;
+            attn.data()[index] = updt;
+        }}}
+    }});
+}
+
+template <int KS, int DKS, int NS, int DNS, typename scalar_t>
+void pointwise_neighborhood_3d_bias( // QK    / A-grad
+    const Tensor6D<scalar_t> query,  // query / d_out
+    const Tensor6D<scalar_t> key,    // key   / value
+    const Tensor4D<scalar_t> bias,   // relative positional bias tensor
+    Tensor6D<scalar_t> attn,         // attn  / d_attn
+    const int depth, 
+    const int height, 
+    const int width,
+    const int heads,
+    const int kernel_size_in,
+    const int kernel_size_d_in,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int batch_size) {
+    const int KERNEL_SIZE = (KS>1) ? KS : kernel_size_in;
+    const int KERNEL_SIZE_D = (DKS>1) ? DKS : kernel_size_d_in;
+    const int NEIGHBORHOOD_SIZE = (NS>0) ? NS : KERNEL_SIZE / 2;
+    const int NEIGHBORHOOD_SIZE_D = (DNS>0) ? DNS : KERNEL_SIZE_D / 2;
+    at::parallel_for(0, batch_size*heads*depth*height*width, GRAIN_SIZE, [&](int start, int end) {
+    for (int x = start; x < end; x++) {
+        int indtmp1 = x/width;
+        const int j = x - indtmp1 * width;
+        int indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+        const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        for (int kk = 0; kk < KERNEL_SIZE_D; kk++) {
+        for (int ki = 0; ki < KERNEL_SIZE; ki++) {
+        for (int kj = 0; kj < KERNEL_SIZE; kj++) {
+            scalar_t updt = scalar_t(0);
+            const int batchHeadOffset = b * query.stride(0) + h * query.stride(1);
+            const int queryOffset = batchHeadOffset + k * query.stride(2) + i * query.stride(3) + j * query.stride(4);
+            const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * key.stride(2) + (ki*dilation+ni) * key.stride(3) + (kj*dilation+nj) * key.stride(4);
+            for (int dimOffset=0; dimOffset < dim; ++dimOffset)
+                updt += query.data()[queryOffset+dimOffset] * key.data()[keyOffset+dimOffset];
+            const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + kk*(KERNEL_SIZE*KERNEL_SIZE)+ki*KERNEL_SIZE+kj;
+            const int biasIndex = h * bias.stride(0) + (pk+kk) * bias.stride(1) + (pi+ki) * bias.stride(2) + (pj+kj) * bias.stride(3);
+            updt += bias.data()[biasIndex];
+            attn.data()[index] = updt;
+        }}}
+    }});
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cpu/rel_pos_bias_3d.cpp b/src/natten/csrc/cpu/rel_pos_bias_3d.cpp
new file mode 100644
index 0000000..6383033
--- /dev/null
+++ b/src/natten/csrc/cpu/rel_pos_bias_3d.cpp
@@ -0,0 +1,79 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Relative positional bias backward pass CPU kernel for 3D data.
+*/
+
+#include <torch/extension.h>
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+
+#include "natten_cpu_commons.h"
+
+namespace natten {
+
+#define GRAIN_SIZE 0
+
+template <int KS, int DKS, int NS, int DNS, typename scalar_t>
+void rel_pos_bias_gradient_3d(
+    at::TensorAccessor<scalar_t, 4> d_bias,
+    const at::TensorAccessor<scalar_t, 6> d_attn,
+    const int depth, 
+    const int height, 
+    const int width,
+    const int heads,
+    const int kernel_size_in,
+    const int kernel_size_d_in,
+    const int dilation,
+    const int dilation_d,
+    const int batch_size) {
+    const int KERNEL_SIZE = (KS>1) ? KS : kernel_size_in;
+    const int KERNEL_SIZE_D = (DKS>1) ? DKS : kernel_size_d_in;
+    const int NEIGHBORHOOD_SIZE = (NS>0) ? NS : KERNEL_SIZE / 2;
+    const int NEIGHBORHOOD_SIZE_D = (DNS>0) ? DNS : KERNEL_SIZE_D / 2;
+    at::parallel_for(0, heads, GRAIN_SIZE, [&](int start, int end) {
+    for (int h = start; h < end; h++) {
+        for (int k = 0; k < depth; k++) {
+        const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+            for (int i = 0; i < height; i++) {
+            const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+            for (int j = 0; j < width; j++) {
+                const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                for (int kk = 0; kk < KERNEL_SIZE_D; kk++) {
+                for (int ki = 0; ki < KERNEL_SIZE; ki++) {
+                for (int kj = 0; kj < KERNEL_SIZE; kj++) {
+                    scalar_t d_bias_update = scalar_t(0);
+                    int attnOffset = h * d_attn.stride(1) + k * d_attn.stride(2) + i * d_attn.stride(3) + j * d_attn.stride(4) + kk*(KERNEL_SIZE*KERNEL_SIZE)+ki*KERNEL_SIZE+kj;
+                    for (int b=0; b < batch_size; ++b){
+                        d_bias_update += d_attn.data()[attnOffset];
+                        attnOffset += d_attn.stride(0);
+                    }
+                    const int index = h * d_bias.stride(0) + (pk+kk) * d_bias.stride(1) + (pi+ki) * d_bias.stride(2) + (pj+kj) * d_bias.stride(3);
+                    d_bias.data()[index] += d_bias_update;
+                }}}
+            }}}
+    }});
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cuda/inverse_neighborhood_3d.cu b/src/natten/csrc/cuda/inverse_neighborhood_3d.cu
new file mode 100644
index 0000000..a934e13
--- /dev/null
+++ b/src/natten/csrc/cuda/inverse_neighborhood_3d.cu
@@ -0,0 +1,166 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Inverse-Neighborhood-Neighborhood kernel for 3D data.
+           Applies inverse neighborhood attention weights to inverse neighborhood values.
+           Used to compute key and value grads.
+*/
+
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/ATen.h>
+#include <ATen/native/cuda/KernelUtils.cuh>
+#include <ATen/AccumulateType.h>
+#include <cuda_fp16.h>
+
+#include "natten_commons.cuh"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor6D = typename torch::PackedTensorAccessor32<scalar_t,6,torch::DefaultPtrTraits>;
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void inverse_neighborhood_3d(           // K-grad / V-grad
+    const Tensor6D<scalar_t> weights,              // d_attn / attn
+    const Tensor6D<scalar_t> values,               // query  / d_out
+    Tensor6D<scalar_t> output,                     // d_key  / d_value
+    const int depth,
+    const int height,
+    const int width,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int output_numel) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < output_numel){
+        int indtmp1 = linearIndex/dim;
+        const int d = linearIndex - indtmp1 * dim;
+        int indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+        const int ni = get_backward_window_start(i, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_backward_window_start(j, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_backward_window_start(k, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        const int ei = get_backward_window_end(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int ej = get_backward_window_end(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int ek = get_backward_window_end(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        const int weightsOffset = b * weights.stride(0) + h * weights.stride(1);
+        const int valuesOffset = b * values.stride(0) + h * values.stride(1) + d;
+        scalar_t output_update = scalar_t(0);
+        #pragma unroll
+        for (int xk=nk; xk < ek; xk+=dilation_d){
+            const int onk = get_window_start(xk, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+            #pragma unroll
+            for (int xi=ni; xi < ei; xi+=dilation){
+                const int oni = get_window_start(xi, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                #pragma unroll
+                for (int xj=nj; xj < ej; xj+=dilation){
+                    const int onj = get_window_start(xj, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                    const int valuesIndex = valuesOffset + xk * values.stride(2) + xi * values.stride(3) + xj * values.stride(4);
+                    const int weightsIndex = weightsOffset + xk * weights.stride(2) + xi * weights.stride(3) + xj * weights.stride(4) + (int((k-onk)/dilation_d)*KERNEL_SIZE*KERNEL_SIZE)+int((i-oni)/dilation)*KERNEL_SIZE+int((j-onj)/dilation);
+                    output_update += values.data()[valuesIndex] * weights.data()[weightsIndex];
+                }
+            }
+        }
+        output.data()[linearIndex] = output_update;
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void inverse_neighborhood_3d_fp16(      // K-grad / V-grad
+    const Tensor6D<scalar_t> weights,              // d_attn / attn
+    const Tensor6D<scalar_t> values,               // query  / d_out
+    Tensor6D<scalar_t> output,                     // d_key  / d_value
+    const int depth,
+    const int height,
+    const int width,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dimhalf,
+    const int output_numel) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < output_numel){
+        __half2* values2 = reinterpret_cast<__half2*>(values.data());
+        __half2* output2 = reinterpret_cast<__half2*>(output.data());
+        int indtmp1 = linearIndex/dimhalf;
+        const int d = linearIndex - indtmp1 * dimhalf;
+        int indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+        const int ni = get_backward_window_start(i, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_backward_window_start(j, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_backward_window_start(k, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        const int ei = get_backward_window_end(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int ej = get_backward_window_end(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int ek = get_backward_window_end(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        const int weightsOffset = b * weights.stride(0) + h * weights.stride(1);
+        const int stride3 = dimhalf * width;
+        const int stride2 = stride3 * height;
+        const int valuesOffset = b * (stride2 * depth * heads) + h * (stride2 * depth) + d;
+        __half2 output_update = __float2half2_rn(0.f);
+        #pragma unroll
+        for (int xk=nk; xk < ek; xk+=dilation_d){
+            const int onk = get_window_start(xk, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+            #pragma unroll
+            for (int xi=ni; xi < ei; xi+=dilation){
+                const int oni = get_window_start(xi, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                #pragma unroll
+                for (int xj=nj; xj < ej; xj+=dilation){
+                    const int onj = get_window_start(xj, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                    const int valuesIndex = valuesOffset + xk * stride2 + xi * stride3 + xj * dimhalf;
+                    const int weightsIndex = weightsOffset + xk * weights.stride(2) + xi * weights.stride(3) + xj * weights.stride(4) + (int((k-onk)/dilation_d)*KERNEL_SIZE*KERNEL_SIZE)+int((i-oni)/dilation)*KERNEL_SIZE+int((j-onj)/dilation);
+                    scalar_t a = weights.data()[weightsIndex];
+                    output_update = __hfma2(values2[valuesIndex], __halves2half2(a, a), output_update);
+                }
+            }
+        }
+        output2[linearIndex] = output_update;
+    }
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cuda/na3d.cu b/src/natten/csrc/cuda/na3d.cu
new file mode 100644
index 0000000..1fa1543
--- /dev/null
+++ b/src/natten/csrc/cuda/na3d.cu
@@ -0,0 +1,495 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood Attention 3D - CUDA interface
+*/
+
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/ATen.h>
+#include <ATen/native/cuda/KernelUtils.cuh>
+#include <ATen/AccumulateType.h>
+#include <cuda_fp16.h>
+
+#include "natten_commons.cuh"
+#include "pointwise_neighborhood_3d.cu"
+#include "neighborhood_neighborhood_3d.cu"
+#include "inverse_neighborhood_3d.cu"
+#include "rel_pos_bias_3d.cu"
+
+#define CUDA_NUM_THREADS_Q 512
+#define CUDA_NUM_THREADS_K 512
+#define CUDA_NUM_THREADS_RPB 64
+#define CUDA_NUM_THREADS_Q16 512
+#define CUDA_NUM_THREADS_K16 256
+#define CUDA_NUM_THREADS_RPB16 64
+#define CUDA_NUM_THREADS_F 512
+#define CUDA_NUM_THREADS_FP16 512
+#define CUDA_NUM_THREADS_V 512
+#define CUDA_NUM_THREADS_V16 256
+
+namespace natten {
+
+torch::Tensor natten3dqkrpb_cuda_forward(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = query.size(0);
+    int64_t heads = query.size(1);
+    int64_t depth = query.size(2);
+    int64_t height = query.size(3);
+    int64_t width = query.size(4);
+    int64_t dim = query.size(5);
+    int ysize = kernel_size_d * kernel_size * kernel_size;
+    int zsize = batch_size * heads;
+    int xsize = depth * height * width;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_forward", kernel_size);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_forward", kernel_size_d);
+    int KERNELTHREADS = min(CUDA_NUM_THREADS, ysize);
+    int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), xsize);
+    int BATCHTHREADS = max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+
+    auto attn = torch::zeros(
+            {batch_size, heads, depth, height, width, ysize}, query.options());
+
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    const dim3 blocks(
+            (xsize + PIXELTHREADS - 1) / PIXELTHREADS,
+            (ysize + KERNELTHREADS - 1) / KERNELTHREADS,
+            (zsize + BATCHTHREADS - 1) / BATCHTHREADS);
+    const dim3 threads(PIXELTHREADS, KERNELTHREADS, BATCHTHREADS);
+    AT_DISPATCH_FLOATING_TYPES(query.scalar_type(), "natten3dqkrpb_cuda_forward", ([&] {
+        const auto query_a = query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto key_a = key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        if (rpb.has_value()) {
+            const auto rpb_a = rpb.value().packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                    pointwise_neighborhood_3d_bias, 
+                    blocks, threads, 0, stream, 
+                    query_a, key_a, rpb_a, attn_a, 
+                    depth, height, width, batch_size, heads, dilation, dilation_d, dim);
+        } else {
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                    pointwise_neighborhood_3d, 
+                    blocks, threads, 0, stream, 
+                    query_a, key_a, attn_a, 
+                    depth, height, width, batch_size, heads, dilation, dilation_d, dim);
+        }
+    }));
+    return attn;
+}
+
+torch::Tensor natten3dqkrpb_cuda_forward_fp16(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = query.size(0);
+    int64_t heads = query.size(1);
+    int64_t depth = query.size(2);
+    int64_t height = query.size(3);
+    int64_t width = query.size(4);
+    int64_t dimhalf = query.size(5) / 2;
+    int ysize = kernel_size_d * kernel_size * kernel_size;
+    int zsize = batch_size * heads;
+    int xsize = depth * height * width;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_forward_fp16", kernel_size);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_forward_fp16", kernel_size_d);
+    TORCH_CHECK(dimhalf*2 == query.size(5), "Dims per head must be an even number in FP16.");
+    int KERNELTHREADS = min(CUDA_NUM_THREADS, ysize);
+    int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), xsize);
+    int BATCHTHREADS = max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+
+    auto attn = torch::zeros(
+            {batch_size, heads, depth, height, width, ysize}, query.options());
+
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    const dim3 blocks(
+            (xsize + PIXELTHREADS - 1) / PIXELTHREADS,
+            (ysize + KERNELTHREADS - 1) / KERNELTHREADS,
+            (zsize + BATCHTHREADS - 1) / BATCHTHREADS);
+    const dim3 threads(PIXELTHREADS, KERNELTHREADS, BATCHTHREADS);
+    AT_DISPATCH_HALF_TYPES(at::kHalf, query.scalar_type(), "natten3dqkrpb_cuda_forward_fp16", ([&] {
+        const auto query_a = query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto key_a = key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        if (rpb.has_value()) {
+            const auto rpb_a = rpb.value().packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                    pointwise_neighborhood_3d_bias_fp16, 
+                    blocks, threads, 0, stream, 
+                    query_a, key_a, rpb_a, attn_a, 
+                    depth, height, width, batch_size, heads, dilation, dilation_d, dimhalf);
+        } else {
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                    pointwise_neighborhood_3d_fp16, 
+                    blocks, threads, 0, stream, 
+                    query_a, key_a, attn_a, 
+                    depth, height, width, batch_size, heads, dilation, dilation_d, dimhalf);
+        }
+    }));
+    return attn;
+}
+
+std::vector<torch::Tensor> natten3dqkrpb_cuda_backward(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = query.size(0);
+    int64_t heads = query.size(1);
+    int64_t depth = query.size(2);
+    int64_t height = query.size(3);
+    int64_t width = query.size(4);
+    int64_t dim = query.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_backward", kernel_size);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_backward", kernel_size_d);
+    int64_t RPB_MAX_D = kernel_size_d * 2 - 1;
+    int64_t RPB_MAX = kernel_size * 2 - 1;
+   
+    auto d_query = torch::zeros_like(query);
+    auto d_key = torch::zeros_like(key);
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX_D, RPB_MAX, RPB_MAX}, d_attn.options());
+
+    int32_t n_rpb = heads * depth * height * width * d_attn.size(5);
+    int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB);
+    dim3 grid_rpb(blocks_rpb);
+    dim3 blockr(CUDA_NUM_THREADS_RPB);
+    int32_t n_query = d_query.numel();
+    int blocks_query = GET_BLOCKS(n_query, CUDA_NUM_THREADS_Q);
+    dim3 grid_query(blocks_query);
+    dim3 blockq(CUDA_NUM_THREADS_Q);
+    int32_t n_key = d_key.numel();
+    int blocks_key = GET_BLOCKS(n_key, CUDA_NUM_THREADS_K);
+    dim3 grid_key(blocks_key);
+    dim3 blockk(CUDA_NUM_THREADS_K);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten3dqkrpb_backward_cuda", ([&] {
+        const auto d_attn_a = d_attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto query_a = query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto key_a = key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_query_a = d_query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                    rel_pos_bias_gradient_3d, 
+                    grid_rpb, blockr, 0, stream, 
+                    d_rpb_a, d_attn_a, depth, height, width, dilation, dilation_d, batch_size, d_rpb.numel(), n_rpb);
+        }
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                neighborhood_neighborhood_3d, 
+                grid_query, blockq, 0, stream,
+                d_attn_a, key_a, d_query_a, 
+                depth, height, width, heads, dilation, dilation_d, dim, n_query);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                inverse_neighborhood_3d, 
+                grid_key, blockk, 0, stream,
+                d_attn_a, query_a, d_key_a, 
+                depth, height, width, heads, dilation, dilation_d, dim, n_key);
+    }));
+    return {d_query, d_key, d_rpb};
+}
+
+std::vector<torch::Tensor> natten3dqkrpb_cuda_backward_fp16(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = query.size(0);
+    int64_t heads = query.size(1);
+    int64_t depth = query.size(2);
+    int64_t height = query.size(3);
+    int64_t width = query.size(4);
+    int64_t dimhalf = query.size(5) / 2;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_backward_fp16", kernel_size);
+    CHECK_KERNELSIZE("natten3dqkrpb_cuda_backward_fp16", kernel_size_d);
+    TORCH_CHECK(dimhalf*2 == query.size(5), "Dims per head must be an even number in FP16.");
+    int64_t RPB_MAX_D = kernel_size_d * 2 - 1;
+    int64_t RPB_MAX = kernel_size * 2 - 1;
+   
+    auto d_query = torch::zeros_like(query);
+    auto d_key = torch::zeros_like(key);
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX_D, RPB_MAX, RPB_MAX}, d_attn.options());
+
+    int32_t n_rpb = heads * depth * height * width * d_attn.size(5);
+    int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB16);
+    dim3 grid_rpb(blocks_rpb);
+    dim3 blockr(CUDA_NUM_THREADS_RPB16);
+    int32_t nhalf_query = d_query.numel() / 2;
+    int blocks_query = GET_BLOCKS(nhalf_query, CUDA_NUM_THREADS_Q16);
+    dim3 grid_query(blocks_query);
+    dim3 blockq(CUDA_NUM_THREADS_Q16);
+    int32_t nhalf_key = d_key.numel() / 2;
+    int blocks_key = GET_BLOCKS(nhalf_key, CUDA_NUM_THREADS_K16);
+    dim3 grid_key(blocks_key);
+    dim3 blockk(CUDA_NUM_THREADS_K16);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_HALF_TYPES(at::kHalf, d_query.scalar_type(), "natten3dqkrpb_backward_cuda_fp16", ([&] {
+        const auto d_attn_a = d_attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto query_a = query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto key_a = key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_query_a = d_query.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+            LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                    rel_pos_bias_gradient_3d_fp16, 
+                    grid_rpb, blockr, 0, stream,
+                    d_rpb_a, d_attn_a, 
+                    depth, height, width, dilation, dilation_d, batch_size, d_rpb.numel(), n_rpb);
+        }
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                neighborhood_neighborhood_3d_fp16, 
+                grid_query, blockq, 0, stream,
+                d_attn_a, key_a, d_query_a, 
+                depth, height, width, heads, dilation, dilation_d, dimhalf, nhalf_query);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                inverse_neighborhood_3d_fp16, 
+                grid_key, blockk, 0, stream,
+                d_attn_a, query_a, d_key_a, 
+                depth, height, width, heads, dilation, dilation_d, dimhalf, nhalf_key);
+    }));
+    return {d_query, d_key, d_rpb};
+}
+
+torch::Tensor natten3dav_cuda_forward(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = value.size(0);
+    int64_t heads = value.size(1);
+    int64_t depth = value.size(2);
+    int64_t height = value.size(3);
+    int64_t width = value.size(4);
+    int64_t dim = value.size(5);
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dav_cuda_forward", kernel_size);
+    CHECK_KERNELSIZE("natten3dav_cuda_forward", kernel_size_d);
+
+    auto out = torch::zeros_like(value);
+
+    int32_t n = out.numel();
+    int blocks = GET_BLOCKS(n, CUDA_NUM_THREADS_F);
+    dim3 grid(blocks);
+    dim3 block(CUDA_NUM_THREADS_F);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "natten3dav_forward_cuda", ([&] {
+        const auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto value_a = value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto out_a = out.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        LAUNCH_NA_KDNDS(
+                kernel_size, kernel_size_d,
+                neighborhood_neighborhood_3d, 
+                grid, block, 0, stream, 
+                attn_a, value_a, out_a, 
+                depth, height, width, heads, dilation, dilation_d, dim, n);
+    }));
+    return out;
+}
+
+torch::Tensor natten3dav_cuda_forward_fp16(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = value.size(0);
+    int64_t heads = value.size(1);
+    int64_t depth = value.size(2);
+    int64_t height = value.size(3);
+    int64_t width = value.size(4);
+    int64_t dimhalf = value.size(5) / 2;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dav_cuda_forward_fp16", kernel_size);
+    CHECK_KERNELSIZE("natten3dav_cuda_forward_fp16", kernel_size_d);
+    TORCH_CHECK(dimhalf*2 == value.size(5), "Dims per head must be an even number in FP16.");
+
+    auto out = torch::zeros_like(value);
+
+    int32_t nhalf = out.numel() / 2;
+    int blocks = GET_BLOCKS(nhalf, CUDA_NUM_THREADS_FP16);
+    dim3 grid(blocks);
+    dim3 block(CUDA_NUM_THREADS_FP16);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_HALF_TYPES(at::kHalf, value.scalar_type(), "natten3dav_forward_cuda_fp16", ([&] {
+        const auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto value_a = value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto out_a = out.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        LAUNCH_NA_KDNDS(
+                kernel_size, kernel_size_d,
+                neighborhood_neighborhood_3d_fp16, 
+                grid, block, 0, stream, 
+                attn_a, value_a, out_a, 
+                depth, height, width, heads, dilation, dilation_d, dimhalf, nhalf);
+    }));
+    return out;
+}
+
+std::vector<torch::Tensor> natten3dav_cuda_backward(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = value.size(0);
+    int64_t heads = value.size(1);
+    int64_t depth = value.size(2);
+    int64_t height = value.size(3);
+    int64_t width = value.size(4);
+    int64_t dim = value.size(5);
+    int attn_size = attn.size(5); 
+    int zsize = batch_size * heads;
+    int xsize = depth * height * width;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dav_cuda_backward", kernel_size);
+    CHECK_KERNELSIZE("natten3dav_cuda_backward", kernel_size_d);
+    int KERNELTHREADS = min(CUDA_NUM_THREADS, attn_size);
+    int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), xsize);
+    int BATCHTHREADS = max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+
+    auto d_attn = torch::zeros_like(attn);
+    auto d_value = torch::zeros_like(value);
+
+    const dim3 attn_blocks(
+            (xsize + PIXELTHREADS - 1) / PIXELTHREADS,
+            (attn_size + KERNELTHREADS - 1) / KERNELTHREADS,
+            (zsize + BATCHTHREADS - 1) / BATCHTHREADS);
+    const dim3 attn_threads(PIXELTHREADS, KERNELTHREADS, BATCHTHREADS);
+    int32_t n_value = d_value.numel();
+    int blocks_value = GET_BLOCKS(n_value);
+    dim3 grid_value(blocks_value);
+    dim3 block(CUDA_NUM_THREADS);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_FLOATING_TYPES(d_attn.scalar_type(), "natten3dav_backward_cuda", ([&] {
+        auto d_attn_a = d_attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_value_a = d_value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto d_out_a = d_out.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto value_a = value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                pointwise_neighborhood_3d, 
+                attn_blocks, attn_threads, 0, stream, 
+                d_out_a, value_a, d_attn_a, 
+                depth, height, width, batch_size, heads, dilation, dilation_d, dim);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d,
+                inverse_neighborhood_3d, grid_value, block, 0, stream, 
+                attn_a, d_out_a, d_value_a, 
+                depth, height, width, heads, dilation, dilation_d, dim, n_value);
+    }));
+    return {d_attn, d_value};
+}
+
+std::vector<torch::Tensor> natten3dav_cuda_backward_fp16(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    int64_t batch_size = value.size(0);
+    int64_t heads = value.size(1);
+    int64_t depth = value.size(2);
+    int64_t height = value.size(3);
+    int64_t width = value.size(4);
+    int64_t dimhalf = value.size(5) / 2;
+    int attn_size = attn.size(5); 
+    int zsize = batch_size * heads;
+    int xsize = depth * height * width;
+    CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d);
+    CHECK_KERNELSIZE("natten3dav_cuda_backward_fp16", kernel_size);
+    CHECK_KERNELSIZE("natten3dav_cuda_backward_fp16", kernel_size_d);
+    TORCH_CHECK(dimhalf*2 == value.size(5), "Dims per head must be an even number in FP16.");
+    int KERNELTHREADS = min(CUDA_NUM_THREADS, attn_size);
+    int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), xsize);
+    int BATCHTHREADS = max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+
+    auto d_attn = torch::zeros_like(attn);
+    auto d_value = torch::zeros_like(value);
+
+    const dim3 attn_blocks(
+            (xsize + PIXELTHREADS - 1) / PIXELTHREADS,
+            (attn_size + KERNELTHREADS - 1) / KERNELTHREADS,
+            (zsize + BATCHTHREADS - 1) / BATCHTHREADS);
+    const dim3 attn_threads(PIXELTHREADS, KERNELTHREADS, BATCHTHREADS);
+    int32_t nhalf_value = d_value.numel() / 2;
+    int blocks_value = GET_BLOCKS(nhalf_value);
+    dim3 grid_value(blocks_value);
+    dim3 block(CUDA_NUM_THREADS);
+    const auto stream = c10::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_HALF_TYPES(at::kHalf, d_attn.scalar_type(), "natten3dav_backward_cuda_fp16", ([&] {
+        auto d_attn_a = d_attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        auto d_value_a = d_value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto d_out_a = d_out.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto value_a = value.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        const auto attn_a = attn.packed_accessor32<scalar_t,6,torch::DefaultPtrTraits>();
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                pointwise_neighborhood_3d_fp16, 
+                attn_blocks, attn_threads, 0, stream, 
+                d_out_a, value_a, d_attn_a,  
+                depth, height, width, batch_size, heads, dilation, dilation_d, dimhalf);
+        LAUNCH_NA_KDNDS(kernel_size, kernel_size_d, 
+                inverse_neighborhood_3d_fp16, 
+                grid_value, block, 0, stream, 
+                d_out_a, d_value_a, attn_a, 
+                depth, height, width, heads, dilation, dilation_d, dimhalf, nhalf_value);
+    }));
+    return {d_attn, d_value};
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cuda/natten_commons.cuh b/src/natten/csrc/cuda/natten_commons.cuh
index a0b43c2..deb742c 100644
--- a/src/natten/csrc/cuda/natten_commons.cuh
+++ b/src/natten/csrc/cuda/natten_commons.cuh
@@ -112,7 +112,7 @@ inline __host__ __device__ int get_pb_start(const int index, const int length, c
 
 #define CHECK_SEQUENCE(length, kernel_size, dilation) TORCH_CHECK(length >= kernel_size*dilation, "Input sequence length must be greater than or equal to kernel size x dilation.")
 #define CHECK_FEATMAP(height, width, kernel_size, dilation) TORCH_CHECK(height >= kernel_size*dilation && width >= kernel_size*dilation, "Input resolution must be greater than or equal to kernel size x dilation.")
-#define CHECK_3DFEATMAP(depth, height, width, kernel_size_d, kernel_size, dilation_d, dilation) TORCH_CHECK(depth >= kernel_size_d*dilation_d && height >= kernel_size*dilation && width >= kernel_size*dilation, "Input resolution must be greater than or equal to kernel size x dilation.")
+#define CHECK_3DFEATMAP(depth, height, width, kernel_size, kernel_size_d, dilation, dilation_d) TORCH_CHECK(depth >= kernel_size_d*dilation_d && height >= kernel_size*dilation && width >= kernel_size*dilation, "Input resolution must be greater than or equal to kernel size x dilation.")
 
 // 2D Neighborhood Attention
 
diff --git a/src/natten/csrc/cuda/neighborhood_neighborhood_3d.cu b/src/natten/csrc/cuda/neighborhood_neighborhood_3d.cu
new file mode 100644
index 0000000..2ce3cb5
--- /dev/null
+++ b/src/natten/csrc/cuda/neighborhood_neighborhood_3d.cu
@@ -0,0 +1,151 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood-Neighborhood kernel for 3D data.
+           Applies neighborhood attention weights to neighborhood values.
+*/
+
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/ATen.h>
+#include <ATen/native/cuda/KernelUtils.cuh>
+#include <ATen/AccumulateType.h>
+#include <cuda_fp16.h>
+
+#include "natten_commons.cuh"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor6D = typename torch::PackedTensorAccessor32<scalar_t,6,torch::DefaultPtrTraits>;
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void neighborhood_neighborhood_3d(           // AV     / Q-grad
+    const Tensor6D<scalar_t> weights,                   // attn   / d_attn
+    const Tensor6D<scalar_t> values,                    // value  / key
+    Tensor6D<scalar_t> output,                          // output / d_query
+    const int depth,
+    const int height,
+    const int width,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dim,
+    const int totalElements) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < totalElements){
+        int indtmp1 = linearIndex/dim;
+        const int d = linearIndex - indtmp1 * dim;
+        int indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+
+        const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        scalar_t updt = scalar_t(0);
+        int weightsOffset = b * weights.stride(0) + h * weights.stride(1) + k * weights.stride(2) + i * weights.stride(3) + j * weights.stride(4);
+        const int valuesOffset = b * values.stride(0) + h * values.stride(1) + d;
+        #pragma unroll
+        for (int xk=nk; xk < nk + KERNEL_SIZE_D * dilation_d; xk+=dilation_d)
+            #pragma unroll
+            for (int xi=ni; xi < ni + KERNEL_SIZE * dilation; xi+=dilation)
+                #pragma unroll
+                for (int xj=nj; xj < nj + KERNEL_SIZE * dilation; xj+=dilation){
+                    const int valuesIndex = valuesOffset + xk * values.stride(2)  + xi * values.stride(3) + xj * values.stride(4);
+                    updt += weights.data()[weightsOffset] * values.data()[valuesIndex];
+                    ++weightsOffset;
+                }
+        output.data()[linearIndex] = updt;
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void neighborhood_neighborhood_3d_fp16(      // AV     / Q-grad
+    const Tensor6D<scalar_t> weights,                   // attn   / d_attn
+    const Tensor6D<scalar_t> values,                    // value  / key
+    Tensor6D<scalar_t> output,                          // output / d_query
+    const int depth,
+    const int height,
+    const int width,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dimhalf,
+    const int totalElements) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < totalElements){
+        __half2* values2 = reinterpret_cast<__half2*>(values.data());
+        __half2* output2 = reinterpret_cast<__half2*>(output.data());
+        int indtmp1 = linearIndex/dimhalf;
+        const int d = linearIndex - indtmp1 * dimhalf;
+        int indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/heads;
+        const int h = indtmp1 - indtmp2 * heads;
+        const int b = indtmp2;
+
+        const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        __half2 updt = __float2half2_rn(0.f);
+        int weightsOffset = b * weights.stride(0) + h * weights.stride(1) + k * weights.stride(2) + i * weights.stride(3) + j * weights.stride(4);
+        const int stride3 = dimhalf * width;
+        const int stride2 = stride3 * height;
+        const int valuesOffset = b * (stride2 * depth * heads) + h * (stride2 * depth) + d;
+        #pragma unroll
+        for (int xk=nk; xk < nk + KERNEL_SIZE_D * dilation_d; xk+=dilation_d)
+            #pragma unroll
+            for (int xi=ni; xi < ni + KERNEL_SIZE * dilation; xi+=dilation)
+                #pragma unroll
+                for (int xj=nj; xj < nj + KERNEL_SIZE * dilation; xj+=dilation){
+                    const int valuesIndex = valuesOffset + xk * stride2 + xi * stride3 + xj * dimhalf;
+                    scalar_t a = weights.data()[weightsOffset];
+                    updt = __hfma2(__halves2half2(a, a), values2[valuesIndex], updt);
+                    ++weightsOffset;
+                }
+        output2[linearIndex] = updt;
+    }
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cuda/pointwise_neighborhood_2d.cu b/src/natten/csrc/cuda/pointwise_neighborhood_2d.cu
index b12bb01..9c8c943 100644
--- a/src/natten/csrc/cuda/pointwise_neighborhood_2d.cu
+++ b/src/natten/csrc/cuda/pointwise_neighborhood_2d.cu
@@ -21,7 +21,7 @@
  *
  **************************************************************************************************/
 /*! \file
-    \brief Pointwise-Neighborhood kernel for 1D data.
+    \brief Pointwise-Neighborhood kernel for 2D data.
            Computes attention weights between query points and their corresponding
            key neighborhood.
            Extra kernel with fused bias (relative positional bias.)
diff --git a/src/natten/csrc/cuda/pointwise_neighborhood_3d.cu b/src/natten/csrc/cuda/pointwise_neighborhood_3d.cu
new file mode 100644
index 0000000..eb98e43
--- /dev/null
+++ b/src/natten/csrc/cuda/pointwise_neighborhood_3d.cu
@@ -0,0 +1,269 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Pointwise-Neighborhood kernel for 3D data.
+           Computes attention weights between query points and their corresponding
+           key neighborhood.
+           Extra kernel with fused bias (relative positional bias.)
+*/
+
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/ATen.h>
+#include <ATen/native/cuda/KernelUtils.cuh>
+#include <ATen/AccumulateType.h>
+#include <cuda_fp16.h>
+
+#include "natten_commons.cuh"
+
+namespace natten {
+
+template<class scalar_t>
+using Tensor4D = typename torch::PackedTensorAccessor32<scalar_t,4,torch::DefaultPtrTraits>;
+template<class scalar_t>
+using Tensor6D = typename torch::PackedTensorAccessor32<scalar_t,6,torch::DefaultPtrTraits>;
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void pointwise_neighborhood_3d_bias( // QK
+    const Tensor6D<scalar_t> query,             // query
+    const Tensor6D<scalar_t> key,               // key  
+    const Tensor4D<scalar_t> bias,              // relative positional bias tensor
+    Tensor6D<scalar_t> attn,                    // attn
+    const int depth,
+    const int height,
+    const int width,
+    const int batch_size,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dim) {
+    const int z = blockIdx.z * blockDim.z + threadIdx.z;
+    if (z < batch_size * heads){
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        if (x < depth * height * width){
+            const int y = blockIdx.y * blockDim.y + threadIdx.y;
+            if (y < KERNEL_SIZE_D * KERNEL_SIZE * KERNEL_SIZE){
+                int indtmp1 = y / KERNEL_SIZE;
+                const int kk = indtmp1 / KERNEL_SIZE;
+                const int kj = y - indtmp1 * KERNEL_SIZE;
+                const int ki = indtmp1 - kk * KERNEL_SIZE;
+
+                indtmp1 = x / width;
+                const int k = indtmp1 / height;
+                const int j = x - indtmp1 * width;
+                const int i = indtmp1 - k * height;
+
+                const int b = z / heads;
+                const int h = z - b * heads;
+
+                const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+                const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+
+                scalar_t updt = scalar_t(0);
+                const int batchHeadOffset = b * query.stride(0) + h * query.stride(1);
+                const int queryOffset = batchHeadOffset + k * query.stride(2) + i * query.stride(3) + j * query.stride(4);
+                const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * key.stride(2) + (ki*dilation+ni) * key.stride(3) + (kj*dilation+nj) * key.stride(4);
+                #pragma unroll
+                for (int dimOffset=0; dimOffset < dim; ++dimOffset)
+                    updt += query.data()[queryOffset+dimOffset] * key.data()[keyOffset+dimOffset];
+                const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + y * attn.stride(5);
+                const int biasIndex = h * bias.stride(0) + (pk+kk) * bias.stride(1) + (pi+ki) * bias.stride(2) + (pj+kj) * bias.stride(3);
+                updt += bias.data()[biasIndex];
+                attn.data()[index] = updt;
+            }
+        }
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void pointwise_neighborhood_3d_bias_fp16( // QK
+    const Tensor6D<scalar_t> query,                  // query
+    const Tensor6D<scalar_t> key,                    // key  
+    const Tensor4D<scalar_t> bias,                   // relative positional bias tensor
+    Tensor6D<scalar_t> attn,                         // attn
+    const int depth,
+    const int height,
+    const int width,
+    const int batch_size,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dimhalf) {
+    const int z = blockIdx.z * blockDim.z + threadIdx.z;
+    if (z < batch_size * heads){
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        if (x < depth* height * width){
+            const int y = blockIdx.y * blockDim.y + threadIdx.y;
+            if (y < KERNEL_SIZE_D * KERNEL_SIZE * KERNEL_SIZE){
+                __half2* query2 = reinterpret_cast<__half2*>(query.data());
+                __half2* key2 = reinterpret_cast<__half2*>(key.data());
+                int indtmp1 = y / KERNEL_SIZE;
+                const int kk = indtmp1 / KERNEL_SIZE;
+                const int kj = y - indtmp1 * KERNEL_SIZE;
+                const int ki = indtmp1 - kk * KERNEL_SIZE;
+
+                indtmp1 = x / width;
+                const int k = indtmp1 / height;
+                const int j = x - indtmp1 * width;
+                const int i = indtmp1 - k * height;
+
+                const int b = z / heads;
+                const int h = z - b * heads;
+
+                const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+                const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+
+                __half2 updt = __float2half2_rn(0.f);
+                const int stride3 = dimhalf * width;
+                const int stride2 = stride3 * height;
+                const int batchHeadOffset = b * (stride2*depth*heads) + h * (stride2*depth);
+                const int queryOffset = batchHeadOffset + k * stride2 + i * stride3 + j * dimhalf;
+                const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * stride2 + (ki*dilation+ni) * stride3 + (kj*dilation+nj) * dimhalf;
+                #pragma unroll
+                for (int dimOffset=0; dimOffset < dimhalf; ++dimOffset)
+                    updt = __hfma2(query2[queryOffset+dimOffset], key2[keyOffset+dimOffset], updt);
+                const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + y * attn.stride(5);
+                const int biasIndex = h * bias.stride(0) + (pk+kk) * bias.stride(1) + (pi+ki) * bias.stride(2) + (pj+kj) * bias.stride(3);
+                attn.data()[index] = static_cast<scalar_t>(__hadd(updt.x, updt.y)) + bias.data()[biasIndex];
+            }
+        }
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void pointwise_neighborhood_3d(      // QK
+    const Tensor6D<scalar_t> query,             // query
+    const Tensor6D<scalar_t> key,               // key  
+    Tensor6D<scalar_t> attn,                    // attn
+    const int depth,
+    const int height,
+    const int width,
+    const int batch_size,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dim) {
+    const int z = blockIdx.z * blockDim.z + threadIdx.z;
+    if (z < batch_size * heads){
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        if (x < depth * height * width){
+            const int y = blockIdx.y * blockDim.y + threadIdx.y;
+            if (y < KERNEL_SIZE_D * KERNEL_SIZE * KERNEL_SIZE){
+                int indtmp1 = y / KERNEL_SIZE;
+                const int kk = indtmp1 / KERNEL_SIZE;
+                const int kj = y - indtmp1 * KERNEL_SIZE;
+                const int ki = indtmp1 - kk * KERNEL_SIZE;
+
+                indtmp1 = x / width;
+                const int k = indtmp1 / height;
+                const int j = x - indtmp1 * width;
+                const int i = indtmp1 - k * height;
+
+                const int b = z / heads;
+                const int h = z - b * heads;
+
+                const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+
+                scalar_t updt = scalar_t(0);
+                const int batchHeadOffset = b * query.stride(0) + h * query.stride(1);
+                const int queryOffset = batchHeadOffset + k * query.stride(2) + i * query.stride(3) + j * query.stride(4);
+                const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * key.stride(2) + (ki*dilation+ni) * key.stride(3) + (kj*dilation+nj) * key.stride(4);
+                #pragma unroll
+                for (int dimOffset=0; dimOffset < dim; ++dimOffset)
+                    updt += query.data()[queryOffset+dimOffset] * key.data()[keyOffset+dimOffset];
+                const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + y * attn.stride(5);
+                attn.data()[index] = updt;
+            }
+        }
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void pointwise_neighborhood_3d_fp16(      // QK
+    const Tensor6D<scalar_t> query,                  // query
+    const Tensor6D<scalar_t> key,                    // key  
+    Tensor6D<scalar_t> attn,                         // attn
+    const int depth,
+    const int height,
+    const int width,
+    const int batch_size,
+    const int heads,
+    const int dilation,
+    const int dilation_d,
+    const int dimhalf) {
+    const int z = blockIdx.z * blockDim.z + threadIdx.z;
+    if (z < batch_size * heads){
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        if (x < depth* height * width){
+            const int y = blockIdx.y * blockDim.y + threadIdx.y;
+            if (y < KERNEL_SIZE_D * KERNEL_SIZE * KERNEL_SIZE){
+                __half2* query2 = reinterpret_cast<__half2*>(query.data());
+                __half2* key2 = reinterpret_cast<__half2*>(key.data());
+                int indtmp1 = y / KERNEL_SIZE;
+                const int kk = indtmp1 / KERNEL_SIZE;
+                const int kj = y - indtmp1 * KERNEL_SIZE;
+                const int ki = indtmp1 - kk * KERNEL_SIZE;
+
+                indtmp1 = x / width;
+                const int k = indtmp1 / height;
+                const int j = x - indtmp1 * width;
+                const int i = indtmp1 - k * height;
+
+                const int b = z / heads;
+                const int h = z - b * heads;
+
+                const int ni = get_window_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nj = get_window_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+                const int nk = get_window_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+
+                __half2 updt = __float2half2_rn(0.f);
+                const int stride3 = dimhalf * width;
+                const int stride2 = stride3 * height;
+                const int batchHeadOffset = b * (stride2*depth*heads) + h * (stride2*depth);
+                const int queryOffset = batchHeadOffset + k * stride2 + i * stride3 + j * dimhalf;
+                const int keyOffset = batchHeadOffset + (kk*dilation_d+nk) * stride2 + (ki*dilation+ni) * stride3 + (kj*dilation+nj) * dimhalf;
+                #pragma unroll
+                for (int dimOffset=0; dimOffset < dimhalf; ++dimOffset)
+                    updt = __hfma2(query2[queryOffset+dimOffset], key2[keyOffset+dimOffset], updt);
+                const int index = b * attn.stride(0) + h * attn.stride(1) + k * attn.stride(2) + i * attn.stride(3) + j * attn.stride(4) + y * attn.stride(5);
+                attn.data()[index] = static_cast<scalar_t>(__hadd(updt.x, updt.y));
+            }
+        }
+    }
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/cuda/rel_pos_bias_3d.cu b/src/natten/csrc/cuda/rel_pos_bias_3d.cu
new file mode 100644
index 0000000..7acac7c
--- /dev/null
+++ b/src/natten/csrc/cuda/rel_pos_bias_3d.cu
@@ -0,0 +1,134 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Relative positional bias backward pass kernel for 3D data.
+*/
+
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/ATen.h>
+#include <ATen/native/cuda/KernelUtils.cuh>
+#include <ATen/AccumulateType.h>
+#include <cuda_fp16.h>
+
+#include "natten_commons.cuh"
+
+namespace natten {
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void rel_pos_bias_gradient_3d(
+    torch::PackedTensorAccessor32<scalar_t,4,torch::DefaultPtrTraits> d_rpb,
+    const torch::PackedTensorAccessor32<scalar_t,6,torch::DefaultPtrTraits> d_attn,
+    const int depth,
+    const int height,
+    const int width,
+    const int dilation,
+    const int dilation_d,
+    const int batch_size,
+    const int d_rpb_numel,
+    const int totalThreads) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < totalThreads){
+        int indtmp1 = linearIndex/KERNEL_SIZE;
+        const int kj = linearIndex - indtmp1 * KERNEL_SIZE;
+        int indtmp2 = indtmp1/KERNEL_SIZE;
+        const int ki = indtmp1 - indtmp2 * KERNEL_SIZE;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/KERNEL_SIZE_D;
+        const int kk = indtmp1 - indtmp2 * KERNEL_SIZE_D;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        const int h = indtmp2;
+        const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        scalar_t d_rpb_update = scalar_t(0);
+        int attnOffset = h * d_attn.stride(1) + k * d_attn.stride(2) + i * d_attn.stride(3) + j * d_attn.stride(4) + ((kk*KERNEL_SIZE*KERNEL_SIZE)+(ki*KERNEL_SIZE)+kj);
+        #pragma unroll
+        for (int b=0; b < batch_size; ++b){
+            d_rpb_update += d_attn.data()[attnOffset];
+            attnOffset += d_attn.stride(0);
+        }
+        const int index = h * d_rpb.stride(0) + (pk+kk) * d_rpb.stride(1) + (pi+ki) * d_rpb.stride(2) + (pj+kj) * d_rpb.stride(3);
+        at::native::fastAtomicAdd(d_rpb.data(), index, d_rpb_numel, static_cast<scalar_t>(d_rpb_update), true);
+    }
+}
+
+template <int KERNEL_SIZE, int KERNEL_SIZE_D, int NEIGHBORHOOD_SIZE, int NEIGHBORHOOD_SIZE_D, typename scalar_t>
+__global__ void rel_pos_bias_gradient_3d_fp16(
+    torch::PackedTensorAccessor32<scalar_t,4,torch::DefaultPtrTraits> d_rpb,
+    const torch::PackedTensorAccessor32<scalar_t,6,torch::DefaultPtrTraits> d_attn,
+    const int depth,
+    const int height,
+    const int width,
+    const int dilation,
+    const int dilation_d,
+    const int batch_size,
+    const int d_rpb_numel,
+    const int totalThreads) {
+    const int linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+    if (linearIndex < totalThreads){
+        int indtmp1 = linearIndex/KERNEL_SIZE;
+        const int kj = linearIndex - indtmp1 * KERNEL_SIZE;
+        int indtmp2 = indtmp1/KERNEL_SIZE;
+        const int ki = indtmp1 - indtmp2 * KERNEL_SIZE;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/KERNEL_SIZE_D;
+        const int kk = indtmp1 - indtmp2 * KERNEL_SIZE_D;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/width;
+        const int j = indtmp1 - indtmp2 * width;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/height;
+        const int i = indtmp1 - indtmp2 * height;
+        indtmp1 = indtmp2;
+        indtmp2 = indtmp1/depth;
+        const int k = indtmp1 - indtmp2 * depth;
+        const int h = indtmp2;
+        const int pi = get_pb_start(i, height, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pj = get_pb_start(j, width, KERNEL_SIZE, NEIGHBORHOOD_SIZE, dilation);
+        const int pk = get_pb_start(k, depth, KERNEL_SIZE_D, NEIGHBORHOOD_SIZE_D, dilation_d);
+        float d_rpb_update = float(0.0);
+        int attnOffset = h * d_attn.stride(1) + k * d_attn.stride(2) + i * d_attn.stride(3) + j * d_attn.stride(4) + ((kk*KERNEL_SIZE*KERNEL_SIZE)+(ki*KERNEL_SIZE)+kj);
+        #pragma unroll
+        for (int b=0; b < batch_size; ++b){
+            d_rpb_update += static_cast<float>(d_attn.data()[attnOffset]);
+            attnOffset += d_attn.stride(0);
+        }
+        const int index = h * d_rpb.stride(0) + (pk+kk) * d_rpb.stride(1) + (pi+ki) * d_rpb.stride(2) + (pj+kj) * d_rpb.stride(3);
+        at::native::fastAtomicAdd(d_rpb.data(), index, d_rpb_numel, static_cast<scalar_t>(d_rpb_update), true);
+    }
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/natten.cpp b/src/natten/csrc/natten.cpp
index 45ccdde..4d7781e 100644
--- a/src/natten/csrc/natten.cpp
+++ b/src/natten/csrc/natten.cpp
@@ -29,6 +29,8 @@
 #include "natten1dqkrpb.h"
 #include "natten2dav.h"
 #include "natten2dqkrpb.h"
+#include "natten3dav.h"
+#include "natten3dqkrpb.h"
 
 namespace natten {
 
@@ -42,6 +44,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("natten2dqkrpb_backward", &natten2dqkrpb_backward, "NATTEN2DQK+RPB backward");
   m.def("natten2dav_forward", &natten2dav_forward, "NATTEN2DAV forward");
   m.def("natten2dav_backward", &natten2dav_backward, "NATTEN2DAV backward");
+
+  m.def("natten3dqkrpb_forward", &natten3dqkrpb_forward, "NATTEN3DQK+RPB forward");
+  m.def("natten3dqkrpb_backward", &natten3dqkrpb_backward, "NATTEN3DQK+RPB backward");
+  m.def("natten3dav_forward", &natten3dav_forward, "NATTEN3DAV forward");
+  m.def("natten3dav_backward", &natten3dav_backward, "NATTEN3DAV backward");
 }
 
 } // namespace natten
diff --git a/src/natten/csrc/natten3dav.h b/src/natten/csrc/natten3dav.h
new file mode 100644
index 0000000..b63d1c7
--- /dev/null
+++ b/src/natten/csrc/natten3dav.h
@@ -0,0 +1,144 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood Attention 3D - AV (attention * value) bindings
+*/
+
+#include <torch/extension.h>
+#include <vector>
+
+namespace natten {
+
+// CPU forward declarations
+torch::Tensor natten3dav_cpu_forward(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+// CPU backward declarations
+std::vector<torch::Tensor> natten3dav_cpu_backward(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+#if defined(WITH_CUDA)
+// CUDA forward declarations
+torch::Tensor natten3dav_cuda_forward(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+torch::Tensor natten3dav_cuda_forward_fp16(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+// CUDA backward declarations
+std::vector<torch::Tensor> natten3dav_cuda_backward(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+std::vector<torch::Tensor> natten3dav_cuda_backward_fp16(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+#endif
+
+// C++ interface
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+torch::Tensor natten3dav_forward(
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    CHECK_CONTIGUOUS(attn);
+    CHECK_CONTIGUOUS(value);
+    assert(int(sqrt(attn.size(5))) == kernel_size);
+    assert(attn.device().is_cuda() == value.device().is_cuda());
+    if (attn.device().is_cuda()) {
+#if defined(WITH_CUDA)
+        bool half = ::detail::scalar_type(value.scalar_type()) == at::ScalarType::Half;
+        if (half)
+            return natten3dav_cuda_forward_fp16(attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+        return natten3dav_cuda_forward(attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+#else
+    AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
+#endif
+    }
+    return natten3dav_cpu_forward(attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+}
+
+std::vector<torch::Tensor> natten3dav_backward(
+    const torch::Tensor &d_out,
+    const torch::Tensor &attn,
+    const torch::Tensor &value,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    CHECK_CONTIGUOUS(d_out);
+    CHECK_CONTIGUOUS(attn);
+    CHECK_CONTIGUOUS(value);
+    assert(int(sqrt(attn.size(5))) == kernel_size);
+    assert(attn.device().is_cuda() == value.device().is_cuda() && d_out.device().is_cuda() == value.device().is_cuda());
+    if (attn.device().is_cuda()) {
+#if defined(WITH_CUDA)
+        bool half = ::detail::scalar_type(value.scalar_type()) == at::ScalarType::Half;
+        if (half)
+            return natten3dav_cuda_backward_fp16(d_out, attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+        return natten3dav_cuda_backward(d_out, attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+#else
+    AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
+#endif
+    }
+    return natten3dav_cpu_backward(d_out, attn, value, kernel_size_d, kernel_size, dilation_d, dilation);
+}
+
+} // namespace natten
diff --git a/src/natten/csrc/natten3dqkrpb.h b/src/natten/csrc/natten3dqkrpb.h
new file mode 100644
index 0000000..5f9db91
--- /dev/null
+++ b/src/natten/csrc/natten3dqkrpb.h
@@ -0,0 +1,156 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 Ali Hassani.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Neighborhood Attention 3D - QK (query * key) bindings
+*/
+
+#include <torch/extension.h>
+#include <vector>
+
+namespace natten {
+
+// CPU forward declarations
+torch::Tensor natten3dqkrpb_cpu_forward(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+// CPU backward declarations
+std::vector<torch::Tensor> natten3dqkrpb_cpu_backward(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+#if defined(WITH_CUDA)
+// CUDA forward declarations
+torch::Tensor natten3dqkrpb_cuda_forward(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+torch::Tensor natten3dqkrpb_cuda_forward_fp16(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+// CUDA backward declarations
+std::vector<torch::Tensor> natten3dqkrpb_cuda_backward(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+std::vector<torch::Tensor> natten3dqkrpb_cuda_backward_fp16(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation);
+
+#endif
+
+// C++ interface
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+torch::Tensor natten3dqkrpb_forward(
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const at::optional<at::Tensor> &rpb,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    CHECK_CONTIGUOUS(query);
+    CHECK_CONTIGUOUS(key);
+    const int heads = query.size(1);
+    assert(query.device().is_cuda() == key.device().is_cuda());
+    if (rpb.has_value()) {
+        assert(rpb.value().device().is_cuda() == key.device().is_cuda());
+        CHECK_CONTIGUOUS(rpb.value());
+        assert(int((rpb.value().size(1) + 1) / 2) == kernel_size);
+    }
+    if (query.device().is_cuda()) {
+#if defined(WITH_CUDA)
+        bool half = ::detail::scalar_type(query.scalar_type()) == at::ScalarType::Half;
+        if (half)
+            return natten3dqkrpb_cuda_forward_fp16(query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation);
+        return natten3dqkrpb_cuda_forward(query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation);
+#else
+    AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
+#endif
+    }
+    return natten3dqkrpb_cpu_forward(query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation);
+}
+
+std::vector<torch::Tensor> natten3dqkrpb_backward(
+    const torch::Tensor &d_attn,
+    const torch::Tensor &query,
+    const torch::Tensor &key,
+    const bool biasEnabled,
+    const int kernel_size_d,
+    const int kernel_size,
+    const int dilation_d,
+    const int dilation) {
+    CHECK_CONTIGUOUS(d_attn);
+    CHECK_CONTIGUOUS(query);
+    CHECK_CONTIGUOUS(key);
+    assert(query.device().is_cuda() == key.device().is_cuda() && d_attn.device().is_cuda() == key.device().is_cuda());
+    if (query.device().is_cuda()) {
+#if defined(WITH_CUDA)
+        bool half = ::detail::scalar_type(query.scalar_type()) == at::ScalarType::Half;
+        if (half)
+            return natten3dqkrpb_cuda_backward_fp16(d_attn, query, key, biasEnabled, kernel_size_d, kernel_size, dilation_d, dilation);
+        return natten3dqkrpb_cuda_backward(d_attn, query, key, biasEnabled, kernel_size_d, kernel_size, dilation_d, dilation);
+#else
+    AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
+#endif
+    }
+    return natten3dqkrpb_cpu_backward(d_attn, query, key, biasEnabled, kernel_size_d, kernel_size, dilation_d, dilation);
+}
+
+} // namespace natten
diff --git a/src/natten/flops.py b/src/natten/flops.py
index 89df7be..00720e6 100644
--- a/src/natten/flops.py
+++ b/src/natten/flops.py
@@ -27,11 +27,70 @@
 from fvcore.nn.jit_handles import get_shape
 
 
-def qk_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def qk_1d_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     """
-    Count flops for the QKRPB kernel.
+    Counts flops for the 1D QK operation.
     """
-    assert len(inputs) >= 3, f"Expected 3 inputs (query, key, rpb), got {len(inputs)}"
+    assert (
+        len(inputs) >= 2
+    ), f"Expected at least 2 inputs (query, key), got {len(inputs)}"
+    assert len(outputs) == 1, f"Expected 1 output (attn), got {len(outputs)}"
+    input_shapes = [get_shape(v) for v in inputs]
+    output_shapes = [get_shape(v) for v in outputs]
+    assert (
+        len(input_shapes[0]) == 4
+    ), f"Query must be a 4-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 4
+    ), f"Key must be a 4-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(output_shapes[0]) == 4
+    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+    assert (
+        input_shapes[0] == input_shapes[1]
+    ), f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
+    batch_size, heads, length, dim = input_shapes[0]
+    batch_size, heads, length, kernel_size = output_shapes[0]
+
+    flops = batch_size * heads * length * dim * kernel_size
+    flops += batch_size * heads * length * kernel_size
+    return flops
+
+
+def av_1d_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+    """
+    Counts flops for the 1D AV operation.
+    """
+    assert len(inputs) == 2, f"Expected 2 inputs (attn and value), got {len(inputs)}"
+    assert len(outputs) == 1, f"Expected 1 output (out), got {len(outputs)}"
+    input_shapes = [get_shape(v) for v in inputs]
+    output_shapes = [get_shape(v) for v in outputs]
+    assert (
+        len(input_shapes[0]) == 4
+    ), f"Attn must be a 4-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 4
+    ), f"Value must be a 4-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(output_shapes[0]) == 4
+    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+    assert output_shapes[0] == input_shapes[1], (
+        f"Out and Value shapes did not match! O: {output_shapes[0]}, V:"
+        f" {input_shapes[1]}"
+    )
+    batch_size, heads, length, kernel_size = input_shapes[0]
+    batch_size, heads, length, dim = output_shapes[0]
+    flops = batch_size * heads * length * dim * kernel_size
+    return flops
+
+
+def qk_2d_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+    """
+    Counts flops for the 2D QK operation.
+    """
+    assert (
+        len(inputs) >= 2
+    ), f"Expected at least 2 inputs (query, key), got {len(inputs)}"
     assert len(outputs) == 1, f"Expected 1 output (attn), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
@@ -41,9 +100,6 @@ def qk_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     assert (
         len(input_shapes[1]) == 5
     ), f"Key must be a 5-dim tensor, got {len(input_shapes[1])}"
-    assert (
-        len(input_shapes[2]) == 3
-    ), f"RelPosBias must be a 3-dim tensor, got {len(input_shapes[2])}"
     assert (
         len(output_shapes[0]) == 5
     ), f"Output must be a 5-dim tensor, got {len(output_shapes[0])}"
@@ -58,11 +114,11 @@ def qk_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     return flops
 
 
-def av_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def av_2d_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     """
-    Count flops for the AV kernel.
+    Counts flops for the 2D AV operation.
     """
-    assert len(inputs) >= 2, f"Expected 2 inputs (attn and value), got {len(inputs)}"
+    assert len(inputs) == 2, f"Expected 2 inputs (attn and value), got {len(inputs)}"
     assert len(outputs) == 1, f"Expected 1 output (out), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
@@ -85,71 +141,72 @@ def av_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     return flops
 
 
-def qk_1d_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def qk_3d_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     """
-    Count flops for the 1D QKRPB kernel.
+    Counts flops for the 3D QK operation.
     """
-    assert len(inputs) >= 3, f"Expected 3 inputs (query, key, rpb), got {len(inputs)}"
+    assert (
+        len(inputs) >= 2
+    ), f"Expected at least 2 inputs (query, key), got {len(inputs)}"
     assert len(outputs) == 1, f"Expected 1 output (attn), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
     assert (
-        len(input_shapes[0]) == 4
-    ), f"Query must be a 4-dim tensor, got {len(input_shapes[0])}"
-    assert (
-        len(input_shapes[1]) == 4
-    ), f"Key must be a 4-dim tensor, got {len(input_shapes[1])}"
+        len(input_shapes[0]) == 6
+    ), f"Query must be a 6-dim tensor, got {len(input_shapes[0])}"
     assert (
-        len(input_shapes[2]) == 2
-    ), f"RelPosBias must be a 2-dim tensor, got {len(input_shapes[2])}"
+        len(input_shapes[1]) == 6
+    ), f"Key must be a 6-dim tensor, got {len(input_shapes[1])}"
     assert (
-        len(output_shapes[0]) == 4
-    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+        len(output_shapes[0]) == 6
+    ), f"Output must be a 6-dim tensor, got {len(output_shapes[0])}"
     assert (
         input_shapes[0] == input_shapes[1]
     ), f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
-    batch_size, heads, length, dim = input_shapes[0]
-    batch_size, heads, length, kernel_size = output_shapes[0]
+    batch_size, heads, depth, height, width, dim = input_shapes[0]
+    batch_size, heads, depth, height, width, kernel_size_cu = output_shapes[0]
 
-    flops = batch_size * heads * length * dim * kernel_size
-    flops += batch_size * heads * length * kernel_size
+    flops = batch_size * heads * depth * height * width * dim * kernel_size_cu
+    flops += batch_size * heads * depth * height * width * kernel_size_cu
     return flops
 
 
-def av_1d_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def av_3d_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     """
-    Count flops for the 1D AV kernel.
+    Counts flops for the 3D AV operation.
     """
-    assert len(inputs) >= 2, f"Expected 2 inputs (attn and value), got {len(inputs)}"
+    assert len(inputs) == 2, f"Expected 2 inputs (attn and value), got {len(inputs)}"
     assert len(outputs) == 1, f"Expected 1 output (out), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
     assert (
-        len(input_shapes[0]) == 4
-    ), f"Attn must be a 4-dim tensor, got {len(input_shapes[0])}"
+        len(input_shapes[0]) == 6
+    ), f"Attn must be a 6-dim tensor, got {len(input_shapes[0])}"
     assert (
-        len(input_shapes[1]) == 4
-    ), f"Value must be a 4-dim tensor, got {len(input_shapes[1])}"
+        len(input_shapes[1]) == 6
+    ), f"Value must be a 6-dim tensor, got {len(input_shapes[1])}"
     assert (
-        len(output_shapes[0]) == 4
-    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+        len(output_shapes[0]) == 6
+    ), f"Output must be a 6-dim tensor, got {len(output_shapes[0])}"
     assert output_shapes[0] == input_shapes[1], (
         f"Out and Value shapes did not match! O: {output_shapes[0]}, V:"
         f" {input_shapes[1]}"
     )
-    batch_size, heads, length, kernel_size = input_shapes[0]
-    batch_size, heads, length, dim = output_shapes[0]
-    flops = batch_size * heads * length * dim * kernel_size
+    batch_size, heads, depth, height, width, kernel_size_cu = input_shapes[0]
+    batch_size, heads, depth, height, width, dim = output_shapes[0]
+    flops = batch_size * heads * depth * height * width * dim * kernel_size_cu
     return flops
 
 
 def add_natten_handle(flop_ctr):
     return flop_ctr.set_op_handle(
         **{
-            "prim::PythonOp.NATTEN2DQKRPBFunction": qk_rpb_flop,
-            "prim::PythonOp.NATTEN2DAVFunction": av_flop,
             "prim::PythonOp.NATTEN1DQKRPBFunction": qk_1d_rpb_flop,
             "prim::PythonOp.NATTEN1DAVFunction": av_1d_flop,
+            "prim::PythonOp.NATTEN2DQKRPBFunction": qk_2d_rpb_flop,
+            "prim::PythonOp.NATTEN2DAVFunction": av_2d_flop,
+            "prim::PythonOp.NATTEN3DQKRPBFunction": qk_3d_rpb_flop,
+            "prim::PythonOp.NATTEN3DAVFunction": av_3d_flop,
         }
     )
 
diff --git a/src/natten/functional.py b/src/natten/functional.py
index 38feccb..bf43f87 100644
--- a/src/natten/functional.py
+++ b/src/natten/functional.py
@@ -170,6 +170,85 @@ def backward(ctx, grad_out):
         return d_attn, d_value, None, None
 
 
+class NATTEN3DQKRPBFunction(Function):
+    """
+    3D QK+RPB autograd function
+    Computes neighborhood attention weights given queries and keys,
+    and adds relative positional biases.
+    This calls the `QKRPB` kernel.
+    """
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float16)
+    def forward(ctx, query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation):
+        query = query.contiguous()
+        key = key.contiguous()
+        attn = _C.natten3dqkrpb_forward(
+            query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation
+        )
+        ctx.save_for_backward(query, key)
+        ctx.kernel_size_d = kernel_size_d
+        ctx.kernel_size = kernel_size
+        ctx.dilation_d = dilation_d
+        ctx.dilation = dilation
+        ctx.bias = rpb is not None
+        return attn
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad_out):
+        outputs = _C.natten3dqkrpb_backward(
+            grad_out.contiguous(),
+            ctx.saved_tensors[0],
+            ctx.saved_tensors[1],
+            ctx.bias,
+            ctx.kernel_size_d,
+            ctx.kernel_size,
+            ctx.dilation_d,
+            ctx.dilation,
+        )
+        d_query, d_key, d_rpb = outputs
+        return d_query, d_key, d_rpb, None, None, None, None
+
+
+class NATTEN3DAVFunction(Function):
+    """
+    3D AV autograd function
+    Computes neighborhood attention outputs given attention weights, and values.
+    This calls the `AV` kernel.
+    """
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float16)
+    def forward(ctx, attn, value, kernel_size_d, kernel_size, dilation_d, dilation):
+        attn = attn.contiguous()
+        value = value.contiguous()
+        out = _C.natten3dav_forward(
+            attn, value, kernel_size_d, kernel_size, dilation_d, dilation
+        )
+        ctx.save_for_backward(attn, value)
+        ctx.kernel_size_d = kernel_size_d
+        ctx.kernel_size = kernel_size
+        ctx.dilation_d = dilation_d
+        ctx.dilation = dilation
+        return out
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad_out):
+        outputs = _C.natten3dav_backward(
+            grad_out.contiguous(),
+            ctx.saved_tensors[0],
+            ctx.saved_tensors[1],
+            ctx.kernel_size_d,
+            ctx.kernel_size,
+            ctx.dilation_d,
+            ctx.dilation,
+        )
+        d_attn, d_value = outputs
+        return d_attn, d_value, None, None, None, None
+
+
 def natten1dqkrpb(query, key, rpb, kernel_size, dilation):
     return NATTEN1DQKRPBFunction.apply(query, key, rpb, kernel_size, dilation)
 
@@ -192,3 +271,21 @@ def natten2dqk(query, key, kernel_size, dilation):
 
 def natten2dav(attn, value, kernel_size, dilation):
     return NATTEN2DAVFunction.apply(attn, value, kernel_size, dilation)
+
+
+def natten3dqkrpb(query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation):
+    return NATTEN3DQKRPBFunction.apply(
+        query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation
+    )
+
+
+def natten3dqk(query, key, kernel_size_d, kernel_size, dilation_d, dilation):
+    return NATTEN3DQKRPBFunction.apply(
+        query, key, None, kernel_size_d, kernel_size, dilation_d, dilation
+    )
+
+
+def natten3dav(attn, value, kernel_size_d, kernel_size, dilation_d, dilation):
+    return NATTEN3DAVFunction.apply(
+        attn, value, kernel_size_d, kernel_size, dilation_d, dilation
+    )
diff --git a/src/natten/natten3d.py b/src/natten/natten3d.py
new file mode 100644
index 0000000..74d5fb4
--- /dev/null
+++ b/src/natten/natten3d.py
@@ -0,0 +1,126 @@
+#################################################################################################
+# Copyright (c) 2023 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import torch
+from torch import nn
+from torch.nn.functional import pad
+from torch.nn.init import trunc_normal_
+
+from .functional import natten3dav, natten3dqkrpb
+
+
+class NeighborhoodAttention3D(nn.Module):
+    """
+    Neighborhood Attention 3D Module
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        kernel_size,
+        kernel_size_d=None,
+        dilation=1,
+        dilation_d=None,
+        bias=True,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // self.num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+        kernel_size_d = kernel_size_d or kernel_size
+        dilation = dilation or 1
+        dilation_d = dilation_d or dilation
+        assert (
+            kernel_size > 1 and kernel_size % 2 == 1
+        ), f"Kernel size must be an odd number greater than 1, got {kernel_size}."
+        assert (
+            dilation >= 1
+        ), f"Dilation must be greater than or equal to 1, got {dilation}."
+        assert (
+            dilation_d >= 1
+        ), f"Dilation (depth) must be greater than or equal to 1, got {dilation_d}."
+        self.kernel_size = kernel_size
+        self.kernel_size_d = kernel_size_d
+        self.dilation = dilation
+        self.dilation_d = dilation_d
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        if bias:
+            self.rpb = nn.Parameter(
+                torch.zeros(
+                    num_heads,
+                    (2 * kernel_size_d - 1),
+                    (2 * kernel_size - 1),
+                    (2 * kernel_size - 1),
+                )
+            )
+            trunc_normal_(self.rpb, std=0.02, mean=0.0, a=-2.0, b=2.0)
+        else:
+            self.register_parameter("rpb", None)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, D, H, W, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, D, H, W, 3, self.num_heads, self.head_dim)
+            .permute(4, 0, 5, 1, 2, 3, 6)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q = q * self.scale
+        attn = natten3dqkrpb(
+            q,
+            k,
+            self.rpb,
+            self.kernel_size_d,
+            self.kernel_size,
+            self.dilation_d,
+            self.dilation,
+        )
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = natten3dav(
+            attn,
+            v,
+            self.kernel_size_d,
+            self.kernel_size,
+            self.dilation_d,
+            self.dilation,
+        )
+        x = x.permute(0, 2, 3, 4, 1, 5).reshape(B, D, H, W, C)
+
+        return self.proj_drop(self.proj(x))
+
+    def extra_repr(self) -> str:
+        return (
+            f"head_dim={self.head_dim}, num_heads={self.num_heads}, "
+            + f"kernel_size_d={self.kernel_size_d}, dilation_d={self.dilation_d}, "
+            + f"kernel_size={self.kernel_size}, dilation={self.dilation}, "
+            + f"rel_pos_bias={self.rpb is not None}"
+        )


Test patch:
diff --git a/tests/test_na3d.py b/tests/test_na3d.py
new file mode 100644
index 0000000..a08de1f
--- /dev/null
+++ b/tests/test_na3d.py
@@ -0,0 +1,439 @@
+#################################################################################################
+# Copyright (c) 2023 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+
+import logging
+import os
+import unittest
+
+import torch
+from torch.autograd import gradcheck
+from torch.utils.cpp_extension import CUDA_HOME
+
+from natten import NeighborhoodAttention3D
+from natten.functional import natten3dav, natten3dqkrpb
+
+HAS_CUDA = (
+    torch.cuda.is_available()
+    and (CUDA_HOME is not None)
+    or os.getenv("FORCE_CUDA", "0") == "1"
+)
+logger = logging.getLogger(__name__)
+
+
+def _priv_test_gradcheck_natten3dqk(
+    batch_size,
+    heads,
+    depth,
+    height,
+    width,
+    dim,
+    kernel_size_d,
+    kernel_size,
+    dilation_d,
+    dilation,
+    dtype,
+    device,
+    eps,
+    atol,
+    rtol,
+    ndtol,
+    fast_mode,
+):
+    torch.manual_seed(42)
+    md = "FAST" if fast_mode else "SLOW"
+    kwargs = {"dtype": dtype, "device": device, "requires_grad": True}
+    query = torch.randn((batch_size, heads, depth, height, width, dim), **kwargs)
+    key = torch.randn((batch_size, heads, depth, height, width, dim), **kwargs)
+    rpb = None
+    variables = [query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation]
+
+    assert gradcheck(
+        natten3dqkrpb,
+        variables,
+        eps=eps,
+        atol=atol,
+        rtol=rtol,
+        nondet_tol=ndtol,
+        fast_mode=fast_mode,
+    ), f"Gradcheck failed for `NATTEN3D-QK` with {dtype} on {device}, {md} MODE."
+
+
+def _priv_test_gradcheck_natten3dqkrpb(
+    batch_size,
+    heads,
+    depth,
+    height,
+    width,
+    dim,
+    kernel_size_d,
+    kernel_size,
+    dilation_d,
+    dilation,
+    dtype,
+    device,
+    eps,
+    atol,
+    rtol,
+    ndtol,
+    fast_mode,
+):
+    torch.manual_seed(42)
+    md = "FAST" if fast_mode else "SLOW"
+    kwargs = {"dtype": dtype, "device": device, "requires_grad": True}
+    query = torch.randn((batch_size, heads, depth, height, width, dim), **kwargs)
+    key = torch.randn((batch_size, heads, depth, height, width, dim), **kwargs)
+    rpb = torch.randn(
+        (heads, 2 * kernel_size_d - 1, 2 * kernel_size - 1, 2 * kernel_size - 1),
+        **kwargs,
+    )
+    variables = [query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation]
+
+    assert gradcheck(
+        natten3dqkrpb,
+        variables,
+        eps=eps,
+        atol=atol,
+        rtol=rtol,
+        nondet_tol=ndtol,
+        fast_mode=fast_mode,
+    ), f"Gradcheck failed for `NATTEN3D-QKRPB` with {dtype} on {device}, {md} MODE."
+
+
+def _priv_test_gradcheck_natten3dav(
+    batch_size,
+    heads,
+    depth,
+    height,
+    width,
+    dim,
+    kernel_size_d,
+    kernel_size,
+    dilation_d,
+    dilation,
+    dtype,
+    device,
+    eps,
+    atol,
+    rtol,
+    ndtol,
+    fast_mode,
+):
+    torch.manual_seed(42)
+    md = "FAST" if fast_mode else "SLOW"
+    kwargs = {"dtype": dtype, "device": device, "requires_grad": True}
+    attn = torch.randn(
+        (
+            batch_size,
+            heads,
+            depth,
+            height,
+            width,
+            kernel_size_d * kernel_size * kernel_size,
+        ),
+        **kwargs,
+    )
+    value = torch.randn((batch_size, heads, depth, height, width, dim), **kwargs)
+    variables = [attn, value, kernel_size_d, kernel_size, dilation_d, dilation]
+
+    assert gradcheck(
+        natten3dav,
+        variables,
+        eps=eps,
+        atol=atol,
+        rtol=rtol,
+        nondet_tol=ndtol,
+        fast_mode=fast_mode,
+    ), f"Gradcheck failed for `NATTEN3D-AV` with {dtype} on {device}, {md} MODE."
+
+
+def _priv_test_allclose_cpu_cuda(
+    batch_size,
+    depth,
+    height,
+    width,
+    depth_kernel_sizes=[3, 5, 7, 9, 11, 13],
+    kernel_sizes=[3, 5, 7, 9, 11, 13],
+    dims=[2, 4, 8],
+    heads=[1, 2, 4],
+    tol=1e-6,
+):
+    for kernel_size_d in depth_kernel_sizes:
+        for kernel_size in kernel_sizes:
+            for dim in dims:
+                for num_heads in heads:
+                    for rpb in [True, False]:
+                        model_kwargs = {
+                            "dim": dim * num_heads,
+                            "kernel_size_d": kernel_size_d,
+                            "kernel_size": kernel_size,
+                            "num_heads": num_heads,
+                            "qkv_bias": True,
+                            "bias": rpb,
+                        }
+
+                        base_state_dict = NeighborhoodAttention3D(
+                            **model_kwargs
+                        ).state_dict()
+
+                        x1 = torch.randn(
+                            (batch_size, depth, height, width, dim * num_heads)
+                        )
+                        x2 = x1.clone().detach().cuda(0)
+
+                        nat1 = NeighborhoodAttention3D(**model_kwargs).eval()
+                        nat1.load_state_dict(base_state_dict, strict=False)
+
+                        nat2 = NeighborhoodAttention3D(**model_kwargs).cuda(0).eval()
+                        nat2.load_state_dict(base_state_dict, strict=False)
+
+                        y1 = nat1(x1)
+                        y2 = nat2(x2)
+
+                        forward_mse = ((y1.data - y2.cpu().data) ** 2).mean()
+
+                        assert forward_mse < tol, (
+                            f"FAIL: Forward MSE ({forward_mse}) was above the specified"
+                            f" tolerance (tol) for heads={num_heads}, dim={dim},"
+                            f" kernel_size_d={kernel_size_d},"
+                            f" kernel_size={kernel_size}, rpb={rpb}."
+                        )
+
+
+class NA3DTest(unittest.TestCase):
+    def test_natten3dqk_gradcheck_cpu_fast(self):
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dqk(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cpu",
+            1e-6,
+            1e-5,
+            1e-3,
+            0,
+            True,
+        )
+
+    def test_natten3dqk_gradcheck_cuda_fast(self):
+        if not HAS_CUDA:
+            self.skipTest("CUDA not available.")
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dqk(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cuda",
+            1e-6,
+            1e-5,
+            1e-3,
+            1e-8,
+            True,
+        )
+
+    def test_natten3dqkrpb_gradcheck_cpu_fast(self):
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dqkrpb(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cpu",
+            1e-6,
+            1e-5,
+            1e-3,
+            0,
+            True,
+        )
+
+    def test_natten3dqkrpb_gradcheck_cuda_fast(self):
+        if not HAS_CUDA:
+            self.skipTest("CUDA not available.")
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dqkrpb(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cuda",
+            1e-6,
+            1e-5,
+            1e-3,
+            1e-8,
+            True,
+        )
+
+    def test_natten3dav_gradcheck_cpu_fast(self):
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dav(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cpu",
+            1e-6,
+            1e-5,
+            1e-3,
+            0,
+            True,
+        )
+
+    def test_natten3dav_gradcheck_cuda_fast(self):
+        if not HAS_CUDA:
+            self.skipTest("CUDA not available.")
+        (
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+        ) = (1, 2, 15, 16, 17, 32, 3, 5, 2, 3)
+        _priv_test_gradcheck_natten3dav(
+            batch,
+            heads,
+            depth,
+            height,
+            width,
+            dim,
+            kernel_d,
+            kernel,
+            dilation_d,
+            dilation,
+            torch.float64,
+            "cuda",
+            1e-6,
+            1e-5,
+            1e-3,
+            1e-8,
+            True,
+        )
+
+    def test_cpu_cuda_allclose(self):
+        if not HAS_CUDA:
+            self.skipTest("CUDA not available.")
+        _priv_test_allclose_cpu_cuda(
+            batch_size=2,
+            depth=18,
+            height=16,
+            width=17,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()