Module neuralops.core.triton_kernels.selective_scan

Triton-accelerated selective scan kernel for Mamba-1 S6 architecture.

Computes the continuous-time discretization state recurrence: dA_t = exp(softplus(delta_t) * A) dB_t = (softplus(delta_t) * u_t) * B_t h_t = dA_t * h_{t-1} + dB_t y_t = sum_N(h_t * C_t) + D * u_t

FIXED: backward pass now computes dA, dB, dC (previously left at zero). Stores full h_t trajectory in forward so backward can access h_{t-1}.

Functions

def selective_scan_fn(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
delta_softplus: bool = True) ‑> torch.Tensor

Expand source code

def selective_scan_fn(
    u: torch.Tensor,
    delta: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor | None = None,
    delta_softplus: bool = True,
) -> torch.Tensor:
    """
    Triton selective scan functional interface.
    Shapes:
        u:     (B, L, D)
        delta: (B, L, D)
        A:     (D, N)
        B:     (B, L, N)
        C:     (B, L, N)
        D:     (D,) optional
    """
    return TritonSelectiveScanFunction.apply(u, delta, A, B, C, D, delta_softplus)

Triton selective scan functional interface.

Shapes -----= u: (B, L, D) delta: (B, L, D) A: (D, N) B: (B, L, N) C: (B, L, N) D: (D,) optional

Classes

class TritonSelectiveScanFunction (*args, **kwargs)

Expand source code

class TritonSelectiveScanFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        u: torch.Tensor,  # (B, L, D)
        delta: torch.Tensor,  # (B, L, D)
        A: torch.Tensor,  # (D, N)
        B: torch.Tensor,  # (B, L, N)
        C: torch.Tensor,  # (B, L, N)
        D: torch.Tensor | None = None,  # (D,)
        delta_softplus: bool = True,
    ) -> torch.Tensor:
        B_size, L_size, D_dim = u.shape
        N_state = A.shape[1]

        u = u.contiguous()
        delta = delta.contiguous()
        A = A.contiguous()
        B = B.contiguous()
        C = C.contiguous()
        if D is not None:
            D = D.contiguous()

        out = torch.empty_like(u)

        # Workspace for full h_t trajectory: (B, D, L, N), fp32
        h_workspace = torch.empty(
            B_size,
            D_dim,
            L_size,
            N_state,
            device=u.device,
            dtype=torch.float32,
        )

        grid = (B_size, D_dim)

        _selective_scan_fwd_kernel[grid](
            u,
            delta,
            A,
            B,
            C,
            D if D is not None else u,
            out,
            h_workspace,
            u.stride(0),
            u.stride(1),
            u.stride(2),
            delta.stride(0),
            delta.stride(1),
            delta.stride(2),
            A.stride(0),
            A.stride(1),
            B.stride(0),
            B.stride(1),
            B.stride(2),
            C.stride(0),
            C.stride(1),
            C.stride(2),
            out.stride(0),
            out.stride(1),
            out.stride(2),
            h_workspace.stride(0),
            h_workspace.stride(1),
            h_workspace.stride(2),
            h_workspace.stride(3),
            B_size,
            L_size,
            D_dim,
            N_STATE=N_state,
            HAS_D=(D is not None),
            DELTA_SOFTPLUS=delta_softplus,
        )

        ctx.save_for_backward(u, delta, A, B, C, D, h_workspace)
        ctx.delta_softplus = delta_softplus
        return out

    @staticmethod
    def backward(ctx, dout: torch.Tensor):
        u, delta, A, B, C, D, h_workspace = ctx.saved_tensors
        B_size, L_size, D_dim = u.shape
        N_state = A.shape[1]

        dout = dout.contiguous()

        du = torch.empty_like(u)
        ddelta = torch.empty_like(delta)
        # dA and dD must be fp32 because we use atomic_add (fp16 atomics not
        # universally supported). Cast back to source dtype at the end.
        dA = torch.zeros_like(A, dtype=torch.float32)
        dB = torch.zeros_like(B)  # atomic_add نیاز به شروع از صفر داره
        dC = torch.zeros_like(C)
        dD = torch.zeros_like(D, dtype=torch.float32) if D is not None else None

        grid = (B_size, D_dim)

        _selective_scan_bwd_kernel[grid](
            u,
            delta,
            A,
            B,
            C,
            D if D is not None else u,
            dout,
            du,
            ddelta,
            dA,
            dB,
            dC,
            dD if dD is not None else u,
            h_workspace,
            u.stride(0),
            u.stride(1),
            u.stride(2),
            delta.stride(0),
            delta.stride(1),
            delta.stride(2),
            A.stride(0),
            A.stride(1),
            B.stride(0),
            B.stride(1),
            B.stride(2),
            C.stride(0),
            C.stride(1),
            C.stride(2),
            dout.stride(0),
            dout.stride(1),
            dout.stride(2),
            h_workspace.stride(0),
            h_workspace.stride(1),
            h_workspace.stride(2),
            h_workspace.stride(3),
            B_size,
            L_size,
            D_dim,
            N_STATE=N_state,
            HAS_D=(D is not None),
            DELTA_SOFTPLUS=ctx.delta_softplus,
        )

        dA = dA.to(A.dtype)
        if dD is not None:
            dD = dD.to(D.dtype)

        return du, ddelta, dA, dB, dC, dD, None

Base class to create custom autograd.Function.

To create a custom autograd.Function, subclass this class and implement the :meth:forward and :meth:backward static methods. Then, to use your custom op in the forward pass, call the class method apply. Do not call :meth:forward directly.

