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

Issue:
Neighborhood-Attention whether there is a requirement for the size of the tensor?
I would like to ask when I have changed the self-attention of transformer to natten and this error occurs.

  File "/root/autodl-tmp/pycharm_project_983/basicsr/archs/my_arch.py", line 131, in forward
    x = self.attn(x)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.virtualenvs/pythonProject/lib/python3.8/site-packages/natten/natten2d.py", line 95, in forward
    attn = natten2dqkrpb(q, k, self.rpb, self.kernel_size, self.dilation)
  File "/root/.virtualenvs/pythonProject/lib/python3.8/site-packages/natten/functional.py", line 186, in natten2dqkrpb
    return NATTEN2DQKRPBFunction.apply(query, key, rpb, kernel_size, dilation)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/root/miniconda3/lib/python3.8/site-packages/torch/cuda/amp/autocast_mode.py", line 106, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/root/.virtualenvs/pythonProject/lib/python3.8/site-packages/natten/functional.py", line 119, in forward
    attn = _C.natten2dqkrpb_forward(query, key, rpb, kernel_size, dilation)
RuntimeError: numel needs to be smaller than int32_t max; otherwise, please use packed_accessor64




Discussion / hints:
Hello and thank you for your interest,

The issue appears to be indexing. At the moment NATTEN does int32 indexing, and that can lead to an overflow with very large inputs.
Could you share your input size?

If that is the issue, we can add int64 indexing to NATTEN's priority list.
Thank you for the response. My configuration is as follows:

 embed_dim: 96
  kernel_size:7
  num_heads: [3, 6, 12, 24 ]
  depths: [2, 2, 6, 2 ]
  dilations: [  [1, 8],
  [1, 4],
  [1, 2, 1, 2, 1, 2],
  [1, 1],
  ]

x_size:(1,1120,1680,96)

torch==2.0.0+cu118
Thanks for sharing that. We'll try and resolve this on our end and get back to you.

Context: because your feature map is quite larger than your average ImageNet or even COCO feature map (1120 x 1680), the last Transformer encoder will have inputs of size 1 x 1120 x 1680 x 768 = 1,445,068,800 which shouldn't really cause an index overflow, but that is the error ATen (backend API) is throwing.
Could you please try this branch and see if it resolves your issue?

https://github.com/alihassanijr/NATTEN-Torch/tree/int64

It avoids the error on our end, but there's still more changes that I'd add before merging into the main branch and pushing a new release.

If you're unfamiliar with building from source, just clone my fork, switch to the branch, and build:
```bash
git clone https://github.com/alihassanijr/NATTEN-Torch

cd NATTEN-Torch

git checkout int64

pip3 install ninja # assuming torch is already installed
pip3 install -r requirements.txt # if torch is not installed -- will install 2.0 + cu118.

make
```
Thank you very much. That solves my problem
I moved this issue here since it's related to NATTEN.
The reason it hasn't been merged is that the change that was made to resolve this issue (switch from accessor32 to accessor64) introduces a very considerable latency in every kernel launch.
As a consequence, we may have to template accessors in every kernel and only conditionally call those with accessor64s.

Since this would affect the ongoing refactoring of NATTEN, I'll keep this issue open until we resolve it.

