Module neuralops.tests.test_triton_gdn_scan

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

Coverage: - Numerical correctness vs PyTorch reference - Multiple dtypes: fp32, fp16, bf16 - Multiple shapes: batch, seq, head, dim variations - Edge cases: head_dim bounds - 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_gdn_scan_pytorch(q, k, v, beta, alpha)

Expand source code

def reference_gdn_scan_pytorch(q, k, v, beta, alpha):
    """
    Pure PyTorch sequential reference for Gated DeltaNet scan.
    Recurrence:
        inter_v_t = k_t @ S_{t-1}
        err_t     = v_t - inter_v_t
        norm_sq_t = sum(k_t^2)
        step_t    = alpha_t / max(1 + norm_sq_t, 1)
        S_t       = beta_t * S_{t-1} + step_t * k_t^T @ err_t
        y_t       = q_t @ S_t
    """
    B, H, L, D = q.shape
    q_f = q.float()
    k_f = k.float()
    v_f = v.float()
    beta_f = beta.float()
    alpha_f = alpha.float()

    y = torch.zeros(B, H, L, D, device=q.device, dtype=torch.float32)

    for b in range(B):
        for h in range(H):
            S = torch.zeros(D, D, device=q.device, dtype=torch.float32)
            for t in range(L):
                qt = q_f[b, h, t]          # (D,)
                kt = k_f[b, h, t]          # (D,)
                vt = v_f[b, h, t]          # (D,)
                bt = beta_f[b, h, t]       # scalar
                at = alpha_f[b, h, t]      # scalar

                inter_v = torch.matmul(kt, S)  # (D,)
                err = vt - inter_v
                norm_sq = torch.dot(kt, kt)
                step = at / max(1.0 + norm_sq.item(), 1.0)

                S = bt * S + step * torch.outer(kt, err)
                y[b, h, t] = torch.matmul(qt, S)

    return y.to(q.dtype)

Pure PyTorch sequential reference for Gated DeltaNet scan.

Recurrence -----= inter_v_t = k_t @ S_{t-1} err_t = v_t - inter_v_t norm_sq_t = sum(k_t^2) step_t = alpha_t / max(1 + norm_sq_t, 1) S_t = beta_t * S_{t-1} + step_t * k_t^T @ err_t y_t = q_t @ S_t

def test_gdn_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_gdn_scan_triton_backward(device):
    """Verify backward gradient computation."""
    B, H, L, D = 1, 2, 16, 32
    q = torch.randn(B, H, L, D, device=device, requires_grad=True)
    k = torch.randn(B, H, L, D, device=device, requires_grad=True)
    v = torch.randn(B, H, L, D, device=device, requires_grad=True)
    beta = torch.sigmoid(torch.randn(B, H, L, device=device, requires_grad=True))
    alpha = torch.sigmoid(torch.randn(B, H, L, device=device, requires_grad=True))

    out = gdn_scan_triton(q, k, v, beta, alpha)
    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()
    assert beta.grad is not None and torch.isfinite(beta.grad).all()
    assert alpha.grad is not None and torch.isfinite(alpha.grad).all()

Verify backward gradient computation.

def test_gdn_scan_triton_correctness(device, dtype, 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("D", [16, 32, 64])
def test_gdn_scan_triton_correctness(device, dtype, D):
    """Verify Triton GDN matches PyTorch reference on CUDA."""
    B, H, L = 2, 2, 16
    q = torch.randn(B, H, L, D, device=device, dtype=dtype)
    k = torch.randn(B, H, L, D, device=device, dtype=dtype)
    v = torch.randn(B, H, L, D, device=device, dtype=dtype)
    beta = torch.sigmoid(torch.randn(B, H, L, device=device, dtype=dtype))
    alpha = torch.sigmoid(torch.randn(B, H, L, device=device, dtype=dtype))

    out = gdn_scan_triton(q, k, v, beta, alpha)
    ref = reference_gdn_scan_pytorch(q, k, v, beta, alpha)

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

Verify Triton GDN matches PyTorch reference on CUDA.

def test_gdn_scan_triton_dim_exceeded(device)

Expand source code

@pytest.mark.skipif(not HAS_TRITON or not torch.cuda.is_available(), reason="Requires Triton and CUDA")
def test_gdn_scan_triton_dim_exceeded(device):
    """Verify exception raised when D > MAX_REGISTER_DIM."""
    B, H, L, D = 1, 1, 8, MAX_REGISTER_DIM + 16
    q = torch.randn(B, H, L, D, device=device)
    k = torch.randn(B, H, L, D, device=device)
    v = torch.randn(B, H, L, D, device=device)
    beta = torch.sigmoid(torch.randn(B, H, L, device=device))
    alpha = torch.sigmoid(torch.randn(B, H, L, device=device))

    with pytest.raises(ValueError, match="head_dim"):
        gdn_scan_triton(q, k, v, beta, alpha)

Verify exception raised when D > MAX_REGISTER_DIM.

def test_reference_gdn_scan_pytorch_basic()

Expand source code

def test_reference_gdn_scan_pytorch_basic():
    """Verify PyTorch reference runs and outputs correct shape on CPU."""
    B, H, L, D = 2, 4, 16, 32
    q = torch.randn(B, H, L, D)
    k = torch.randn(B, H, L, D)
    v = torch.randn(B, H, L, D)
    beta = torch.sigmoid(torch.randn(B, H, L))
    alpha = torch.sigmoid(torch.randn(B, H, L))

    out = reference_gdn_scan_pytorch(q, k, v, beta, alpha)
    assert out.shape == (B, H, L, D)
    assert torch.isfinite(out).all()

Verify PyTorch reference runs and outputs correct shape on CPU.