To ensure correctness and best performance, make sure you are calling the correct methods on ctx and validating your backward function using :func:torch.autograd.gradcheck.

See :ref:extending-autograd for more details on how to use this class.

Examples::

>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
>>> class Exp(Function):
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result
>>>
>>> # Use it by calling the apply method:
>>> # xdoctest: +SKIP
>>> output = Exp.apply(input)

Ancestors

• torch.autograd.function.Function
• torch.autograd.function._SingleLevelFunction
• torch._C._FunctionBase
• torch.autograd.function.FunctionCtx
• torch.autograd.function._HookMixin

Static methods

def backward(ctx, dout: torch.Tensor)

Expand source code

@staticmethod
def backward(ctx, dout: torch.Tensor):
    u, delta, A, B, C, D, h_workspace = ctx.saved_tensors
    B_size, L_size, D_dim = u.shape
    N_state = A.shape[1]

    dout = dout.contiguous()

    du = torch.empty_like(u)
    ddelta = torch.empty_like(delta)
    # dA and dD must be fp32 because we use atomic_add (fp16 atomics not
    # universally supported). Cast back to source dtype at the end.
    dA = torch.zeros_like(A, dtype=torch.float32)
    dB = torch.zeros_like(B)  # atomic_add نیاز به شروع از صفر داره
    dC = torch.zeros_like(C)
    dD = torch.zeros_like(D, dtype=torch.float32) if D is not None else None

    grid = (B_size, D_dim)

    _selective_scan_bwd_kernel[grid](
        u,
        delta,
        A,
        B,
        C,
        D if D is not None else u,
        dout,
        du,
        ddelta,
        dA,
        dB,
        dC,
        dD if dD is not None else u,
        h_workspace,
        u.stride(0),
        u.stride(1),
        u.stride(2),
        delta.stride(0),
        delta.stride(1),
        delta.stride(2),
        A.stride(0),
        A.stride(1),
        B.stride(0),
        B.stride(1),
        B.stride(2),
        C.stride(0),
        C.stride(1),
        C.stride(2),
        dout.stride(0),
        dout.stride(1),
        dout.stride(2),
        h_workspace.stride(0),
        h_workspace.stride(1),
        h_workspace.stride(2),
        h_workspace.stride(3),
        B_size,
        L_size,
        D_dim,
        N_STATE=N_state,
        HAS_D=(D is not None),
        DELTA_SOFTPLUS=ctx.delta_softplus,
    )

    dA = dA.to(A.dtype)
    if dD is not None:
        dD = dD.to(D.dtype)

    return du, ddelta, dA, dB, dC, dD, None

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context :attr:ctx as the first argument, followed by as many outputs as the :func:forward returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to :func:forward. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute :attr:ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., :func:backward will have ctx.needs_input_grad[0] = True if the first input to :func:forward needs gradient computed w.r.t. the output.

def forward(ctx,
u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
delta_softplus: bool = True) ‑> torch.Tensor

Expand source code

@staticmethod
def forward(
    ctx,
    u: torch.Tensor,  # (B, L, D)
    delta: torch.Tensor,  # (B, L, D)
    A: torch.Tensor,  # (D, N)
    B: torch.Tensor,  # (B, L, N)
    C: torch.Tensor,  # (B, L, N)
    D: torch.Tensor | None = None,  # (D,)
    delta_softplus: bool = True,
) -> torch.Tensor:
    B_size, L_size, D_dim = u.shape
    N_state = A.shape[1]

    u = u.contiguous()
    delta = delta.contiguous()
    A = A.contiguous()
    B = B.contiguous()
    C = C.contiguous()
    if D is not None:
        D = D.contiguous()

    out = torch.empty_like(u)

    # Workspace for full h_t trajectory: (B, D, L, N), fp32
    h_workspace = torch.empty(
        B_size,
        D_dim,
        L_size,
        N_state,
        device=u.device,
        dtype=torch.float32,
    )

    grid = (B_size, D_dim)

    _selective_scan_fwd_kernel[grid](
        u,
        delta,
        A,
        B,
        C,
        D if D is not None else u,
        out,
        h_workspace,
        u.stride(0),
        u.stride(1),
        u.stride(2),
        delta.stride(0),
        delta.stride(1),
        delta.stride(2),
        A.stride(0),
        A.stride(1),
        B.stride(0),
        B.stride(1),
        B.stride(2),
        C.stride(0),
        C.stride(1),
        C.stride(2),
        out.stride(0),
        out.stride(1),
        out.stride(2),
        h_workspace.stride(0),
        h_workspace.stride(1),
        h_workspace.stride(2),
        h_workspace.stride(3),
        B_size,
        L_size,
        D_dim,
        N_STATE=N_state,
        HAS_D=(D is not None),
        DELTA_SOFTPLUS=delta_softplus,
    )

    ctx.save_for_backward(u, delta, A, B, C, D, h_workspace)
    ctx.delta_softplus = delta_softplus
    return out

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
• See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.
• Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
• See :ref:extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with :func:ctx.save_for_backward if they are intended to be used in backward (equivalently, vjp) or :func:ctx.save_for_forward if they are intended to be used for in jvp.

class triton

Methods

def autotune(*a, **kw)

Expand source code

"autotune": lambda *a, **kw: lambda f: f,

def heuristics(*a, **kw)

Expand source code

"heuristics": lambda *a, **kw: lambda f: f,

def jit(f)

Expand source code

"jit": lambda f: f,