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

Issue:
Nested tensors
Hi, thanks for the library; we're using NATTEN currently for image generation and it's really helpful.

I was wondering whether you had any plan to (or already do) support [nested tensors](https://pytorch.org/docs/stable/nested.html)?  
this would enable image training to use batches with varying dimensions (you wouldn't need aspect ratio bucketing any more), and would enable languge model training to use batches with varying lengths (you wouldn't need padding or packing any more).

torch scaled_dot_product_attention [supports nested tensors](https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial#nestedtensor-and-dense-tensor-support) nowadays, perhaps the code provides a clue as to how they approached it:  
- forward support: https://github.com/pytorch/pytorch/commit/5fb687182dba781d9c95388d19f4784b98cb8b20
- backward support: https://github.com/pytorch/pytorch/commit/5183760ca52ea7dbd892ebc88cabdd43ebc81cbe
- flash attn support: https://github.com/pytorch/pytorch/commit/80614783e3333bd19eeec5856c60dd2eeaa16431
CUDA error: out of memory
When I run following code
```python
import torch
from natten import NeighborhoodAttention2D

device = torch.device("cuda:5")
na2d = NeighborhoodAttention2D(dim=128, kernel_size=7, dilation=2, num_heads=4).to(device)
x = torch.rand((1, 64, 64, 128)).to(device)
y = na2d(x)
```

But OOM:
```shell
attn = attn.softmax(dim=-1)
RuntimeError: CUDA error: out of memory
```

Here is response of `nvidia-smi`:

```shell
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.141.03   Driver Version: 470.141.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:1A:00.0 Off |                  N/A |
| 78%   68C    P2   321W / 350W |  23294MiB / 24268MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  On   | 00000000:1B:00.0 Off |                  N/A |
|100%   81C    P2   326W / 350W |   8768MiB / 24268MiB |     65%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce ...  On   | 00000000:1C:00.0 Off |                  N/A |
| 90%   74C    P2   311W / 350W |  21104MiB / 24268MiB |     73%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA GeForce ...  On   | 00000000:1D:00.0 Off |                  N/A |
|100%   80C    P2   324W / 350W |  21104MiB / 24268MiB |     73%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   4  NVIDIA GeForce ...  On   | 00000000:1E:00.0 Off |                  N/A |
| 94%   76C    P2   318W / 350W |  21104MiB / 24268MiB |     98%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   5  NVIDIA GeForce ...  On   | 00000000:3D:00.0 Off |                  N/A |
| 31%   39C    P8    22W / 350W |      3MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   6  NVIDIA GeForce ...  On   | 00000000:3F:00.0 Off |                  N/A |
| 90%   73C    P2   318W / 350W |  21104MiB / 24268MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   7  NVIDIA GeForce ...  On   | 00000000:40:00.0 Off |                  N/A |
| 87%   74C    P2   316W / 350W |  21104MiB / 24268MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   8  NVIDIA GeForce ...  On   | 00000000:41:00.0 Off |                  N/A |
| 82%   71C    P2   314W / 350W |  11930MiB / 24268MiB |     60%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
```


Discussion / hints:
Hello and thanks for your interest; I'm glad you're finding it useful.

I think it should be relatively easy to add support for nested tensors; I'll take a look at that later today.
The changes I made in the summer to the library structure separates out a lot of torch tensor API dependencies into a different layer, which means we could in theory have the API layer do multiple kernel calls on the different tensors within the nested tensor.
that's fantastic news! thanks; I hope it turns out to be as straightforward as you're anticipating. 🙂 
Okay; it looks like adding a C++ API will be difficult, mostly because PyTorch doesn't expose all the nested utils that are present in their source code (i.e. ATen/native/nested sources don't appear to ship even with the latest nightly build.)

I doubt any other plugin / framework with PyTorch integration supports nested tensors yet, but if you happen to know any, please let us know.

This unfortunately leaves us with only one choice: unpacking the nested tensors in the python interface and calling the C++ API on the buffers individually. It might be suboptimal, but it will get the job done. It'll just require a minor change in the API that I've been wanting to make for quite some time. I'll try to get it done by the end of the week.
I have a bit of bad news. I think PyTorch still doesn't support custom autograd functions taking in nested tensors, regardless of what the underlying operations are.

I even tried a simple identity op:

```python3
import torch
from torch.autograd import Function

class Na1dQKNested(Function):
    @staticmethod
    def forward(ctx, x):
        return x

    @staticmethod
    def jvp(ctx, x):
        return x

    @staticmethod
    def backward(ctx, x):
        return x
```

Forward pass is of course working fine; but passing in any tensors with `requires_grad=True`, or with the forward mode AD context will just fail. I tried a couple of the most recent torch releases, and the latest nightly.

But if you happen to be interested more in inference than training at the moment, I can still push that out today.
> only one choice: unpacking the nested tensors in the python interface and calling the C++ API on the buffers individually

this is _probably_ fine, because for image training the microbatch sizes are pretty small anyway. perhaps there's a way to dispatch all the work asynchronously, then wait for all results to come back? so progress is made on all batch items simultaneously?

> PyTorch still doesn't support custom autograd functions taking in nested tensors
> Forward pass is of course working fine

hmmm darn, thanks for implementing. maybe nested mode forward can still be useful for things like inference services, where multiple users request images of different sizes simultaneously.

I wonder if the pytorch team know that nested tensors for custom autograd functions is needed. not just for simplifying multi-aspect image training, but also for simplifying the training of Mistral-style language models. it's also interesting to see that utilities that could help for invoking C++ on nested tensors, are not exposed. maybe we're too early. but if I get a chance to talk to any pytorch people on Discord, I'll let them know there's a use-case for this.
>this is probably fine, because for image training the microbatch sizes are pretty small anyway. perhaps there's a way to dispatch all the work asynchronously, then wait for all results to come back? so progress is made on all batch items simultaneously?

It's all perfectly possible; the only issue is that if we want to use autograd, then our only real choice is to stick to the torch API. 

But yes, the change that I made to the NATTEN C++ API is that instead of the C++ API handling the creation of output tensors, we're now handling that in Python. This is generally the preferred solution, and I recall a very hard-to-reproduce bug where the C++ API for new tensors started a separate cuda stream on device index 0, regardless of the device all the input tensors were on.

So this means we can easily preallocate all the memory that's necessary for our output tensors, and then just call the NATTEN C++ handles, and the kernel calls are asynchronous, so the only thing running synchronously would be the checks.

With this change, I got one op (1D - QK) to support nested tensors, and forward pass appeared to work, but when I toggled `requires_grad`, it hits the following error, and I'm assuming somewhere in the Autograd C++ API, they're calling an API that's only implemented for non-nested tensors:
```
RuntimeError: Internal error: NestedTensorImpl doesn't support sizes. Please file an issue.
```

Forward mode AD isn't nearly as nice, it just spits out tens of lines of error traces in the C++ library that are difficult to even comprehend unless one is very familiar with the torch C++ APIs.

>I wonder if the pytorch team know that nested tensors for custom autograd functions is needed. not just for simplifying multi-aspect image training, but also for simplifying the training of Mistral-style language models. it's also interesting to see that utilities that could help for invoking C++ on nested tensors, are not exposed. maybe we're too early. but if I get a chance to talk to any pytorch people on Discord, I'll let them know there's a use-case for this.

I think I'm also at a point that I'm convinced it's not something I'm doing wrong, and I haven't seen a single pytorch issue, open or closed, refer to this exact scenario, so I'll probably open an issue there and see how that goes.

But anyway, I'll try and push support for forward pass only today.
But strangely, I'm training model with `NeighborhoodAttention`, it's work. When I testing my model, it show `out of memory`, the training is running at this time.
Hi,

That should not happen. Are you sure your SMI is configured correctly? SMI GPU IDs don't necessarily correspond to the cuda id. Can you confirm that, and make sure it's not a change you made?

I just tried it on a 2080 Ti with only 11 gigs of memory and it works just fine:
```
import torch
from natten import NeighborhoodAttention2D

device = torch.device("cuda:5")
na2d = NeighborhoodAttention2D(dim=128, kernel_size=7, dilation=2, num_heads=4).to(device)
x = torch.rand((1, 64, 64, 128)).to(device)
y = na2d(x)
```
After `x` is allocated:
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.63.01    Driver Version: 470.63.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   5  NVIDIA GeForce ...  On   | 00000000:89:00.0 Off |                  N/A |
| 18%   31C    P2    49W / 250W |    662MiB / 11019MiB |      1%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    5   N/A  N/A    327991      C   python3                           659MiB |
+-----------------------------------------------------------------------------+
```
After `y = na2d(x)`:
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.63.01    Driver Version: 470.63.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   5  NVIDIA GeForce ...  On   | 00000000:89:00.0 Off |                  N/A |
| 18%   32C    P2    49W / 250W |    682MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    5   N/A  N/A    327991      C   python3                           679MiB |
+-----------------------------------------------------------------------------+
```
By the way, `to` is not a reliable way of limiting all of torch's CUDA ops to a single GPU.
`with torch.device(device):` and `torch.cuda.set_device(device)` are better choices.
`CUDA_VISIBLE_DEVICES` can be used as a stronger alternative.
GPU 2,3,4,6,7 is my training process with DDP, and model contain `NeighborhoodAttention`, it's works just fine.
When I ran the snippet you sent I noticed torch reserves memory on GPU 0 as well. Not sure why, but not too surprising either. Can you try adding `torch.cuda.set_device(device)` or doing `CUDA_VISIBLE_DEVICES=5 python3` and trying it?
thanks your reply, now I running this code with `python demo.py`

```python
import torch
from natten import NeighborhoodAttention2D

device = torch.device("cuda:5")
torch.cuda.set_device(device)

na2d = NeighborhoodAttention2D(dim=128, kernel_size=7, dilation=2, num_heads=4).to(device)
fc = torch.nn.Linear(128, 1024*1024).to(device)

x = torch.rand((1, 64, 64, 128)).to(device)
y = fc(x)
print('fc ok')

y = na2d(x)
print('na ok')
```

result is:
```
fc ok
na ok
```

when I comment `torch.cuda.set_device(device)`, only `fc` ok

but why?
So what's (most likely) happening is that when you don't set the torch cuda device, the API call to the cuda backend reserves memory on cuda:0. I'm not sure why that's happening, but after thoroughly testing it, that appears to be the issue.
The only tensors that our kernels reserve are on the same device (I specifically checked that it is true just now).
And I tested the call to make sure it's not something like the autograd function.

The only temporary fix I can think of is adding the device wrapper to the functions (in natten.functional; just thinking out loud) -- which is not a solution I like, because it only fixes the issue, and doesn't explain it.

[This](https://discuss.pytorch.org/t/memory-leak-due-to-a-tensor-created-with-thcudatensor-newwithsize2d-in-an-extension/22430) seems related.
Thank very much for your reply. In any case, this is a very very great work. It prefectly introduces inductive bias of convolution to self-attention.
Thank you very much for your comment; glad you find it useful.

As far as this issue goes, I'll mark this as a bug so we resolve it in a future release.
It is probably because we're creating the output tensors within our CPP interface, and that for some reason reserves memory on the default CUDA device (cuda:0).
We'll test taking that part out and if it doesn't result in a performance hit, we'll merge, and if it does, we'll figure something else out.

Fix patch:
diff --git a/csrc/include/natten/pytorch/cpu/na1d.h b/csrc/include/natten/pytorch/cpu/na1d.h
index 0d062ec..c0b8392 100644
--- a/csrc/include/natten/pytorch/cpu/na1d.h
+++ b/csrc/include/natten/pytorch/cpu/na1d.h
@@ -49,7 +49,7 @@ void na1d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int length,
diff --git a/csrc/include/natten/pytorch/cpu/na2d.h b/csrc/include/natten/pytorch/cpu/na2d.h
index 381083a..83ca7d0 100644
--- a/csrc/include/natten/pytorch/cpu/na2d.h
+++ b/csrc/include/natten/pytorch/cpu/na2d.h
@@ -50,7 +50,7 @@ void na2d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int height,
diff --git a/csrc/include/natten/pytorch/cpu/na3d.h b/csrc/include/natten/pytorch/cpu/na3d.h
index b2606c3..0ce8d95 100644
--- a/csrc/include/natten/pytorch/cpu/na3d.h
+++ b/csrc/include/natten/pytorch/cpu/na3d.h
@@ -53,7 +53,7 @@ void na3d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int depth,
diff --git a/csrc/include/natten/pytorch/cuda/na1d.cuh b/csrc/include/natten/pytorch/cuda/na1d.cuh
index 34e3ac7..510f757 100644
--- a/csrc/include/natten/pytorch/cuda/na1d.cuh
+++ b/csrc/include/natten/pytorch/cuda/na1d.cuh
@@ -49,7 +49,7 @@ void na1d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int length,
diff --git a/csrc/include/natten/pytorch/cuda/na2d.cuh b/csrc/include/natten/pytorch/cuda/na2d.cuh
index 8129d65..ae3f966 100644
--- a/csrc/include/natten/pytorch/cuda/na2d.cuh
+++ b/csrc/include/natten/pytorch/cuda/na2d.cuh
@@ -50,7 +50,7 @@ void na2d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int height,
diff --git a/csrc/include/natten/pytorch/cuda/na3d.cuh b/csrc/include/natten/pytorch/cuda/na3d.cuh
index 9e97ce7..1685fac 100644
--- a/csrc/include/natten/pytorch/cuda/na3d.cuh
+++ b/csrc/include/natten/pytorch/cuda/na3d.cuh
@@ -53,7 +53,7 @@ void na3d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int depth,
diff --git a/csrc/include/natten/pytorch/helpers.h b/csrc/include/natten/pytorch/helpers.h
index 7531f1e..e6e9340 100644
--- a/csrc/include/natten/pytorch/helpers.h
+++ b/csrc/include/natten/pytorch/helpers.h
@@ -62,3 +62,101 @@
     }                                                                                              \
   }()
 #endif
