liger-kernel 0.5.10__py3-none-any.whl → 0.6.0__py3-none-any.whl

liger_kernel/ops/softmax.py

@@ -0,0 +1,201 @@
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+@triton.jit
+def _softmax_single_block_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < n_cols
+
+    x = tl.load(X_ptr + row_id * X_row_stride + offs, mask=mask, other=-float("inf"), cache_modifier=".ca")
+    m = tl.max(x, axis=0)
+    e = tl.exp(x - m)
+    d = tl.sum(e, axis=0)
+    y = e / d
+    tl.store(Y_ptr + row_id * Y_row_stride + offs, y, mask=mask, cache_modifier=".cs")
+
+
+@triton.jit
+def _softmax_multi_block_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+
+    m = tl.float32(-float("inf"))
+    d = tl.float32(0.0)
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        xblk = tl.load(X_ptr + row_id * X_row_stride + idx, mask=mask, other=-float("inf"), cache_modifier=".ca")
+        blk_max = tl.max(xblk, axis=0)
+        new_m = tl.max(m, blk_max)
+        d = d * tl.exp(m - new_m) + tl.sum(tl.exp(xblk - new_m), axis=0)
+        m = new_m
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        xblk = tl.load(X_ptr + row_id * X_row_stride + idx, mask=mask, other=-float("inf"), cache_modifier=".ca")
+        yblk = tl.exp(xblk - m) / d
+        tl.store(Y_ptr + row_id * Y_row_stride + idx, yblk, mask=mask, cache_modifier=".cs")
+
+
+@triton.jit
+def _softmax_single_block_backward_kernel(
+    dy_ptr,
+    dy_stride,
+    y_ptr,
+    y_stride,
+    dx_ptr,
+    dx_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < n_cols
+
+    dy = tl.load(dy_ptr + row_id * dy_stride + offs, mask=mask, other=0.0)
+    y = tl.load(y_ptr + row_id * y_stride + offs, mask=mask, other=0.0, cache_modifier=".ca")
+    dot = tl.sum(dy * y, axis=0)
+    dx = y * (dy - dot)
+    tl.store(dx_ptr + row_id * dx_stride + offs, dx, mask=mask, cache_modifier=".wb")
+
+
+@triton.jit
+def _softmax_multi_block_backward_kernel(
+    dy_ptr,
+    dy_stride,
+    y_ptr,
+    y_stride,
+    dx_ptr,
+    dx_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    acc = tl.float32(0.0)
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        dy_blk = tl.load(dy_ptr + row_id * dy_stride + idx, mask=mask, other=0.0)
+        y_blk = tl.load(y_ptr + row_id * y_stride + idx, mask=mask, other=0.0, cache_modifier=".ca")
+        acc += tl.sum(dy_blk * y_blk, axis=0)
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        dy_blk = tl.load(dy_ptr + row_id * dy_stride + idx, mask=mask, other=0.0)
+        y_blk = tl.load(y_ptr + row_id * y_stride + idx, mask=mask, other=0.0, cache_modifier=".ca")
+        dx_blk = y_blk * (dy_blk - acc)
+        tl.store(dx_ptr + row_id * dx_stride + idx, dx_blk, mask=mask, cache_modifier=".wb")
+
+
+def _softmax_forward(x: torch.Tensor) -> Tuple[torch.Tensor, int, int, bool]:
+    *batch, n_cols = x.shape
+    x2d = x.contiguous().view(-1, n_cols)
+    n_rows = x2d.shape[0]
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    y2d = torch.empty_like(x2d)
+
+    if n_cols <= BLOCK_SIZE:
+        _softmax_single_block_forward_kernel[(n_rows,)](
+            y2d, y2d.stride(0), x2d, x2d.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
+        )
+        multi_block_launch = False
+    else:
+        _softmax_multi_block_forward_kernel[(n_rows,)](
+            y2d, y2d.stride(0), x2d, x2d.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
+        )
+        multi_block_launch = True
+
+    return y2d.view(*batch, n_cols), BLOCK_SIZE, num_warps, multi_block_launch
+
+
+def _softmax_backward(
+    dy: torch.Tensor,
+    y: torch.Tensor,
+    BLOCK_SIZE: int,
+    num_warps: int,
+    multi_block_launch: bool,
+) -> torch.Tensor:
+    *batch, n_cols = dy.shape
+    dy2d = dy.contiguous().view(-1, n_cols)
+    y2d = y.contiguous().view(-1, n_cols)
+    n_rows = dy2d.shape[0]
+    dx2d = torch.empty_like(dy2d)
+
+    if not multi_block_launch and n_cols <= BLOCK_SIZE:
+        _softmax_single_block_backward_kernel[(n_rows,)](
+            dy2d,
+            dy2d.stride(0),
+            y2d,
+            y2d.stride(0),
+            dx2d,
+            dx2d.stride(0),
+            n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+    else:
+        _softmax_multi_block_backward_kernel[(n_rows,)](
+            dy2d,
+            dy2d.stride(0),
+            y2d,
+            y2d.stride(0),
+            dx2d,
+            dx2d.stride(0),
+            n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+
+    return dx2d.view(*batch, n_cols)
+
+
+class LigerSoftmaxFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, input_: torch.Tensor):
+        y, BLOCK_SIZE, num_warps, multi_block_launch = _softmax_forward(input_)
+        ctx.save_for_backward(y)
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        ctx.num_warps = num_warps
+        ctx.multi_block_launch = multi_block_launch
+        return y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        (y,) = ctx.saved_tensors
+        dx = _softmax_backward(
+            grad_output,
+            y,
+            ctx.BLOCK_SIZE,
+            ctx.num_warps,
+            ctx.multi_block_launch,
+        )
+        return dx

liger_kernel/ops/sparsemax.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 import triton
 import triton.language as tl
@@ -105,63 +107,73 @@ def _sparsemax_backward_kernel(
         tl.store(gi_row + offs_iter, gi_val.to(gi_row.dtype.element_ty), mask=mask_iter, cache_modifier=".wb")
 
 
+def _sparsemax_forward(x: torch.Tensor, dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    if dim < 0:
+        dim += x.dim()
+    x_sw = x.transpose(dim, -1).contiguous()
+    n_cols = x_sw.size(-1)
+    n_rows = x_sw.numel() // n_cols
+    x_flat = x_sw.view(n_rows, n_cols)
+    x_sorted_flat = torch.sort(x_flat.float(), dim=-1, descending=True).values
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    out_flat = torch.empty_like(x_flat)
+    grid = (n_rows,)
+    _sparsemax_forward_kernel[grid](
+        x_flat,
+        x_flat.stride(0),
+        x_sorted_flat,
+        x_sorted_flat.stride(0),
+        out_flat,
+        out_flat.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    y = out_flat.view_as(x_sw).transpose(dim, -1)
+    return y, out_flat
+
+
+def _sparsemax_backward(
+    grad_out: torch.Tensor,
+    out_flat: torch.Tensor,
+    dim: int,
+) -> torch.Tensor:
+    grad_sw = grad_out.transpose(dim, -1).contiguous()
+    n_cols = grad_sw.size(-1)
+    n_rows = grad_sw.numel() // n_cols
+    go_flat = grad_sw.view(n_rows, n_cols)
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    dx_flat = torch.empty_like(go_flat)
+    grid = (n_rows,)
+    _sparsemax_backward_kernel[grid](
+        out_flat,
+        go_flat,
+        dx_flat,
+        out_flat.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    dx = dx_flat.view_as(grad_sw).transpose(dim, -1)
+    return dx
+
+
 class LigerSparsemaxFunction(torch.autograd.Function):
     @staticmethod
     @ensure_contiguous
     def forward(ctx, x: torch.Tensor, dim: int):
-        if dim < 0:
-            dim += x.dim()
-        ctx.dim = dim
-
-        x_sw = x.transpose(dim, -1).contiguous()
-        n_cols = x_sw.size(-1)
-        n_rows = x_sw.numel() // n_cols
-        x_flat = x_sw.view(n_rows, n_cols)
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-        out_flat = torch.empty_like(x_flat)
-        grid = (n_rows,)
-
-        x_sorted_flat = torch.sort(x_flat.float(), dim=-1, descending=True).values
-
-        _sparsemax_forward_kernel[grid](
-            x_flat,
-            x_flat.stride(0),
-            x_sorted_flat,
-            x_sorted_flat.stride(0),
-            out_flat,
-            out_flat.stride(0),
-            n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
+        y, out_flat = _sparsemax_forward(x, dim)
         ctx.save_for_backward(out_flat)
-        return out_flat.view_as(x_sw).transpose(dim, -1)
+        ctx.dim = dim
+        return y
 
     @staticmethod
     @ensure_contiguous
     def backward(ctx, grad_out: torch.Tensor):
         (out_flat,) = ctx.saved_tensors
-        dim = ctx.dim
-
-        go_sw = grad_out.transpose(dim, -1).contiguous()
-        n_cols = go_sw.size(-1)
-        n_rows = go_sw.numel() // n_cols
-        go_flat = go_sw.view(n_rows, n_cols)
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-        gi_flat = torch.empty_like(go_flat)
-        grid = (n_rows,)
-
-        _sparsemax_backward_kernel[grid](
-            out_flat,
-            go_flat,
-            gi_flat,
-            out_flat.stride(0),
-            n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
-        return gi_flat.view_as(go_sw).transpose(dim, -1), None
+        dx = _sparsemax_backward(grad_out, out_flat, ctx.dim)
+        return dx, None

liger_kernel/ops/swiglu.py

@@ -26,7 +26,7 @@ def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BL
     # sigmoid requires type float32
     a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
     b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
-    c_row = silu(a_row) * b_row
+    c_row = silu(a_row).cast(b_row.dtype) * b_row
     tl.store(c_ptr + col_offsets, c_row, mask=mask)
 
 

liger_kernel/transformers/__init__.py

@@ -30,6 +30,7 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
@@ -87,6 +88,7 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_granite",
         "apply_liger_kernel_to_llama",
         "apply_liger_kernel_to_llava",
+        "apply_liger_kernel_to_llama4",
         "apply_liger_kernel_to_mistral",
         "apply_liger_kernel_to_mixtral",
         "apply_liger_kernel_to_mllama",
@@ -141,6 +143,7 @@ if _TRANSFORMERS_AVAILABLE:
             "apply_liger_kernel_to_granite",
             "apply_liger_kernel_to_llama",
             "apply_liger_kernel_to_llava",
+            "apply_liger_kernel_to_llama4",
             "apply_liger_kernel_to_mistral",
             "apply_liger_kernel_to_mixtral",
             "apply_liger_kernel_to_mllama",

liger_kernel/transformers/functional.py

@@ -4,14 +4,17 @@ from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
 from liger_kernel.ops.dyt import LigerDyTFunction
 from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
+from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
 from liger_kernel.ops.geglu import LigerGELUMulFunction
 from liger_kernel.ops.group_norm import LigerGroupNormFunction
 from liger_kernel.ops.jsd import LigerJSDFunction
 from liger_kernel.ops.kl_div import LigerKLDivLossFunction
 from liger_kernel.ops.layer_norm import LigerLayerNormFunction
+from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
 from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
 from liger_kernel.ops.rms_norm import LigerRMSNormFunction
 from liger_kernel.ops.rope import LigerRopeFunction
+from liger_kernel.ops.softmax import LigerSoftmaxFunction
 from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
 from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.tvd import LigerTVDLossFunction
@@ -167,6 +170,61 @@ def liger_sparsemax(
     return LigerSparsemaxFunction.apply(input, dim)
 
 
+def liger_multi_token_attention(
+    scores,
+    weight,
+    bias=None,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    sparse: bool = False,
+):
+    """
+    Functional interface for multi-token attention.
+
+    Args:
+        scores: Input tensor of shape (B, C_in, L, L)
+        weight: Convolution weight tensor of shape (C_out, C_in // groups, K, K)
+        bias: Optional bias tensor of shape (C_out,)
+        stride: Stride for the convolution (default: 1)
+        padding: Padding for the convolution (default: 0)
+        dilation: Dilation factor for the convolution (default: 1)
+        groups: Number of groups for the convolution (default: 1)
+        sparse: Specifies if input tensors are expected to be sparse (default: False)
+    Returns:
+        Output tensor after applying multi-token attention.
+    """
+    return LigerMultiTokenAttentionFunction.apply(scores, weight, bias, stride, padding, dilation, groups, sparse)
+
+
+def liger_fused_neighborhood_attention(
+    query,
+    key,
+    value,
+    kernel_size: int = 7,
+    dilation: int = 1,
+    scale: float = None,
+):
+    """
+    Liger fused neighborhood attention.
+
+    paper: https://arxiv.org/pdf/2504.16922
+
+    Args:
+        query: Query tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        key: Key tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        value: Value tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        kernel_size: Size of the neighborhood window (default: 7)
+        dilation: Dilation factor for the neighborhood (default: 1)
+        scale: Scaling factor for attention scores (default: rsqrt(head_dim))
+
+    Returns:
+        Output tensor of shape [batch_size, num_heads, seq_len, head_dim]
+    """
+    return LigerFusedNeighborhoodAttentionFunction.apply(query, key, value, kernel_size, dilation, scale)
+
+
 def liger_tvd(
     input,
     target,
@@ -203,5 +261,9 @@ def liger_swiglu(a, b):
     return LigerSiLUMulFunction.apply(a, b)
 
 
+def liger_softmax(x):
+    return LigerSoftmaxFunction.apply(x)
+
+
 def liger_dyt(x, alpha, gamma, beta):
     return LigerDyTFunction.apply(x, alpha, gamma, beta)

liger_kernel/transformers/fused_neighborhood_attention.py

@@ -0,0 +1,234 @@
+import math
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
+
+
+class LigerFusedNeighborhoodAttention(nn.Module):
+    """
+    Liger Fused Neighborhood Attention Module.
+
+    Paper: https://arxiv.org/pdf/2504.16922
+
+    Fused Neighborhood attention restricts the attention mechanism to a local neighborhood
+    around each position, reducing computational complexity from O(n²) to O(n*k)
+    where k is the neighborhood size.
+
+    Args:
+        hidden_size (int): The hidden dimension size
+        num_heads (int): Number of attention heads
+        kernel_size (int): Size of the neighborhood window (default: 7)
+        dilation (int): Dilation factor for the neighborhood (default: 1)
+        bias (bool): Whether to use bias in linear projections (default: True)
+        dropout (float): Dropout probability (default: 0.0)
+        scale (Optional[float]): Scaling factor for attention scores.
+                                If None, uses 1/sqrt(head_dim) (default: None)
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        kernel_size: int = 7,
+        dilation: int = 1,
+        bias: bool = True,
+        dropout: float = 0.0,
+        scale: Optional[float] = None,
+    ):
+        super().__init__()
+
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})")
+
+        if kernel_size <= 0:
+            raise ValueError(f"kernel_size ({kernel_size}) must be positive")
+
+        if kernel_size % 2 == 0:
+            raise ValueError(f"kernel_size ({kernel_size}) must be odd")
+
+        if dilation < 1:
+            raise ValueError(f"dilation ({dilation}) must be positive")
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.scale = scale if scale is not None else 1.0 / math.sqrt(self.head_dim)
+        self.dropout_p = dropout
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+        if dropout > 0.0:
+            self.dropout = nn.Dropout(dropout)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the fused neighborhood attention module.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor of shape [batch_size, seq_len, hidden_size]
+            attention_mask (Optional[torch.Tensor]): Attention mask (currently not supported)
+
+        Returns:
+            torch.Tensor: Output tensor of shape [batch_size, seq_len, hidden_size]
+        """
+        if attention_mask is not None:
+            raise NotImplementedError("Attention mask is not yet supported in LigerFusedNeighborhoodAttention")
+
+        batch_size, seq_len, hidden_size = hidden_states.shape
+
+        query = self.q_proj(hidden_states)
+        key = self.k_proj(hidden_states)
+        value = self.v_proj(hidden_states)
+
+        query = query.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn_output = LigerFusedNeighborhoodAttentionFunction.apply(
+            query, key, value, self.kernel_size, self.dilation, self.scale
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)
+
+        if self.dropout is not None:
+            attn_output = self.dropout(attn_output)
+
+        output = self.out_proj(attn_output)
+
+        return output
+
+    def extra_repr(self) -> str:
+        return (
+            f"hidden_size={self.hidden_size}, num_heads={self.num_heads}, "
+            f"head_dim={self.head_dim}, kernel_size={self.kernel_size}, "
+            f"dilation={self.dilation}, scale={self.scale}, dropout={self.dropout_p}"
+        )
+
+
+class LigerFusedNeighborhoodAttentionLayer(nn.Module):
+    """
+    A complete neighborhood attention layer with layer norm and residual connection.
+
+    Args:
+        hidden_size (int): The hidden dimension size
+        num_heads (int): Number of attention heads
+        kernel_size (int): Size of the neighborhood window (default: 7)
+        dilation (int): Dilation factor for the neighborhood (default: 1)
+        bias (bool): Whether to use bias in linear projections (default: True)
+        dropout (float): Dropout probability (default: 0.0)
+        layer_norm_eps (float): Epsilon for layer normalization (default: 1e-5)
+        scale (Optional[float]): Scaling factor for attention scores (default: None)
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        kernel_size: int = 7,
+        dilation: int = 1,
+        bias: bool = True,
+        dropout: float = 0.0,
+        layer_norm_eps: float = 1e-5,
+        scale: Optional[float] = None,
+    ):
+        super().__init__()
+
+        self.attention = LigerFusedNeighborhoodAttention(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            bias=bias,
+            dropout=dropout,
+            scale=scale,
+        )
+
+        self.layer_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+        if dropout > 0.0:
+            self.dropout = nn.Dropout(dropout)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass with residual connection and layer normalization.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor of shape [batch_size, seq_len, hidden_size]
+            attention_mask (Optional[torch.Tensor]): Attention mask (currently not supported)
+
+        Returns:
+            torch.Tensor: Output tensor of shape [batch_size, seq_len, hidden_size]
+        """
+        normed_hidden_states = self.layer_norm(hidden_states)
+
+        attn_output = self.attention(normed_hidden_states, attention_mask)
+
+        if self.dropout is not None:
+            attn_output = self.dropout(attn_output)
+
+        output = hidden_states + attn_output
+
+        return output
+
+
+class LigerFusedNeighborhoodAttentionConfig:
+    """
+    Configuration class for Fused Neighborhood Attention.
+
+    This can be used to easily configure neighborhood attention parameters
+    for different model architectures.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        num_heads: int = 12,
+        kernel_size: int = 7,
+        dilation: int = 1,
+        bias: bool = True,
+        dropout: float = 0.0,
+        layer_norm_eps: float = 1e-5,
+        scale: Optional[float] = None,
+    ):
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.bias = bias
+        self.dropout = dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.scale = scale
+
+    def to_dict(self):
+        return {
+            "hidden_size": self.hidden_size,
+            "num_heads": self.num_heads,
+            "kernel_size": self.kernel_size,
+            "dilation": self.dilation,
+            "bias": self.bias,
+            "dropout": self.dropout,
+            "layer_norm_eps": self.layer_norm_eps,
+            "scale": self.scale,
+        }