Upload luminars/ssm.py

luminars/ssm.py

@@ -0,0 +1,42 @@
+"""
+Selective State Space (Mamba2) cell + SelectiveScanKernel.
+No dependencies on mamba_ssm -- pure PyTorch.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, einsum
+
+def selective_scan_oneshot(x, delta, A, B, C, D):
+    """
+    x: (B, L, N)  -- input tokens
+    delta: (B, L, N) -- time-step, elementwise
+    A: (N,)           -- diagonal S4D real part
+    B, C: (B, L, N)   -- input-dependent
+    D: (N,)           -- skip connection
+    Returns y: (B, L, N)
+    """
+    B_, L, N = x.shape
+    # discretize: A_bar = exp(delta * A), B_bar = delta * B
+    # A is negative (stable), delta > 0
+    A = -torch.abs(A)                 # force stability
+    delta = F.softplus(delta)           # >0
+    A_bar = torch.exp(delta.unsqueeze(-1) * A)          # (B, L, N, N)?? No, A is (N,)
+    A_bar = torch.exp(delta * A)         # (B, L, N)
+    B_x = delta * B * x                # (B, L, N)
+
+    # recurrent scan
+    h = torch.zeros(B_, N, device=x.device, dtype=x.dtype)
+    ys = []
+    for t in range(L):
+        h = A_bar[:, t] * h + B_x[:, t]
+        y = einsum(h, C[:, t], 'b n, b n -> b')
+        ys.append(y)
+    # Actually y = (C_t * h).sum(-1) gives scalar per token... reshape needed.
+    # Let's do it vectorised:
+    # We actually need y_t = sum_n C_{b,t,n} * h_{b,n} = inner product in N dim
+    y = torch.stack(ys, dim=1).unsqueeze(-1) * C  # no, this is wrong dimension
+    # FIX: h is (B,N), output is (B,N) from h*C where C is (B,L,N)
+    # Let me rewrite properly
+    pass