torch-rocm/functional.py

"""Triton-backed implementations of the Mamba-2 causal-conv1d ABI.

Layout (channel-first, Mamba-2 convention):
    x      : (B, D, L)
    weight : (D, WIDTH)
    bias   : (D,)
    y      : (B, D, L)

All supported activations (silu, swish, relu, identity) run through Triton.
No PyTorch ``F.conv1d`` fallback remains in the production path.
"""

from __future__ import annotations

from typing import Optional

import torch
import triton

from .triton_kernels import (
    causal_conv1d_fwd_kernel,
    causal_conv1d_update_kernel,
    causal_conv1d_bwd_dx_kernel,
    causal_conv1d_bwd_dw_kernel,
    causal_conv1d_bwd_db_kernel,
)

BLOCK_L = 256
BLOCK_D = 64


def _resolve_activation(activation: Optional[str]):
    if activation in ("silu", "swish"):
        return True, False
    if activation == "relu":
        return False, True
    return False, False


class _CausalConv1dFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, activation, orig_len):
        b, d, l_full = x.shape
        width = weight.shape[-1]
        apply_silu, apply_relu = _resolve_activation(activation)

        y = torch.empty_like(x)
        grid = (b, triton.cdiv(l_full, BLOCK_L), triton.cdiv(d, BLOCK_D))
        causal_conv1d_fwd_kernel[grid](
            x,
            weight,
            bias,
            y,
            b,
            l_full,
            d,
            width,
            x.stride(0),
            x.stride(1),
            x.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,
        )

        needs_grad = ctx.needs_input_grad[0] or ctx.needs_input_grad[1] or ctx.needs_input_grad[2]
        pre_act = None
        if needs_grad and (apply_silu or apply_relu):
            pre_act = torch.empty_like(x)
            causal_conv1d_fwd_kernel[grid](
                x, weight, bias, pre_act,
                b, l_full, d, width,
                x.stride(0), x.stride(1), x.stride(2),
                weight.stride(0), weight.stride(1),
                pre_act.stride(0), pre_act.stride(1), pre_act.stride(2),
                False, apply_relu, BLOCK_L, BLOCK_D,
            )

        ctx.save_for_backward(x, weight, bias, pre_act)
        ctx.activation = activation
        ctx.orig_len = orig_len
        return y

    @staticmethod
    def backward(ctx, dy):
        x, weight, bias, pre_act = ctx.saved_tensors
        activation = ctx.activation
        l_full = x.shape[-1]

        if pre_act is not None:
            if activation in ("silu", "swish"):
                sig = torch.sigmoid(pre_act)
                d_pre = dy * sig * (1.0 + pre_act * (1.0 - sig))
            else:
                d_pre = dy * (pre_act > 0).to(dy.dtype)
        else:
            d_pre = dy

        dx = torch.empty_like(x)
        b, d, l = d_pre.shape
        width = weight.shape[-1]
        grid_dx = (b, triton.cdiv(l, BLOCK_L), triton.cdiv(d, BLOCK_D))
        causal_conv1d_bwd_dx_kernel[grid_dx](
            d_pre,
            weight,
            dx,
            b,
            l,
            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 = None
        if ctx.needs_input_grad[1]:
            dw = torch.empty(d, width, dtype=x.dtype, device=x.device)
            grid_dw = (triton.cdiv(d, BLOCK_D), width)
            causal_conv1d_bwd_dw_kernel[grid_dw](
                d_pre,
                x,
                dw,
                b,
                l,
                d,
                width,
                d_pre.stride(0),
                d_pre.stride(1),
                d_pre.stride(2),
                x.stride(0),
                x.stride(1),
                x.stride(2),
                dw.stride(0),
                dw.stride(1),
                BLOCK_D,
            )

        db = None
        if bias is not None and ctx.needs_input_grad[2]:
            db = torch.empty(d, dtype=x.dtype, device=x.device)
            grid_db = (triton.cdiv(d, BLOCK_D),)
            causal_conv1d_bwd_db_kernel[grid_db](
                d_pre,
                db,
                b,
                l,
                d,
                d_pre.stride(0),
                d_pre.stride(1),
                d_pre.stride(2),
                db.stride(0),
                BLOCK_D,
            )

        orig_len = ctx.orig_len
        dx = dx[..., -orig_len:] if orig_len is not None else dx
        return dx, dw, db, None, None


def causal_conv1d_fn(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
    seq_idx: Optional[torch.Tensor] = None,
    initial_states: Optional[torch.Tensor] = None,
    return_final_states: bool = False,
    final_states_out: Optional[torch.Tensor] = None,
    activation: Optional[str] = "silu",
):
    """Causal 1D depthwise convolution (prefill path) -- Triton forward.

    ABI-compatible with Dao-AILab/causal-conv1d ``causal_conv1d_fn``.
    """
    if activation not in ("silu", "swish", "relu", "identity", None):
        raise NotImplementedError(f"activation {activation!r} not supported by causal_conv1d")

    x = x.contiguous()
    b, d, l = x.shape
    width = weight.shape[-1]
    weight = weight.contiguous()
    bias_t = bias.contiguous() if bias is not None else torch.zeros(d, dtype=x.dtype, device=x.device)

    x_in = x
    orig_len = l
    if initial_states is not None:
        x_in = torch.cat([initial_states.to(x.dtype), x], dim=-1).contiguous()
        orig_len = l

    y = _CausalConv1dFunction.apply(x_in, weight, bias_t, activation, orig_len)

    if return_final_states:
        final_states = x[..., -width + 1 :].contiguous()
        if final_states_out is not None:
            final_states_out.copy_(final_states)
            return y, final_states_out
        return y, final_states
    return y


def causal_conv1d_update(
    x: torch.Tensor,
    conv_state: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    activation: Optional[str] = "silu",
    cache_seqlens: Optional[torch.Tensor] = None,
    conv_state_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Causal 1D depthwise convolution for a single decode step -- Triton.

    ABI-compatible with Dao-AILab/causal-conv1d ``causal_conv1d_update``.

    Shapes:
        x: (batch, dim)            single token
        conv_state: (batch, dim, width - 1)   cached input window
        weight: (dim, width)
        bias: (dim,)
        returns: (batch, dim)

    Side-effect: ``conv_state`` is mutated in place (shift left, append x).
    """
    x = x.contiguous()
    b, d = x.shape
    width = weight.shape[-1]
    weight = weight.contiguous()
    has_bias = bias is not None
    bias_t = bias.contiguous() if has_bias else torch.zeros(d, dtype=x.dtype, device=x.device)

    apply_silu, apply_relu = _resolve_activation(activation)
    out = torch.empty(b, d, dtype=x.dtype, device=x.device)

    grid = (b, triton.cdiv(d, BLOCK_D))
    causal_conv1d_update_kernel[grid](
        conv_state,
        x,
        weight,
        bias_t,
        out,
        b,
        d,
        width,
        conv_state.stride(0),
        conv_state.stride(1),
        conv_state.stride(2),
        x.stride(0),
        x.stride(1),
        weight.stride(0),
        weight.stride(1),
        bias_t.stride(0) if has_bias else 0,
        out.stride(0),
        out.stride(1),
        has_bias,
        apply_silu,
        apply_relu,
        BLOCK_D,
    )
    return out