Module neuralops.tests.test_triton_state_update

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

Coverage: - Numerical correctness vs PyTorch SSM recurrence reference - Multiple dtypes: fp32, fp16, bf16 - Multiple shapes: batch, seq, state dimension - Edge cases: initial state provided vs None, D feedthrough present vs None - Reset behavior and sequential carry-over - Fallback path when Triton unavailable - Determinism

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_ssm_recurrence_pytorch(u, delta, A, B, C, D=None, initial_state=None)

Expand source code

def reference_ssm_recurrence_pytorch(u, delta, A, B, C, D=None, initial_state=None):
    """Pure PyTorch reference implementation of continuous SSM recurrence."""
    B_sz, L_sz, D_sz = u.shape
    N_sz = A.shape[-1]
    
    if initial_state is not None:
        state = initial_state.clone()
    else:
        state = torch.zeros(B_sz, D_sz, N_sz, device=u.device, dtype=torch.float32)

    delta_c = delta.expand(B_sz, L_sz, D_sz) if delta.dim() == 2 or delta.shape[-1] != D_sz else delta
    b_c = B.expand(B_sz, L_sz, N_sz) if B.dim() == 2 else B
    c_c = C.expand(B_sz, L_sz, N_sz) if C.dim() == 2 else C

    y = torch.zeros(B_sz, L_sz, D_sz, device=u.device, dtype=torch.float32)
    u_f = u.float()
    delta_f = delta_c.float()
    A_f = A.float()
    B_f = b_c.float()
    C_f = c_c.float()

    for t in range(L_sz):
        # discrete dA = exp(delta_t * A)
        d_t = delta_f[:, t, :, None]  # [B, D, 1]
        dA = torch.exp(d_t * A_f)     # [B, D, N]
        dB = d_t * B_f[:, t, None, :] # [B, D, N]
        u_t = u_f[:, t, :, None]      # [B, D, 1]

        state = state * dA + dB * u_t
        # y_t = state @ C_t
        y[:, t, :] = torch.sum(state * C_f[:, t, None, :], dim=-1)

    if D is not None:
        d_val = D.float().view(1, 1, D_sz) if D.dim() <= 2 else D.float()
        y = y + u_f * d_val

    return y.to(u.dtype), state.to(u.dtype)

Pure PyTorch reference implementation of continuous SSM recurrence.

def test_state_carryover(device)

Expand source code

def test_state_carryover(device):
    """Running step 1 then step 2 with state carry-over matches full sequence."""
    B, L, D_sz, N_sz = 2, 16, 16, 8
    u = torch.randn(B, L, D_sz, device=device, dtype=torch.float32)
    delta = torch.randn(B, L, D_sz, device=device, dtype=torch.float32).sigmoid() * 0.1
    A = -torch.rand(D_sz, N_sz, device=device, dtype=torch.float32)
    B_t = torch.randn(B, L, N_sz, device=device, dtype=torch.float32)
    C_t = torch.randn(B, L, N_sz, device=device, dtype=torch.float32)

    # Full sequence run
    y_full, state_full = fwd(u, delta, A, B_t, C_t)

    # Split into 2 chunks of 8
    y_part1, state_mid = fwd(u[:, :8], delta[:, :8], A, B_t[:, :8], C_t[:, :8])
    y_part2, state_end = fwd(u[:, 8:], delta[:, 8:], A, B_t[:, 8:], C_t[:, 8:], state=state_mid)

    assert torch.allclose(y_full[:, :8], y_part1, atol=1e-4)
    assert torch.allclose(y_full[:, 8:], y_part2, atol=1e-4)
    assert torch.allclose(state_full, state_end, atol=1e-4)

Running step 1 then step 2 with state carry-over matches full sequence.

def test_state_update_correctness(device, dtype, batch, seq_len, d_dim, n_dim)

Expand source code

@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
@pytest.mark.parametrize("batch", [1, 2])
@pytest.mark.parametrize("seq_len", [8, 32])
@pytest.mark.parametrize("d_dim", [16, 32])
@pytest.mark.parametrize("n_dim", [8, 16])
def test_state_update_correctness(device, dtype, batch, seq_len, d_dim, n_dim):
    """Verify state update recurrence matches PyTorch reference."""
    u = torch.randn(batch, seq_len, d_dim, device=device, dtype=dtype)
    delta = torch.randn(batch, seq_len, d_dim, device=device, dtype=dtype).sigmoid() * 0.1
    A = -torch.rand(d_dim, n_dim, device=device, dtype=dtype)
    B = torch.randn(batch, seq_len, n_dim, device=device, dtype=dtype)
    C = torch.randn(batch, seq_len, n_dim, device=device, dtype=dtype)
    D = torch.randn(d_dim, device=device, dtype=dtype)

    y, next_state = fwd(u, delta, A, B, C, D=D)
    y_ref, state_ref = reference_ssm_recurrence_pytorch(u, delta, A, B, C, D=D)

    assert y.shape == (batch, seq_len, d_dim)
    assert next_state.shape == (batch, d_dim, n_dim)
    assert torch.allclose(y, y_ref, atol=1e-2, rtol=1e-2)
    assert torch.allclose(next_state, state_ref, atol=1e-2, rtol=1e-2)

Verify state update recurrence matches PyTorch reference.

def test_state_update_cpu_fallback()

Expand source code

def test_state_update_cpu_fallback():
    """Verify PyTorch fallback works seamlessly on CPU."""
    u = torch.randn(2, 8, 16, device="cpu")
    delta = torch.randn(2, 8, 16, device="cpu").sigmoid() * 0.1
    A = -torch.rand(16, 8, device="cpu")
    B = torch.randn(2, 8, 8, device="cpu")
    C = torch.randn(2, 8, 8, device="cpu")

    y, state = state_update_pytorch_fallback(u, delta, A, B, C)
    assert y.shape == (2, 8, 16)
    assert state.shape == (2, 16, 8)
    assert torch.isfinite(y).all()
    assert torch.isfinite(state).all()

Verify PyTorch fallback works seamlessly on CPU.

def test_state_update_single_token(device)

Expand source code

def test_state_update_single_token(device):
    """Single token sequence (seq_len = 1)."""
    u = torch.randn(1, 1, 16, device=device)
    delta = torch.randn(1, 1, 16, device=device).sigmoid()
    A = -torch.rand(16, 8, device=device)
    B = torch.randn(1, 1, 8, device=device)
    C = torch.randn(1, 1, 8, device=device)
    y, state = fwd(u, delta, A, B, C)
    assert y.shape == (1, 1, 16)
    assert state.shape == (1, 16, 8)

Single token sequence (seq_len = 1).