Fix patch:
diff --git a/csrc/include/natten/cuda/naive/pointwise_neighborhood_1d.cuh b/csrc/include/natten/cuda/naive/pointwise_neighborhood_1d.cuh
index eadc148..2901367 100644
--- a/csrc/include/natten/cuda/naive/pointwise_neighborhood_1d.cuh
+++ b/csrc/include/natten/cuda/naive/pointwise_neighborhood_1d.cuh
@@ -130,7 +130,7 @@ struct PointwiseNeighborhood1DBase {
     int KERNELTHREADS = min(CUDA_NUM_THREADS, kernel_size);
     int TOKENTHREADS = min(CUDA_NUM_THREADS / KERNELTHREADS, length);
     int BATCHTHREADS =
-        max(1, CUDA_NUM_THREADS / (TOKENTHREADS * KERNELTHREADS));
+        min(64, max(1, CUDA_NUM_THREADS / (TOKENTHREADS * KERNELTHREADS)));
     dim3 grid(
         (length + TOKENTHREADS - 1) / TOKENTHREADS,
         (kernel_size + KERNELTHREADS - 1) / KERNELTHREADS,
diff --git a/csrc/include/natten/cuda/naive/pointwise_neighborhood_3d.cuh b/csrc/include/natten/cuda/naive/pointwise_neighborhood_3d.cuh
index 535bfd7..bf87106 100644
--- a/csrc/include/natten/cuda/naive/pointwise_neighborhood_3d.cuh
+++ b/csrc/include/natten/cuda/naive/pointwise_neighborhood_3d.cuh
@@ -171,7 +171,7 @@ struct PointwiseNeighborhood3DBase {
         min(CUDA_NUM_THREADS, attention_span /* == kernel_size^3 */);
     int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), spatial_size);
     int BATCHTHREADS =
-        max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+        min(64, max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS)));
     dim3 grid(
         (spatial_size + PIXELTHREADS - 1) / PIXELTHREADS,
         (attention_span + KERNELTHREADS - 1) / KERNELTHREADS,
diff --git a/csrc/include/natten/cuda/naive/tiled/base.cuh b/csrc/include/natten/cuda/naive/tiled/base.cuh
index 8946be6..5ed4ad0 100644
--- a/csrc/include/natten/cuda/naive/tiled/base.cuh
+++ b/csrc/include/natten/cuda/naive/tiled/base.cuh
@@ -148,7 +148,7 @@ struct PointwiseNeighborhood2DBase {
         min(CUDA_NUM_THREADS, attention_span /* == kernel_size^2 */);
     int PIXELTHREADS = min(int(CUDA_NUM_THREADS / KERNELTHREADS), spatial_size);
     int BATCHTHREADS =
-        max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS));
+        min(64, max(1, CUDA_NUM_THREADS / (PIXELTHREADS * KERNELTHREADS)));
     dim3 grid(
         (spatial_size + PIXELTHREADS - 1) / PIXELTHREADS,
         (attention_span + KERNELTHREADS - 1) / KERNELTHREADS,


Test patch:
diff --git a/tests/test_na1d.py b/tests/test_na1d.py
index aab914a..b272e86 100644
--- a/tests/test_na1d.py
+++ b/tests/test_na1d.py
@@ -51,44 +51,47 @@
 logger = logging.getLogger(__name__)
 
 
+def init_cpu_ref(B, H, L, D, kernel_size, dilation, has_bias):
+    with torch.no_grad():
+        q, k, v = (
+            torch.randn((B, H, L, D)) * (D**-0.5),
+            torch.randn((B, H, L, D)),
+            torch.randn((B, H, L, D)),
+        )
+        rpb = None if not has_bias else torch.randn(H, 2 * kernel_size - 1)
+        q_, k_, v_ = (
+            q.clone().cuda(),
+            k.clone().cuda(),
+            v.clone().cuda(),
+        )
+        rpb_ = None if rpb is None else rpb.clone().cuda()
+
+        attn_ref = na1d_qk_with_bias(q, k, rpb, kernel_size, dilation)
+        attn_ref = attn_ref.softmax(dim=-1)
+        out_ref = na1d_av(attn_ref, v, kernel_size, dilation)
+    return (q_, k_, v_, rpb_, kernel_size, dilation), (attn_ref.cuda(), out_ref.cuda())
+
+
 class NA1DTests(unittest.TestCase):
-    def _test_against_cpu(
-        self, B, H, L, D, kernel_size, dilation, has_bias, dtype, eps
-    ):
+    def _test_against_cpu(self, inputs, reference, eps, dtype):
+        q, k, v, rpb, kernel_size, dilation = inputs
+        attn_ref, out_ref = reference
         with torch.no_grad():
-            q, k, v = (
-                torch.randn((B, H, L, D)) * (D**-0.5),
-                torch.randn((B, H, L, D)),
-                torch.randn((B, H, L, D)),
-            )
-            rpb = None
-            if has_bias:
-                rpb = torch.randn(H, 2 * kernel_size - 1)
-            q_, k_, v_ = (
-                q.clone().cuda().to(dtype),
-                k.clone().cuda().to(dtype),
-                v.clone().cuda().to(dtype),
-            )
-            rpb_ = None if rpb is None else rpb.clone().cuda().to(dtype)
-
-            attn_ref = na1d_qk_with_bias(q, k, rpb, kernel_size, dilation)
-            attn_ref = attn_ref.softmax(dim=-1)
-            out_ref = na1d_av(attn_ref, v, kernel_size, dilation)
+            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
+            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
+            assert q_.is_cuda and k_.is_cuda and v_.is_cuda
 
             attn = na1d_qk_with_bias(q_, k_, rpb_, kernel_size, dilation)
             attn = attn.softmax(dim=-1)
             out = na1d_av(attn, v_, kernel_size, dilation)
 
-            torch.testing.assert_close(attn.float().cpu(), attn_ref, atol=eps, rtol=0)
-            torch.testing.assert_close(out.float().cpu(), out_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
 
     def _test_all_dtypes_against_cpu(
         self, B, H, L, D, kernel_size, dilation, has_bias=False
     ):
-        # Test naive kernels
-        disable_gemm_na()
-        disable_tf32()
-        self._test_against_cpu(
+        inputs, reference = init_cpu_ref(
             B=B,
             H=H,
             L=L,
@@ -96,57 +99,44 @@ def _test_all_dtypes_against_cpu(
             kernel_size=kernel_size,
             dilation=dilation,
             has_bias=has_bias,
+        )
+        # Test naive kernels
+        disable_gemm_na()
+        disable_tf32()
+        self._test_against_cpu(
+            inputs=inputs,
+            reference=reference,
             dtype=torch.float32,
             eps=1e-4,
         )
         if HAS_HALF:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                L=L,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float16,
                 eps=1e-1,
             )
         if HAS_BFLOAT:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                L=L,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.bfloat16,
                 eps=1e-1,
             )
+        # Test GEMM-based kernels
         if HAS_GEMM:
-            # Test GEMM-based kernels
             enable_gemm_na()
             if HAS_FLOAT_GEMM:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    L=L,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.float32,
                     eps=1e-2,
                 )
                 enable_tf32()
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    L=L,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.float32,
                     eps=1e-2,
                 )
@@ -154,32 +144,21 @@ def _test_all_dtypes_against_cpu(
                 HAS_HALF
             ), "GEMM kernels must support FP16 across on all supported architectures."
             self._test_against_cpu(
-                B=B,
-                H=H,
-                L=L,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float16,
                 eps=1e-1,
             )
             if HAS_BFLOAT:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    L=L,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.bfloat16,
                     eps=1e-1,
                 )
 
     @skip_if_cuda_is_not_supported()
     def test_cpu_vs_cuda(self):
-        torch.manual_seed(42)
         self._test_all_dtypes_against_cpu(
             B=1, H=1, L=16, D=32, kernel_size=7, dilation=1
         )
@@ -209,7 +188,6 @@ def _test_autograd_qk(
         self, B, H, L, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, L, D), **kwargs)
         key = torch.randn((B, H, L, D), **kwargs)
@@ -231,7 +209,6 @@ def _test_autograd_av(
         self, B, H, L, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, L, kernel_size + L_extra), **kwargs)
         value = torch.randn((B, H, L, D), **kwargs)
@@ -345,7 +322,6 @@ def test_invalid_dilation(self):
         )
 
     def _test_fwad_qk(self, B, H, L, D, kernel_size, dilation, device, L_extra=0):
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, L, D), **kwargs)
         key = torch.randn((B, H, L, D), **kwargs)
@@ -363,7 +339,6 @@ def _test_fwad_qk(self, B, H, L, D, kernel_size, dilation, device, L_extra=0):
         ), "Forward mode autograd check failed for NA1D: QK."
 
     def _test_fwad_av(self, B, H, L, D, kernel_size, dilation, device, L_extra=0):
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, L, kernel_size + L_extra), **kwargs)
         value = torch.randn((B, H, L, D), **kwargs)
