Add op_check tests

Author: bhack

mamba_ssm/ops/selective_scan_interface_ compilable.py

@@ -1,15 +1,15 @@
 import torch
-import torch.nn.functional as F
 from einops import rearrange
 from typing import Optional, Tuple
 
-import selective_scan_cuda
+from mamba_ssm.ops.selective_scan_interface import selective_scan_cuda
 
 
 @torch.library.custom_op(
     "custom_ops::selective_scan_fwd",
     device_types=["cuda"],
     mutates_args=(),
+    schema="(Tensor u, Tensor delta, Tensor A, Tensor B, Tensor C, Tensor? D, Tensor? z, Tensor? delta_bias, bool delta_softplus, bool return_last_state) -> (Tensor, Tensor, Tensor, Tensor, bool, bool, bool)",
 )
 def custom_selective_scan_fwd(
     u: torch.Tensor,
@@ -22,28 +22,33 @@ def custom_selective_scan_fwd(
     delta_bias: Optional[torch.Tensor],
     delta_softplus: bool,
     return_last_state: bool,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool, bool, bool]:
+):
     pass
 
+
 @torch.library.register_fake("custom_ops::selective_scan_fwd")
 def custom_selective_scan_fwd_fake(
-    u,
-    delta,
-    A,
-    B,
-    C,
-    D,
-    z,
-    delta_bias,
-    delta_softplus,
-    return_last_state,
+    u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state
 ):
-    final_out = torch.empty_like(u)
     dstate = A.size(1) * (2 if A.is_complex() else 1)
