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

Issue:
✨ feat(NeighborhoodAttention2D and NeighborhoodAttention1D):
Support to remove RPB without side-effects by setting rpb to a buffer tensor.

---

In some special cases, we may not need rpb.

Considering that the implementation of NATTEN depends on rpb, I provide some simple modifications to support this feature.

@alihassanijr 


Discussion / hints:
Thanks for the PR @lartpang .
Give me a bit to go through this and get back to you.
@alihassanijr 
Thanks for your reply and the valuable advice.
About the above changes, I'm actually a little confused here.
Besides the more concise form, are there some other differences between the two versions?

i.e.:

`self.register_buffer(tensor)` vs `nn.Parameter(tensor, requires_grad=False)`, are they really equivalent??

It is just the preferred way of doing it.
Buffers are exactly that, buffers.
Bias terms are never buffers, they are parameters.
In other words, no they are not equivalent in all cases. Buffers never get stored into state dict, parameters do. A good example of that is the running mean and variance in normalization layers.
Now the "correct" way of doing it would be to set rpb to `None`, similar to standard torch practice; however that would not work until we modify the backend to support NoneTypes for rpb.

@alihassanijr 

> Bias terms are never buffers, they are parameters.

I do agree with this.

> Buffers never get stored into state dict, parameters do. 

The `persistent=True` of `register_buffer` will make the buffer be part of the module’s state_dict. 😸 

https://pytorch.org/docs/1.8.0/generated/torch.nn.Module.html?highlight=buffer#torch.nn.Module.register_buffer

Fix patch:
diff --git a/natten/__init__.py b/natten/__init__.py
index 52cab78..0f25d5e 100644
--- a/natten/__init__.py
+++ b/natten/__init__.py
@@ -5,6 +5,6 @@
 LICENSE file in the root directory of this source tree.
 """
 from .natten1d import NeighborhoodAttention1D
-from .natten2d import NeighborhoodAttention2D, NeighborhoodAttention
+from .natten2d import NeighborhoodAttention2D
 
-__version__ = "0.14.4"
+__version__ = "0.14.5.dev0"
diff --git a/natten/flops.py b/natten/flops.py
index e9e5f97..d3e5076 100644
--- a/natten/flops.py
+++ b/natten/flops.py
@@ -20,11 +20,21 @@ def qk_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     assert len(outputs) == 1, f"Expected 1 output (attn), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
-    assert len(input_shapes[0]) == 5, f"Query must be a 5-dim tensor, got {len(input_shapes[0])}"
-    assert len(input_shapes[1]) == 5, f"Key must be a 5-dim tensor, got {len(input_shapes[1])}"
-    assert len(input_shapes[2]) == 3, f"RelPosBias must be a 3-dim tensor, got {len(input_shapes[2])}"
-    assert len(output_shapes[0]) == 5, f"Output must be a 5-dim tensor, got {len(output_shapes[0])}"
-    assert input_shapes[0] == input_shapes[1], f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
+    assert (
+        len(input_shapes[0]) == 5
+    ), f"Query must be a 5-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 5
+    ), f"Key must be a 5-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(input_shapes[2]) == 3
+    ), f"RelPosBias must be a 3-dim tensor, got {len(input_shapes[2])}"
+    assert (
+        len(output_shapes[0]) == 5
+    ), f"Output must be a 5-dim tensor, got {len(output_shapes[0])}"
+    assert (
+        input_shapes[0] == input_shapes[1]
+    ), f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
     batch_size, heads, height, width, dim = input_shapes[0]
     batch_size, heads, height, width, kernel_size_sq = output_shapes[0]
 
@@ -41,10 +51,19 @@ def av_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     assert len(outputs) == 1, f"Expected 1 output (out), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
-    assert len(input_shapes[0]) == 5, f"Attn must be a 5-dim tensor, got {len(input_shapes[0])}"
-    assert len(input_shapes[1]) == 5, f"Value must be a 5-dim tensor, got {len(input_shapes[1])}"
-    assert len(output_shapes[0]) == 5, f"Output must be a 5-dim tensor, got {len(output_shapes[0])}"
-    assert output_shapes[0] == input_shapes[1], f"Out and Value shapes did not match! O: {output_shapes[0]}, V: {input_shapes[1]}"
+    assert (
+        len(input_shapes[0]) == 5
+    ), f"Attn must be a 5-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 5
+    ), f"Value must be a 5-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(output_shapes[0]) == 5
+    ), f"Output must be a 5-dim tensor, got {len(output_shapes[0])}"
+    assert output_shapes[0] == input_shapes[1], (
+        f"Out and Value shapes did not match! O: {output_shapes[0]}, V:"
+        f" {input_shapes[1]}"
+    )
     batch_size, heads, height, width, kernel_size_sq = input_shapes[0]
     batch_size, heads, height, width, dim = output_shapes[0]
     flops = batch_size * heads * height * width * dim * kernel_size_sq
@@ -59,11 +78,21 @@ def qk_1d_rpb_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     assert len(outputs) == 1, f"Expected 1 output (attn), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
-    assert len(input_shapes[0]) == 4, f"Query must be a 4-dim tensor, got {len(input_shapes[0])}"
-    assert len(input_shapes[1]) == 4, f"Key must be a 4-dim tensor, got {len(input_shapes[1])}"
-    assert len(input_shapes[2]) == 2, f"RelPosBias must be a 2-dim tensor, got {len(input_shapes[2])}"
-    assert len(output_shapes[0]) == 4, f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
-    assert input_shapes[0] == input_shapes[1], f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
+    assert (
+        len(input_shapes[0]) == 4
+    ), f"Query must be a 4-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 4
+    ), f"Key must be a 4-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(input_shapes[2]) == 2
+    ), f"RelPosBias must be a 2-dim tensor, got {len(input_shapes[2])}"
+    assert (
+        len(output_shapes[0]) == 4
+    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+    assert (
+        input_shapes[0] == input_shapes[1]
+    ), f"Query and Key shapes did not match! Q: {input_shapes[0]}, K: {input_shapes[1]}"
     batch_size, heads, length, dim = input_shapes[0]
     batch_size, heads, length, kernel_size = output_shapes[0]
 
@@ -80,10 +109,19 @@ def av_1d_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     assert len(outputs) == 1, f"Expected 1 output (out), got {len(outputs)}"
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
-    assert len(input_shapes[0]) == 4, f"Attn must be a 4-dim tensor, got {len(input_shapes[0])}"
-    assert len(input_shapes[1]) == 4, f"Value must be a 4-dim tensor, got {len(input_shapes[1])}"
-    assert len(output_shapes[0]) == 4, f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
-    assert output_shapes[0] == input_shapes[1], f"Out and Value shapes did not match! O: {output_shapes[0]}, V: {input_shapes[1]}"
+    assert (
+        len(input_shapes[0]) == 4
+    ), f"Attn must be a 4-dim tensor, got {len(input_shapes[0])}"
+    assert (
+        len(input_shapes[1]) == 4
+    ), f"Value must be a 4-dim tensor, got {len(input_shapes[1])}"
+    assert (
+        len(output_shapes[0]) == 4
+    ), f"Output must be a 4-dim tensor, got {len(output_shapes[0])}"
+    assert output_shapes[0] == input_shapes[1], (
+        f"Out and Value shapes did not match! O: {output_shapes[0]}, V:"
+        f" {input_shapes[1]}"
+    )
     batch_size, heads, length, kernel_size = input_shapes[0]
     batch_size, heads, length, dim = output_shapes[0]
     flops = batch_size * heads * length * dim * kernel_size
@@ -97,7 +135,8 @@ def add_natten_handle(flop_ctr):
             "prim::PythonOp.NATTEN2DAVFunction": av_flop,
             "prim::PythonOp.NATTEN1DQKRPBFunction": qk_1d_rpb_flop,
             "prim::PythonOp.NATTEN1DAVFunction": av_1d_flop,
-        })
+        }
+    )
 
 
 def get_flops(model, input, disable_warnings=False):
@@ -106,4 +145,3 @@ def get_flops(model, input, disable_warnings=False):
     if disable_warnings:
         flop_ctr = flop_ctr.unsupported_ops_warnings(False)
     return flop_ctr.total()
-
diff --git a/natten/functional.py b/natten/functional.py
index 9ee0627..0470e5c 100644
--- a/natten/functional.py
+++ b/natten/functional.py
@@ -7,14 +7,17 @@
 import torch
 from torch.autograd import Function
 from torch.cuda.amp import custom_fwd, custom_bwd
+
 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 correct torch build: " + \
-                      f"natten.shi-labs.com."
-                      )
+    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"natten.shi-labs.com."
+    )
 
 
 class NATTEN1DQKRPBFunction(Function):
@@ -24,27 +27,30 @@ class NATTEN1DQKRPBFunction(Function):
     and adds relative positional biases.
     This calls the `QKRPB` kernel.
     """
