File size: 5,081 Bytes
45eef7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """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,
)
|