Module neuralops.tests.test_triton_fms_scan

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

Coverage: - Numerical correctness vs PyTorch reference - Multiple dtypes: fp32, fp16, bf16 - Multiple shapes: batch, seq, head, dim variations - Edge cases: single-token, seq_len > chunk_size, 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_fms_scan_pytorch(q, k, v, decay)

Expand source code

def reference_fms_scan_pytorch(q, k, v, decay):
    """
    Pure PyTorch sequential reference for Factorised Matrix State scan.
    Recurrence:
        q_feat = ELU(q) + 1.0
        k_feat = (ELU(k) + 1.0) / sqrt(D)
        S_t    = decay_h * S_{t-1} + k_feat_t^T @ v_t    # (D, D) matrix state
        Z_t    = decay_h * Z_{t-1} + k_feat_t            # (D,) vector normalizer
        num_t  = q_feat_t @ S_t                          # (D,)
        den_t  = q_feat_t @ Z_t + 1e-5                   # scalar
        y_t    = num_t / den_t                           # (D,)
    """
    B, H, L, D = q.shape
    q_f = (F.elu(q.float()) + 1.0)
    k_f = (F.elu(k.float()) + 1.0) / (D ** 0.5)
    v_f = v.float()
    decay_f = decay.float()

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

    for b in range(B):
        for h in range(H):
            d_h = decay_f[h]
            S = torch.zeros(D, D, device=q.device, dtype=torch.float32)
            Z = torch.zeros(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,)

                S = d_h * S + torch.outer(kt, vt)
                Z = d_h * Z + kt

                num = torch.matmul(qt, S)
                den = torch.dot(qt, Z) + 1e-5
                y[b, h, t] = num / den

    return y.to(q.dtype)

Pure PyTorch sequential reference for Factorised Matrix State scan.

Recurrence -----= q_feat = ELU(q) + 1.0 k_feat = (ELU(k) + 1.0) / sqrt(D) S_t = decay_h * S_{t-1} + k_feat_t^T @ v_t # (D, D) matrix state Z_t = decay_h * Z_{t-1} + k_feat_t # (D,) vector normalizer num_t = q_feat_t @ S_t # (D,) den_t = q_feat_t @ Z_t + 1e-5 # scalar y_t = num_t / den_t # (D,)

def test_fms_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_fms_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)
    decay = torch.sigmoid(torch.randn(H, device=device, requires_grad=True))

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

Verify backward gradient computation.

def test_fms_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_fms_scan_triton_correctness(device, dtype, D):
    """Verify Triton FMS matches PyTorch reference on CUDA."""
    B, H, L = 2, 2, 32
    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)
    decay = torch.sigmoid(torch.randn(H, device=device, dtype=dtype))

    out = fms_scan_triton(q, k, v, decay)
    ref = reference_fms_scan_pytorch(q, k, v, decay)

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

Verify Triton FMS matches PyTorch reference on CUDA.

def test_fms_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_fms_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)
    decay = torch.sigmoid(torch.randn(H, device=device))

    with pytest.raises(ValueError, match="head_dim"):
        fms_scan_triton(q, k, v, decay)

Verify exception raised when D > MAX_REGISTER_DIM.

def test_reference_fms_scan_pytorch_basic()

Expand source code

def test_reference_fms_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)
    decay = torch.sigmoid(torch.randn(H))

    out = reference_fms_scan_pytorch(q, k, v, decay)
    assert out.shape == (B, H, L, D)
    assert torch.isfinite(out).all()

Verify PyTorch reference runs and outputs correct shape on CPU.