@@ -443,7 +418,6 @@ def test_fwad_cuda_gemm(self):
         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, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
         query = torch.nested.nested_tensor(
@@ -493,7 +467,6 @@ def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
                 torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
 
     def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
         attn = torch.nested.nested_tensor(
@@ -561,6 +534,428 @@ def test_nested_forward_cuda(self):
         self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
         self._test_nested_av_forward(dtype=torch.float16, device="cuda")
 
+    def _test_with_extra_tokens(
+        self,
+        B,
+        H,
+        L,
+        D,
+        kernel_size,
+        dilation,
+        has_bias,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+        L_extra=9,
+    ):
+        assert L_extra > 0
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, L, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, L, D), **kwargs),
+                torch.randn((B, H, L, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            rpb, rpb_ref = None, None
+            if has_bias:
+                rpb = torch.randn(H, 2 * kernel_size - 1, **kwargs)
+                rpb_ref = rpb.clone()
+
+            extra_k = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_v = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_k_ref, extra_v_ref = extra_k.clone(), extra_v.clone()
+
+            # Reference implementation
+            attn_extra_ref = (
+                (
+                    q_ref.view(B * H, L, D).contiguous()
+                    @ extra_k_ref.view(B * H, L_extra, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, L, L_extra)
+                .contiguous()
+            )
+            attn_na_ref = na1d_qk_with_bias(
+                q_ref, k_ref, rpb_ref, kernel_size, dilation
+            )
+            attn_ref = torch.cat([attn_na_ref, attn_extra_ref], dim=-1)
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            attn_na_softmax_ref, attn_extra_softmax_ref = attn_ref_softmax.split(
+                [kernel_size, L_extra], dim=-1
+            )
+            attn_na_softmax_ref, attn_extra_softmax_ref = (
+                attn_na_softmax_ref.contiguous(),
+                attn_extra_softmax_ref.contiguous(),
+            )
+            out_na_ref = na1d_av(attn_na_softmax_ref, v_ref, kernel_size, dilation)
+            out_extra_ref = (
+                (
+                    attn_extra_softmax_ref.view(B * H, L, L_extra).contiguous()
+                    @ extra_v_ref.view(B * H, L_extra, D).contiguous()
+                )
+                .view(B, H, L, D)
+                .contiguous()
+            )
+            out_ref = out_extra_ref + out_na_ref
+
+            # Op
+            attn = na1d_qk_with_bias(
+                q, k, rpb, kernel_size, dilation, additional_keys=extra_k
+            )
+            attn_na, attn_extra = attn.split([kernel_size, L_extra], dim=-1)
+            attn_na, attn_extra = attn_na.contiguous(), attn_extra.contiguous()
+            attn_softmax = attn.softmax(dim=-1)
+            attn_na_softmax, attn_extra_softmax = attn_softmax.split(
+                [kernel_size, L_extra], dim=-1
+            )
+            attn_na_softmax, attn_extra_softmax = (
+                attn_na_softmax.contiguous(),
+                attn_extra_softmax.contiguous(),
+            )
+            out = na1d_av(
+                attn_softmax, v, kernel_size, dilation, additional_values=extra_v
+            )
+
+            # Elementwise checks
+            torch.testing.assert_close(attn_na, attn_na_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_na_softmax, attn_na_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn_extra, attn_extra_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_extra_softmax, attn_extra_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn, attn_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_with_extra_tokens(
+        self, B, H, L, D, kernel_size, dilation, has_bias=False
+    ):
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            L=L,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_with_extra_tokens(
+        self, B, H, L, D, kernel_size, dilation, has_bias=False
+    ):
+        # Test naive kernels
+        disable_gemm_na()
+        disable_tf32()
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            L=L,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cuda",
+        )
+        if HAS_HALF:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float16,
+                device="cuda",
+            )
+        if HAS_BFLOAT:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+
+        # Test GEMM-based kernels
+        if HAS_GEMM:
+            enable_gemm_na()
+            if HAS_FLOAT_GEMM:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.float32,
+                    device="cuda",
+                )
+                enable_tf32()
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.float32,
+                    device="cuda",
+                )
+            assert (
+                HAS_HALF
+            ), "GEMM kernels must support FP16 across on all supported architectures."
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float16,
+                device="cuda",
+            )
+            if HAS_BFLOAT:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                )
+
+    def test_cpu_with_extra_tokens(self):
+        self._test_cpu_with_extra_tokens(
+            B=1, H=1, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=1, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=2, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=4, H=4, L=100, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=4, H=8, L=100, D=32, kernel_size=3, dilation=4
+        )
+        self._test_cpu_with_extra_tokens(
+            B=4, H=4, L=100, D=32, kernel_size=13, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=4, H=8, L=100, D=32, kernel_size=13, dilation=2
+        )
+        self._test_cpu_with_extra_tokens(
+            B=4, H=8, L=100, D=32, has_bias=True, kernel_size=13, dilation=2
+        )
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_with_extra_tokens(self):
+        self._test_cuda_with_extra_tokens(
+            B=1, H=1, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=1, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=2, L=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=4, H=4, L=100, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=4, H=8, L=100, D=32, kernel_size=3, dilation=4
+        )
+        self._test_cuda_with_extra_tokens(
+            B=4, H=4, L=100, D=32, kernel_size=13, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=4, H=8, L=100, D=32, kernel_size=13, dilation=2
+        )
+        self._test_cuda_with_extra_tokens(
+            B=4, H=8, L=100, D=32, has_bias=True, kernel_size=13, dilation=2
+        )
+
+    def _test_against_self_attention(
+        self,
+        B,
+        H,
+        L,
+        D,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+    ):
+        # Only odd-sized kernels are valid in NA,
+        # and NA == SA when kernel size equals
+        # input size. NATTEN only allows square-sized
+        # kernels, so we need to make sure that the
+        # input size is an odd value across every
+        # spatial dim.
+        assert L % 2 == 1
+        kernel_size = L
+        dilation = 1
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, L, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, L, D), **kwargs),
+                torch.randn((B, H, L, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            # Reference implementation
+            attn_ref = (
+                (
+                    q_ref.view(B * H, L, D).contiguous()
+                    @ k_ref.view(B * H, L, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, L, L)
+                .contiguous()
+            )
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            out_ref = (
+                (
+                    attn_ref_softmax.view(B * H, L, L).contiguous()
+                    @ v_ref.view(B * H, L, D).contiguous()
+                )
+                .view(B, H, L, D)
+                .contiguous()
+            )
+
+            # Op
+            attn = na1d_qk(q, k, kernel_size, dilation)
+            attn_softmax = attn.softmax(dim=-1)
+            out = na1d_av(attn_softmax, v, kernel_size, dilation)
+
+            # We can only check the outputs against each other, and
+            # not attention weights, because they are stored in a
+            # different order with NATTEN.
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_against_self_attention(self, B, H, L, D):
+        assert L % 2 == 1
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            L=L,
+            D=D,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_against_self_attention(self, B, H, L, D):
+        assert L % 2 == 1
+        # Test naive kernels
+        disable_gemm_na()
+        disable_tf32()
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            L=L,
+            D=D,
+            dtype=torch.float32,
+            device="cuda",
+            eps=1e-4,
+        )
+        if HAS_HALF:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                dtype=torch.float16,
+                device="cuda",
+                eps=1e-1,
+            )
+        if HAS_BFLOAT:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                dtype=torch.bfloat16,
+                device="cuda",
+                eps=1e-1,
+            )
+        # Test GEMM-based kernels
+        if HAS_GEMM:
+            enable_gemm_na()
+            if HAS_FLOAT_GEMM:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    dtype=torch.float32,
+                    device="cuda",
+                    eps=1e-2,
+                )
+                enable_tf32()
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    dtype=torch.float32,
+                    device="cuda",
+                    eps=1e-2,
+                )
+            assert (
+                HAS_HALF
+            ), "GEMM kernels must support FP16 across on all supported architectures."
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                L=L,
+                D=D,
+                dtype=torch.float16,
+                device="cuda",
+                eps=1e-1,
+            )
+            if HAS_BFLOAT:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    L=L,
+                    D=D,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                    eps=1e-1,
+                )
+
+    def test_cpu_against_self_attention(self):
+        self._test_cpu_against_self_attention(B=1, H=2, L=7, D=32)
+        self._test_cpu_against_self_attention(B=1, H=2, L=15, D=32)
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_against_self_attention(self):
+        self._test_cuda_against_self_attention(B=2, H=4, L=3, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=5, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=7, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=9, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=11, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=13, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, L=15, D=32)
+
 
 if __name__ == "__main__":
