| """Fused Triton scan for the PSSA selective diagonal state-space recurrence. |
| |
| h_t = exp(delta_t * A) * h_{t-1} + delta_t * b_t * x_t |
| y_t = <h_t, c_t> |
| |
| `A` is (D, DS) negative, `delta` is (B, T, D) positive, `b`/`c` are (B, T, DS). |
| The whole (forward + backward) runs as one kernel launch per chunk, with the |
| per-token state `h_t` materialised so the backward pass is a single reverse |
| sweep. Carries `h0` across chunks and returns `hT`. |
| """ |
| import torch |
| import triton |
| import triton.language as tl |
|
|
| __all__ = ["ScanFn", "scan_reference", "selftest"] |
|
|
|
|
| @triton.jit |
| def scan_fwd(x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, h0_ptr, y_ptr, hT_ptr, H_ptr, |
| T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr): |
| b = tl.program_id(0) |
| pd = tl.program_id(1) |
| od = pd * BD + tl.arange(0, BD) |
| os_ = tl.arange(0, DS) |
| A = tl.load(A_ptr + od[:, None] * DS + os_[None, :]) |
| h = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :]) |
| for t in range(T): |
| x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32) |
| dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32) |
| bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32) |
| cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32) |
| h = tl.exp(dl[:, None] * A) * h + dl[:, None] * bp[None, :] * x[:, None] |
| tl.store(y_ptr + b * T * D + t * D + od, tl.sum(h * cp[None, :], axis=1)) |
| tl.store(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :], h) |
| tl.store(hT_ptr + b * D * DS + od[:, None] * DS + os_[None, :], h) |
|
|
|
|
| @triton.jit |
| def scan_bwd(dy_ptr, x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, H_ptr, h0_ptr, |
| gh0_ptr, gx_ptr, gd_ptr, gbp_ptr, gcp_ptr, gA_ptr, |
| T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr): |
| b = tl.program_id(0) |
| pd = tl.program_id(1) |
| od = pd * BD + tl.arange(0, BD) |
| os_ = tl.arange(0, DS) |
| A = tl.load(A_ptr + od[:, None] * DS + os_[None, :]) |
| g = tl.zeros((BD, DS), tl.float32) |
| gA = tl.zeros((BD, DS), tl.float32) |
| for t in range(T - 1, -1, -1): |
| dy = tl.load(dy_ptr + b * T * D + t * D + od).to(tl.float32) |
| x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32) |
| dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32) |
| bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32) |
| cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32) |
| h = tl.load(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :]) |
| if t < T - 1: |
| dln = tl.load(d_ptr + b * T * D + (t + 1) * D + od).to(tl.float32) |
| g = dy[:, None] * cp[None, :] + tl.exp(dln[:, None] * A) * g |
| else: |
| g = dy[:, None] * cp[None, :] |
| if t > 0: |
| hp = tl.load(H_ptr + b * T * D * DS + (t - 1) * D * DS + od[:, None] * DS + os_[None, :]) |
| else: |
| hp = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :]) |
| a = tl.exp(dl[:, None] * A) |
| ga = g * hp |
| gd = tl.sum(ga * a * A, axis=1) + tl.sum(g * bp[None, :] * x[:, None], axis=1) |
| gbp = tl.sum(g * dl[:, None] * x[:, None], axis=0) |
| gcp = tl.sum(dy[:, None] * h, axis=0) |
| gx = tl.sum(g * dl[:, None] * bp[None, :], axis=1) |
| gA += ga * a * dl[:, None] |
| if t == 0: |
| tl.store(gh0_ptr + b * D * DS + od[:, None] * DS + os_[None, :], a * g) |
| tl.store(gx_ptr + b * T * D + t * D + od, gx) |
| tl.store(gd_ptr + b * T * D + t * D + od, gd) |
| tl.atomic_add(gbp_ptr + b * T * DS + t * DS + os_, gbp) |
| tl.atomic_add(gcp_ptr + b * T * DS + t * DS + os_, gcp) |
| tl.atomic_add(gA_ptr + od[:, None] * DS + os_[None, :], gA) |
|
|
|
|
| class ScanFn(torch.autograd.Function): |
| """Differentiable fused scan. Returns (y, hT); hT is not differentiated.""" |
|
|
| @staticmethod |
| def forward(ctx, x, delta, bp, cp, A, h0, bd=16): |
| B, T, D = x.shape |
| DS = A.shape[1] |
| y = torch.empty(B, T, D, device=x.device, dtype=torch.float32) |
| hT = torch.empty(B, D, DS, device=x.device, dtype=torch.float32) |
| H = torch.empty(B, T, D, DS, device=x.device, dtype=torch.float32) |
| scan_fwd[(B, triton.cdiv(D, bd))]( |
| x, delta, bp, cp, A, h0, y, hT, H, T=T, D=D, DS=DS, BD=bd) |
| ctx.save_for_backward(x, delta, bp, cp, A, H, h0) |
| ctx.bd = bd |
| ctx.mark_non_differentiable(hT) |
| return y, hT |
|
|
| @staticmethod |
| def backward(ctx, dy, dhT): |
| x, delta, bp, cp, A, H, h0 = ctx.saved_tensors |
| B, T, D = x.shape |
| DS = A.shape[1] |
| bd = ctx.bd |
| gx = torch.empty(B, T, D, device=x.device, dtype=torch.float32) |
| gd = torch.empty_like(delta, dtype=torch.float32) |
| gbp = torch.zeros_like(bp, dtype=torch.float32) |
| gcp = torch.zeros_like(cp, dtype=torch.float32) |
| gA = torch.zeros_like(A, dtype=torch.float32) |
| gh0 = torch.zeros(B, D, DS, device=x.device, dtype=torch.float32) |
| scan_bwd[(B, triton.cdiv(D, bd))]( |
| dy.contiguous(), x, delta, bp, cp, A, H, h0, gh0, |
| gx, gd, gbp, gcp, gA, T=T, D=D, DS=DS, BD=bd) |
| return gx.to(x.dtype), gd, gbp, gcp, gA, gh0.to(h0.dtype), None |
|
|
|
|
| def scan_reference(x, delta, bp, cp, A, h0): |
| """Plain PyTorch loop, for tests. Returns (y, H, hT).""" |
| B, T, D = x.shape |
| h = h0 |
| ys, hs = [], [] |
| for t in range(T): |
| h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None] |
| ys.append((h * cp[:, t, None, :]).sum(-1)) |
| hs.append(h) |
| return torch.stack(ys, 1), torch.stack(hs, 1), h |
|
|
|
|
| def selftest(ds=16, d=64, b=2, t=7, verbose=True): |
| torch.manual_seed(1) |
| dev = "cuda" |
| x = torch.randn(b, t, d, device=dev, requires_grad=True) |
| delta = torch.nn.functional.softplus(torch.randn(b, t, d, device=dev)).requires_grad_(True) |
| bp = torch.randn(b, t, ds, device=dev, requires_grad=True) |
| cp = torch.randn(b, t, ds, device=dev, requires_grad=True) |
| A = (-torch.nn.functional.softplus(torch.randn(d, ds, device=dev))).requires_grad_(True) |
| h0 = (torch.randn(b, d, ds, device=dev) * 0.5).requires_grad_(True) |
| yr, _, _ = scan_reference(x, delta, bp, cp, A, h0) |
| y, hT = ScanFn.apply(x, delta, bp, cp, A, h0) |
| gy = torch.randn_like(yr) |
| gref = torch.autograd.grad(yr, (x, delta, bp, cp, A, h0), gy) |
| gtri = torch.autograd.grad(y, (x, delta, bp, cp, A, h0), gy) |
| out = {"fwd": ((y - yr).abs().max() / yr.abs().max()).item()} |
| for nm, a, bb in zip(["dx", "dd", "dbp", "dcp", "dA", "dh0"], gtri, gref): |
| out[nm] = ((a - bb).abs().max() / bb.abs().max()).item() |
| if verbose: |
| print("scan selftest rel errs: " + " ".join(f"{k}={v:.1e}" for k, v in out.items())) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| selftest() |
|
|
