torch-rocm/cpp_functions.py

"""C++ ABI wrappers for the Chakra Triton causal-conv1d kernels.

Upstream ``mamba_ssm`` imports these symbols from ``causal_conv1d.cpp_functions``.
We provide thin wrappers that delegate to our Triton implementations so that
``selective_scan_interface`` and ``ssd_combined`` can resolve the fast path
without requiring the upstream C++/CUDA extension.
"""

from __future__ import annotations

from typing import Optional

import torch

from .functional import causal_conv1d_fn, causal_conv1d_update


def causal_conv1d_fwd_function(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor],
    seq_idx: Optional[torch.Tensor],
    cu_seqlens: Optional[torch.Tensor],
    max_seqlen: Optional[int],
    activation: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """C++ ABI wrapper for ``causal_conv1d_fn``.

    Upstream signature:
        (x, weight, bias, seq_idx, cu_seqlens, max_seqlen, activation)
        -> (y, final_states)
    """
    activation_str = "silu" if activation else "identity"
    y, final_states = causal_conv1d_fn(
        x,
        weight,
        bias=bias,
        seq_idx=seq_idx,
        initial_states=None,
        return_final_states=True,
        activation=activation_str,
    )
    return y, final_states


def causal_conv1d_bwd_function(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor],
    seq_idx: Optional[torch.Tensor],
    cu_seqlens: Optional[torch.Tensor],
    max_seqlen: Optional[int],
    dy: torch.Tensor,
    dconv_state: Optional[torch.Tensor],
    dinitial_states: Optional[torch.Tensor],
    activation: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """C++ ABI wrapper for the Triton backward path.

    Returns ``(dx, dweight, dbias, dinitial_states, None)`` to match the
    upstream C extension tuple layout.
    """
    activation_str = "silu" if activation else "identity"

    x_in = x
    if dinitial_states is not None and dinitial_states.numel() > 0:
        x_in = torch.cat([dinitial_states.to(x.dtype), x], dim=-1).contiguous()

    y = torch.empty_like(x_in)
    b, d, l_full = x_in.shape
    width = weight.shape[-1]
    from .triton_kernels import causal_conv1d_fwd_kernel
    BLOCK_L = 256
    BLOCK_D = 64
    apply_silu = activation in ("silu", True)
    apply_relu = activation == "relu"
    grid = (b, __import__("triton").cdiv(l_full, BLOCK_L), __import__("triton").cdiv(d, BLOCK_D))
    causal_conv1d_fwd_kernel[grid](
        x_in, weight, bias if bias is not None else torch.zeros(d, device=x.device, dtype=x.dtype), y,
        b, l_full, d, width,
        x_in.stride(0), x_in.stride(1), x_in.stride(2),
        weight.stride(0), weight.stride(1),
        y.stride(0), y.stride(1), y.stride(2),
        apply_silu, apply_relu, BLOCK_L, BLOCK_D,
    )

    pre_act = y if (apply_silu or apply_relu) else None
    if apply_silu:
        sig = torch.sigmoid(pre_act)
        d_pre = dy * sig * (1.0 + pre_act * (1.0 - sig))
    elif apply_relu:
        d_pre = dy * (pre_act > 0).to(dy.dtype)
    else:
        d_pre = dy

    from .triton_kernels import causal_conv1d_bwd_dx_kernel, causal_conv1d_bwd_dw_kernel, causal_conv1d_bwd_db_kernel

    dx = torch.empty_like(x_in)
    grid_dx = (b, __import__("triton").cdiv(l_full, BLOCK_L), __import__("triton").cdiv(d, BLOCK_D))
    causal_conv1d_bwd_dx_kernel[grid_dx](
        d_pre, weight, dx,
        b, l_full, d, width,
        d_pre.stride(0), d_pre.stride(1), d_pre.stride(2),
        weight.stride(0), weight.stride(1),
        dx.stride(0), dx.stride(1), dx.stride(2),
        BLOCK_L, BLOCK_D,
    )

    dw = torch.zeros(d, width, dtype=x.dtype, device=x.device)
    grid_dw = (__import__("triton").cdiv(d, BLOCK_D), width)
    causal_conv1d_bwd_dw_kernel[grid_dw](
        d_pre, x_in, dw,
        b, l_full, d, width,
        d_pre.stride(0), d_pre.stride(1), d_pre.stride(2),
        x_in.stride(0), x_in.stride(1), x_in.stride(2),
        dw.stride(0), dw.stride(1),
        BLOCK_D,
    )

    db = torch.zeros(d, dtype=x.dtype, device=x.device)
    grid_db = (__import__("triton").cdiv(d, BLOCK_D),)
    causal_conv1d_bwd_db_kernel[grid_db](
        d_pre, db,
        b, l_full, d,
        d_pre.stride(0), d_pre.stride(1), d_pre.stride(2),
        db.stride(0),
        BLOCK_D,
    )

    orig_len = l_full - (x_in.shape[-1] - x.shape[-1])
    dx = dx[..., -orig_len:] if orig_len is not None and orig_len < dx.shape[-1] else dx
    dinitial_states_out = torch.zeros(0)
    return dx, dw, db, dinitial_states_out, None


def causal_conv1d_update_function(
    x: torch.Tensor,
    conv_state: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor],
    activation: bool = True,
) -> torch.Tensor:
    """C++ ABI wrapper for ``causal_conv1d_update``."""
    activation_str = "silu" if activation else "identity"
    return causal_conv1d_update(
        x, conv_state, weight, bias if bias is not None else torch.zeros(x.shape[1], device=x.device, dtype=x.dtype),
        activation=activation_str,
    )