+    torch.manual_seed(42)
     unittest.main()
diff --git a/tests/test_na2d.py b/tests/test_na2d.py
index 180535b..b57a0b0 100644
--- a/tests/test_na2d.py
+++ b/tests/test_na2d.py
@@ -53,45 +53,51 @@
 logger = logging.getLogger(__name__)
 
 
+def init_cpu_ref(B, H, X, Y, D, kernel_size, dilation, has_bias):
+    with torch.no_grad():
+        q, k, v = (
+            torch.randn((B, H, X, Y, D)) * (D**-0.5),
+            torch.randn((B, H, X, Y, D)),
+            torch.randn((B, H, X, Y, D)),
+        )
+        rpb = (
+            None
+            if not has_bias
+            else torch.randn(H, 2 * kernel_size - 1, 2 * kernel_size - 1)
+        )
+        q_, k_, v_ = (
+            q.clone().cuda(),
+            k.clone().cuda(),
+            v.clone().cuda(),
+        )
+        rpb_ = None if rpb is None else rpb.clone().cuda()
+
+        attn_ref = na2d_qk_with_bias(q, k, rpb, kernel_size, dilation)
+        attn_ref = attn_ref.softmax(dim=-1)
+        out_ref = na2d_av(attn_ref, v, kernel_size, dilation)
+    return (q_, k_, v_, rpb_, kernel_size, dilation), (attn_ref.cuda(), out_ref.cuda())
+
+
 class NA2DTests(unittest.TestCase):
-    def _test_against_cpu(
-        self, B, H, X, Y, D, kernel_size, dilation, has_bias, dtype, eps
-    ):
+    def _test_against_cpu(self, inputs, reference, eps, dtype):
+        q, k, v, rpb, kernel_size, dilation = inputs
+        attn_ref, out_ref = reference
         with torch.no_grad():
-            q, k, v = (
-                torch.randn((B, H, X, Y, D)) * (D**-0.5),
-                torch.randn((B, H, X, Y, D)),
-                torch.randn((B, H, X, Y, D)),
-            )
-            rpb = None
-            if has_bias:
-                rpb = torch.randn(H, 2 * kernel_size - 1, 2 * kernel_size - 1)
-            q_, k_, v_ = (
-                q.clone().cuda().to(dtype),
-                k.clone().cuda().to(dtype),
-                v.clone().cuda().to(dtype),
-            )
-            rpb_ = None if rpb is None else rpb.clone().cuda().to(dtype)
-
-            attn_ref = na2d_qk_with_bias(q, k, rpb, kernel_size, dilation)
-            attn_ref = attn_ref.softmax(dim=-1)
-            out_ref = na2d_av(attn_ref, v, kernel_size, dilation)
+            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
+            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
+            assert q_.is_cuda and k_.is_cuda and v_.is_cuda
 
             attn = na2d_qk_with_bias(q_, k_, rpb_, kernel_size, dilation)
             attn = attn.softmax(dim=-1)
             out = na2d_av(attn, v_, kernel_size, dilation)
 
-            torch.testing.assert_close(attn.float().cpu(), attn_ref, atol=eps, rtol=0)
-            torch.testing.assert_close(out.float().cpu(), out_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
 
     def _test_all_dtypes_against_cpu(
         self, B, H, X, Y, D, kernel_size, dilation, has_bias=False
     ):
-        # Test naive kernels
-        disable_tiled_na()
-        disable_gemm_na()
-        disable_tf32()
-        self._test_against_cpu(
+        inputs, reference = init_cpu_ref(
             B=B,
             H=H,
             X=X,
@@ -100,73 +106,51 @@ def _test_all_dtypes_against_cpu(
             kernel_size=kernel_size,
             dilation=dilation,
             has_bias=has_bias,
+        )
+        # Test naive kernels
+        disable_tiled_na()
+        disable_gemm_na()
+        disable_tf32()
+        self._test_against_cpu(
+            inputs=inputs,
+            reference=reference,
             dtype=torch.float32,
             eps=1e-4,
         )
         if HAS_HALF:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float16,
                 eps=1e-1,
             )
         if HAS_BFLOAT:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.bfloat16,
                 eps=1e-1,
             )
-        if kernel_size < 15:
+        if kernel_size < 15 and D == 32:
             # Test tiled kernels
             enable_tiled_na()
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                D=32,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float32,
                 eps=1e-4,
             )
             if HAS_HALF:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    X=X,
-                    Y=Y,
-                    D=32,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.float16,
                     eps=1e-1,
                 )
             if HAS_BFLOAT:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    X=X,
-                    Y=Y,
-                    D=32,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.bfloat16,
                     eps=1e-1,
                 )
@@ -175,27 +159,15 @@ def _test_all_dtypes_against_cpu(
             enable_gemm_na()
             if HAS_FLOAT_GEMM:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    X=X,
-                    Y=Y,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.float32,
                     eps=1e-2,
                 )
                 enable_tf32()
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    X=X,
-                    Y=Y,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.float32,
                     eps=1e-2,
                 )
@@ -203,34 +175,21 @@ def _test_all_dtypes_against_cpu(
                 HAS_HALF
             ), "GEMM kernels must support FP16 across on all supported architectures."
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float16,
                 eps=1e-1,
             )
             if HAS_BFLOAT:
                 self._test_against_cpu(
-                    B=B,
-                    H=H,
-                    X=X,
-                    Y=Y,
-                    D=D,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                    has_bias=has_bias,
+                    inputs=inputs,
+                    reference=reference,
                     dtype=torch.bfloat16,
                     eps=1e-1,
                 )
 
     @skip_if_cuda_is_not_supported()
     def test_cpu_vs_cuda(self):
-        torch.manual_seed(42)
         self._test_all_dtypes_against_cpu(
             B=1, H=1, X=16, Y=16, D=32, kernel_size=7, dilation=1
         )
@@ -255,12 +214,17 @@ def test_cpu_vs_cuda(self):
         self._test_all_dtypes_against_cpu(
             B=1, H=2, X=28, Y=28, D=32, has_bias=True, kernel_size=13, dilation=2
         )
