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

        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]
        
src/liger_kernel/transformers/fused_neighborhood_attention.py