+
     @staticmethod
     @custom_fwd(cast_inputs=torch.float16)
-    def forward(ctx, query, key, rpb, dilation):
+    def forward(ctx, query, key, rpb, kernel_size, dilation):
         query = query.contiguous()
         key = key.contiguous()
-        attn = _C.natten1dqkrpb_forward(
-                query,
-                key,
-                rpb, 
-                dilation)
+        attn = _C.natten1dqkrpb_forward(query, key, rpb, kernel_size, dilation)
         ctx.save_for_backward(query, key)
         ctx.dilation = dilation
+        ctx.bias = rpb is not None
         return attn
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
         outputs = _C.natten1dqkrpb_backward(
-            grad_out.contiguous(), ctx.saved_variables[0], ctx.saved_variables[1], ctx.dilation)
+            grad_out.contiguous(),
+            ctx.saved_variables[0],
+            ctx.saved_variables[1],
+            ctx.bias,
+            ctx.dilation,
+        )
         d_query, d_key, d_rpb = outputs
-        return d_query, d_key, d_rpb, None
+        return d_query, d_key, d_rpb, None, None
 
 
 class NATTEN1DAVFunction(Function):
@@ -53,15 +59,13 @@ class NATTEN1DAVFunction(Function):
     Computes neighborhood attention outputs given attention weights, and values.
     This calls the `AV` kernel.
     """
+
     @staticmethod
     @custom_fwd(cast_inputs=torch.float16)
     def forward(ctx, attn, value, dilation):
         attn = attn.contiguous()
         value = value.contiguous()
-        out = _C.natten1dav_forward(
-                attn, 
-                value, 
-                dilation)
+        out = _C.natten1dav_forward(attn, value, dilation)
         ctx.save_for_backward(attn, value)
         ctx.dilation = dilation
         return out
@@ -70,7 +74,11 @@ def forward(ctx, attn, value, dilation):
     @custom_bwd
     def backward(ctx, grad_out):
         outputs = _C.natten1dav_backward(
-            grad_out.contiguous(), ctx.saved_variables[0], ctx.saved_variables[1], ctx.dilation)
+            grad_out.contiguous(),
+            ctx.saved_variables[0],
+            ctx.saved_variables[1],
+            ctx.dilation,
+        )
         d_attn, d_value = outputs
         return d_attn, d_value, None
 
@@ -82,27 +90,30 @@ class NATTEN2DQKRPBFunction(Function):
     and adds relative positional biases.
     This calls the `QKRPB` kernel.
     """