-    last_state_fake = u.new_empty((u.size(0), u.size(1), dstate)) if return_last_state else u.new_empty(0)
-    out_fake = torch.empty_like(u)
-    x_fake = u.new_empty((u.size(0), u.size(1), u.size(2), 2 * dstate))
-    return final_out, last_state_fake, out_fake, x_fake, False, False, z is not None
+    seqlen = u.size(2)
+    n_chunks = (seqlen + 2048 - 1) // 2048
+
+    squeeze_B = B.dim() == 3
+    squeeze_C = C.dim() == 3
+    has_z = z is not None
+
+    final_out = torch.empty_like(delta)
+    out_fake = torch.empty_like(delta)
+    last_state_fake = (
+        u.new_empty((u.size(0), u.size(1), dstate))
+        if return_last_state
+        else u.new_empty(0)
+    )
+    x_fake = u.new_empty((u.size(0), u.size(1), n_chunks, 2 * A.size(1)), dtype=A.dtype)
+
+    return final_out, last_state_fake, out_fake, x_fake, squeeze_B, squeeze_C, has_z
+
 
 @torch.library.register_kernel("custom_ops::selective_scan_fwd", "cuda")
 def custom_selective_scan_fwd_cuda(
@@ -81,16 +86,23 @@ def custom_selective_scan_fwd_cuda(
         C = rearrange(C, "b dstate l -> b 1 dstate l").contiguous()
         squeeze_C = True
 
-    out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, z, delta_bias, delta_softplus)
+    out, x, *rest = selective_scan_cuda.fwd(
+        u, delta, A, B, C, D, z, delta_bias, delta_softplus
+    )
     has_z = z is not None
-    final_out = rest[0].clone() if has_z else out.clone()
+    if has_z:
+        final_out = rest[0].clone()
+    else:
+        final_out = out.clone()
     last_state = x[:, :, -1, 1::2].clone() if return_last_state else u.new_empty(0)
     return final_out, last_state, out, x, squeeze_B, squeeze_C, has_z
 
+
 @torch.library.custom_op(
     "custom_ops::selective_scan_bwd",
     device_types=["cuda"],
     mutates_args=(),
+    schema="(Tensor dout, Tensor u, Tensor delta, Tensor A, Tensor B, Tensor C, Tensor? D, Tensor? z, Tensor? delta_bias, bool delta_softplus, Tensor out, Tensor x, bool squeeze_B, bool squeeze_C, bool recompute_out_z) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor?, Tensor?, Tensor?)",
 )
 def custom_selective_scan_bwd(
     dout: torch.Tensor,
@@ -107,9 +119,11 @@ def custom_selective_scan_bwd(
     x: torch.Tensor,
     squeeze_B: bool,
     squeeze_C: bool,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    recompute_out_z: bool,
+):
     pass
 
+
 @torch.library.register_fake("custom_ops::selective_scan_bwd")
 def custom_selective_scan_bwd_fake(
     dout,
@@ -126,16 +140,33 @@ def custom_selective_scan_bwd_fake(
     x,
     squeeze_B,
     squeeze_C,
+    recompute_out_z,
 ):
+    # Here we just return shape-compatible fake tensors
     du = torch.empty_like(u)
     ddelta = torch.empty_like(delta)
     dA = torch.empty_like(A)
-    dB = torch.empty_like(B)
-    dC = torch.empty_like(C)
-    dD = torch.empty_like(D) if (D is not None and D.numel() > 0) else u.new_empty(0)
-    dz = torch.empty_like(z) if (z is not None and z.numel() > 0) else u.new_empty(0)
-    ddelta_bias = torch.empty_like(delta_bias) if (delta_bias is not None and delta_bias.numel() > 0) else u.new_empty(0)
-    return du, ddelta, dA, dB, dC, dD, dz, ddelta_bias
+
+    # Decide if variable B/C
+    is_variable_B = B.dim() > 3
+    is_variable_C = C.dim() > 3
+
+    dB = torch.empty_like(
+        B, dtype=B.dtype
+    )  # If variable_B, still float32 is okay for fake
+    dC = torch.empty_like(C, dtype=C.dtype)
+
+    dD = torch.empty_like(D) if (D is not None) else None
+    ddelta_bias_out = torch.empty_like(delta_bias) if (delta_bias is not None) else None
+    dz = torch.empty_like(z) if (z is not None) else None
+
+    if squeeze_B and dB.numel() > 0:
+        dB = dB.squeeze(1)
+    if squeeze_C and dC.numel() > 0:
+        dC = dC.squeeze(1)
+
+    return du, ddelta, dA, dB, dC, dD, ddelta_bias_out, dz
+
 
 @torch.library.register_kernel("custom_ops::selective_scan_bwd", "cuda")
 def custom_selective_scan_bwd_cuda(
@@ -153,68 +184,101 @@ def custom_selective_scan_bwd_cuda(
     x: torch.Tensor,
     squeeze_B: bool,
     squeeze_C: bool,
+    recompute_out_z: bool,
 ):
     if dout.stride(-1) != 1:
         dout = dout.contiguous()
-    B = B.contiguous()
-    C = C.contiguous()
 
     results = selective_scan_cuda.bwd(
-        u, delta, A, B, C, D, z, delta_bias, dout, x, out, None, delta_softplus, False
+        u,
+        delta,
+        A,
+        B,
+        C,
+        D,
+        z,
+        delta_bias,
+        dout,
+        x,
+        out,
+        None,
+        delta_softplus,
+        recompute_out_z,
     )
+
     has_z = z is not None
     if has_z:
-        du, ddelta, dA, dB, dC, dD, ddelta_bias, dz = results
+        du, ddelta, dA, dB, dC, dD, ddelta_bias_out, dz = results
     else:
-        du, ddelta, dA, dB, dC, dD, ddelta_bias = results
-        dz = u.new_empty(0)
+        du, ddelta, dA, dB, dC, dD, ddelta_bias_out = results
+        dz = None
 
     if squeeze_B and dB.numel() > 0:
         dB = dB.squeeze(1)
     if squeeze_C and dC.numel() > 0:
         dC = dC.squeeze(1)
 
-    return du, ddelta, dA, dB, dC, dD, dz, ddelta_bias
+    return du, ddelta, dA, dB, dC, dD, ddelta_bias_out, dz
+
 
 def custom_bridge(ctx, *grads):
     dout = grads[0] if grads else ctx.saved_tensors[0].new_empty(0)
     saved = ctx.saved_tensors
+
     if not ctx.has_z:
         u, delta, A, B, C, D, delta_bias, x, out = saved
         z = None
     else:
         u, delta, A, B, C, D, z, delta_bias, x, out = saved
 
-    du, ddelta, dA, dB, dC, dD, dz, ddelta_bias = torch.ops.custom_ops.selective_scan_bwd(
-        dout,
-        u,
-        delta,
-        A,
-        B,
-        C,
-        D,
-        z,
-        delta_bias,
-        ctx.delta_softplus,
-        out,
-        x,
-        ctx.squeeze_B,
-        ctx.squeeze_C
+    du, ddelta, dA, dB, dC, dD, ddelta_bias_out, dz = (
+        torch.ops.custom_ops.selective_scan_bwd(
+            dout,
+            u,
+            delta,
+            A,
+            B,
+            C,
+            D,
+            z,
+            delta_bias,
+            ctx.delta_softplus,
+            out,
+            x,
+            ctx.squeeze_B,
+            ctx.squeeze_C,
+            False,
+        )
     )
 
+    # For optional inputs, return None if not provided in forward
+    if D is None:
+        dD = None
+    if z is None:
+        dz = None
+    if delta_bias is None:
+        ddelta_bias_out = None
+
+    # Return gradients in the order of forward inputs:
+    # (u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
+    # `delta_softplus` and `return_last_state` are bools -> gradient = None
+    d_delta_softplus = None
+    d_return_last_state = None
+
     return (
         du,
         ddelta,
         dA,
         dB,
         dC,
-        dD if D is not None else None,
-        dz if z is not None else None,
-        ddelta_bias if delta_bias is not None else None,
-        None,
-        None,
+        dD,
+        dz,
+        ddelta_bias_out,
+        d_delta_softplus,
+        d_return_last_state,
     )
 
+
 def custom_setup_context(ctx, inputs, output):
     (u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state) = inputs
     (final_out, last_state, out, x, squeeze_B, squeeze_C, has_z) = output
@@ -236,10 +300,12 @@ def custom_setup_context(ctx, inputs, output):
     else:
         ctx.save_for_backward(u, delta, A, B, C, D, z, delta_bias, x, out)
 
+
 torch.library.register_autograd(
     "custom_ops::selective_scan_fwd", custom_bridge, setup_context=custom_setup_context
 )
 
+
 def selective_scan_fn_custom_op(
     u: torch.Tensor,
     delta: torch.Tensor,
@@ -252,20 +318,9 @@ def selective_scan_fn_custom_op(
     delta_softplus: bool,
     return_last_state: bool,
 ) -> torch.Tensor:
-    # Pass all arguments positionally, exactly in schema order:
     final_out, last_state, _, _, _, _, _ = torch.ops.custom_ops.selective_scan_fwd(
-        u,
-        delta,
-        A,
-        B,
-        C,
-        D,
-        z,
-        delta_bias,
-        delta_softplus,
-        return_last_state
+        u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state
     )
-
     if return_last_state:
         return final_out, last_state
     else:

tests/ops/test_selective_scan.py

@@ -254,3 +254,28 @@ def test_mamba_inner_fn(is_variable_B, is_variable_C, seqlen, itype, wtype):
     #                       atol=atolw if not is_variable_C else atol)
     # assert torch.allclose(D.grad, D_ref.grad, rtol=rtolw, atol=atolw)
     # assert torch.allclose(delta_bias.grad, delta_bias_ref.grad, rtol=rtolw, atol=atolw)
+
+def test_selective_scan_opcheck():
+    from torch.library import opcheck
+    
+    device = "cuda"
+    # small inputs for opcheck
+    u = torch.randn(1, 2, 8, device=device, requires_grad=True)
+    delta = torch.randn(1, 2, 8, device=device, requires_grad=True)
+    A = torch.randn(2, 8, device=device, requires_grad=True)
+    B = torch.randn(1, 1, 8, 8, device=device, requires_grad=True)
+    C = torch.randn(1, 1, 8, 8, device=device, requires_grad=True)
+    D = torch.randn(2, device=device, requires_grad=True)
+    z = torch.randn(1, 2, 8, device=device, requires_grad=True)
+    delta_bias = torch.randn(2, device=device, requires_grad=True)
+    delta_softplus = False
+    return_last_state = False
+
+    # Run opcheck
+    result = opcheck(
+        torch.ops.custom_ops.selective_scan_fwd,
+        (u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state),
+        test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic", "test_autograd_registration"),
+        raise_exception=True
+    )
+    print("Opcheck result:", result)