+
+namespace natten {
+namespace pytorch {
+
+inline void CheckArgs(int kernel_size, int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1, got ", kernel_size, ".");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer, got ", dilation, ".");
+}
+
+inline void CheckArgsAgainstDim(int dim, int kernel_size, int dilation) {
+    TORCH_CHECK(kernel_size * dilation <= dim, "Input axes must be less than or equal to the product of kernel size and dilation. "
+            "Got kernel size ", kernel_size, ", dilation ", dilation, ", but dimension size was ", dim, ".");
+}
+
+inline void CheckIfPropertiesMatch(const at::Tensor& a, const at::Tensor& b) {
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(b);
+    TORCH_CHECK(a.device().is_cuda() == b.device().is_cuda(), "Expected all tensors to be on the same device.");
+    TORCH_CHECK(a.scalar_type() == b.scalar_type(), "Input tensors must match in dtype!");
+}
+
+inline void CheckIfPropertiesMatch(const at::Tensor& a, const at::Tensor& b, const at::Tensor& c) {
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(b);
+    CHECK_CONTIGUOUS(c);
+    TORCH_CHECK(a.device().is_cuda() == b.device().is_cuda() && b.device().is_cuda() == c.device().is_cuda(), "Expected all tensors to be on the same device.");
+    TORCH_CHECK(a.scalar_type() == b.scalar_type() && b.scalar_type() == c.scalar_type(), "Input tensors must match in dtype!");
+}
+
+template <size_t NaDim>
+void CheckIfTensorShapesMatch(const at::Tensor& a, const at::Tensor& b) {
+    static_assert(NaDim >= 1 && NaDim < 4);
+    static constexpr size_t Rank = NaDim + 3;
+    TORCH_CHECK(a.dim() == b.dim() && a.dim() == Rank, "Expected ", Rank, "-D tensors.");
+    for (size_t i=0; i < Rank; ++i) {
+        TORCH_CHECK(a.size(i) == b.size(i), "Tensor shape mismatch at dimension ", i, ": ", a.size(i), " != ", b.size(i));
+    }
+}
+
+template <size_t NaDim>
+void CheckAttnShape(const at::Tensor& input, const at::Tensor& attn, int kernel_size) {
+    static_assert(NaDim >= 1 && NaDim < 4);
+    TORCH_CHECK(attn.dim() == NaDim + 3, "Expected ", NaDim + 3, "-D tensors.");
+    for (size_t i=0; i < NaDim + 2; ++i) {
+        TORCH_CHECK(input.size(i) == attn.size(i), "Tensor shape mismatch at dimension ", i, ": ", input.size(i), " != ", input.size(i));
+    }
+    auto expected_kernel_size = std::pow(kernel_size, NaDim);
+    TORCH_CHECK(attn.size(NaDim + 2) == expected_kernel_size, "Expected attention dim was ", expected_kernel_size, ", got ", attn.size(NaDim + 2));
+}
+
+template <size_t NaDim>
+void CheckBias(const at::Tensor& input, const at::Tensor& bias, int kernel_size) {
+    static_assert(NaDim >= 1 && NaDim < 4);
+    TORCH_CHECK(input.scalar_type() == bias.scalar_type(), "Inputs and bias must match in dtype.");
+    TORCH_CHECK(bias.device().is_cuda() == input.device().is_cuda(), 
+            "Expected positional bias to be on the same device as the inputs.");
+    CHECK_CONTIGUOUS(bias);
+    TORCH_CHECK(bias.size(0) == input.size(1), "Expected bias.shape[0] == input.shape[1] == heads.");
+    for (size_t i=0; i < NaDim; ++i) {
+        auto expected_bias_dim = kernel_size * 2 - 1;
+        TORCH_CHECK(bias.size(i + 1) == expected_bias_dim, "Invalid bias shape at dim ", i + 1, "; "
+                "expected ", expected_bias_dim, ", got ", bias.size(i + 1), ".");
+    }
+}
+
+// TODO: I resent this; please do it the right way.
+template <size_t NaDim>
+void CheckAttnShape(const at::Tensor& input, const at::Tensor& attn, int kernel_size, int kernel_size_d) {
+    static_assert(NaDim == 3);
+    TORCH_CHECK(attn.dim() == NaDim + 3, "Expected ", NaDim + 3, "-D tensors.");
+    for (size_t i=0; i < NaDim + 2; ++i) {
+        TORCH_CHECK(input.size(i) == attn.size(i), "Tensor shape mismatch at dimension ", i, ": ", input.size(i), " != ", input.size(i));
+    }
+    auto expected_kernel_size = kernel_size * kernel_size * kernel_size_d;
+    TORCH_CHECK(attn.size(NaDim + 2) == expected_kernel_size, "Expected attention dim was ", expected_kernel_size, ", got ", attn.size(NaDim + 2));
+}
+
+template <size_t NaDim>
+void CheckBias(const at::Tensor& input, const at::Tensor& bias, int kernel_size, int kernel_size_d) {
+    static_assert(NaDim ==3);
+    TORCH_CHECK(input.scalar_type() == bias.scalar_type(), "Inputs and bias must match in dtype.");
+    TORCH_CHECK(bias.device().is_cuda() == input.device().is_cuda(), 
+            "Expected positional bias to be on the same device as the inputs.");
+    CHECK_CONTIGUOUS(bias);
+    TORCH_CHECK(bias.size(0) == input.size(1), "Expected bias.shape[0] == input.shape[1] == heads.");
+
+    auto expected_bias_dim_0 = kernel_size_d * 2 - 1;
+    TORCH_CHECK(bias.size(1) == expected_bias_dim_0, "Invalid bias shape at dim 1; expected ", 
+        expected_bias_dim_0, ", got ", bias.size(1), ".");
+    for (size_t i=1; i < NaDim; ++i) {
+        auto expected_bias_dim = kernel_size * 2 - 1;
+        TORCH_CHECK(bias.size(i + 1) == expected_bias_dim, "Invalid bias shape at dim ", i + 1, "; "
+                "expected ", expected_bias_dim, ", got ", bias.size(i + 1), ".");
+    }
+}
+
+} // namespace pytorch
+} // namespace natten
diff --git a/csrc/include/natten/pytorch/na1d.h b/csrc/include/natten/pytorch/na1d.h
index 579210d..1690a61 100644
--- a/csrc/include/natten/pytorch/na1d.h
+++ b/csrc/include/natten/pytorch/na1d.h
@@ -26,33 +26,38 @@
 
 #pragma once
 #include <ATen/ATen.h>
