| """ |
| Parallel Scan (Blelloch algorithm) in pure PyTorch for ARG selective SSM recurrence. |
| |
| Eliminates all external dependencies (mamba-ssm, selective_scan_cuda). |
| Runs in 2*log2(T) parallel steps (11 steps for T=2048) on any GPU. |
| 100% mathematically exact equivalence to sequential SSM recurrence. |
| """ |
|
|
| from __future__ import annotations |
| import math |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def npo2(length: int) -> int: |
| """Returns the next power of 2 above length.""" |
| return 2 ** math.ceil(math.log2(length)) |
|
|
|
|
| def pad_npo2(x: torch.Tensor) -> torch.Tensor: |
| """Pads input length dimension (dim 1) to the next power of 2.""" |
| length = x.size(1) |
| len_p2 = npo2(length) |
| if len_p2 == length: |
| return x |
| pad_tuple = (0, 0, 0, 0, 0, len_p2 - length) |
| return F.pad(x, pad_tuple, "constant", 0) |
|
|
|
|
| class PScan(torch.autograd.Function): |
| """ |
| Computes the parallel associative scan: |
| H[t] = A[t] * H[t-1] + X[t] with H[-1] = 0 |
| in O(log2(T)) parallel steps via Blelloch tree reduction. |
| """ |
|
|
| @staticmethod |
| def pscan(A: torch.Tensor, X: torch.Tensor) -> None: |
| |
| |
| B, D, L, _ = A.size() |
| num_steps = int(math.log2(L)) |
|
|
| |
| Aa = A |
| Xa = X |
| for _ in range(num_steps - 2): |
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T // 2, 2, -1) |
| Xa = Xa.view(B, D, T // 2, 2, -1) |
|
|
| Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0])) |
| Aa[:, :, :, 1].mul_(Aa[:, :, :, 0]) |
|
|
| Aa = Aa[:, :, :, 1] |
| Xa = Xa[:, :, :, 1] |
|
|
| if Xa.size(2) == 4: |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0])) |
| Aa[:, :, 1].mul_(Aa[:, :, 0]) |
| Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1]))) |
| elif Xa.size(2) == 2: |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0])) |
| return |
| else: |
| return |
|
|
| |
| Aa = A[:, :, 2 ** (num_steps - 2) - 1 : L : 2 ** (num_steps - 2)] |
| Xa = X[:, :, 2 ** (num_steps - 2) - 1 : L : 2 ** (num_steps - 2)] |
| Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1])) |
| Aa[:, :, 2].mul_(Aa[:, :, 1]) |
|
|
| for k in range(num_steps - 3, -1, -1): |
| stride = 2 ** k |
| Aa = A[:, :, stride - 1 : L : stride] |
| Xa = X[:, :, stride - 1 : L : stride] |
|
|
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T // 2, 2, -1) |
| Xa = Xa.view(B, D, T // 2, 2, -1) |
|
|
| Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1])) |
| Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1]) |
|
|
| @staticmethod |
| def pscan_rev(A: torch.Tensor, X: torch.Tensor) -> None: |
| |
| B, D, L, _ = A.size() |
| num_steps = int(math.log2(L)) |
|
|
| |
| Aa = A |
| Xa = X |
| for _ in range(num_steps - 2): |
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T // 2, 2, -1) |
| Xa = Xa.view(B, D, T // 2, 2, -1) |
|
|
| Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1])) |
| Aa[:, :, :, 0].mul_(Aa[:, :, :, 1]) |
|
|
| Aa = Aa[:, :, :, 0] |
| Xa = Xa[:, :, :, 0] |
|
|
| if Xa.size(2) == 4: |
| Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3])) |
| Aa[:, :, 2].mul_(Aa[:, :, 3]) |
| Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2])))) |
| elif Xa.size(2) == 2: |
| Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1])) |
| return |
| else: |
| return |
|
|
| |
| stride_top = 2 ** (num_steps - 2) |
| Aa = A[:, :, 0 : L : stride_top] |
| Xa = X[:, :, 0 : L : stride_top] |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2])) |
| Aa[:, :, 1].mul_(Aa[:, :, 2]) |
|
|
| for k in range(num_steps - 3, -1, -1): |
| stride = 2 ** k |
| Aa = A[:, :, 0 : L : stride] |
| Xa = X[:, :, 0 : L : stride] |
|
|
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T // 2, 2, -1) |
| Xa = Xa.view(B, D, T // 2, 2, -1) |
|
|
| Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0])) |
| Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0]) |
|
|
| @staticmethod |
| def forward(ctx, A_in: torch.Tensor, X_in: torch.Tensor) -> torch.Tensor: |
| |
| |
| L = X_in.size(1) |
|
|
| if L == npo2(L): |
| A = A_in.clone() |
| X = X_in.clone() |
| else: |
| A = pad_npo2(A_in).clone() |
| X = pad_npo2(X_in).clone() |
|
|
| A = A.transpose(2, 1).contiguous() |
| X = X.transpose(2, 1).contiguous() |
|
|
| PScan.pscan(A, X) |
|
|
| ctx.save_for_backward(A_in, X) |
| return X.transpose(2, 1)[:, :L].contiguous() |
|
|
| @staticmethod |
| def backward(ctx, grad_output_in: torch.Tensor): |
| A_in, X = ctx.saved_tensors |
| L = grad_output_in.size(1) |
|
|
| if L == npo2(L): |
| grad_output = grad_output_in.clone() |
| else: |
| grad_output = pad_npo2(grad_output_in).clone() |
| A_in = pad_npo2(A_in) |
|
|
| grad_output = grad_output.transpose(2, 1).contiguous() |
| A_in_t = A_in.transpose(2, 1).contiguous() |
|
|
| |
| A_shifted = F.pad(A_in_t[:, :, 1:], (0, 0, 0, 1)) |
|
|
| PScan.pscan_rev(A_shifted, grad_output) |
|
|
| grad_A = torch.zeros_like(X) |
| grad_A[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:]) |
|
|
| return grad_A.transpose(2, 1)[:, :L].contiguous(), grad_output.transpose(2, 1)[:, :L].contiguous() |
|
|
|
|
| pscan = PScan.apply |
|
|