+
     @staticmethod
     @custom_fwd(cast_inputs=torch.float16)
-    def forward(ctx, query, key, rpb, dilation):
+    def forward(ctx, query, key, rpb, kernel_size, dilation):
         query = query.contiguous()
         key = key.contiguous()
-        attn = _C.natten2dqkrpb_forward(
-                query,
-                key,
-                rpb, 
-                dilation)
+        attn = _C.natten2dqkrpb_forward(query, key, rpb, kernel_size, dilation)
         ctx.save_for_backward(query, key)
         ctx.dilation = dilation
+        ctx.bias = rpb is not None
         return attn
 
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_out):
         outputs = _C.natten2dqkrpb_backward(
-            grad_out.contiguous(), ctx.saved_variables[0], ctx.saved_variables[1], ctx.dilation)
+            grad_out.contiguous(),
+            ctx.saved_variables[0],
+            ctx.saved_variables[1],
+            ctx.bias,
+            ctx.dilation,
+        )
         d_query, d_key, d_rpb = outputs
-        return d_query, d_key, d_rpb, None
+        return d_query, d_key, d_rpb, None, None
 
 
 class NATTEN2DAVFunction(Function):
@@ -111,15 +122,13 @@ class NATTEN2DAVFunction(Function):
     Computes neighborhood attention outputs given attention weights, and values.
     This calls the `AV` kernel.
     """
+
     @staticmethod
     @custom_fwd(cast_inputs=torch.float16)
     def forward(ctx, attn, value, dilation):
         attn = attn.contiguous()
         value = value.contiguous()
-        out = _C.natten2dav_forward(
-                attn, 
-                value, 
-                dilation)
+        out = _C.natten2dav_forward(attn, value, dilation)
         ctx.save_for_backward(attn, value)
         ctx.dilation = dilation
         return out
@@ -128,23 +137,26 @@ def forward(ctx, attn, value, dilation):
     @custom_bwd
     def backward(ctx, grad_out):
         outputs = _C.natten2dav_backward(
-            grad_out.contiguous(), ctx.saved_variables[0], ctx.saved_variables[1], ctx.dilation)
+            grad_out.contiguous(),
+            ctx.saved_variables[0],
+            ctx.saved_variables[1],
+            ctx.dilation,
+        )
         d_attn, d_value = outputs
         return d_attn, d_value, None
 
 
-def natten1dqkrpb(query, key, rpb, dilation):
-    return NATTEN1DQKRPBFunction.apply(query, key, rpb, dilation)
+def natten1dqkrpb(query, key, rpb, kernel_size, dilation):
+    return NATTEN1DQKRPBFunction.apply(query, key, rpb, kernel_size, dilation)
 
 
 def natten1dav(attn, value, dilation):
     return NATTEN1DAVFunction.apply(attn, value, dilation)
 
 
-def natten2dqkrpb(query, key, rpb, dilation):
-    return NATTEN2DQKRPBFunction.apply(query, key, rpb, dilation)
+def natten2dqkrpb(query, key, rpb, kernel_size, dilation):
+    return NATTEN2DQKRPBFunction.apply(query, key, rpb, kernel_size, dilation)
 
 
 def natten2dav(attn, value, dilation):
     return NATTEN2DAVFunction.apply(attn, value, dilation)
-
diff --git a/natten/natten1d.py b/natten/natten1d.py
index 3ff235a..f74eeb8 100644
--- a/natten/natten1d.py
+++ b/natten/natten1d.py
@@ -15,23 +15,39 @@ class NeighborhoodAttention1D(nn.Module):
     """
     Neighborhood Attention 1D Module
     """
-    def __init__(self, dim, kernel_size, num_heads,
-                 qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.,
-                 dilation=None, bias=True):
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        kernel_size,
+        dilation=1,
+        bias=True,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = dim // self.num_heads
-        self.scale = qk_scale or self.head_dim ** -0.5
-        assert kernel_size > 1 and kernel_size % 2 == 1, \
-            f"Kernel size must be an odd number greater than 1, got {kernel_size}."
+        self.scale = qk_scale or self.head_dim**-0.5
+        assert (
+            kernel_size > 1 and kernel_size % 2 == 1
+        ), f"Kernel size must be an odd number greater than 1, got {kernel_size}."
         self.kernel_size = kernel_size
+        assert (
+            dilation is None or dilation >= 1
+        ), f"Dilation must be greater than or equal to 1, got {dilation}."
         self.dilation = dilation or 1
         self.window_size = self.kernel_size * self.dilation
 
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.rpb = nn.Parameter(torch.zeros(num_heads, (2 * kernel_size - 1)), requires_grad=bias)
         if bias:
-            trunc_normal_(self.rpb, std=.02, mean=0., a=-2., b=2.)
+            self.rpb = nn.Parameter(torch.zeros(num_heads, (2 * kernel_size - 1)))
+            trunc_normal_(self.rpb, std=0.02, mean=0.0, a=-2.0, b=2.0)
+        else:
+            self.register_parameter("rpb", None)
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
@@ -44,10 +60,14 @@ def forward(self, x):
             pad_r = max(0, self.window_size - L)
             x = pad(x, (0, 0, pad_l, pad_r))
             _, L, _ = x.shape
-        qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, L, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
         q = q * self.scale
-        attn = natten1dqkrpb(q, k, self.rpb, self.dilation)
+        attn = natten1dqkrpb(q, k, self.rpb, self.kernel_size, self.dilation)
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
         x = natten1dav(attn, v, self.dilation)
@@ -58,4 +78,8 @@ def forward(self, x):
         return self.proj_drop(self.proj(x))
 
     def extra_repr(self) -> str:
-        return f'kernel_size={self.kernel_size}, dilation={self.dilation}, head_dim={self.head_dim}, num_heads={self.num_heads}, bias={self.rpb.requires_grad}'
+        return (
+            f"head_dim={self.head_dim}, num_heads={self.num_heads}, "
+            + f"kernel_size={self.kernel_size}, dilation={self.dilation}, "
+            + f"rel_pos_bias={self.rpb is not None}"
+        )
diff --git a/natten/natten2d.py b/natten/natten2d.py
index 3031bd1..bc26ecf 100644
--- a/natten/natten2d.py
+++ b/natten/natten2d.py
@@ -9,36 +9,47 @@
 from torch.nn.functional import pad
 from torch.nn.init import trunc_normal_
 from .functional import natten2dqkrpb, natten2dav
-import warnings
 
 
 class NeighborhoodAttention2D(nn.Module):
     """
     Neighborhood Attention 2D Module
     """
-    def __init__(self, dim, kernel_size, num_heads,
-                 qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.,
-                 dilation=None, bias=True):
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        kernel_size,
+        dilation=1,
+        bias=True,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = dim // self.num_heads
-        self.scale = qk_scale or self.head_dim ** -0.5
-        assert kernel_size > 1 and kernel_size % 2 == 1, \
-            f"Kernel size must be an odd number greater than 1, got {kernel_size}."
+        self.scale = qk_scale or self.head_dim**-0.5
+        assert (
+            kernel_size > 1 and kernel_size % 2 == 1
+        ), f"Kernel size must be an odd number greater than 1, got {kernel_size}."
         self.kernel_size = kernel_size
-        if type(dilation) is str:
-            self.dilation = None
-            self.window_size = None
-        else:
-            assert dilation is None or dilation >= 1, \
-                f"Dilation must be greater than or equal to 1, got {dilation}."
-            self.dilation = dilation or 1
-            self.window_size = self.kernel_size * self.dilation
+        assert (
+            dilation is None or dilation >= 1
+        ), f"Dilation must be greater than or equal to 1, got {dilation}."
+        self.dilation = dilation or 1
+        self.window_size = self.kernel_size * self.dilation
 
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.rpb = nn.Parameter(torch.zeros(num_heads, (2 * kernel_size - 1), (2 * kernel_size - 1)), requires_grad=bias)
         if bias:
-            trunc_normal_(self.rpb, std=.02, mean=0., a=-2., b=2.)
+            self.rpb = nn.Parameter(
+                torch.zeros(num_heads, (2 * kernel_size - 1), (2 * kernel_size - 1))
+            )
+            trunc_normal_(self.rpb, std=0.02, mean=0.0, a=-2.0, b=2.0)
+        else:
+            self.register_parameter("rpb", None)
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
@@ -47,24 +58,23 @@ def forward(self, x):
         B, Hp, Wp, C = x.shape
         H, W = int(Hp), int(Wp)
         pad_l = pad_t = pad_r = pad_b = 0
-        dilation = self.dilation
-        window_size = self.window_size
-        if window_size is None:
-            dilation = max(min(H, W) // self.kernel_size, 1)
-            window_size = dilation * self.kernel_size
-        if H < window_size or W < window_size:
+        if H < self.window_size or W < self.window_size:
             pad_l = pad_t = 0
-            pad_r = max(0, window_size - W)
-            pad_b = max(0, window_size - H)
+            pad_r = max(0, self.window_size - W)
+            pad_b = max(0, self.window_size - H)
             x = pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
             _, H, W, _ = x.shape
-        qkv = self.qkv(x).reshape(B, H, W, 3, self.num_heads, self.head_dim).permute(3, 0, 4, 1, 2, 5)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, H, W, 3, self.num_heads, self.head_dim)
+            .permute(3, 0, 4, 1, 2, 5)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
         q = q * self.scale
-        attn = natten2dqkrpb(q, k, self.rpb, dilation)
+        attn = natten2dqkrpb(q, k, self.rpb, self.kernel_size, self.dilation)
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
-        x = natten2dav(attn, v, dilation)
+        x = natten2dav(attn, v, self.dilation)
         x = x.permute(0, 2, 3, 1, 4).reshape(B, H, W, C)
         if pad_r or pad_b:
             x = x[:, :Hp, :Wp, :]
@@ -72,18 +82,8 @@ def forward(self, x):
         return self.proj_drop(self.proj(x))
 
     def extra_repr(self) -> str:
-        return f'kernel_size={self.kernel_size}, dilation={self.dilation}, head_dim={self.head_dim}, num_heads={self.num_heads}, bias={self.rpb.requires_grad}'
-
-
-class NeighborhoodAttention(NeighborhoodAttention2D):
-    """
-    Neighborhood Attention 2D Module
-    """
-    def __init__(self, dim, kernel_size, num_heads,
-                 qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.,
-                 dilation=None):
-        super().__init__(dim=dim, kernel_size=kernel_size, num_heads=num_heads,
-                         qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop,
-                         dilation=dilation)
-        warnings.warn('Using NeighborhoodAttention has been deprecated since natten v0.13. ' +
-                      'Please consider using NeighborhoodAttention2D instead.', DeprecationWarning, stacklevel=2)
+        return (
+            f"head_dim={self.head_dim}, num_heads={self.num_heads}, "
+            + f"kernel_size={self.kernel_size}, dilation={self.dilation}, "
+            + f"rel_pos_bias={self.rpb is not None}"
+        )
diff --git a/natten/src/cpp/natten1dqkrpb_cpu_kernel.cpp b/natten/src/cpp/natten1dqkrpb_cpu_kernel.cpp
index 2261935..ea06bb5 100644
--- a/natten/src/cpp/natten1dqkrpb_cpu_kernel.cpp
+++ b/natten/src/cpp/natten1dqkrpb_cpu_kernel.cpp
@@ -224,6 +224,7 @@ std::vector<torch::Tensor> natten1dqkrpb_cpu_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int batch_size = query.size(0);
     int heads = query.size(1);
@@ -235,7 +236,9 @@ std::vector<torch::Tensor> natten1dqkrpb_cpu_backward(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
 
     AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten1dqkrpb_backward_cpu", ([&] {
         const auto d_attn_a = d_attn.accessor<scalar_t, 4>();
@@ -243,9 +246,11 @@ std::vector<torch::Tensor> natten1dqkrpb_cpu_backward(
         const auto key_a = key.accessor<scalar_t, 4>();
         auto d_query_a = d_query.accessor<scalar_t, 4>();
         auto d_key_a = d_key.accessor<scalar_t, 4>();
-        auto d_rpb_a = d_rpb.accessor<scalar_t, 2>();
-        LAUNCH_DNA_KNS(kernel_size, dilation, natten1drpb_cpu_backward_kernel, 
-                d_rpb_a, d_attn_a, length, heads, kernel_size, dilation, batch_size);
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.accessor<scalar_t, 2>();
+            LAUNCH_DNA_KNS(kernel_size, dilation, natten1drpb_cpu_backward_kernel, 
+                    d_rpb_a, d_attn_a, length, heads, kernel_size, dilation, batch_size);
+        }
         LAUNCH_DNA_KNS(kernel_size, dilation, natten1dq_cpu_backward_kernel, 
                 d_query_a, d_attn_a, key_a, length, heads, kernel_size, dilation, dim, batch_size);
         LAUNCH_DNA_KNS(kernel_size, dilation, natten1dk_cpu_backward_kernel, 
diff --git a/natten/src/cpp/natten2dqkrpb_cpu_kernel.cpp b/natten/src/cpp/natten2dqkrpb_cpu_kernel.cpp
index 216aebf..3c2c80f 100644
--- a/natten/src/cpp/natten2dqkrpb_cpu_kernel.cpp
+++ b/natten/src/cpp/natten2dqkrpb_cpu_kernel.cpp
@@ -253,6 +253,7 @@ std::vector<torch::Tensor> natten2dqkrpb_cpu_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int batch_size = query.size(0);
     int heads = query.size(1);
@@ -266,7 +267,9 @@ std::vector<torch::Tensor> natten2dqkrpb_cpu_backward(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
 
     AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten2dqkrpb_backward_cpu", ([&] {
         const auto d_attn_a = d_attn.accessor<scalar_t, 5>();
@@ -274,9 +277,11 @@ std::vector<torch::Tensor> natten2dqkrpb_cpu_backward(
         const auto key_a = key.accessor<scalar_t, 5>();
         auto d_query_a = d_query.accessor<scalar_t, 5>();
         auto d_key_a = d_key.accessor<scalar_t, 5>();
-        auto d_rpb_a = d_rpb.accessor<scalar_t, 3>();
-        LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cpu_backward_kernel, 
-                d_rpb_a, d_attn_a, height, width, heads, kernel_size, dilation, batch_size);
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.accessor<scalar_t, 3>();
+            LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cpu_backward_kernel, 
+                    d_rpb_a, d_attn_a, height, width, heads, kernel_size, dilation, batch_size);
+        }
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dq_cpu_backward_kernel, 
                 d_query_a, d_attn_a, key_a, height, width, heads, kernel_size, dilation, dim, batch_size);
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dk_cpu_backward_kernel, 
diff --git a/natten/src/cuda/natten1dqkrpb_cuda_kernel.cu b/natten/src/cuda/natten1dqkrpb_cuda_kernel.cu
index 7607eae..0ec627f 100644
--- a/natten/src/cuda/natten1dqkrpb_cuda_kernel.cu
+++ b/natten/src/cuda/natten1dqkrpb_cuda_kernel.cu
@@ -417,6 +417,7 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int64_t batch_size = query.size(0);
     int64_t heads = query.size(1);
@@ -428,8 +429,9 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
-
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
     int32_t n_rpb = heads * length * kernel_size;
     int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB);
     dim3 grid_rpb(blocks_rpb);
@@ -444,14 +446,16 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward(
     dim3 blockk(CUDA_NUM_THREADS_K);
     const auto stream = c10::cuda::getCurrentCUDAStream();
     AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten1dqkrpb_backward_cuda", ([&] {
-        auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,2,torch::DefaultPtrTraits>();
-        auto d_query_a = d_query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
-        auto d_key_a = d_key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto d_attn_a = d_attn.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto query_a = query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto key_a = key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
-        LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1drpb_cuda_backward_kernel, grid_rpb, blockr, 0, stream, 
-                d_rpb_a, d_attn_a, length, kernel_size, dilation, batch_size, d_rpb.numel(), n_rpb);
+        auto d_query_a = d_query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,2,torch::DefaultPtrTraits>();
+            LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1drpb_cuda_backward_kernel, grid_rpb, blockr, 0, stream, 
+                    d_rpb_a, d_attn_a, length, kernel_size, dilation, batch_size, d_rpb.numel(), n_rpb);
+        }
         LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1dq_cuda_backward_kernel_fp32, grid_query, blockq, 0, stream, 
                 d_query_a, d_attn_a, key_a, length, heads, kernel_size, dilation, dim, n_query);
         LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1dk_cuda_backward_kernel_fp32, grid_key, blockk, 0, stream, 
@@ -464,6 +468,7 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward_fp16(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int64_t batch_size = query.size(0);
     int64_t heads = query.size(1);
@@ -476,7 +481,9 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward_fp16(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX}, d_attn.options());
 
     int32_t n_rpb = heads * length * kernel_size;
     int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB16);
@@ -492,14 +499,16 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward_fp16(
     dim3 blockk(CUDA_NUM_THREADS_K16);
     const auto stream = c10::cuda::getCurrentCUDAStream();
     AT_DISPATCH_HALF_TYPES(at::kHalf, d_query.scalar_type(), "natten1dqkrpb_backward_cuda_fp16", ([&] {
-        auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,2,torch::DefaultPtrTraits>();
-        auto d_query_a = d_query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
-        auto d_key_a = d_key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto d_attn_a = d_attn.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto query_a = query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
         const auto key_a = key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
-        LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1drpb_cuda_backward_kernel_fp16, grid_rpb, blockr, 0, stream, 
-                d_rpb_a, d_attn_a, length, kernel_size, dilation, batch_size, d_rpb.numel(), n_rpb);
+        auto d_query_a = d_query.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,4,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,2,torch::DefaultPtrTraits>();
+            LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1drpb_cuda_backward_kernel_fp16, grid_rpb, blockr, 0, stream, 
+                    d_rpb_a, d_attn_a, length, kernel_size, dilation, batch_size, d_rpb.numel(), n_rpb);
+        }
         LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1dq_cuda_backward_kernel_fp16, grid_query, blockq, 0, stream, 
                 d_query_a, d_attn_a, key_a, length, heads, kernel_size, dilation, dimhalf, nhalf_query);
         LAUNCH_DNA_KNS_1D(kernel_size, dilation, natten1dk_cuda_backward_kernel_fp16, grid_key, blockk, 0, stream, 
diff --git a/natten/src/cuda/natten2dqkrpb_cuda_kernel.cu b/natten/src/cuda/natten2dqkrpb_cuda_kernel.cu
index 8e17ea1..6c40904 100644
--- a/natten/src/cuda/natten2dqkrpb_cuda_kernel.cu
+++ b/natten/src/cuda/natten2dqkrpb_cuda_kernel.cu
@@ -1457,6 +1457,7 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int64_t batch_size = query.size(0);
     int64_t heads = query.size(1);
@@ -1471,7 +1472,9 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
 
     int32_t n_rpb = heads * height * width * kernel_size_sq;
     int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB);
@@ -1487,14 +1490,16 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward(
     dim3 blockk(CUDA_NUM_THREADS_K);
     const auto stream = c10::cuda::getCurrentCUDAStream();
     AT_DISPATCH_FLOATING_TYPES(d_query.scalar_type(), "natten2dqkrpb_backward_cuda", ([&] {
-        auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,3,torch::DefaultPtrTraits>();
-        auto d_query_a = d_query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
-        auto d_key_a = d_key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto d_attn_a = d_attn.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto query_a = query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto key_a = key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
-        LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cuda_backward_kernel, grid_rpb, blockr, 0, stream,
-                d_rpb_a, d_attn_a, height, width, dilation, batch_size, d_rpb.numel(), n_rpb);
+        auto d_query_a = d_query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,3,torch::DefaultPtrTraits>();
+            LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cuda_backward_kernel, grid_rpb, blockr, 0, stream,
+                    d_rpb_a, d_attn_a, height, width, dilation, batch_size, d_rpb.numel(), n_rpb);
+        }
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dq_cuda_backward_kernel_fp32, grid_query, blockq, 0, stream,
                 d_query_a, d_attn_a, key_a, height, width, heads, dilation, dim, n_query);
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dk_cuda_backward_kernel_fp32, grid_key, blockk, 0, stream,
@@ -1507,6 +1512,7 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward_fp16(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     int64_t batch_size = query.size(0);
     int64_t heads = query.size(1);
@@ -1522,7 +1528,9 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward_fp16(
    
     auto d_query = torch::zeros_like(query);
     auto d_key = torch::zeros_like(key);
-    auto d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
+    at::Tensor d_rpb;
+    if (biasEnabled)
+        d_rpb = torch::zeros({heads, RPB_MAX, RPB_MAX}, d_attn.options());
 
     int32_t n_rpb = heads * height * width * kernel_size_sq;
     int blocks_rpb = GET_BLOCKS(n_rpb, CUDA_NUM_THREADS_RPB16);
@@ -1538,14 +1546,16 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward_fp16(
     dim3 blockk(CUDA_NUM_THREADS_K16);
     const auto stream = c10::cuda::getCurrentCUDAStream();
     AT_DISPATCH_HALF_TYPES(at::kHalf, d_query.scalar_type(), "natten2dqkrpb_backward_cuda_fp16", ([&] {
-        auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,3,torch::DefaultPtrTraits>();
-        auto d_query_a = d_query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
-        auto d_key_a = d_key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto d_attn_a = d_attn.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto query_a = query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
         const auto key_a = key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
-        LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cuda_backward_kernel_fp16, grid_rpb, blockr, 0, stream,
-                d_rpb_a, d_attn_a, height, width, dilation, batch_size, d_rpb.numel(), n_rpb);
+        auto d_query_a = d_query.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
+        auto d_key_a = d_key.packed_accessor32<scalar_t,5,torch::DefaultPtrTraits>();
+        if (biasEnabled) {
+            auto d_rpb_a = d_rpb.packed_accessor32<scalar_t,3,torch::DefaultPtrTraits>();
+            LAUNCH_DNA_KNS(kernel_size, dilation, natten2drpb_cuda_backward_kernel_fp16, grid_rpb, blockr, 0, stream,
+                    d_rpb_a, d_attn_a, height, width, dilation, batch_size, d_rpb.numel(), n_rpb);
+        }
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dq_cuda_backward_kernel_fp16, grid_query, blockq, 0, stream,
                 d_query_a, d_attn_a, key_a, height, width, heads, dilation, dimhalf, nhalf_query);
         LAUNCH_DNA_KNS(kernel_size, dilation, natten2dk_cuda_backward_kernel_fp16, grid_key, blockk, 0, stream,
diff --git a/natten/src/natten1dqkrpb.h b/natten/src/natten1dqkrpb.h
index 266c294..1312d0f 100644
--- a/natten/src/natten1dqkrpb.h
+++ b/natten/src/natten1dqkrpb.h
@@ -21,6 +21,7 @@ std::vector<torch::Tensor> natten1dqkrpb_cpu_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 #if defined(WITH_CUDA)
@@ -42,12 +43,14 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 std::vector<torch::Tensor> natten1dqkrpb_cuda_backward_fp16(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 #endif
@@ -59,10 +62,13 @@ std::vector<torch::Tensor> natten1dqkrpb_cuda_backward_fp16(
 torch::Tensor natten1dqkrpb_forward(
     const torch::Tensor &query,
     const torch::Tensor &key,
-    const torch::Tensor &rpb,
+    const at::optional<at::Tensor> &rpb_opt,
+    const int kernel_size,
     const int dilation) {
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
+    const int heads = query.size(1);
+    auto rpb = rpb_opt.has_value() ? rpb_opt.value() : torch::zeros({heads, kernel_size}, query.options()); 
     CHECK_CONTIGUOUS(rpb);
     assert(query.device().is_cuda() == key.device().is_cuda() && rpb.device().is_cuda() == key.device().is_cuda());
     if (query.device().is_cuda()) {
@@ -82,6 +88,7 @@ std::vector<torch::Tensor> natten1dqkrpb_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     CHECK_CONTIGUOUS(d_attn);
     CHECK_CONTIGUOUS(query);
@@ -91,12 +98,12 @@ std::vector<torch::Tensor> natten1dqkrpb_backward(
 #if defined(WITH_CUDA)
         bool half = ::detail::scalar_type(query.scalar_type()) == at::ScalarType::Half;
         if (half)
-            return natten1dqkrpb_cuda_backward_fp16(d_attn, query, key, dilation);
-        return natten1dqkrpb_cuda_backward(d_attn, query, key, dilation);
+            return natten1dqkrpb_cuda_backward_fp16(d_attn, query, key, biasEnabled, dilation);
+        return natten1dqkrpb_cuda_backward(d_attn, query, key, biasEnabled, dilation);
 #else
     AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
 #endif
     }
-    return natten1dqkrpb_cpu_backward(d_attn, query, key, dilation);
+    return natten1dqkrpb_cpu_backward(d_attn, query, key, biasEnabled, dilation);
 }
 } // namespace natten
diff --git a/natten/src/natten2dqkrpb.h b/natten/src/natten2dqkrpb.h
index 1c32ce9..5081291 100644
--- a/natten/src/natten2dqkrpb.h
+++ b/natten/src/natten2dqkrpb.h
@@ -21,6 +21,7 @@ std::vector<torch::Tensor> natten2dqkrpb_cpu_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 #if defined(WITH_CUDA)
@@ -54,12 +55,14 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 std::vector<torch::Tensor> natten2dqkrpb_cuda_backward_fp16(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation);
 
 #endif
@@ -71,10 +74,13 @@ std::vector<torch::Tensor> natten2dqkrpb_cuda_backward_fp16(
 torch::Tensor natten2dqkrpb_forward(
     const torch::Tensor &query,
     const torch::Tensor &key,
-    const torch::Tensor &rpb,
+    const at::optional<at::Tensor> &rpb_opt,
+    const int kernel_size,
     const int dilation) {
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
+    const int heads = query.size(1);
+    auto rpb = rpb_opt.has_value() ? rpb_opt.value() : torch::zeros({heads, kernel_size, kernel_size}, query.options()); 
     CHECK_CONTIGUOUS(rpb);
     assert(query.device().is_cuda() == key.device().is_cuda() && rpb.device().is_cuda() == key.device().is_cuda());
     if (query.device().is_cuda()) {
@@ -104,6 +110,7 @@ std::vector<torch::Tensor> natten2dqkrpb_backward(
     const torch::Tensor &d_attn,
     const torch::Tensor &query,
     const torch::Tensor &key,
+    const bool biasEnabled,
     const int dilation) {
     CHECK_CONTIGUOUS(d_attn);
     CHECK_CONTIGUOUS(query);
@@ -113,12 +120,12 @@ std::vector<torch::Tensor> natten2dqkrpb_backward(
 #if defined(WITH_CUDA)
         bool half = ::detail::scalar_type(query.scalar_type()) == at::ScalarType::Half;
         if (half)
-            return natten2dqkrpb_cuda_backward_fp16(d_attn, query, key, dilation);
-        return natten2dqkrpb_cuda_backward(d_attn, query, key, dilation);
+            return natten2dqkrpb_cuda_backward_fp16(d_attn, query, key, biasEnabled, dilation);
+        return natten2dqkrpb_cuda_backward(d_attn, query, key, biasEnabled, dilation);
 #else
     AT_ERROR("NATTEN is not compiled with CUDA! Please make sure you installed correctly by referring to shi-labs.com/natten.");
 #endif
     }
-    return natten2dqkrpb_cpu_backward(d_attn, query, key, dilation);
+    return natten2dqkrpb_cpu_backward(d_attn, query, key, biasEnabled, dilation);
 }
 } // namespace natten


Test patch:
diff --git a/tests/test_na1d.py b/tests/test_na1d.py
index 7b295ed..a062b0f 100644
--- a/tests/test_na1d.py
+++ b/tests/test_na1d.py
@@ -10,8 +10,8 @@
 import torch
 from torch.utils.cpp_extension import CUDA_HOME
 from torch.autograd import gradcheck
-from natten import NeighborhoodAttention1D as NeighborhoodAttention
-from natten.functional import NATTEN1DAVFunction, NATTEN1DQKRPBFunction
+from natten import NeighborhoodAttention1D
+from natten.functional import natten1dqkrpb, natten1dav
 
 
 HAS_CUDA = (torch.cuda.is_available() and (CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1")
@@ -31,10 +31,10 @@ def _priv_test_gradcheck_natten1dqkrpb(
     query = torch.randn((batch_size, heads, length, dim), **kwargs)
     key = torch.randn((batch_size, heads, length, dim), **kwargs)
     rpb = torch.randn((heads, 2 * kernel_size - 1), **kwargs)
-    variables = [query, key, rpb, dilation]
+    variables = [query, key, rpb, kernel_size, dilation]
 
     assert gradcheck(
-            NATTEN1DQKRPBFunction.apply, 
+            natten1dqkrpb, 
             variables, 
             eps=eps, 
             atol=atol, 
@@ -58,7 +58,7 @@ def _priv_test_gradcheck_natten1dav(
     variables = [attn, value, dilation]
 
     assert gradcheck(
-            NATTEN1DAVFunction.apply, 
+            natten1dav, 
             variables, 
             eps=eps, 
             atol=atol, 
@@ -82,15 +82,15 @@ def _priv_test_allclose_cpu_cuda(batch_size, length,
                     'qkv_bias': True,
                 }
 
-                base_state_dict = NeighborhoodAttention(**model_kwargs).state_dict()
+                base_state_dict = NeighborhoodAttention1D(**model_kwargs).state_dict()
 
                 x1 = torch.randn((batch_size, length, dim * num_heads))
                 x2 = x1.clone().detach().cuda(0)
 
-                nat1 = NeighborhoodAttention(**model_kwargs).eval()
+                nat1 = NeighborhoodAttention1D(**model_kwargs).eval()
                 nat1.load_state_dict(base_state_dict, strict=False)
 
-                nat2 = NeighborhoodAttention(**model_kwargs).cuda(0).eval()
+                nat2 = NeighborhoodAttention1D(**model_kwargs).cuda(0).eval()
                 nat2.load_state_dict(base_state_dict, strict=False)
 
                 y1 = nat1(x1)
diff --git a/tests/test_na2d.py b/tests/test_na2d.py
index d05925e..735990f 100644
--- a/tests/test_na2d.py
+++ b/tests/test_na2d.py
@@ -10,8 +10,8 @@
 import torch
 from torch.utils.cpp_extension import CUDA_HOME
 from torch.autograd import gradcheck
-from natten import NeighborhoodAttention2D as NeighborhoodAttention
-from natten.functional import NATTEN2DAVFunction, NATTEN2DQKRPBFunction
+from natten import NeighborhoodAttention2D
+from natten.functional import natten2dqkrpb, natten2dav
 
 
 HAS_CUDA = (torch.cuda.is_available() and (CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1")
@@ -31,10 +31,10 @@ def _priv_test_gradcheck_natten2dqkrpb(
     query = torch.randn((batch_size, heads, height, width, dim), **kwargs)
     key = torch.randn((batch_size, heads, height, width, dim), **kwargs)
     rpb = torch.randn((heads, 2 * kernel_size - 1, 2 * kernel_size - 1), **kwargs)
-    variables = [query, key, rpb, dilation]
+    variables = [query, key, rpb, kernel_size, dilation]
 
     assert gradcheck(
-            NATTEN2DQKRPBFunction.apply, 
+            natten2dqkrpb, 
             variables, 
             eps=eps, 
             atol=atol, 
@@ -58,7 +58,7 @@ def _priv_test_gradcheck_natten2dav(
     variables = [attn, value, dilation]
 
     assert gradcheck(
-            NATTEN2DAVFunction.apply, 
+            natten2dav, 
             variables, 
             eps=eps, 
             atol=atol, 
@@ -82,15 +82,15 @@ def _priv_test_allclose_cpu_cuda(batch_size, height, width,
                     'qkv_bias': True,
                 }
 
-                base_state_dict = NeighborhoodAttention(**model_kwargs).state_dict()
+                base_state_dict = NeighborhoodAttention2D(**model_kwargs).state_dict()
 
                 x1 = torch.randn((batch_size, height, width, dim * num_heads))
                 x2 = x1.clone().detach().cuda(0)
 
-                nat1 = NeighborhoodAttention(**model_kwargs).eval()
+                nat1 = NeighborhoodAttention2D(**model_kwargs).eval()
                 nat1.load_state_dict(base_state_dict, strict=False)
 
-                nat2 = NeighborhoodAttention(**model_kwargs).cuda(0).eval()
+                nat2 = NeighborhoodAttention2D(**model_kwargs).cuda(0).eval()
                 nat2.load_state_dict(base_state_dict, strict=False)
 
                 y1 = nat1(x1)