-#include <vector>
 
 namespace natten {
 namespace pytorch {
 
-at::Tensor na1d_qk_forward(
+void na1d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation);
 
-std::vector<at::Tensor> na1d_qk_backward(
+void na1d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation);
 
-at::Tensor na1d_av_forward(
+void na1d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation);
 
-std::vector<at::Tensor> na1d_av_backward(
+void na1d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
diff --git a/csrc/include/natten/pytorch/na2d.h b/csrc/include/natten/pytorch/na2d.h
index 06a3732..4be94c4 100644
--- a/csrc/include/natten/pytorch/na2d.h
+++ b/csrc/include/natten/pytorch/na2d.h
@@ -26,33 +26,38 @@
 
 #pragma once
 #include <ATen/ATen.h>
-#include <vector>
 
 namespace natten {
 namespace pytorch {
 
-at::Tensor na2d_qk_forward(
+void na2d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation);
 
-std::vector<at::Tensor> na2d_qk_backward(
+void na2d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation);
 
-at::Tensor na2d_av_forward(
+void na2d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation);
 
-std::vector<at::Tensor> na2d_av_backward(
+void na2d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
diff --git a/csrc/include/natten/pytorch/na3d.h b/csrc/include/natten/pytorch/na3d.h
index 248d628..73e19a0 100644
--- a/csrc/include/natten/pytorch/na3d.h
+++ b/csrc/include/natten/pytorch/na3d.h
@@ -26,12 +26,12 @@
 
 #pragma once
 #include <ATen/ATen.h>
-#include <vector>
 
 namespace natten {
 namespace pytorch {
 
-at::Tensor na3d_qk_forward(
+void na3d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
@@ -40,17 +40,20 @@ at::Tensor na3d_qk_forward(
     const int depth_kernel_size,
     const int depth_dilation);
 
-std::vector<at::Tensor> na3d_qk_backward(
+void na3d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation,
     const int depth_kernel_size,
     const int depth_dilation);
 
-at::Tensor na3d_av_forward(
+void na3d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
@@ -58,7 +61,9 @@ at::Tensor na3d_av_forward(
     const int depth_kernel_size,
     const int depth_dilation);
 
-std::vector<at::Tensor> na3d_av_backward(
+void na3d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
@@ -69,4 +74,3 @@ std::vector<at::Tensor> na3d_av_backward(
 
 } // namespace pytorch
 } // namespace natten
-
diff --git a/csrc/src/pytorch/cpu/na1d.cpp b/csrc/src/pytorch/cpu/na1d.cpp
index c1aeea1..3f9f23b 100644
--- a/csrc/src/pytorch/cpu/na1d.cpp
+++ b/csrc/src/pytorch/cpu/na1d.cpp
@@ -61,7 +61,7 @@ void na1d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int length,
@@ -74,7 +74,7 @@ void na1d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, length, dim,
             kernel_size, dilation);
 }
diff --git a/csrc/src/pytorch/cpu/na2d.cpp b/csrc/src/pytorch/cpu/na2d.cpp
index f306a67..d7cdc36 100644
--- a/csrc/src/pytorch/cpu/na2d.cpp
+++ b/csrc/src/pytorch/cpu/na2d.cpp
@@ -62,7 +62,7 @@ void na2d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int height,
@@ -76,7 +76,7 @@ void na2d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
 }
diff --git a/csrc/src/pytorch/cpu/na3d.cpp b/csrc/src/pytorch/cpu/na3d.cpp
index ccb7656..25ad84c 100644
--- a/csrc/src/pytorch/cpu/na3d.cpp
+++ b/csrc/src/pytorch/cpu/na3d.cpp
@@ -65,7 +65,7 @@ void na3d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int depth,
@@ -82,7 +82,7 @@ void na3d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, depth, height, width, dim,
             kernel_size, dilation, depth_kernel_size, depth_dilation);
 }
diff --git a/csrc/src/pytorch/cuda/na1d.cu b/csrc/src/pytorch/cuda/na1d.cu
index b4016db..d3e7316 100644
--- a/csrc/src/pytorch/cuda/na1d.cu
+++ b/csrc/src/pytorch/cuda/na1d.cu
@@ -62,7 +62,7 @@ void na1d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int length,
@@ -75,7 +75,7 @@ void na1d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, length, dim,
             kernel_size, dilation);
 }
diff --git a/csrc/src/pytorch/cuda/na2d.cu b/csrc/src/pytorch/cuda/na2d.cu
index 711e7f3..9f68d51 100644
--- a/csrc/src/pytorch/cuda/na2d.cu
+++ b/csrc/src/pytorch/cuda/na2d.cu
@@ -63,7 +63,7 @@ void na2d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int height,
@@ -77,7 +77,7 @@ void na2d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
 }
diff --git a/csrc/src/pytorch/cuda/na3d.cu b/csrc/src/pytorch/cuda/na3d.cu
index 43fd7f0..ab7d86f 100644
--- a/csrc/src/pytorch/cuda/na3d.cu
+++ b/csrc/src/pytorch/cuda/na3d.cu
@@ -66,7 +66,7 @@ void na3d_qk_backward(
     const at::Tensor &key,
     at::Tensor &d_query,
     at::Tensor &d_key,
-    at::Tensor &d_bias,
+    at::optional<at::Tensor> &d_bias,
     const int batch_size,
     const int heads,
     const int depth,
@@ -83,7 +83,7 @@ void na3d_qk_backward(
             static_cast<void *>(d_attn.data_ptr()),
             static_cast<void *>(d_query.data_ptr()),
             static_cast<void *>(d_key.data_ptr()),
-            d_bias.has_storage() ? static_cast<void *>(d_bias.data_ptr()) : nullptr,
+            d_bias.has_value() ? static_cast<void *>(d_bias.value().data_ptr()) : nullptr,
             batch_size, heads, depth, height, width, dim,
             kernel_size, dilation, depth_kernel_size, depth_dilation);
 }
diff --git a/csrc/src/pytorch/na1d.cpp b/csrc/src/pytorch/na1d.cpp
index 6a8cbcc..bb8acd4 100644
--- a/csrc/src/pytorch/na1d.cpp
+++ b/csrc/src/pytorch/na1d.cpp
@@ -26,7 +26,6 @@
 
 #include <torch/extension.h>
 #include <ATen/ATen.h>
-#include <vector>
 #include "natten/pytorch/cpu/na1d.h"
 #ifdef NATTEN_WITH_CUDA
 #include "natten/pytorch/cuda/na1d.cuh"
@@ -37,38 +36,25 @@
 namespace natten {
 namespace pytorch {
 
-at::Tensor na1d_qk_forward(
+void na1d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == 4, "Expected query and key to be two 4-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda(), "Expected both query and key to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(query, key, attn);
+    CheckIfTensorShapesMatch<1>(query, key);
+    CheckAttnShape<1>(query, attn, kernel_size);
     if (bias.has_value()) {
-        TORCH_CHECK(query.scalar_type() == bias.value().scalar_type(), "Query, key, and bias tensors must match in dtype.");
-        TORCH_CHECK(bias.value().device().is_cuda() == key.device().is_cuda(), 
-                "Expected positional bias to be on the same device as the query and key tensors.");
-        CHECK_CONTIGUOUS(bias.value());
-        TORCH_CHECK(bias.value().size(0) == query.size(1), "Expected bias.shape[0] == query.shape[1] == heads.");
-        TORCH_CHECK(int((bias.value().size(1) + 1) / 2) == kernel_size, "Invalid bias shape.");
+        CheckBias<1>(query, bias.value(), kernel_size);
     }
     int batch_size = query.size(0);
     int heads      = query.size(1);
     int length     = query.size(2);
     int dim        = query.size(3);
-    TORCH_CHECK(kernel_size * dilation <= length, "Kernel size * dilation must be less than or equal to sequence length.");
-    auto attn = torch::empty({batch_size, heads, length, kernel_size}, query.options());
+    CheckArgsAgainstDim(length, kernel_size, dilation);
     DISPATCH_DEVICE(query.device(), na1d_qk_forward,
             query,
             key,
@@ -76,50 +62,33 @@ at::Tensor na1d_qk_forward(
             attn,
             batch_size, heads, length, dim,
             kernel_size, dilation);
-    return attn;
 }
 
 
-std::vector<at::Tensor> na1d_qk_backward(
+void na1d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    CHECK_CONTIGUOUS(d_attn);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == d_attn.dim() && query.dim() == 4, 
-            "Expected query, key, and d_attn to be 4-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(
-            query.size(0) == d_attn.size(0) && 
-            query.size(1) == d_attn.size(1) && 
-            query.size(2) == d_attn.size(2) && 
-            kernel_size == d_attn.size(3), "Wrong shape for d_attn.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda() && query.device().is_cuda() == d_attn.device().is_cuda(), 
-            "Expected query, key, and d_attn to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(query, key);
+    CheckIfPropertiesMatch(d_query, d_key, d_attn);
+    CheckIfTensorShapesMatch<1>(query, key);
+    CheckIfTensorShapesMatch<1>(d_query, d_key);
+    CheckIfTensorShapesMatch<1>(query, d_key);
+    CheckAttnShape<1>(query, d_attn, kernel_size);
+    if (d_bias.has_value()) {
+        CheckBias<1>(query, d_bias.value(), kernel_size);
+    }
     int batch_size = query.size(0);
     int heads      = query.size(1);
     int length     = query.size(2);
     int dim        = query.size(3);
-    TORCH_CHECK(kernel_size * dilation <= length, "Kernel size * dilation must be less than or equal to sequence length.");
-    auto d_query = torch::empty_like(query);
-    auto d_key   = torch::empty_like(key);
-    at::Tensor d_bias;
-    if (has_bias) {
-        auto rpb_dtype = query.scalar_type() == torch::kFloat64 ? query.scalar_type() : torch::kFloat32;
-        auto options = query.options().dtype(rpb_dtype);
-        d_bias = torch::zeros({heads, 2 * kernel_size - 1}, options);
-    }
+    CheckArgsAgainstDim(length, kernel_size, dilation);
     DISPATCH_DEVICE(d_attn.device(), na1d_qk_backward,
             d_attn,
             query,
@@ -129,74 +98,51 @@ std::vector<at::Tensor> na1d_qk_backward(
             d_bias,
             batch_size, heads, length, dim,
             kernel_size, dilation);
-    return {d_query, d_key, d_bias};
 }
 
-at::Tensor na1d_av_forward(
+void na1d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
-    TORCH_CHECK(attn.dim() == value.dim() && attn.dim() == 4, "Expected attention and value to be two 4-D tensors.");
-    TORCH_CHECK(
-            attn.size(0) == value.size(0) && 
-            attn.size(1) == value.size(1) && 
-            attn.size(2) == value.size(2), "Expected attention and value to match in batch size, heads, and length.");
-    TORCH_CHECK(kernel_size == attn.size(3), "Attention weights per token do not match kernel size.");
-    TORCH_CHECK(attn.device().is_cuda() == value.device().is_cuda(), "Expected both attention and value to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(out, value, attn);
+    CheckIfTensorShapesMatch<1>(out, value);
+    CheckAttnShape<1>(value, attn, kernel_size);
     int batch_size = value.size(0);
     int heads      = value.size(1);
     int length     = value.size(2);
     int dim        = value.size(3);
-    TORCH_CHECK(kernel_size * dilation <= length, "Kernel size * dilation must be less than or equal to sequence length.");
-    auto output = torch::empty_like(value);
+    CheckArgsAgainstDim(length, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na1d_av_forward,
             attn,
             value,
-            output,
+            out,
             batch_size, heads, length, dim,
             kernel_size, dilation);
-    return output;
 }
 
-std::vector<at::Tensor> na1d_av_backward(
+void na1d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    CHECK_CONTIGUOUS(d_out);
-    TORCH_CHECK(d_out.scalar_type() == value.scalar_type(), "d_out and value tensors must match in dtype.");
-    TORCH_CHECK(d_out.dim() == value.dim() && d_out.dim() == attn.dim() && d_out.dim() == 4, 
-            "Expected d_out, value, and attn to be 4-D tensors.");
-    TORCH_CHECK(
-            d_out.size(0) == value.size(0) && 
-            d_out.size(1) == value.size(1) && 
-            d_out.size(2) == value.size(2) && 
-            d_out.size(3) == value.size(3), "Expected d_out and value to be of the same shape.");
-    TORCH_CHECK(
-            d_out.size(0) == attn.size(0) && 
-            d_out.size(1) == attn.size(1) && 
-            d_out.size(2) == attn.size(2) && 
-            kernel_size == attn.size(3), "Wrong shape for attn.");
-    TORCH_CHECK(d_out.device().is_cuda() == value.device().is_cuda() && d_out.device().is_cuda() == attn.device().is_cuda(), 
-            "Expected d_out, value, and attn to be on the same device.");
-    int batch_size = d_out.size(0);
-    int heads      = d_out.size(1);
-    int length     = d_out.size(2);
-    int dim        = d_out.size(3);
-    TORCH_CHECK(kernel_size * dilation <= length, "Kernel size * dilation must be less than or equal to sequence length.");
-    auto d_attn  = torch::empty_like(attn);
-    auto d_value = torch::empty_like(value);
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(attn, value);
+    CheckIfPropertiesMatch(d_attn, d_value, d_out);
+    CheckIfTensorShapesMatch<1>(value, d_value);
+    CheckIfTensorShapesMatch<1>(attn, d_attn);
+    CheckIfTensorShapesMatch<1>(value, d_out);
+    CheckAttnShape<1>(value, attn, kernel_size);
+    int batch_size = value.size(0);
+    int heads      = value.size(1);
+    int length     = value.size(2);
+    int dim        = value.size(3);
+    CheckArgsAgainstDim(length, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na1d_av_backward,
             d_out,
             attn,
@@ -205,7 +151,6 @@ std::vector<at::Tensor> na1d_av_backward(
             d_value,
             batch_size, heads, length, dim,
             kernel_size, dilation);
-    return {d_attn, d_value};
 }
 
 } // namespace pytorch
diff --git a/csrc/src/pytorch/na2d.cpp b/csrc/src/pytorch/na2d.cpp
index 4bc832d..4e7ec5a 100644
--- a/csrc/src/pytorch/na2d.cpp
+++ b/csrc/src/pytorch/na2d.cpp
@@ -26,7 +26,6 @@
 
 #include <torch/extension.h>
 #include <ATen/ATen.h>
-#include <vector>
 #include "natten/pytorch/cpu/na2d.h"
 #ifdef NATTEN_WITH_CUDA
 #include "natten/pytorch/cuda/na2d.cuh"
@@ -37,42 +36,27 @@
 namespace natten {
 namespace pytorch {
 
-at::Tensor na2d_qk_forward(
+void na2d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == 5, "Expected query and key to be two 5-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3) && 
-            query.size(4) == key.size(4), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda(), "Expected both query and key to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(query, key, attn);
+    CheckIfTensorShapesMatch<2>(query, key);
+    CheckAttnShape<2>(query, attn, kernel_size);
     if (bias.has_value()) {
-        TORCH_CHECK(query.scalar_type() == bias.value().scalar_type(), "Query, key, and bias tensors must match in dtype.");
-        TORCH_CHECK(bias.value().device().is_cuda() == key.device().is_cuda(), 
-                "Expected positional bias to be on the same device as the query and key tensors.");
-        CHECK_CONTIGUOUS(bias.value());
-        TORCH_CHECK(bias.value().size(0) == query.size(1), "Expected bias.shape[0] == query.shape[1] == heads.");
-        TORCH_CHECK(int((bias.value().size(1) + 1) / 2) == kernel_size, "Invalid bias shape.");
-        TORCH_CHECK(int((bias.value().size(2) + 1) / 2) == kernel_size, "Invalid bias shape.");
+        CheckBias<2>(query, bias.value(), kernel_size);
     }
     int batch_size = query.size(0);
     int heads      = query.size(1);
     int height     = query.size(2);
     int width      = query.size(3);
     int dim        = query.size(4);
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto attn = torch::empty({batch_size, heads, height, width, kernel_size * kernel_size}, query.options());
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(query.device(), na2d_qk_forward,
             query,
             key,
@@ -80,54 +64,35 @@ at::Tensor na2d_qk_forward(
             attn,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
-    return attn;
 }
 
 
-std::vector<at::Tensor> na2d_qk_backward(
+void na2d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    CHECK_CONTIGUOUS(d_attn);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == d_attn.dim() && query.dim() == 5, 
-            "Expected query, key, and d_attn to be 5-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3) && 
-            query.size(4) == key.size(4), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(
-            query.size(0) == d_attn.size(0) && 
-            query.size(1) == d_attn.size(1) && 
-            query.size(2) == d_attn.size(2) && 
-            query.size(3) == d_attn.size(3) && 
-            kernel_size * kernel_size == d_attn.size(4), "Wrong shape for d_attn.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda() && query.device().is_cuda() == d_attn.device().is_cuda(), 
-            "Expected query, key, and d_attn to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(query, key);
+    CheckIfPropertiesMatch(d_query, d_key, d_attn);
+    CheckIfTensorShapesMatch<2>(query, key);
+    CheckIfTensorShapesMatch<2>(d_query, d_key);
+    CheckIfTensorShapesMatch<2>(query, d_key);
+    CheckAttnShape<2>(query, d_attn, kernel_size);
+    if (d_bias.has_value()) {
+        CheckBias<2>(query, d_bias.value(), kernel_size);
+    }
     int batch_size = query.size(0);
     int heads      = query.size(1);
     int height     = query.size(2);
     int width      = query.size(3);
     int dim        = query.size(4);
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto d_query = torch::empty_like(query);
-    auto d_key   = torch::empty_like(key);
-    at::Tensor d_bias;
-    if (has_bias) {
-        auto rpb_dtype = query.scalar_type() == torch::kFloat64 ? query.scalar_type() : torch::kFloat32;
-        auto options = query.options().dtype(rpb_dtype);
-        d_bias = torch::zeros({heads, 2 * kernel_size - 1, 2 * kernel_size - 1}, options);
-    }
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(d_attn.device(), na2d_qk_backward,
             d_attn,
             query,
@@ -137,81 +102,55 @@ std::vector<at::Tensor> na2d_qk_backward(
             d_bias,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
-    return {d_query, d_key, d_bias};
 }
 
-at::Tensor na2d_av_forward(
+void na2d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
-    TORCH_CHECK(attn.dim() == value.dim() && attn.dim() == 5, "Expected attention and value to be two 5-D tensors.");
-    TORCH_CHECK(
-            attn.size(0) == value.size(0) && 
-            attn.size(1) == value.size(1) && 
-            attn.size(2) == value.size(2) && 
-            attn.size(3) == value.size(3), "Expected attention and value to match in batch size, heads, height, and width.");
-    TORCH_CHECK(kernel_size * kernel_size == attn.size(4), "Attention weights per token do not match kernel size.");
-    TORCH_CHECK(attn.device().is_cuda() == value.device().is_cuda(), "Expected both attention and value to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(out, value, attn);
+    CheckIfTensorShapesMatch<2>(out, value);
+    CheckAttnShape<2>(value, attn, kernel_size);
     int batch_size = value.size(0);
     int heads      = value.size(1);
     int height     = value.size(2);
     int width      = value.size(3);
     int dim        = value.size(4);
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto output = torch::empty_like(value);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na2d_av_forward,
             attn,
             value,
-            output,
+            out,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
-    return output;
 }
 
-std::vector<at::Tensor> na2d_av_backward(
+void na2d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    CHECK_CONTIGUOUS(d_out);
-    TORCH_CHECK(d_out.scalar_type() == value.scalar_type(), "d_out and value tensors must match in dtype.");
-    TORCH_CHECK(d_out.dim() == value.dim() && d_out.dim() == attn.dim() && d_out.dim() == 5, 
-            "Expected d_out, value, and attn to be 5-D tensors.");
-    TORCH_CHECK(
-            d_out.size(0) == value.size(0) && 
-            d_out.size(1) == value.size(1) && 
-            d_out.size(2) == value.size(2) && 
-            d_out.size(3) == value.size(3) && 
-            d_out.size(4) == value.size(4), "Expected d_out and value to be of the same shape.");
-    TORCH_CHECK(
-            d_out.size(0) == attn.size(0) && 
-            d_out.size(1) == attn.size(1) && 
-            d_out.size(2) == attn.size(2) && 
-            d_out.size(3) == attn.size(3) && 
-            kernel_size * kernel_size == attn.size(4), "Wrong shape for attn.");
-    TORCH_CHECK(d_out.device().is_cuda() == value.device().is_cuda() && d_out.device().is_cuda() == attn.device().is_cuda(), 
-            "Expected d_out, value, and attn to be on the same device.");
+    CheckArgs(kernel_size, dilation);
+    CheckIfPropertiesMatch(attn, value);
+    CheckIfPropertiesMatch(d_attn, d_value, d_out);
+    CheckIfTensorShapesMatch<2>(value, d_value);
+    CheckIfTensorShapesMatch<2>(attn, d_attn);
+    CheckIfTensorShapesMatch<2>(value, d_out);
+    CheckAttnShape<2>(value, attn, kernel_size);
     int batch_size = d_out.size(0);
     int heads      = d_out.size(1);
     int height     = d_out.size(2);
     int width      = d_out.size(3);
     int dim        = d_out.size(4);
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto d_attn  = torch::empty_like(attn);
-    auto d_value = torch::empty_like(value);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na2d_av_backward,
             d_out,
             attn,
@@ -220,7 +159,6 @@ std::vector<at::Tensor> na2d_av_backward(
             d_value,
             batch_size, heads, height, width, dim,
             kernel_size, dilation);
-    return {d_attn, d_value};
 }
 
 } // namespace pytorch
diff --git a/csrc/src/pytorch/na3d.cpp b/csrc/src/pytorch/na3d.cpp
index 996219e..01fce18 100644
--- a/csrc/src/pytorch/na3d.cpp
+++ b/csrc/src/pytorch/na3d.cpp
@@ -26,7 +26,6 @@
 
 #include <torch/extension.h>
 #include <ATen/ATen.h>
-#include <vector>
 #include "natten/pytorch/cpu/na3d.h"
 #ifdef NATTEN_WITH_CUDA
 #include "natten/pytorch/cuda/na3d.cuh"
@@ -37,38 +36,22 @@
 namespace natten {
 namespace pytorch {
 
-at::Tensor na3d_qk_forward(
+void na3d_qk_forward(
+    at::Tensor &attn,
     const at::Tensor &query,
     const at::Tensor &key,
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation,
-    const int depth_kernel_size,
-    const int depth_dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == 6, "Expected query and key to be two 6-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3) && 
-            query.size(4) == key.size(4) && 
-            query.size(5) == key.size(5), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda(), "Expected both query and key to be on the same device.");
+    const int kernel_size_d,
+    const int dilation_d) {
+    CheckArgs(kernel_size, dilation);
+    CheckArgs(kernel_size_d, dilation_d);
+    CheckIfPropertiesMatch(query, key, attn);
+    CheckIfTensorShapesMatch<3>(query, key);
+    CheckAttnShape<3>(query, attn, kernel_size, kernel_size_d);
     if (bias.has_value()) {
-        TORCH_CHECK(query.scalar_type() == bias.value().scalar_type(), "Query, key, and bias tensors must match in dtype.");
-        TORCH_CHECK(bias.value().device().is_cuda() == key.device().is_cuda(), 
-                "Expected positional bias to be on the same device as the query and key tensors.");
-        CHECK_CONTIGUOUS(bias.value());
-        TORCH_CHECK(bias.value().size(0) == query.size(1), "Expected bias.shape[0] == query.shape[1] == heads.");
-        TORCH_CHECK(int((bias.value().size(1) + 1) / 2) == depth_kernel_size, "Invalid bias shape.");
-        TORCH_CHECK(int((bias.value().size(2) + 1) / 2) == kernel_size, "Invalid bias shape.");
-        TORCH_CHECK(int((bias.value().size(3) + 1) / 2) == kernel_size, "Invalid bias shape.");
+        CheckBias<3>(query, bias.value(), kernel_size, kernel_size_d);
     }
     int batch_size = query.size(0);
     int heads      = query.size(1);
@@ -76,72 +59,50 @@ at::Tensor na3d_qk_forward(
     int height     = query.size(3);
     int width      = query.size(4);
     int dim        = query.size(5);
-    TORCH_CHECK(depth_kernel_size * depth_dilation <= depth, "Kernel size * dilation must be less than or equal to depth.");
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto attn = torch::empty({batch_size, heads, depth, height, width, depth_kernel_size * kernel_size * kernel_size}, query.options());
+    CheckArgsAgainstDim(depth, kernel_size_d, dilation_d);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(query.device(), na3d_qk_forward,
             query,
             key,
             bias,
             attn,
             batch_size, heads, depth, height, width, dim,
-            kernel_size, dilation, depth_kernel_size, depth_dilation);
-    return attn;
+            kernel_size, dilation, kernel_size_d, dilation_d);
 }
 
 
-std::vector<at::Tensor> na3d_qk_backward(
+void na3d_qk_backward(
+    at::Tensor &d_query,
+    at::Tensor &d_key,
+    at::optional<at::Tensor> &d_bias,
     const at::Tensor &d_attn,
     const at::Tensor &query,
     const at::Tensor &key,
-    const bool has_bias,
     const int kernel_size,
     const int dilation,
-    const int depth_kernel_size,
-    const int depth_dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(query);
-    CHECK_CONTIGUOUS(key);
-    CHECK_CONTIGUOUS(d_attn);
-    TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
-    TORCH_CHECK(query.dim() == key.dim() && query.dim() == d_attn.dim() && query.dim() == 6, 
-            "Expected query, key, and d_attn to be 6-D tensors.");
-    TORCH_CHECK(
-            query.size(0) == key.size(0) && 
-            query.size(1) == key.size(1) && 
-            query.size(2) == key.size(2) && 
-            query.size(3) == key.size(3) && 
-            query.size(4) == key.size(4) && 
-            query.size(5) == key.size(5), "Expected query and key to be of the same shape.");
-    TORCH_CHECK(
-            query.size(0) == d_attn.size(0) && 
-            query.size(1) == d_attn.size(1) && 
-            query.size(2) == d_attn.size(2) && 
-            query.size(3) == d_attn.size(3) && 
-            query.size(4) == d_attn.size(4) && 
-            depth_kernel_size * kernel_size * kernel_size == d_attn.size(5), "Wrong shape for d_attn.");
-    TORCH_CHECK(query.device().is_cuda() == key.device().is_cuda() && query.device().is_cuda() == d_attn.device().is_cuda(), 
-            "Expected query, key, and d_attn to be on the same device.");
+    const int kernel_size_d,
+    const int dilation_d) {
+    CheckArgs(kernel_size, dilation);
+    CheckArgs(kernel_size_d, dilation_d);
+    CheckIfPropertiesMatch(query, key);
+    CheckIfPropertiesMatch(d_query, d_key, d_attn);
+    CheckIfTensorShapesMatch<3>(query, key);
+    CheckIfTensorShapesMatch<3>(d_query, d_key);
+    CheckIfTensorShapesMatch<3>(query, d_key);
+    CheckAttnShape<3>(query, d_attn, kernel_size, kernel_size_d);
+    if (d_bias.has_value()) {
+        CheckBias<3>(query, d_bias.value(), kernel_size, kernel_size_d);
+    }
     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);
-    TORCH_CHECK(depth_kernel_size * depth_dilation <= depth, "Kernel size * dilation must be less than or equal to depth.");
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto d_query = torch::empty_like(query);
-    auto d_key   = torch::empty_like(key);
-    at::Tensor d_bias;
-    if (has_bias) {
-        auto rpb_dtype = query.scalar_type() == torch::kFloat64 ? query.scalar_type() : torch::kFloat32;
-        auto options = query.options().dtype(rpb_dtype);
-        d_bias = torch::zeros({heads, 2 * depth_kernel_size - 1, 2 * kernel_size - 1, 2 * kernel_size - 1}, options);
-    }
+    CheckArgsAgainstDim(depth, kernel_size_d, dilation_d);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(d_attn.device(), na3d_qk_backward,
             d_attn,
             query,
@@ -150,94 +111,66 @@ std::vector<at::Tensor> na3d_qk_backward(
             d_key,
             d_bias,
             batch_size, heads, depth, height, width, dim,
-            kernel_size, dilation, depth_kernel_size, depth_dilation);
-    return {d_query, d_key, d_bias};
+            kernel_size, dilation, kernel_size_d, dilation_d);
 }
 
-at::Tensor na3d_av_forward(
+void na3d_av_forward(
+    at::Tensor &out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation,
-    const int depth_kernel_size,
-    const int depth_dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
-    TORCH_CHECK(attn.dim() == value.dim() && attn.dim() == 6, "Expected attention and value to be two 6-D tensors.");
-    TORCH_CHECK(
-            attn.size(0) == value.size(0) && 
-            attn.size(1) == value.size(1) && 
-            attn.size(2) == value.size(2) && 
-            attn.size(3) == value.size(3), "Expected attention and value to match in batch size, heads, height, and width.");
-    TORCH_CHECK(depth_kernel_size * kernel_size * kernel_size == attn.size(5), "Attention weights per token do not match kernel size.");
-    TORCH_CHECK(attn.device().is_cuda() == value.device().is_cuda(), "Expected both attention and value to be on the same device.");
+    const int kernel_size_d,
+    const int dilation_d) {
+    CheckArgs(kernel_size, dilation);
+    CheckArgs(kernel_size_d, dilation_d);
+    CheckIfPropertiesMatch(out, value, attn);
+    CheckIfTensorShapesMatch<3>(out, value);
+    CheckAttnShape<3>(value, attn, kernel_size, kernel_size_d);
     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);
-    TORCH_CHECK(depth_kernel_size * depth_dilation <= depth, "Kernel size * dilation must be less than or equal to depth.");
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto output = torch::empty_like(value);
+    CheckArgsAgainstDim(depth, kernel_size_d, dilation_d);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na3d_av_forward,
             attn,
             value,
-            output,
+            out,
             batch_size, heads, depth, height, width, dim,
-            kernel_size, dilation, depth_kernel_size, depth_dilation);
-    return output;
+            kernel_size, dilation, kernel_size_d, dilation_d);
 }
 
-std::vector<at::Tensor> na3d_av_backward(
+void na3d_av_backward(
+    at::Tensor &d_attn,
+    at::Tensor &d_value,
     const at::Tensor &d_out,
     const at::Tensor &attn,
     const at::Tensor &value,
     const int kernel_size,
     const int dilation,
-    const int depth_kernel_size,
-    const int depth_dilation) {
-    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
-    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
-    CHECK_CONTIGUOUS(attn);
-    CHECK_CONTIGUOUS(value);
-    CHECK_CONTIGUOUS(d_out);
-    TORCH_CHECK(d_out.scalar_type() == value.scalar_type(), "d_out and value tensors must match in dtype.");
-    TORCH_CHECK(d_out.dim() == value.dim() && d_out.dim() == attn.dim() && d_out.dim() == 6, 
-            "Expected d_out, value, and attn to be 6-D tensors.");
-    TORCH_CHECK(
-            d_out.size(0) == value.size(0) && 
-            d_out.size(1) == value.size(1) && 
-            d_out.size(2) == value.size(2) && 
-            d_out.size(3) == value.size(3) && 
-            d_out.size(4) == value.size(4) && 
-            d_out.size(5) == value.size(5), "Expected d_out and value to be of the same shape.");
-    TORCH_CHECK(
-            d_out.size(0) == attn.size(0) && 
-            d_out.size(1) == attn.size(1) && 
-            d_out.size(2) == attn.size(2) && 
-            d_out.size(3) == attn.size(3) && 
-            d_out.size(4) == attn.size(4) && 
-            depth_kernel_size * kernel_size * kernel_size == attn.size(5), "Wrong shape for attn.");
-    TORCH_CHECK(d_out.device().is_cuda() == value.device().is_cuda() && d_out.device().is_cuda() == attn.device().is_cuda(), 
-            "Expected d_out, value, and attn to be on the same device.");
+    const int kernel_size_d,
+    const int dilation_d) {
+    CheckArgs(kernel_size, dilation);
+    CheckArgs(kernel_size_d, dilation_d);
+    CheckIfPropertiesMatch(attn, value);
+    CheckIfPropertiesMatch(d_attn, d_value, d_out);
+    CheckIfTensorShapesMatch<3>(value, d_value);
+    CheckIfTensorShapesMatch<3>(attn, d_attn);
+    CheckIfTensorShapesMatch<3>(value, d_out);
+    CheckAttnShape<3>(value, attn, kernel_size, kernel_size_d);
     int batch_size = d_out.size(0);
     int heads      = d_out.size(1);
     int depth      = d_out.size(2);
     int height     = d_out.size(3);
     int width      = d_out.size(4);
     int dim        = d_out.size(5);
-    TORCH_CHECK(depth_kernel_size * depth_dilation <= depth, "Kernel size * dilation must be less than or equal to depth.");
-    TORCH_CHECK(kernel_size * dilation <= height, "Kernel size * dilation must be less than or equal to height.");
-    TORCH_CHECK(kernel_size * dilation <= width, "Kernel size * dilation must be less than or equal to width.");
-    auto d_attn  = torch::empty_like(attn);
-    auto d_value = torch::empty_like(value);
+    CheckArgsAgainstDim(depth, kernel_size_d, dilation_d);
+    CheckArgsAgainstDim(height, kernel_size, dilation);
+    CheckArgsAgainstDim(width, kernel_size, dilation);
     DISPATCH_DEVICE(attn.device(), na3d_av_backward,
             d_out,
             attn,
@@ -245,8 +178,7 @@ std::vector<at::Tensor> na3d_av_backward(
             d_attn,
             d_value,
             batch_size, heads, depth, height, width, dim,
-            kernel_size, dilation, depth_kernel_size, depth_dilation);
-    return {d_attn, d_value};
+            kernel_size, dilation, kernel_size_d, dilation_d);
 }
 
 } // namespace pytorch
diff --git a/src/natten/functional.py b/src/natten/functional.py
index 5339603..4286875 100644
--- a/src/natten/functional.py
+++ b/src/natten/functional.py
@@ -21,6 +21,7 @@
 #
 #################################################################################################
 import torch
+from torch import Tensor
 from torch.autograd import Function
 from torch.cuda.amp import custom_bwd, custom_fwd
 
@@ -35,6 +36,10 @@
         + f"shi-labs.com/natten"
     )
 
+from .nested import (na1d_av_nested, na1d_qk_nested, na2d_av_nested,
+                     na2d_qk_nested, na3d_av_nested, na3d_qk_nested)
+from .utils import make_attn_tensor_from_input
+
 
 def has_cuda():
     return _C.has_cuda()
@@ -82,11 +87,13 @@ class NeighborhoodAttention1DQKAutogradFunction(Function):
     def forward(ctx, query, key, rpb, kernel_size, dilation):
         query = query.contiguous()
         key = key.contiguous()
-        attn = _C.na1d_qk_forward(query, key, rpb, kernel_size, dilation)
-        ctx.save_for_backward(query, key)
+        if rpb is not None:
+            rpb = rpb.to(key.dtype)
+        attn = make_attn_tensor_from_input(query, kernel_size)
+        _C.na1d_qk_forward(attn, query, key, rpb, kernel_size, dilation)
+        ctx.save_for_backward(query, key, rpb)
         ctx.kernel_size = kernel_size
         ctx.dilation = dilation
-        ctx.bias = rpb is not None
         return attn
 
     @staticmethod
@@ -101,25 +108,33 @@ def jvp(ctx, query_t, key_t, rpb, kernel_size, dilation):
                 "Positional biases are currently not supported "
                 "in forward mode autodiff."
             )
-        query_p, key_p = ctx.to_save
+        query_p, key_p, _ = ctx.to_save
         query_t = query_t.contiguous()
         key_t = key_t.contiguous()
-        return _C.na1d_qk_forward(
-            query_t, key_p, None, ctx.kernel_size, ctx.dilation
-        ) + _C.na1d_qk_forward(query_p, key_t, None, ctx.kernel_size, ctx.dilation)
+        attn_0 = make_attn_tensor_from_input(query_t, ctx.kernel_size)
+        attn_1 = torch.empty_like(attn_0)
+        _C.na1d_qk_forward(attn_0, query_t, key_p, None, ctx.kernel_size, ctx.dilation)
+        _C.na1d_qk_forward(attn_1, query_p, key_t, None, ctx.kernel_size, ctx.dilation)
+        return attn_0 + attn_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na1d_qk_backward(
+        query, key, rpb = ctx.saved_tensors
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        # dRPB has to be zero filled
+        d_rpb = None if rpb is None else torch.zeros_like(rpb)
+        _C.na1d_qk_backward(
+            d_query,
+            d_key,
+            d_rpb,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
-            ctx.bias,
+            query,
+            key,
             ctx.kernel_size,
             ctx.dilation,
         )
-        d_query, d_key, d_rpb = outputs
         return d_query, d_key, d_rpb, None, None
 
 
@@ -129,7 +144,8 @@ class NeighborhoodAttention1DAVAutogradFunction(Function):
     def forward(ctx, attn, value, kernel_size, dilation):
         attn = attn.contiguous()
         value = value.contiguous()
-        out = _C.na1d_av_forward(attn, value, kernel_size, dilation)
+        out = torch.empty_like(value)
+        _C.na1d_av_forward(out, attn, value, kernel_size, dilation)
         ctx.save_for_backward(attn, value)
         ctx.kernel_size = kernel_size
         ctx.dilation = dilation
@@ -145,21 +161,27 @@ def jvp(ctx, attn_t, value_t, kernel_size, dilation):
         attn_p, value_p = ctx.to_save
         attn_t = attn_t.contiguous()
         value_t = value_t.contiguous()
-        return _C.na1d_av_forward(
-            attn_t, value_p, ctx.kernel_size, ctx.dilation
-        ) + _C.na1d_av_forward(attn_p, value_t, ctx.kernel_size, ctx.dilation)
+        out_0 = torch.empty_like(value_p)
+        out_1 = torch.empty_like(out_0)
+        _C.na1d_av_forward(out_0, attn_t, value_p, ctx.kernel_size, ctx.dilation)
+        _C.na1d_av_forward(out_1, attn_p, value_t, ctx.kernel_size, ctx.dilation)
+        return out_0 + out_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na1d_av_backward(
+        attn, value = ctx.saved_tensors
+        d_attn = torch.empty_like(attn)
+        d_value = torch.empty_like(value)
+        _C.na1d_av_backward(
+            d_attn,
+            d_value,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
+            attn,
+            value,
             ctx.kernel_size,
             ctx.dilation,
         )
-        d_attn, d_value = outputs
         return d_attn, d_value, None, None
 
 
@@ -171,11 +193,11 @@ def forward(ctx, query, key, rpb, kernel_size, dilation):
         key = key.contiguous()
         if rpb is not None:
             rpb = rpb.to(key.dtype)
-        attn = _C.na2d_qk_forward(query, key, rpb, kernel_size, dilation)
-        ctx.save_for_backward(query, key)
+        attn = make_attn_tensor_from_input(query, kernel_size**2)
+        _C.na2d_qk_forward(attn, query, key, rpb, kernel_size, dilation)
+        ctx.save_for_backward(query, key, rpb)
         ctx.kernel_size = kernel_size
         ctx.dilation = dilation
-        ctx.bias = rpb is not None
         return attn
 
     @staticmethod
@@ -190,25 +212,33 @@ def jvp(ctx, query_t, key_t, rpb, kernel_size, dilation):
                 "Positional biases are currently not supported "
                 "in forward mode autodiff."
             )
-        query_p, key_p = ctx.to_save
+        query_p, key_p, _ = ctx.to_save
         query_t = query_t.contiguous()
         key_t = key_t.contiguous()
-        return _C.na2d_qk_forward(
-            query_t, key_p, None, ctx.kernel_size, ctx.dilation
-        ) + _C.na2d_qk_forward(query_p, key_t, None, ctx.kernel_size, ctx.dilation)
+        attn_0 = make_attn_tensor_from_input(query_t, ctx.kernel_size**2)
+        attn_1 = torch.empty_like(attn_0)
+        _C.na2d_qk_forward(attn_0, query_t, key_p, None, ctx.kernel_size, ctx.dilation)
+        _C.na2d_qk_forward(attn_1, query_p, key_t, None, ctx.kernel_size, ctx.dilation)
+        return attn_0 + attn_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na2d_qk_backward(
+        query, key, rpb = ctx.saved_tensors
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        # dRPB has to be zero filled
+        d_rpb = None if rpb is None else torch.zeros_like(rpb)
+        _C.na2d_qk_backward(
+            d_query,
+            d_key,
+            d_rpb,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
-            ctx.bias,
+            query,
+            key,
             ctx.kernel_size,
             ctx.dilation,
         )
-        d_query, d_key, d_rpb = outputs
         return d_query, d_key, d_rpb, None, None
 
 
@@ -216,9 +246,10 @@ class NeighborhoodAttention2DAVAutogradFunction(Function):
     @staticmethod
     @custom_fwd
     def forward(ctx, attn, value, kernel_size, dilation):
-        attn = attn.contiguous().to(value.dtype)
+        attn = attn.contiguous()
         value = value.contiguous()
-        out = _C.na2d_av_forward(attn, value, kernel_size, dilation)
+        out = torch.empty_like(value)
+        _C.na2d_av_forward(out, attn, value, kernel_size, dilation)
         ctx.save_for_backward(attn, value)
         ctx.kernel_size = kernel_size
         ctx.dilation = dilation
@@ -234,21 +265,27 @@ def jvp(ctx, attn_t, value_t, kernel_size, dilation):
         attn_p, value_p = ctx.to_save
         attn_t = attn_t.contiguous()
         value_t = value_t.contiguous()
-        return _C.na2d_av_forward(
-            attn_t, value_p, ctx.kernel_size, ctx.dilation
-        ) + _C.na2d_av_forward(attn_p, value_t, ctx.kernel_size, ctx.dilation)
+        out_0 = torch.empty_like(value_p)
+        out_1 = torch.empty_like(out_0)
+        _C.na2d_av_forward(out_0, attn_t, value_p, ctx.kernel_size, ctx.dilation)
+        _C.na2d_av_forward(out_1, attn_p, value_t, ctx.kernel_size, ctx.dilation)
+        return out_0 + out_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na2d_av_backward(
+        attn, value = ctx.saved_tensors
+        d_attn = torch.empty_like(attn)
+        d_value = torch.empty_like(value)
+        _C.na2d_av_backward(
+            d_attn,
+            d_value,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
+            attn,
+            value,
             ctx.kernel_size,
             ctx.dilation,
         )
-        d_attn, d_value = outputs
         return d_attn, d_value, None, None
 
 
@@ -258,15 +295,19 @@ class NeighborhoodAttention3DQKAutogradFunction(Function):
     def forward(ctx, query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation):
         query = query.contiguous()
         key = key.contiguous()
-        attn = _C.na3d_qk_forward(
-            query, key, rpb, kernel_size, dilation, kernel_size_d, dilation_d
+        if rpb is not None:
+            rpb = rpb.to(key.dtype)
+        attn = make_attn_tensor_from_input(
+            query, kernel_size * kernel_size * kernel_size_d
+        )
+        _C.na3d_qk_forward(
+            attn, query, key, rpb, kernel_size, dilation, kernel_size_d, dilation_d
         )
-        ctx.save_for_backward(query, key)
+        ctx.save_for_backward(query, key, rpb)
         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
@@ -281,10 +322,15 @@ def jvp(ctx, query_t, key_t, rpb, kernel_size_d, kernel_size, dilation_d, dilati
                 "Positional biases are currently not supported "
                 "in forward mode autodiff."
             )
-        query_p, key_p = ctx.to_save
+        query_p, key_p, _ = ctx.to_save
         query_t = query_t.contiguous()
         key_t = key_t.contiguous()
-        return _C.na3d_qk_forward(
+        attn_0 = make_attn_tensor_from_input(
+            query_t, ctx.kernel_size * ctx.kernel_size * ctx.kernel_size_d
+        )
+        attn_1 = torch.empty_like(attn_0)
+        _C.na3d_qk_forward(
+            attn_0,
             query_t,
             key_p,
             None,
@@ -292,7 +338,9 @@ def jvp(ctx, query_t, key_t, rpb, kernel_size_d, kernel_size, dilation_d, dilati
             ctx.dilation,
             ctx.kernel_size_d,
             ctx.dilation_d,
-        ) + _C.na3d_qk_forward(
+        )
+        _C.na3d_qk_forward(
+            attn_1,
             query_p,
             key_t,
             None,
@@ -301,21 +349,28 @@ def jvp(ctx, query_t, key_t, rpb, kernel_size_d, kernel_size, dilation_d, dilati
             ctx.kernel_size_d,
             ctx.dilation_d,
         )
+        return attn_0 + attn_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na3d_qk_backward(
+        query, key, rpb = ctx.saved_tensors
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        # dRPB has to be zero filled
+        d_rpb = None if rpb is None else torch.zeros_like(rpb)
+        _C.na3d_qk_backward(
+            d_query,
+            d_key,
+            d_rpb,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
-            ctx.bias,
+            query,
+            key,
             ctx.kernel_size,
             ctx.dilation,
             ctx.kernel_size_d,
             ctx.dilation_d,
         )
-        d_query, d_key, d_rpb = outputs
         return d_query, d_key, d_rpb, None, None, None, None
 
 
@@ -325,8 +380,9 @@ class NeighborhoodAttention3DAVAutogradFunction(Function):
     def forward(ctx, attn, value, kernel_size_d, kernel_size, dilation_d, dilation):
         attn = attn.contiguous()
         value = value.contiguous()
-        out = _C.na3d_av_forward(
-            attn, value, kernel_size, dilation, kernel_size_d, dilation_d
+        out = torch.empty_like(value)
+        _C.na3d_av_forward(
+            out, attn, value, kernel_size, dilation, kernel_size_d, dilation_d
         )
         ctx.save_for_backward(attn, value)
         ctx.kernel_size_d = kernel_size_d
@@ -345,14 +401,19 @@ def jvp(ctx, attn_t, value_t, kernel_size_d, kernel_size, dilation_d, dilation):
         attn_p, value_p = ctx.to_save
         attn_t = attn_t.contiguous()
         value_t = value_t.contiguous()
-        return _C.na3d_av_forward(
+        out_0 = torch.empty_like(value_p)
+        out_1 = torch.empty_like(out_0)
+        _C.na3d_av_forward(
+            out_0,
             attn_t,
             value_p,
             ctx.kernel_size,
             ctx.dilation,
             ctx.kernel_size_d,
             ctx.dilation_d,
-        ) + _C.na3d_av_forward(
+        )
+        _C.na3d_av_forward(
+            out_1,
             attn_p,
             value_t,
             ctx.kernel_size,
@@ -360,72 +421,101 @@ def jvp(ctx, attn_t, value_t, kernel_size_d, kernel_size, dilation_d, dilation):
             ctx.kernel_size_d,
             ctx.dilation_d,
         )
+        return out_0 + out_1
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
-        outputs = _C.na3d_av_backward(
+        attn, value = ctx.saved_tensors
+        d_attn = torch.empty_like(attn)
+        d_value = torch.empty_like(value)
+        _C.na3d_av_backward(
+            d_attn,
+            d_value,
             grad_out.contiguous(),
-            ctx.saved_tensors[0],
-            ctx.saved_tensors[1],
+            attn,
+            value,
             ctx.kernel_size,
             ctx.dilation,
             ctx.kernel_size_d,
             ctx.dilation_d,
         )
-        d_attn, d_value = outputs
         return d_attn, d_value, None, None, None, None
 
 
 def natten1dqkrpb(query, key, rpb, kernel_size, dilation):
+    if query.is_nested or key.is_nested:
+        return na1d_qk_nested(query, key, rpb, kernel_size, dilation)
     return NeighborhoodAttention1DQKAutogradFunction.apply(
         query, key, rpb, kernel_size, dilation
     )
 
 
 def natten1dqk(query, key, kernel_size, dilation):
+    if query.is_nested or key.is_nested:
+        return na1d_qk_nested(query, key, None, kernel_size, dilation)
     return NeighborhoodAttention1DQKAutogradFunction.apply(
         query, key, None, kernel_size, dilation
     )
 
 
 def natten1dav(attn, value, kernel_size, dilation):
+    if attn.is_nested or value.is_nested:
+        return na1d_av_nested(attn, value, kernel_size, dilation)
     return NeighborhoodAttention1DAVAutogradFunction.apply(
         attn, value, kernel_size, dilation
     )
 
 
 def natten2dqkrpb(query, key, rpb, kernel_size, dilation):
+    if query.is_nested or key.is_nested:
+        return na2d_qk_nested(query, key, rpb, kernel_size, dilation)
     return NeighborhoodAttention2DQKAutogradFunction.apply(
         query, key, rpb, kernel_size, dilation
     )
 
 
 def natten2dqk(query, key, kernel_size, dilation):
+    if query.is_nested or key.is_nested:
+        return na2d_qk_nested(query, key, None, kernel_size, dilation)
     return NeighborhoodAttention2DQKAutogradFunction.apply(
         query, key, None, kernel_size, dilation
     )
 
 
 def natten2dav(attn, value, kernel_size, dilation):
+    if attn.is_nested or value.is_nested:
+        return na2d_av_nested(attn, value, kernel_size, dilation)
     return NeighborhoodAttention2DAVAutogradFunction.apply(
         attn, value, kernel_size, dilation
     )
 
 
 def natten3dqkrpb(query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation):
+    if query.is_nested or key.is_nested:
+        return na3d_qk_nested(
+            query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation
+        )
     return NeighborhoodAttention3DQKAutogradFunction.apply(
         query, key, rpb, kernel_size_d, kernel_size, dilation_d, dilation
     )
 
 
 def natten3dqk(query, key, kernel_size_d, kernel_size, dilation_d, dilation):
+    if query.is_nested or key.is_nested:
+        return na3d_qk_nested(
+            query, key, None, kernel_size_d, kernel_size, dilation_d, dilation
+        )
     return NeighborhoodAttention3DQKAutogradFunction.apply(
         query, key, None, kernel_size_d, kernel_size, dilation_d, dilation
     )
 
 
 def natten3dav(attn, value, kernel_size_d, kernel_size, dilation_d, dilation):
+    if attn.is_nested or value.is_nested:
+        return na3d_av_nested(
+            attn, value, kernel_size_d, kernel_size, dilation_d, dilation
+        )
     return NeighborhoodAttention3DAVAutogradFunction.apply(
         attn, value, kernel_size_d, kernel_size, dilation_d, dilation
     )
diff --git a/src/natten/nested.py b/src/natten/nested.py
new file mode 100644
index 0000000..619dafb
--- /dev/null
+++ b/src/natten/nested.py
@@ -0,0 +1,192 @@
+#################################################################################################
+# 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.
+#
+#################################################################################################
+from typing import Optional
+
+import torch
+from torch import Tensor
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+try:
+    from natten import _C
+except ImportError:
+    raise ImportError(
+        f"Failed to import NATTEN's CPP backend. "
+        + f"This could be due to an invalid/incomplete install. "
+        + f"Please uninstall NATTEN (pip uninstall natten) and re-install with the"
+        f" correct torch build: "
+        + f"shi-labs.com/natten"
+    )
+
+from .utils import make_attn_tensor_from_input
+
+
+def na1d_qk_nested(
+    query: Tensor, key: Tensor, rpb: Optional[Tensor], kernel_size: int, dilation: int
+) -> Tensor:
+    if not query.is_nested or not key.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not query.is_leaf or not key.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if query.requires_grad or key.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+    if rpb is not None and rpb.is_nested:
+        raise ValueError("Positional biases cannot be nested.")
+
+    if rpb is not None:
+        rpb = rpb.contiguous().to(key.dtype)
+
+    if query.size(0) != key.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    attn = torch.nested.nested_tensor(
+        [make_attn_tensor_from_input(q, kernel_size) for q in query]
+    )
+    for q, k, a in zip(query, key, attn):
+        _C.na1d_qk_forward(a, q, k, rpb, kernel_size, dilation)
+
+    return attn
+
+
+def na1d_av_nested(attn: Tensor, value: Tensor, kernel_size: int, dilation: int):
+    if not attn.is_nested or not value.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not attn.is_leaf or not value.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if attn.requires_grad or value.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+
+    if attn.size(0) != value.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    out = torch.empty_like(value)
+    for a, v, o in zip(attn, value, out):
+        _C.na1d_av_forward(o, a, v, kernel_size, dilation)
+
+    return out
+
+
+def na2d_qk_nested(
+    query: Tensor, key: Tensor, rpb: Optional[Tensor], kernel_size: int, dilation: int
+) -> Tensor:
+    if not query.is_nested or not key.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not query.is_leaf or not key.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if query.requires_grad or key.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+    if rpb is not None and rpb.is_nested:
+        raise ValueError("Positional biases cannot be nested.")
+
+    if rpb is not None:
+        rpb = rpb.contiguous().to(key.dtype)
+
+    if query.size(0) != key.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    attn = torch.nested.nested_tensor(
+        [make_attn_tensor_from_input(q, kernel_size**2) for q in query]
+    )
+    for q, k, a in zip(query, key, attn):
+        _C.na2d_qk_forward(a, q, k, rpb, kernel_size, dilation)
+
+    return attn
+
+
+def na2d_av_nested(attn: Tensor, value: Tensor, kernel_size: int, dilation: int):
+    if not attn.is_nested or not value.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not attn.is_leaf or not value.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if attn.requires_grad or value.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+
+    if attn.size(0) != value.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    out = torch.empty_like(value)
+    for a, v, o in zip(attn, value, out):
+        _C.na2d_av_forward(o, a, v, kernel_size, dilation)
+
+    return out
+
+
+def na3d_qk_nested(
+    query: Tensor,
+    key: Tensor,
+    rpb: Optional[Tensor],
+    kernel_size_d: int,
+    kernel_size: int,
+    dilation_d: int,
+    dilation: int,
+) -> Tensor:
+    if not query.is_nested or not key.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not query.is_leaf or not key.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if query.requires_grad or key.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+    if rpb is not None and rpb.is_nested:
+        raise ValueError("Positional biases cannot be nested.")
+
+    if rpb is not None:
+        rpb = rpb.contiguous().to(key.dtype)
+
+    if query.size(0) != key.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    attn = torch.nested.nested_tensor([
+        make_attn_tensor_from_input(q, kernel_size * kernel_size * kernel_size_d)
+        for q in query
+    ])
+    for q, k, a in zip(query, key, attn):
+        _C.na3d_qk_forward(
+            a, q, k, rpb, kernel_size, dilation, kernel_size_d, dilation_d
+        )
+
+    return attn
+
+
+def na3d_av_nested(
+    attn: Tensor,
+    value: Tensor,
+    kernel_size_d: int,
+    kernel_size: int,
+    dilation_d: int,
+    dilation: int,
+):
+    if not attn.is_nested or not value.is_nested:
+        raise ValueError("Expected all inputs to be nested.")
+    if not attn.is_leaf or not value.is_leaf:
+        raise ValueError("Only one level of nested tensors is supported at the moment.")
+    if attn.requires_grad or value.requires_grad:
+        raise ValueError("Autograd is not supported for nested tensors.")
+
+    if attn.size(0) != value.size(0):
+        raise ValueError("Got nested inputs, but they don't match in size.")
+
+    out = torch.empty_like(value)
+    for a, v, o in zip(attn, value, out):
+        _C.na3d_av_forward(o, a, v, kernel_size, dilation, kernel_size_d, dilation_d)
+
+    return out
diff --git a/src/natten/utils.py b/src/natten/utils.py
new file mode 100644
index 0000000..c6f2ac7
--- /dev/null
+++ b/src/natten/utils.py
@@ -0,0 +1,35 @@
+#################################################################################################
+# 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 Tensor
+
+
+def make_attn_tensor_from_input(input_tensor: Tensor, attention_dim: int) -> Tensor:
+    shape = [x for x in input_tensor.shape[:-1]] + [attention_dim]
+    return torch.empty(
+        shape,
+        device=input_tensor.device,
+        dtype=input_tensor.dtype,
+        requires_grad=input_tensor.requires_grad,
+    )


Test patch:
diff --git a/tests/test_na1d.py b/tests/test_na1d.py
index 95c94ec..1a30108 100644
--- a/tests/test_na1d.py
+++ b/tests/test_na1d.py
@@ -34,6 +34,7 @@
                     has_cuda, has_gemm, has_half)
 from natten.functional import natten1dav, natten1dqkrpb
 
+SKIP_NESTED_TESTS = [int(x) for x in torch.__version__.split(".")[:2]] < [2, 1]
 HAS_CUDA = torch.cuda.is_available() and (CUDA_HOME is not None) and has_cuda()
 HAS_GEMM = has_gemm()
 HAS_HALF = has_half()
@@ -384,6 +385,76 @@ def test_fwad_cuda_gemm(self):
         self._test_fwad(B=1, H=4, L=64, D=16, kernel_size=21, dilation=1, device="cuda")
         self._test_fwad(B=1, H=2, L=64, D=16, kernel_size=21, dilation=2, device="cuda")
 
+    def _test_nested_qk_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        query = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16),
+                torch.randn(2, 8, 16, 32),
+                torch.randn(4, 1, 32, 16),
+            ],
+            **kwargs,
+        )
+        key = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16),
+                torch.randn(2, 8, 16, 32),
+                torch.randn(4, 1, 32, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten1dqkrpb(query, key, None, kernel_size, dilation)
+        out_ref = []
+        for q, k in zip(query, key):
+            out_ref.append(natten1dqkrpb(q, k, None, kernel_size, dilation))
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def _test_nested_av_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        attn = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, kernel_size),
+                torch.randn(2, 8, 16, kernel_size),
+                torch.randn(4, 1, 32, kernel_size),
+            ],
+            **kwargs,
+        )
+        value = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16),
+                torch.randn(2, 8, 16, 32),
+                torch.randn(4, 1, 32, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten1dav(attn, value, kernel_size, dilation)
+        out_ref = []
+        for a, v in zip(attn, value):
+            out_ref.append(natten1dav(a, v, kernel_size, dilation))
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def test_nested_forward_cpu(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        self._test_nested_qk_forward(dtype=torch.float32, device="cpu")
+        self._test_nested_av_forward(dtype=torch.float32, device="cpu")
+
+    def test_nested_forward_cuda(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        if not HAS_CUDA:
+            self.skipTest("NATTEN not compiled with CUDA.")
+        self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
+        self._test_nested_av_forward(dtype=torch.float16, device="cuda")
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_na2d.py b/tests/test_na2d.py
index a321c20..15807d7 100644
--- a/tests/test_na2d.py
+++ b/tests/test_na2d.py
@@ -34,6 +34,7 @@
                     has_cuda, has_gemm, has_half)
 from natten.functional import natten2dav, natten2dqkrpb
 
+SKIP_NESTED_TESTS = [int(x) for x in torch.__version__.split(".")[:2]] < [2, 1]
 HAS_CUDA = torch.cuda.is_available() and (CUDA_HOME is not None) and has_cuda()
 HAS_GEMM = has_gemm()
 HAS_HALF = has_half()
@@ -509,6 +510,76 @@ def test_fwad_cuda_gemm(self):
             B=1, H=1, X=7, Y=6, D=8, kernel_size=3, dilation=2, device="cuda"
         )
 
+    def _test_nested_qk_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        query = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16, 16),
+                torch.randn(2, 8, 16, 18, 32),
+                torch.randn(4, 1, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        key = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16, 16),
+                torch.randn(2, 8, 16, 18, 32),
+                torch.randn(4, 1, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten2dqkrpb(query, key, None, kernel_size, dilation)
+        out_ref = []
+        for q, k in zip(query, key):
+            out_ref.append(natten2dqkrpb(q, k, None, kernel_size, dilation))
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def _test_nested_av_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        attn = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16, kernel_size**2),
+                torch.randn(2, 8, 16, 18, kernel_size**2),
+                torch.randn(4, 1, 32, 20, kernel_size**2),
+            ],
+            **kwargs,
+        )
+        value = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 14, 16, 16),
+                torch.randn(2, 8, 16, 18, 32),
+                torch.randn(4, 1, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten2dav(attn, value, kernel_size, dilation)
+        out_ref = []
+        for a, v in zip(attn, value):
+            out_ref.append(natten2dav(a, v, kernel_size, dilation))
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def test_nested_forward_cpu(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        self._test_nested_qk_forward(dtype=torch.float32, device="cpu")
+        self._test_nested_av_forward(dtype=torch.float32, device="cpu")
+
+    def test_nested_forward_cuda(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        if not HAS_CUDA:
+            self.skipTest("NATTEN not compiled with CUDA.")
+        self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
+        self._test_nested_av_forward(dtype=torch.float16, device="cuda")
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_na3d.py b/tests/test_na3d.py
index 44e5db5..8497142 100644
--- a/tests/test_na3d.py
+++ b/tests/test_na3d.py
@@ -32,6 +32,7 @@
 from natten import has_bfloat, has_cuda, has_half
 from natten.functional import natten3dav, natten3dqkrpb
 
+SKIP_NESTED_TESTS = [int(x) for x in torch.__version__.split(".")[:2]] < [2, 1]
 HAS_CUDA = torch.cuda.is_available() and (CUDA_HOME is not None) and has_cuda()
 HAS_HALF = has_half()
 HAS_BFLOAT = has_bfloat()
@@ -351,6 +352,88 @@ def test_fwad_cuda(self):
             device="cuda",
         )
 
+    def _test_nested_qk_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kernel_size_d, dilation_d = 3, 3
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        query = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 9, 14, 16, 16),
+                torch.randn(2, 8, 9, 16, 18, 32),
+                torch.randn(4, 1, 9, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        key = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 9, 14, 16, 16),
+                torch.randn(2, 8, 9, 16, 18, 32),
+                torch.randn(4, 1, 9, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten3dqkrpb(
+            query, key, None, kernel_size_d, kernel_size, dilation_d, dilation
+        )
+        out_ref = []
+        for q, k in zip(query, key):
+            out_ref.append(
+                natten3dqkrpb(
+                    q, k, None, kernel_size_d, kernel_size, dilation_d, dilation
+                )
+            )
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def _test_nested_av_forward(self, dtype, device):
+        torch.manual_seed(42)
+        kernel_size, dilation = 7, 2
+        kernel_size_d, dilation_d = 3, 3
+        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
+        attn = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 9, 14, 16, kernel_size_d * kernel_size * kernel_size),
+                torch.randn(2, 8, 9, 16, 18, kernel_size_d * kernel_size * kernel_size),
+                torch.randn(4, 1, 9, 32, 20, kernel_size_d * kernel_size * kernel_size),
+            ],
+            **kwargs,
+        )
+        value = torch.nested.nested_tensor(
+            [
+                torch.randn(1, 2, 9, 14, 16, 16),
+                torch.randn(2, 8, 9, 16, 18, 32),
+                torch.randn(4, 1, 9, 32, 20, 16),
+            ],
+            **kwargs,
+        )
+        out_nested = natten3dav(
+            attn, value, kernel_size_d, kernel_size, dilation_d, dilation
+        )
+        out_ref = []
+        for a, v in zip(attn, value):
+            out_ref.append(
+                natten3dav(a, v, kernel_size_d, kernel_size, dilation_d, dilation)
+            )
+
+        for o, o_ref in zip(out_nested, out_ref):
+            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
+
+    def test_nested_forward_cpu(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        self._test_nested_qk_forward(dtype=torch.float32, device="cpu")
+        self._test_nested_av_forward(dtype=torch.float32, device="cpu")
+
+    def test_nested_forward_cuda(self):
+        if SKIP_NESTED_TESTS:
+            self.skipTest("Nested tensors are only supported with torch >= 2.1.")
+        if not HAS_CUDA:
+            self.skipTest("NATTEN not compiled with CUDA.")
+        self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
+        self._test_nested_av_forward(dtype=torch.float16, device="cuda")
+
 
 if __name__ == "__main__":
     unittest.main()