+        # See issue #32
+        # Fixed since switching to int64 indexing, and throwing
+        # away torch accessors.
+        self._test_all_dtypes_against_cpu(
+            B=1, H=1, X=1280, Y=1920, D=8, kernel_size=5, dilation=1
+        )
 
     def _test_autograd_qk(
         self, B, H, X, Y, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, X, Y, D), **kwargs)
         key = torch.randn((B, H, X, Y, D), **kwargs)
@@ -282,7 +246,6 @@ def _test_autograd_av(
         self, B, H, X, Y, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, X, Y, kernel_size**2 + L_extra), **kwargs)
         value = torch.randn((B, H, X, Y, D), **kwargs)
@@ -427,7 +390,6 @@ def test_invalid_dilation(self):
 
     def _test_fwad_qk(self, B, H, X, Y, D, kernel_size, dilation, device, L_extra=0):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, X, Y, D), **kwargs)
         key = torch.randn((B, H, X, Y, D), **kwargs)
@@ -446,7 +408,6 @@ def _test_fwad_qk(self, B, H, X, Y, D, kernel_size, dilation, device, L_extra=0)
 
     def _test_fwad_av(self, B, H, X, Y, D, kernel_size, dilation, device, L_extra=0):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, X, Y, kernel_size**2 + L_extra), **kwargs)
         value = torch.randn((B, H, X, Y, D), **kwargs)
@@ -572,7 +533,6 @@ def test_fwad_cuda_gemm(self):
         )
 
     def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
         query = torch.nested.nested_tensor(
@@ -622,7 +582,6 @@ def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
                 torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
 
     def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
         attn = torch.nested.nested_tensor(
@@ -689,3 +648,525 @@ def test_nested_forward_cpu(self):
     def test_nested_forward_cuda(self):
         self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
         self._test_nested_av_forward(dtype=torch.float16, device="cuda")
+
+    def _test_with_extra_tokens(
+        self,
+        B,
+        H,
+        X,
+        Y,
+        D,
+        kernel_size,
+        dilation,
+        has_bias,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+        L_extra=9,
+    ):
+        assert L_extra > 0
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, X, Y, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, X, Y, D), **kwargs),
+                torch.randn((B, H, X, Y, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            rpb, rpb_ref = None, None
+            if has_bias:
+                rpb = torch.randn(H, 2 * kernel_size - 1, 2 * kernel_size - 1, **kwargs)
+                rpb_ref = rpb.clone()
+
+            extra_k = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_v = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_k_ref, extra_v_ref = extra_k.clone(), extra_v.clone()
+
+            # Reference implementation
+            attn_extra_ref = (
+                (
+                    q_ref.view(B * H, X * Y, D).contiguous()
+                    @ extra_k_ref.view(B * H, L_extra, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, X, Y, L_extra)
+                .contiguous()
+            )
+            attn_na_ref = na2d_qk_with_bias(
+                q_ref, k_ref, rpb_ref, kernel_size, dilation
+            )
+            attn_ref = torch.cat([attn_na_ref, attn_extra_ref], dim=-1)
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            attn_na_softmax_ref, attn_extra_softmax_ref = attn_ref_softmax.split(
+                [kernel_size**2, L_extra], dim=-1
+            )
+            attn_na_softmax_ref, attn_extra_softmax_ref = (
+                attn_na_softmax_ref.contiguous(),
+                attn_extra_softmax_ref.contiguous(),
+            )
+            out_na_ref = na2d_av(attn_na_softmax_ref, v_ref, kernel_size, dilation)
+            out_extra_ref = (
+                (
+                    attn_extra_softmax_ref.view(B * H, X * Y, L_extra).contiguous()
+                    @ extra_v_ref.view(B * H, L_extra, D).contiguous()
+                )
+                .view(B, H, X, Y, D)
+                .contiguous()
+            )
+            out_ref = out_extra_ref + out_na_ref
+
+            # Op
+            attn = na2d_qk_with_bias(
+                q, k, rpb, kernel_size, dilation, additional_keys=extra_k
+            )
+            attn_na, attn_extra = attn.split([kernel_size**2, L_extra], dim=-1)
+            attn_na, attn_extra = attn_na.contiguous(), attn_extra.contiguous()
+            attn_softmax = attn.softmax(dim=-1)
+            attn_na_softmax, attn_extra_softmax = attn_softmax.split(
+                [kernel_size**2, L_extra], dim=-1
+            )
+            attn_na_softmax, attn_extra_softmax = (
+                attn_na_softmax.contiguous(),
+                attn_extra_softmax.contiguous(),
+            )
+            out = na2d_av(
+                attn_softmax, v, kernel_size, dilation, additional_values=extra_v
+            )
+
+            # Elementwise checks
+            torch.testing.assert_close(attn_na, attn_na_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_na_softmax, attn_na_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn_extra, attn_extra_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_extra_softmax, attn_extra_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn, attn_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_with_extra_tokens(
+        self, B, H, X, Y, D, kernel_size, dilation, has_bias=False
+    ):
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_with_extra_tokens(
+        self, B, H, X, Y, D, kernel_size, dilation, has_bias=False
+    ):
+        # Test naive kernels
+        disable_tiled_na()
+        disable_gemm_na()
+        disable_tf32()
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cuda",
+        )
+        if HAS_HALF:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float16,
+                device="cuda",
+            )
+        if HAS_BFLOAT:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+        if kernel_size < 15:
+            # Test tiled kernels
+            enable_tiled_na()
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=32,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float32,
+                device="cuda",
+            )
+            if HAS_HALF:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=32,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.float16,
+                    device="cuda",
+                )
+            if HAS_BFLOAT:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=32,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                )
+        # Test GEMM-based kernels
+        if HAS_GEMM:
+            enable_gemm_na()
+            if HAS_FLOAT_GEMM:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.float32,
+                    device="cuda",
+                )
+                enable_tf32()
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.float32,
+                    device="cuda",
+                )
+            assert (
+                HAS_HALF
+            ), "GEMM kernels must support FP16 across on all supported architectures."
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float16,
+                device="cuda",
+            )
+            if HAS_BFLOAT:
+                self._test_with_extra_tokens(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    has_bias=has_bias,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                )
+
+    def test_cpu_with_extra_tokens(self):
+        self._test_cpu_with_extra_tokens(
+            B=1, H=1, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=1, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=2, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=3, dilation=4
+        )
+        self._test_cpu_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=13, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=13, dilation=2
+        )
+        self._test_cpu_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, has_bias=True, kernel_size=13, dilation=2
+        )
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_with_extra_tokens(self):
+        self._test_cuda_with_extra_tokens(
+            B=1, H=1, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=1, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=2, X=16, Y=16, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=3, dilation=4
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=13, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, kernel_size=13, dilation=2
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, D=32, has_bias=True, kernel_size=13, dilation=2
+        )
+
+    def _test_against_self_attention(
+        self,
+        B,
+        H,
+        X,
+        Y,
+        D,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+    ):
+        # Only odd-sized kernels are valid in NA,
+        # and NA == SA when kernel size equals
+        # input size. NATTEN only allows square-sized
+        # kernels, so we need to make sure that the
+        # input size is an odd value across every
+        # spatial dim.
+        assert X == Y and X % 2 == 1
+        kernel_size = X
+        dilation = 1
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, X, Y, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, X, Y, D), **kwargs),
+                torch.randn((B, H, X, Y, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            # Reference implementation
+            attn_ref = (
+                (
+                    q_ref.view(B * H, X * Y, D).contiguous()
+                    @ k_ref.view(B * H, X * Y, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, X, Y, X * Y)
+                .contiguous()
+            )
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            out_ref = (
+                (
+                    attn_ref_softmax.view(B * H, X * Y, X * Y).contiguous()
+                    @ v_ref.view(B * H, X * Y, D).contiguous()
+                )
+                .view(B, H, X, Y, D)
+                .contiguous()
+            )
+
+            # Op
+            attn = na2d_qk(q, k, kernel_size, dilation)
+            attn_softmax = attn.softmax(dim=-1)
+            out = na2d_av(attn_softmax, v, kernel_size, dilation)
+
+            # We can only check the outputs against each other, and
+            # not attention weights, because they are stored in a
+            # different order with NATTEN.
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_against_self_attention(self, B, H, X, Y, D):
+        assert X == Y and X % 2 == 1
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            D=D,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_against_self_attention(self, B, H, X, Y, D):
+        assert X == Y and X % 2 == 1
+        kernel_size = X
+        # Test naive kernels
+        disable_tiled_na()
+        disable_gemm_na()
+        disable_tf32()
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            D=D,
+            dtype=torch.float32,
+            device="cuda",
+            eps=1e-4,
+        )
+        if HAS_HALF:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                dtype=torch.float16,
+                device="cuda",
+                eps=1e-1,
+            )
+        if HAS_BFLOAT:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                dtype=torch.bfloat16,
+                device="cuda",
+                eps=1e-1,
+            )
+        if kernel_size < 15:
+            # Test tiled kernels
+            enable_tiled_na()
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=32,
+                dtype=torch.float32,
+                device="cuda",
+                eps=1e-4,
+            )
+            if HAS_HALF:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=32,
+                    dtype=torch.float16,
+                    device="cuda",
+                    eps=1e-1,
+                )
+            if HAS_BFLOAT:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=32,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                    eps=1e-1,
+                )
+        # Test GEMM-based kernels
+        if HAS_GEMM:
+            enable_gemm_na()
+            if HAS_FLOAT_GEMM:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    dtype=torch.float32,
+                    device="cuda",
+                    eps=1e-2,
+                )
+                enable_tf32()
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    dtype=torch.float32,
+                    device="cuda",
+                    eps=1e-2,
+                )
+            assert (
+                HAS_HALF
+            ), "GEMM kernels must support FP16 across on all supported architectures."
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                D=D,
+                dtype=torch.float16,
+                device="cuda",
+                eps=1e-1,
+            )
+            if HAS_BFLOAT:
+                self._test_against_self_attention(
+                    B=B,
+                    H=H,
+                    X=X,
+                    Y=Y,
+                    D=D,
+                    dtype=torch.bfloat16,
+                    device="cuda",
+                    eps=1e-1,
+                )
+
+    def test_cpu_against_self_attention(self):
+        self._test_cpu_against_self_attention(B=1, H=2, X=7, Y=7, D=32)
+        self._test_cpu_against_self_attention(B=1, H=2, X=15, Y=15, D=32)
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_against_self_attention(self):
+        self._test_cuda_against_self_attention(B=2, H=4, X=3, Y=3, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=5, Y=5, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=7, Y=7, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=9, Y=9, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=11, Y=11, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=13, Y=13, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=15, Y=15, D=32)
+
+
+if __name__ == "__main__":
+    torch.manual_seed(42)
+    unittest.main()
diff --git a/tests/test_na3d.py b/tests/test_na3d.py
index 31b2446..b13def0 100644
--- a/tests/test_na3d.py
+++ b/tests/test_na3d.py
@@ -39,33 +39,43 @@
 logger = logging.getLogger(__name__)
 
 
