def fused_neighborhood_attention_forward(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    kernel_size: int = 7,
    dilation: int = 1,
    scale: float = None,
    return_lse: bool = False,
) -> tuple:
    """
    Fused neighborhood attention forward pass.

    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
        dilation: Dilation factor for the neighborhood
        scale: Scaling factor for attention scores (default: rsqrt(head_dim))
        return_lse: Whether to return log-sum-exp values

    Returns:
        Tuple of (output tensor, softmax parameters for backward)
    """
    batch_size, num_heads, seq_len, head_dim = query.shape

    if scale is None:
        scale = 1.0 / math.sqrt(head_dim)

    query = query.contiguous()
    key = key.contiguous()
    value = value.contiguous()

    output = torch.empty_like(query)
    qk_scores = torch.empty(batch_size, num_heads, seq_len, seq_len, device=query.device, dtype=query.dtype)

    mask = torch.zeros(seq_len, seq_len, device=query.device, dtype=torch.float32)

    BLOCK_SIZE, num_warps = calculate_settings(seq_len)
    BLOCK_SIZE_M = min(64, triton.next_power_of_2(seq_len))
    BLOCK_SIZE_N = min(64, triton.next_power_of_2(seq_len))
    BLOCK_SIZE_K = max(16, triton.next_power_of_2(head_dim))

    num_stages = 4 if seq_len >= 512 else 2

    grid_mask = (seq_len,)
    _neighborhood_mask_kernel[grid_mask](
        mask,
        seq_len,
        kernel_size,
        dilation,
        BLOCK_SIZE,
        num_stages,
        num_warps,
    )

    grid_qk = (batch_size * num_heads, triton.cdiv(seq_len, BLOCK_SIZE_M), triton.cdiv(seq_len, BLOCK_SIZE_N))
    _fused_neighborhood_attention_qk_kernel[grid_qk](
        query,
        key,
        qk_scores,
        mask,
        query.stride(0),
        query.stride(1),
        query.stride(2),
        query.stride(3),
        key.stride(0),
        key.stride(1),
        key.stride(2),
        key.stride(3),
        qk_scores.stride(0),
        qk_scores.stride(1),
        qk_scores.stride(2),
        qk_scores.stride(3),
        batch_size,
        num_heads,
        seq_len,
        head_dim,
        scale,
        kernel_size,
        dilation,
        BLOCK_SIZE_M,
        BLOCK_SIZE_N,
        BLOCK_SIZE_K,
        num_stages,
        num_warps,
    )

    qk_reshaped = qk_scores.view(batch_size * num_heads * seq_len, seq_len)
    attn_reshaped, BLOCK_SIZE_softmax, num_warps_softmax, multi_block_launch = _softmax_forward(qk_reshaped)
    attn_weights = attn_reshaped.view(batch_size, num_heads, seq_len, seq_len)

    grid_av = (batch_size * num_heads, triton.cdiv(seq_len, BLOCK_SIZE_M), triton.cdiv(head_dim, BLOCK_SIZE_N))
    _fused_neighborhood_attention_av_kernel[grid_av](
        attn_weights,
        value,
        output,
        attn_weights.stride(0),
        attn_weights.stride(1),
        attn_weights.stride(2),
        attn_weights.stride(3),
        value.stride(0),
        value.stride(1),
        value.stride(2),
        value.stride(3),
        output.stride(0),
        output.stride(1),
        output.stride(2),
        output.stride(3),
        batch_size,
        num_heads,
        seq_len,
        head_dim,
        BLOCK_SIZE_M,
        BLOCK_SIZE_N,
        BLOCK_SIZE_K,
        num_stages,
        num_warps,
    )

    if return_lse:
        raise NotImplementedError("return_lse=True is not supported yet.")

    softmax_params = (BLOCK_SIZE_softmax, num_warps_softmax, multi_block_launch)
    return output, attn_weights, softmax_params

    Fused neighborhood attention forward pass.

    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
        dilation: Dilation factor for the neighborhood
        scale: Scaling factor for attention scores (default: rsqrt(head_dim))
        return_lse: Whether to return log-sum-exp values

    Returns:
        Tuple of (output tensor, softmax parameters for backward)
    
src/liger_kernel/ops/fused_neighborhood_attention.py