def test_liger_fused_neighborhood_attention_functional_parameters(kernel_size, dilation):
    """Test the functional interface with different kernel sizes and dilations."""
    batch_size, num_heads, seq_len, head_dim = 1, 2, 24, 16

    query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)
    key = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)
    value = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)

    output = liger_fused_neighborhood_attention(query, key, value, kernel_size=kernel_size, dilation=dilation)

    expected_shape = (batch_size, num_heads, seq_len, head_dim)
    assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}"

    assert not torch.isnan(output).any(), "Output contains NaN values"
    assert not torch.isinf(output).any(), "Output contains Inf values"
Test the functional interface with different kernel sizes and dilations.
test/transformers/test_fused_neighborhood_attention.py