-class NA3DTests(unittest.TestCase):
-    def _test_against_cpu(
-        self, B, H, X, Y, Z, D, kernel_size, dilation, has_bias, dtype, eps
-    ):
-        with torch.no_grad():
-            q, k, v = (
-                torch.randn((B, H, X, Y, Z, D)) * (D**-0.5),
-                torch.randn((B, H, X, Y, Z, D)),
-                torch.randn((B, H, X, Y, Z, D)),
-            )
-            rpb = None
-            if has_bias:
-                rpb = torch.randn(
-                    H, 2 * kernel_size - 1, 2 * kernel_size - 1, 2 * kernel_size - 1
-                )
-            q_, k_, v_ = (
-                q.clone().cuda().to(dtype),
-                k.clone().cuda().to(dtype),
-                v.clone().cuda().to(dtype),
+def init_cpu_ref(B, H, X, Y, Z, D, kernel_size, dilation, has_bias):
+    with torch.no_grad():
+        q, k, v = (
+            torch.randn((B, H, X, Y, Z, D)) * (D**-0.5),
+            torch.randn((B, H, X, Y, Z, D)),
+            torch.randn((B, H, X, Y, Z, D)),
+        )
+        rpb = (
+            None
+            if not has_bias
+            else torch.randn(
+                H, 2 * kernel_size - 1, 2 * kernel_size - 1, 2 * kernel_size - 1
             )
-            rpb_ = None if rpb is None else rpb.clone().cuda().to(dtype)
+        )
+        q_, k_, v_ = (
+            q.clone().cuda(),
+            k.clone().cuda(),
+            v.clone().cuda(),
+        )
+        rpb_ = None if rpb is None else rpb.clone().cuda()
 
-            attn_ref = na3d_qk_with_bias(
-                q, k, rpb, kernel_size, kernel_size, dilation, dilation
-            )
-            attn_ref = attn_ref.softmax(dim=-1)
-            out_ref = na3d_av(attn_ref, v, kernel_size, kernel_size, dilation, dilation)
+        attn_ref = na3d_qk_with_bias(
+            q, k, rpb, kernel_size, kernel_size, dilation, dilation
+        )
+        attn_ref = attn_ref.softmax(dim=-1)
+        out_ref = na3d_av(attn_ref, v, kernel_size, kernel_size, dilation, dilation)
+    return (q_, k_, v_, rpb_, kernel_size, dilation), (attn_ref.cuda(), out_ref.cuda())
+
+
+class NA3DTests(unittest.TestCase):
+    def _test_against_cpu(self, inputs, reference, eps, dtype):
+        q, k, v, rpb, kernel_size, dilation = inputs
+        attn_ref, out_ref = reference
+        with torch.no_grad():
+            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
+            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
+            assert q_.is_cuda and k_.is_cuda and v_.is_cuda
 
             attn = na3d_qk_with_bias(
                 q_, k_, rpb_, kernel_size, kernel_size, dilation, dilation
@@ -73,13 +83,13 @@ def _test_against_cpu(
             attn = attn.softmax(dim=-1)
             out = na3d_av(attn, v_, kernel_size, kernel_size, dilation, dilation)
 
-            torch.testing.assert_close(attn.float().cpu(), attn_ref, atol=eps, rtol=0)
-            torch.testing.assert_close(out.float().cpu(), out_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
+            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
 
     def _test_all_dtypes_against_cpu(
         self, B, H, X, Y, Z, D, kernel_size, dilation, has_bias=False
     ):
-        self._test_against_cpu(
+        inputs, reference = init_cpu_ref(
             B=B,
             H=H,
             X=X,
@@ -89,41 +99,31 @@ def _test_all_dtypes_against_cpu(
             kernel_size=kernel_size,
             dilation=dilation,
             has_bias=has_bias,
+        )
+
+        self._test_against_cpu(
+            inputs=inputs,
+            reference=reference,
             dtype=torch.float32,
             eps=1e-4,
         )
         if HAS_HALF:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                Z=Z,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.float16,
                 eps=1e-1,
             )
         if HAS_BFLOAT:
             self._test_against_cpu(
-                B=B,
-                H=H,
-                X=X,
-                Y=Y,
-                Z=Z,
-                D=D,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                has_bias=has_bias,
+                inputs=inputs,
+                reference=reference,
                 dtype=torch.bfloat16,
                 eps=1e-1,
             )
 
     @skip_if_cuda_is_not_supported()
     def test_cpu_vs_cuda(self):
-        torch.manual_seed(42)
         self._test_all_dtypes_against_cpu(
             B=2, H=3, X=9, Y=10, Z=11, D=32, kernel_size=3, dilation=1
         )
@@ -138,7 +138,6 @@ def _test_autograd_qk(
         self, B, H, X, Y, Z, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, X, Y, Z, D), **kwargs)
         key = torch.randn((B, H, X, Y, Z, D), **kwargs)
@@ -171,7 +170,6 @@ def _test_autograd_av(
         self, B, H, X, Y, Z, D, kernel_size, dilation, eps, device, L_extra=0
     ):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, X, Y, Z, kernel_size**3 + L_extra), **kwargs)
         value = torch.randn((B, H, X, Y, Z, D), **kwargs)
@@ -311,7 +309,6 @@ def test_invalid_dilation(self):
 
     def _test_fwad_qk(self, B, H, X, Y, Z, D, kernel_size, dilation, device, L_extra=0):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         query = torch.randn((B, H, X, Y, Z, D), **kwargs)
         key = torch.randn((B, H, X, Y, Z, D), **kwargs)
@@ -338,7 +335,6 @@ def _test_fwad_qk(self, B, H, X, Y, Z, D, kernel_size, dilation, device, L_extra
 
     def _test_fwad_av(self, B, H, X, Y, Z, D, kernel_size, dilation, device, L_extra=0):
         assert L_extra >= 0
-        torch.manual_seed(42)
         kwargs = {"dtype": torch.float64, "device": device, "requires_grad": True}
         attn = torch.randn((B, H, X, Y, Z, kernel_size**3 + L_extra), **kwargs)
         value = torch.randn((B, H, X, Y, Z, D), **kwargs)
@@ -441,7 +437,6 @@ def test_fwad_cuda(self):
         )
 
     def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kernel_size_d, dilation_d = 3, 3
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
@@ -510,7 +505,6 @@ def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
                 torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
 
     def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
-        torch.manual_seed(42)
         kernel_size, dilation = 7, 2
         kernel_size_d, dilation_d = 3, 3
         kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
@@ -618,6 +612,333 @@ def test_nested_forward_cuda(self):
         self._test_nested_qk_forward(dtype=torch.float16, device="cuda")
         self._test_nested_av_forward(dtype=torch.float16, device="cuda")
 
+    def _test_with_extra_tokens(
+        self,
+        B,
+        H,
+        X,
+        Y,
+        Z,
+        D,
+        kernel_size,
+        dilation,
+        has_bias,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+        L_extra=9,
+    ):
+        assert L_extra > 0
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, X, Y, Z, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, X, Y, Z, D), **kwargs),
+                torch.randn((B, H, X, Y, Z, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            rpb, rpb_ref = None, None
+            if has_bias:
+                rpb = torch.randn(
+                    H,
+                    2 * kernel_size - 1,
+                    2 * kernel_size - 1,
+                    2 * kernel_size - 1,
+                    **kwargs,
+                )
+                rpb_ref = rpb.clone()
+
+            extra_k = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_v = torch.randn((B, H, L_extra, D), **kwargs)
+            extra_k_ref, extra_v_ref = extra_k.clone(), extra_v.clone()
+
+            # Reference implementation
+            attn_extra_ref = (
+                (
+                    q_ref.view(B * H, X * Y * Z, D).contiguous()
+                    @ extra_k_ref.view(B * H, L_extra, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, X, Y, Z, L_extra)
+                .contiguous()
+            )
+            attn_na_ref = na3d_qk_with_bias(
+                q_ref, k_ref, rpb_ref, kernel_size, kernel_size, dilation, dilation
+            )
+            attn_ref = torch.cat([attn_na_ref, attn_extra_ref], dim=-1)
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            attn_na_softmax_ref, attn_extra_softmax_ref = attn_ref_softmax.split(
+                [kernel_size**3, L_extra], dim=-1
+            )
+            attn_na_softmax_ref, attn_extra_softmax_ref = (
+                attn_na_softmax_ref.contiguous(),
+                attn_extra_softmax_ref.contiguous(),
+            )
+            out_na_ref = na3d_av(
+                attn_na_softmax_ref, v_ref, kernel_size, kernel_size, dilation, dilation
+            )
+            out_extra_ref = (
+                (
+                    attn_extra_softmax_ref.view(B * H, X * Y * Z, L_extra).contiguous()
+                    @ extra_v_ref.view(B * H, L_extra, D).contiguous()
+                )
+                .view(B, H, X, Y, Z, D)
+                .contiguous()
+            )
+            out_ref = out_extra_ref + out_na_ref
+
+            # Op
+            attn = na3d_qk_with_bias(
+                q,
+                k,
+                rpb,
+                kernel_size,
+                kernel_size,
+                dilation,
+                dilation,
+                additional_keys=extra_k,
+            )
+            attn_na, attn_extra = attn.split([kernel_size**3, L_extra], dim=-1)
+            attn_na, attn_extra = attn_na.contiguous(), attn_extra.contiguous()
+            attn_softmax = attn.softmax(dim=-1)
+            attn_na_softmax, attn_extra_softmax = attn_softmax.split(
+                [kernel_size**3, L_extra], dim=-1
+            )
+            attn_na_softmax, attn_extra_softmax = (
+                attn_na_softmax.contiguous(),
+                attn_extra_softmax.contiguous(),
+            )
+            out = na3d_av(
+                attn_softmax,
+                v,
+                kernel_size,
+                kernel_size,
+                dilation,
+                dilation,
+                additional_values=extra_v,
+            )
+
+            # Elementwise checks
+            torch.testing.assert_close(attn_na, attn_na_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_na_softmax, attn_na_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn_extra, attn_extra_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(
+                attn_extra_softmax, attn_extra_softmax_ref, atol=eps, rtol=eps
+            )
+            torch.testing.assert_close(attn, attn_ref, atol=eps, rtol=eps)
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_with_extra_tokens(
+        self, B, H, X, Y, Z, D, kernel_size, dilation, has_bias=False
+    ):
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            Z=Z,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_with_extra_tokens(
+        self, B, H, X, Y, Z, D, kernel_size, dilation, has_bias=False
+    ):
+        self._test_with_extra_tokens(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            Z=Z,
+            D=D,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            has_bias=has_bias,
+            dtype=torch.float32,
+            device="cuda",
+        )
+        if HAS_HALF:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                Z=Z,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.float16,
+                device="cuda",
+            )
+        if HAS_BFLOAT:
+            self._test_with_extra_tokens(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                Z=Z,
+                D=D,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                has_bias=has_bias,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+
+    def test_cpu_with_extra_tokens(self):
+        self._test_cpu_with_extra_tokens(
+            B=2, H=3, X=9, Y=10, Z=11, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=2, X=9, Y=10, Z=11, D=32, kernel_size=3, dilation=3
+        )
+        self._test_cpu_with_extra_tokens(
+            B=2, H=1, X=9, Y=10, Z=11, D=32, kernel_size=7, dilation=1
+        )
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_with_extra_tokens(self):
+        self._test_cuda_with_extra_tokens(
+            B=2, H=3, X=9, Y=10, Z=11, D=32, kernel_size=3, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=2, X=9, Y=10, Z=11, D=32, kernel_size=3, dilation=3
+        )
+        self._test_cuda_with_extra_tokens(
+            B=2, H=1, X=9, Y=10, Z=11, D=32, kernel_size=7, dilation=1
+        )
+        self._test_cuda_with_extra_tokens(
+            B=1, H=2, X=28, Y=28, Z=28, D=32, has_bias=True, kernel_size=13, dilation=2
+        )
+
+    def _test_against_self_attention(
+        self,
+        B,
+        H,
+        X,
+        Y,
+        Z,
+        D,
+        dtype,
+        device="cuda",
+        eps=1e-6,
+    ):
+        # Only odd-sized kernels are valid in NA,
+        # and NA == SA when kernel size equals
+        # input size. NATTEN only allows square-sized
+        # kernels, so we need to make sure that the
+        # input size is an odd value across every
+        # spatial dim.
+        assert X == Y == Z and X % 2 == 1
+        kernel_size = X
+        dilation = 1
+        kwargs = {"device": device, "dtype": dtype}
+        with torch.no_grad():
+            q, k, v = (
+                torch.randn((B, H, X, Y, Z, D), **kwargs) * (D**-0.5),
+                torch.randn((B, H, X, Y, Z, D), **kwargs),
+                torch.randn((B, H, X, Y, Z, D), **kwargs),
+            )
+            q_ref, k_ref, v_ref = q.clone(), k.clone(), v.clone()
+
+            # Reference implementation
+            attn_ref = (
+                (
+                    q_ref.view(B * H, X * Y * Z, D).contiguous()
+                    @ k_ref.view(B * H, X * Y * Z, D).transpose(-2, -1).contiguous()
+                )
+                .view(B, H, X, Y, Z, X * Y * Z)
+                .contiguous()
+            )
+            attn_ref_softmax = attn_ref.softmax(dim=-1)
+            out_ref = (
+                (
+                    attn_ref_softmax.view(B * H, X * Y * Z, X * Y * Z).contiguous()
+                    @ v_ref.view(B * H, X * Y * Z, D).contiguous()
+                )
+                .view(B, H, X, Y, Z, D)
+                .contiguous()
+            )
+
+            # Op
+            attn = na3d_qk(q, k, kernel_size, kernel_size, dilation, dilation)
+            attn_softmax = attn.softmax(dim=-1)
+            out = na3d_av(attn_softmax, v, kernel_size, kernel_size, dilation, dilation)
+
+            # We can only check the outputs against each other, and
+            # not attention weights, because they are stored in a
+            # different order with NATTEN.
+            torch.testing.assert_close(out, out_ref, atol=eps, rtol=eps)
+
+    def _test_cpu_against_self_attention(self, B, H, X, Y, Z, D):
+        assert X == Y == Z and X % 2 == 1
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            Z=Z,
+            D=D,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def _test_cuda_against_self_attention(self, B, H, X, Y, Z, D):
+        assert X == Y == Z and X % 2 == 1
+        self._test_against_self_attention(
+            B=B,
+            H=H,
+            X=X,
+            Y=Y,
+            Z=Z,
+            D=D,
+            dtype=torch.float32,
+            device="cuda",
+            eps=1e-4,
+        )
+        if HAS_HALF:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                Z=Z,
+                D=D,
+                dtype=torch.float16,
+                device="cuda",
+                eps=1e-1,
+            )
+        if HAS_BFLOAT:
+            self._test_against_self_attention(
+                B=B,
+                H=H,
+                X=X,
+                Y=Y,
+                Z=Z,
+                D=D,
+                dtype=torch.bfloat16,
+                device="cuda",
+                eps=1e-1,
+            )
+
+    def test_cpu_against_self_attention(self):
+        self._test_cpu_against_self_attention(B=1, H=2, X=7, Y=7, Z=7, D=32)
+
+    @skip_if_cuda_is_not_supported()
+    def test_cuda_against_self_attention(self):
+        self._test_cuda_against_self_attention(B=2, H=4, X=3, Y=3, Z=3, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=5, Y=5, Z=5, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=7, Y=7, Z=7, D=32)
+        self._test_cuda_against_self_attention(B=2, H=4, X=9, Y=9, Z=9, D=32)
+
 
 if __name__ == "__main__":
+    torch.manual_seed(42)
     unittest.main()