Module neuralops.tests.test_triton_selective_scan

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

Coverage: - Numerical correctness vs PyTorch reference - Multiple dtypes: fp32, fp16, bf16 - Multiple shapes: batch, seq, head, dim variations - Edge cases: with and without D skip connection, delta_softplus - Fallback path & error handling - Determinism - Gradient correctness

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_selective_scan_pytorch(u, delta, A, B, C, D=None, delta_softplus=True)

Expand source code

def reference_selective_scan_pytorch(u, delta, A, B, C, D=None, delta_softplus=True):
    """
    Pure PyTorch sequential reference for Mamba selective scan.
    Recurrence:
        delta_val = softplus(delta_t) if delta_softplus else delta_t
        dA_t = exp(delta_val * A)
        dB_t = (delta_val * u_t) * B_t
        h_t  = dA_t * h_{t-1} + dB_t
        y_t  = sum_N(h_t * C_t) + (D * u_t if D is not None else 0)
    """
    B_size, L, D_dim = u.shape
    N = A.shape[-1]

    u_f = u.float()
    delta_f = F.softplus(delta.float()) if delta_softplus else delta.float()
    A_f = A.float()
    B_f = B.float()
    C_f = C.float()
    D_f = D.float() if D is not None else None

    y = torch.zeros(B_size, L, D_dim, device=u.device, dtype=torch.float32)

    for b in range(B_size):
        h = torch.zeros(D_dim, N, device=u.device, dtype=torch.float32)
        for t in range(L):
            u_t = u_f[b, t, :]       # (D,)
            d_t = delta_f[b, t, :]   # (D,)
            B_t = B_f[b, t, :]       # (N,)
            C_t = C_f[b, t, :]       # (N,)

            dA = torch.exp(d_t.unsqueeze(-1) * A_f)  # (D, N)
            dB = (d_t * u_t).unsqueeze(-1) * B_t.unsqueeze(0)  # (D, N)
            h = dA * h + dB

            yt = (h * C_t.unsqueeze(0)).sum(dim=-1)  # (D,)
            if D_f is not None:
                yt = yt + D_f * u_t
            y[b, t, :] = yt

    return y.to(u.dtype)

Pure PyTorch sequential reference for Mamba selective scan.

Recurrence -----= delta_val = softplus(delta_t) if delta_softplus else delta_t dA_t = exp(delta_val * A) dB_t = (delta_val * u_t) * B_t h_t = dA_t * h_{t-1} + dB_t y_t = sum_N(h_t * C_t) + (D * u_t if D is not None else 0)

def test_reference_selective_scan_pytorch_basic()

Expand source code

def test_reference_selective_scan_pytorch_basic():
    """Verify PyTorch reference runs and outputs correct shape on CPU."""
    B, L, D_dim, N = 2, 16, 16, 8
    u = torch.randn(B, L, D_dim)
    delta = torch.randn(B, L, D_dim)
    A = -torch.rand(D_dim, N)
    B_t = torch.randn(B, L, N)
    C_t = torch.randn(B, L, N)
    D = torch.randn(D_dim)

    out = reference_selective_scan_pytorch(u, delta, A, B_t, C_t, D=D)
    assert out.shape == (B, L, D_dim)
    assert torch.isfinite(out).all()

Verify PyTorch reference runs and outputs correct shape on CPU.

def test_selective_scan_triton_backward(device)

Expand source code

@pytest.mark.skipif(not HAS_TRITON or not torch.cuda.is_available(), reason="Requires Triton and CUDA")
def test_selective_scan_triton_backward(device):
    """Verify backward gradient computation."""
    B, L, D_dim, N = 1, 16, 16, 16
    u = torch.randn(B, L, D_dim, device=device, requires_grad=True)
    delta = torch.randn(B, L, D_dim, device=device, requires_grad=True)
    A = -torch.rand(D_dim, N, device=device, requires_grad=True)
    B_t = torch.randn(B, L, N, device=device, requires_grad=True)
    C_t = torch.randn(B, L, N, device=device, requires_grad=True)
    D = torch.randn(D_dim, device=device, requires_grad=True)

    out = selective_scan_fn(u, delta, A, B_t, C_t, D=D)
    loss = out.sum()
    loss.backward()

    assert u.grad is not None and torch.isfinite(u.grad).all()
    assert delta.grad is not None and torch.isfinite(delta.grad).all()
    assert A.grad is not None and torch.isfinite(A.grad).all()
    assert B_t.grad is not None and torch.isfinite(B_t.grad).all()
    assert C_t.grad is not None and torch.isfinite(C_t.grad).all()
    assert D.grad is not None and torch.isfinite(D.grad).all()

Verify backward gradient computation.

def test_selective_scan_triton_correctness(device, dtype, has_D)

Expand source code

@pytest.mark.skipif(not HAS_TRITON or not torch.cuda.is_available(), reason="Requires Triton and CUDA")
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("has_D", [True, False])
def test_selective_scan_triton_correctness(device, dtype, has_D):
    """Verify Triton selective scan matches PyTorch reference on CUDA."""
    B, L, D_dim, N = 2, 16, 16, 16
    u = torch.randn(B, L, D_dim, device=device, dtype=dtype)
    delta = torch.randn(B, L, D_dim, device=device, dtype=dtype)
    A = -torch.rand(D_dim, N, device=device, dtype=dtype)
    B_t = torch.randn(B, L, N, device=device, dtype=dtype)
    C_t = torch.randn(B, L, N, device=device, dtype=dtype)
    D = torch.randn(D_dim, device=device, dtype=dtype) if has_D else None

    out = selective_scan_fn(u, delta, A, B_t, C_t, D=D)
    ref = reference_selective_scan_pytorch(u, delta, A, B_t, C_t, D=D)

    tol = 1e-2 if dtype != torch.float32 else 1e-4
    assert torch.allclose(out, ref, atol=tol, rtol=tol)

Verify Triton selective scan matches PyTorch reference on CUDA.