Module neuralops.core.triton_kernels.state_update

Triton State-Space Model (SSM) & RWKV State Update Kernel.

Computes discrete sequential recurrence: h_t = exp(delta_t * A) * h_{t-1} + (delta_t * B_t) * u_t y_t = (h_t @ C_t) + D * u_t

Used by: - Mamba, S4, S6, RWKV, HGRN2, Jamba - Hybrid SSM-Attention architectures

Functions

def bwd(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
grad_out: torch.Tensor,
D: torch.Tensor | None = None,
state: torch.Tensor | None = None,
eps: float = 1e-06,
**kwargs: Any) ‑> tuple[torch.Tensor, ...]

Expand source code

def bwd(
    u: torch.Tensor,
    delta: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    grad_out: torch.Tensor,
    D: torch.Tensor | None = None,
    state: torch.Tensor | None = None,
    eps: float = 1e-6,
    **kwargs: Any,
) -> tuple[torch.Tensor, ...]:
    """Backward entrypoint computing gradients for SSM inputs."""
    u_req = u.detach().requires_grad_(True)
    delta_req = delta.detach().requires_grad_(True)
    A_req = A.detach().requires_grad_(True)
    B_req = B.detach().requires_grad_(True)
    C_req = C.detach().requires_grad_(True)
    with torch.enable_grad():
        out, _ = fwd(u_req, delta_req, A_req, B_req, C_req, D=D, state=state, eps=eps, **kwargs)
        grads = torch.autograd.grad(out, [u_req, delta_req, A_req, B_req, C_req], grad_out)
    return grads

Backward entrypoint computing gradients for SSM inputs.

def fwd(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
state: torch.Tensor | None = None,
eps: float = 1e-06,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def fwd(
    u: torch.Tensor,
    delta: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor | None = None,
    state: torch.Tensor | None = None,
    eps: float = 1e-6,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Forward entrypoint for fused Triton SSM state updates: (y, new_state)."""
    if not is_supported(u, delta, A, B, C, D=D, **kwargs):
        return state_update_pytorch_fallback(u, delta, A, B, C, D=D, state=state, eps=eps, **kwargs)
    return _TritonStateUpdateAutogradFunction.apply(u, delta, A, B, C, D, state, eps)

Forward entrypoint for fused Triton SSM state updates: (y, new_state).

def is_supported(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
**kwargs: Any) ‑> bool

Expand source code

def is_supported(
    u: torch.Tensor,
    delta: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor | None = None,
    **kwargs: Any,
) -> bool:
    """Check whether SSM state update can run on the Triton kernel."""
    suitable, _ = check_triton_suitability(u, "state_update")
    if not suitable:
        return False
    N_sz = A.shape[-1] if A.dim() >= 2 else 0
    return N_sz in (8, 16, 32, 64) and A.is_cuda and B.is_cuda and C.is_cuda

Check whether SSM state update can run on the Triton kernel.