diff --git a/build/torch212-cxx11-cu126-x86_64-linux/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..44546b7d92ad639c9d9eb6b06d3c7c4d00d8fbfd --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f174415bc0621baf2c673f90cbd4a9087212e2a8d93e9d4ab661c27506ed9119 +size 78812288 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch212-cxx11-cu126-x86_64-linux/cpp_functions.py b/build/torch212-cxx11-cu126-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch212-cxx11-cu126-x86_64-linux/layers.py b/build/torch212-cxx11-cu126-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a2c37ad73d76675bdfd56b83a591e02b0cdec5bc --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json @@ -0,0 +1,44 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "7.0", + "7.2", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0+PTX" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "8XRBW8BiG68sZz+Qy9SpCHIS4qjZPp1KtmHCdQbtkRk=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..bc19c45fb9e1823f236a8af749bd244ccbf4cdd0 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUF5npDg5ylF6hUslnonJpk4E0ZtIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjMxWhcNMjYwOTA0MTQzMjMxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEXDOEz9luhruHBgZlSts8E6zBKO/sb/tXD3hLKeihYiDI1OEo99vKUtMVh1lWu/q6DP3KYDlsU8yptzBfisxgW6OCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUR339NiMbzWtkDa5hYeoUjjDqb/gwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzMu3MAAAQDAEcwRQIhAPDwxhx6LitdPUJzBKQ3UpWgf0wdfF9U7TlG33vujN5NAiBul6DyVJub05Hz9zDp82qAJufawD/LMyHhV8XSfjNVSzAKBggqhkjOPQQDAwNnADBkAjAM2bbW6fKK4DbtGgaoMkwWzc5SJ73cv93qW/acrcQBDlDaYO6tS3FEjFlHPbMFf5UCMGawfyiBCp5JH3ZTHZW23zpghZgQbeCd1nlnkbFLfrSosyanJB9NliNLt7bzyhIvbA=="}, "tlogEntries":[{"logIndex":"2712272739", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531754", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQD8srDq+MHTdVtlET9nlBIxRovCp6nnbUlVOrydHYo7uQIhAIIyi8uZvABuK/caUZg+9I0KXDkHBoSMO31q07qA+xwS"}, "inclusionProof":{"logIndex":"2590368477", "rootHash":"UWIza+5vwcjkcoaWAVFoxGRGfH+FSU17UEQhwACPatU=", "treeSize":"2590369486", "hashes":["1U11euJ9tOhPEFratAmpG5jvfAjickDCreYmBijU5SQ=", "S3uT5c7zKfDKFgc1qsSo0M3vUXas9d8UH43unGonvDE=", "C8FX0x1l7O6RRBo4mtSrbRlHJB4sOKBFv2Qrbf7id1I=", "LFXwUgL1WUbdVBfWa1ec7Br0lk9bhYuAch5UztUJjRI=", "SqQQwp8ftKUS8E7LXVt4vaGDNwjlwKYIyrPNauQ90X0=", "cFnFw51DhaJRsQ6ylnnpBRrayhUXP2lXJymcHizb/MI=", "zlx3RivqVtxl2pUJb0/AvPue3POksFPLxSkFoj/oBHA=", "L53shECZ2Dt19SXshsVZXgKPbyn8xh7JujI6k9YM4ls=", "zHyvrUWl+DArck+G7I0w6gtAV9tiYJWAASXiOc1UaNg=", "cZDpUl+6xSnn8yG0dI4YWAwhtkD8K/yUZrm0JkW1ys8=", "VI3m9U5kLhiizApcU6yjOQD/h+HNInVZln8XgxSvZQs=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590369486\nUWIza+5vwcjkcoaWAVFoxGRGfH+FSU17UEQhwACPatU=\n\n— rekor.sigstore.dev wNI9ajBFAiAcQLtXU1hYSx6+jWcqZy3yEgfRRnjCnG1DNntiC6xSzwIhAK6C8508tsFmNK5xfhFHX5TE/Mtea7zsV7XyuUpXLcY7\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5NmM3YWU1MjNhNWJiOWRjY2I5YTJlMGM4YjU3MjU1MzcwZTQ3NTU2Mjc5ZDdmYWQyOGMxMTc5NjlhMzRjZmVjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQWFudk4xRVovd0xhMXpJb1ZHN1BJN3NLcm9ZWFdlSzQ2RWM1aW9LNUlHakFpQmpRTjFwUjVEclU4MGVoVTRxM0hUbTJrZnQxNGdTdmtTY1hnM1J2Z3U5Y3c9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlJqVnVjRVJuTlhsc1JqWm9WWE5zYm05dVNuQnJORVV3V25SSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BOZUZkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWsxNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZZUkU5RmVqbHNkV2h5ZFVoQ1oxcHNVM1J6T0VVMmVrSkxUeTl6WWk5MFdFUXphRXdLUzJWcGFGbHBSRWt4VDBWdk9UbDJTMVYwVFZab01XeFhkUzl4TmtSUU0wdFpSR3h6VlRoNWNIUjZRbVpwYzNoblZ6WlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZTTXpNNUNrNXBUV0o2VjNSclJHRTFhRmxsYjFWcWFrUnhZaTluZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazExTTAxQlFVRlJSQXBCUldOM1VsRkphRUZRUkhkNGFIZzJUR2wwWkZCVlNucENTMUV6VlhCWFoyWXdkMlJtUmpsVk4xUnNSek16ZG5WcVRqVk9RV2xDZFd3MlJIbFdTblZpQ2pBMVNIbzVla1J3T0RKeFFVcDFabUYzUkM5TVRYbElhRlk0V0ZObWFrNVdVM3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJUVEppWWxjS05tWkxTelJFWW5SSFoyRnZUV3QzVjNwak5WTktOek5qZGpremNWY3ZZV055WTFGQ1JHeEVZVmxQTm5SVE0wWkZha1pzU0ZCaVRVWm1OVlZEVFVkaGR3cG1lV2xDUTNBMVNrZ3pXbFJJV2xjeU0zcHdaMmhhWjFGaVpVTmtNVzVzYm10aVJreG1jbE52YzNsaGJrcENPVTVzYVU1TWREZGllbmxvU1haaVFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgCRTYR69fqs0jKWpKdbKh2xfIFEFtOw+N1x9C+RhaPZACFQD6yLaZFDoJWgk9llhuCJDQq2RlXxgPMjAyNjA5MDQxNDIyMzFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MjIzMVowLwYJKoZIhvcNAQkEMSIEILLQs9CfWuJfGbLxMw0ym/UC0qjM5HvvQz8XIxHJK7HLMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMHquz3uZXN7inPbI/ErpH4YRtusV2g+CYm2hpssYL7HD1fLz8IOWLDufOyTUcRPEEwIwLO/4uX5Zz9IR+8w0EGkA917kLaML7sWIyF/7F4byny+x4p3xhfO13fgN5nlTEeHv"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"lseuUjpbudzLmi4Mi1clU3DkdVYnnX+tKMEXlpo0z+w="}, "signature":"MEQCIAanvN1EZ/wLa1zIoVG7PI7sKroYXWeK46Ec5ioK5IGjAiBjQN1pR5DrU80ehU4q3HTm2kft14gSvkScXg3Rvgu9cw=="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..75ad11199f2a3f9fbda907472fe4f53f1655ca82 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:da62650066ff7531150c70726936f53e2587bfd22926ffd9052d144bede61a87 +size 74854784 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch212-cxx11-cu130-x86_64-linux/cpp_functions.py b/build/torch212-cxx11-cu130-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch212-cxx11-cu130-x86_64-linux/layers.py b/build/torch212-cxx11-cu130-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b592fd2abb99e34231b2aa44863773b612b7af9c --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "2mJlAGb/dTEVDHByaTb1PiWHv9IpJv/ZBS0US+3mGoc=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..0974ae0e5ab487a9fe28b7038313897e85585d79 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUEm/GUtOGnQ/TZx10GxzB6EJLhn4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjQwWhcNMjYwOTA0MTQzMjQwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEqJV5VOwl8yW6JQjfZ/PG77lcXEzXKU3v1v/MT5FVdMLL6GWrexBOKmJ8u3QCe4bXPNLqaK04sCyJzQ+2NTD2eKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUDvPsQoq5tfz4Hk/yGeUGewaKap0wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzM3G8AAAQDAEcwRQIgPqENJP9g8cTnR2wpD3WpcRYI/IuuOXF9eH3ZLwk2ZNoCIQChMj7f1woiEMFJ4BJHVKOBEIV7wVtToyhj4+knPNK1/zAKBggqhkjOPQQDAwNoADBlAjB2OlyPNZvGUEQZ6qVBm3pqaN5NnFbBT1NQqhlczUa+AWMTyizLMHHqig7YczCWLdgCMQDLQ3spQAiP7XnXstYs020/wFiqP28TBx6+/Y3LOn6UOvoXxHnQXeA67Wl1+oovCV4="}, "tlogEntries":[{"logIndex":"2712273997", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531760", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDtinQAMuOlw9xe/C5kA0zoK7RWJH7h6y+pwV5tcfrj4wIgB3x5Gplwj4Vzk7I7PR2Vh3YVbMn6iHnxDQX9VO4ujq8="}, "inclusionProof":{"logIndex":"2590369735", "rootHash":"GwN1lo+P762wCRUXBMgV+ZZ5pDyrIrcZ+NBBYiM6cu0=", "treeSize":"2590370072", "hashes":["NUQR60N3WmIklku4sDxfu8N+4EYGonKDVS21uwhSGM4=", "Nw5/2KB+QECeFVCB/fxBEgS3AfftGXw+bFhIqLaf0I4=", "8KQIDhFJUjJhzm798P/qWaGMwv8/M4EK7msTwVsLgxk=", "1upZntVx11ut1j4XOdwqQgsho+new3/n94Kd/Btu8Rk=", "iuqQjnrr2PYBw3HqoMMoz4ln6ORKfOOvET3+JuSG31g=", "Pl+olzORMn0kCQFecRdVGG6owx5O5wD2YTpOnjrnHo4=", "ae1yX0AUsNipoQGV7sbYYvLK9AWsrYLwEtaSoNCVdMI=", "oMkpIme9WSTQOT5qYhUV4Cd2Pi1lskhQRzjL4AaBDZQ=", "BpRRvj4YPNwp78pQYHVU3OoA0akcy+APi1IcoBkSSSE=", "9sXHSOooXpYfC0khLe96eESh1it1z6+Gv4HmQVQZSJY=", "H7BK9I9wa5CBbs3MI2doHIqZy8u92kQ0N3kdMehlcUk=", "QhJJ+LR/jN7V5M/WjcsQUqO2uxxa26oJxajbA4whHIc=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590370072\nGwN1lo+P762wCRUXBMgV+ZZ5pDyrIrcZ+NBBYiM6cu0=\n\n— rekor.sigstore.dev wNI9ajBFAiEA9Sbu9RMrgP0ZXGqxBYElJ2m1361QAWYhEx9VOabnYP8CIHn8neeo3gM7zuqDOc2o6pPr1+FlwKled7+yeMsfbVd9\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmNjJlZTM2MzE2ZmQxN2EyM2MzOGEzOGI3YzE3NWM1NzRmNjllMzQxMTVlYzgyOGUwZTA5YmI3MDVhMWJmZTViIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSFFqaTRBUXBLME10SFpjSHFpQ05EcmhLT25PZU80cHU4OGc1LzJDRWhXSkFpRUEyUG5LR2tFc0RHYmkvdlpDbjM3eVd4em9PbnNlQ1h4a1A1QWd1bitnNG1nPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlJXMHZSMVYwVDBkdVVTOVVXbmd4TUVkNGVrSTJSVXBNYUc0MGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BSZDFkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxGM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ4U2xZMVZrOTNiRGg1VnpaS1VXcG1XaTlRUnpjM2JHTllSWHBZUzFVemRqRjJMMDBLVkRWR1ZtUk5URXcyUjFkeVpYaENUMHR0U2poMU0xRkRaVFJpV0ZCT1RIRmhTekEwYzBONVNucFJLekpPVkVReVpVdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZFZGxCekNsRnZjVFYwWm5vMFNHc3ZlVWRsVlVkbGQyRkxZWEF3ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazB6UnpoQlFVRlJSQXBCUldOM1VsRkpaMUJ4UlU1S1VEbG5PR05VYmxJeWQzQkVNMWR3WTFKWlNTOUpkWFZQV0VZNVpVZ3pXa3gzYXpKYVRtOURTVkZEYUUxcU4yWXhkMjlwQ2tWTlJrbzBRa3BJVmt0UFFrVkpWamQzVm5SVWIzbG9halFyYTI1UVRrc3hMM3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENNazlzZVZBS1RscDJSMVZGVVZvMmNWWkNiVE53Y1dGT05VNXVSbUpDVkRGT1VYRm9iR042VldFclFWZE5WSGxwZWt4TlNFaHhhV2MzV1dONlExZE1aR2REVFZGRVRBcFJNM053VVVGcFVEZFlibGh6ZEZsek1ESXdMM2RHYVhGUU1qaFVRbmcyS3k5Wk0weFBialpWVDNadldIaElibEZZWlVFMk4xZHNNU3R2YjNaRFZqUTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgZLzw1sr6OonoTWOxTUsJoj+Y4aEOLaYjc6JYcy7JWU4CFQDHzUiSknjVBIftSGRUQXUIBGRcgRgPMjAyNjA5MDQxNDIyNDBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MjI0MFowLwYJKoZIhvcNAQkEMSIEIO8hguSCb3HE8nk2UiPn1/7zp9KzWU7R4YnFivqwW0uPMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMDZB7TMxI2hDlLSAGpKWxI4CysIzTLCdos7ODJFr5Yp7gIXsLNqUUKIHbQd6/YvUAgIwB7TttD6Ug2MjT7d0aGVHQxy29IRBxCRyJ8GSwF3CsXd5iHd6KPaCZfJ6Wh0Ll9Yn"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"9i7jYxb9F6I8OKOLfBdcV09p40EV7IKODgm7cFob/ls="}, "signature":"MEUCIHQji4AQpK0MtHZcHqiCNDrhKOnOeO4pu88g5/2CEhWJAiEA2PnKGkEsDGbi/vZCn37yWxzoOnseCXxkP5Agun+g4mg="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..61d6f45bfd89ff3eae7f76b4a164900b66d9c1c7 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4166b274afb2be743b59de6eaaea5a5eae5cb11db779a6a105a3af403df8c70d +size 75006088 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch212-cxx11-cu132-x86_64-linux/cpp_functions.py b/build/torch212-cxx11-cu132-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch212-cxx11-cu132-x86_64-linux/layers.py b/build/torch212-cxx11-cu132-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..19c61b1806bac131c2dc0b0fce13698b2c9c47ec --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "QWaydK+yvnQ7Wd5uqupaXq5csR23eaahBaOvQD34xw0=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..8b790d7c3b906580c7566595bd28758e24a770f6 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUYAZG70VALPMX6ABwAPwtXzkyJCowCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjQzWhcNMjYwOTA0MTQzMjQzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEk+NZmSY+bKzK5t3gD9vxLLWxY1Bt2t1KRedfw7vaxS/VPuHBGvARLd+njOKIINBe1GqMERLSuvi1QT1xIcEbXaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUashmpWljW3TDR86zpOepwGqRPu0wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzM6ZIAAAQDAEcwRQIhAK38m0kCX8dDW3bBriXofhEqtovhwWhnfCZqxFHt0SZ+AiBx6E4EeJrEeRHMzE21jJBIGZEGnUdTLbP4/VZfx/n59DAKBggqhkjOPQQDAwNoADBlAjB0SKkBxIco3oOeP2dTLwhRhEI/Tv3Wmqr4WngHGOPTdnXFHZblXk8JNyRfMjm1i6kCMQCYwQOyI5rNPAThDjWDbwKAZ4tx5kmaQajddXnkY+Q2to8xhg2zJdC+LTBvfkPEZ4Q="}, "tlogEntries":[{"logIndex":"2712274654", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531764", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDh6OJ6lpnWXKyO2xxBP0wqWIQRNznencR+8W2K6+h1KgIhAO2bafJAxXX7PRPjT+EJwol93nImtyXHXgjZ24ISjCXQ"}, "inclusionProof":{"logIndex":"2590370392", "rootHash":"Kg2lr4U+FftmnFCjjSRwCDlOQ+YUraf143fYI61JyoQ=", "treeSize":"2590370495", "hashes":["jZa0UWh6rc6JyxF1zh12P0TrSojcA44zX3yDmkvuz7Y=", "EMQs3fR0TZWFn5t9WJ7Kw5o2oC8hfWHIILJE9rrX1Ng=", "T/Rr/B6UND0XNePfS+0zV58kmUYmgor574+TaKQZ9fw=", "8hGwLt/JJMJxVAf2SAwvKiBwwqKwihIARALT30tNgMs=", "RoR53FZk/O1jwxxwztQHdrsPDMz1EG1reQClkAScJRY=", "CtpWe866rpU7oBxd3Ga4PZi5HXcc+KARAj3CmxL4dPs=", "wQuI88fXxNVhR/6N0uw9GpZGNPUC/f7GaNLqwBisoY4=", "CGgrayrnC7lQqGXXevaCxlDgrvMfOraDd0LwueLmMZ8=", "A7DmwgbRSrc2LD36Ce5iqpi/zjfceiCQMFKHmZVjcEA=", "an7xDmqrJ5PN89Dsx6EqYZmD0qEyAoXB9zuf+Z0A9kc=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590370495\nKg2lr4U+FftmnFCjjSRwCDlOQ+YUraf143fYI61JyoQ=\n\n— rekor.sigstore.dev wNI9ajBEAiB/9ZvkITaTL4fEegzDXDnWv56LUPRH5OkA0FPcFVhgjgIgCJ1G7gkQ9X0zRB6omTjlqB3q09z0Yna7wRO2D+PXMF8=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwMDFhYzEzNTA0MzBkYzhjZTJhM2IyMDU5Y2JhNzVmNDAyNGEyOGMyZTYyNjdjMmE1Y2I4OWU5NTIyZGNiODNlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQmVBSVFDbmVWeTJZYkE4Q0x1SkFaZ1Y5UEkwaDNSOFpMTFpManNrREttbEFpQllEV1lkM2lvV05RYWtKOUNmR0p6ZmNLN3FuNzgrdjk2VmFObW9BQjM1enc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVldVRmFSemN3VmtGTVVFMVlOa0ZDZDBGUWQzUlllbXQ1U2tOdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BSZWxkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxGNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZySzA1YWJWTlpLMkpMZWtzMWRETm5SRGwyZUV4TVYzaFpNVUowTW5ReFMxSmxaR1lLZHpkMllYaFRMMVpRZFVoQ1IzWkJVa3hrSzI1cVQwdEpTVTVDWlRGSGNVMUZVa3hUZFhacE1WRlVNWGhKWTBWaVdHRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZoYzJodENuQlhiR3BYTTFSRVVqZzJlbkJQWlhCM1IzRlNVSFV3ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazAyV2tsQlFVRlJSQXBCUldOM1VsRkphRUZMTXpodE1HdERXRGhrUkZjellrSnlhVmh2Wm1oRmNYUnZkbWgzVjJodVprTmFjWGhHU0hRd1Uxb3JRV2xDZURaRk5FVmxTbkpGQ21WU1NFMTZSVEl4YWtwQ1NVZGFSVWR1VldSVVRHSlFOQzlXV21aNEwyNDFPVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENNRk5MYTBJS2VFbGpiek52VDJWUU1tUlVUSGRvVW1oRlNTOVVkak5YYlhGeU5GZHVaMGhIVDFCVVpHNVlSa2hhWW14WWF6aEtUbmxTWmsxcWJURnBObXREVFZGRFdRcDNVVTk1U1RWeVRsQkJWR2hFYWxkRVluZExRVm8wZEhnMWEyMWhVV0ZxWkdSWWJtdFpLMUV5ZEc4NGVHaG5NbnBLWkVNclRGUkNkbVpyVUVWYU5GRTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgQGuRQe4S5L5o66JTzZb9S22kViR9nUmjNl8O9DFf31MCFGsq9EH4XsEikIrO5kCSSqjlIrdUGA8yMDI2MDkwNDE0MjI0M1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQyMjQzWjAvBgkqhkiG9w0BCQQxIgQg4rcO9bre5zbK0ZFnSTKvLz8gHa0WvtQHaiIRk5uRhPAwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwObp51WML0oHu6IxSfinqzEJWJi0jL+CvJhlCodrBoxq9F2PZxkOaHIV9HbrH9NUUAjEAmu6IzvLNUrBuXyBEfzTRi4qFlpoJPlRMc+VVeqGEAOxXVFZSdNpr6lecwC72FnrM"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"ABrBNQQw3Izio7IFnLp19AJKKMLmJnwqXLielSLcuD4="}, "signature":"MEQCIBeAIQCneVy2YbA8CLuJAZgV9PI0h3R8ZLLZLjskDKmlAiBYDWYd3ioWNQakJ9CfGJzfcK7qn78+v96VaNmoAB35zw=="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-x86_64-linux/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..9626e4c3dbcfb86d33f4e88928b0d141557a2f66 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f21eb7a4c65dd97d54d45d014ae2c122b8f026cc186b83da7418157ae421dd34 +size 78812136 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_ops.py b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch213-cxx11-cu126-x86_64-linux/cpp_functions.py b/build/torch213-cxx11-cu126-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch213-cxx11-cu126-x86_64-linux/layers.py b/build/torch213-cxx11-cu126-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c82a1b556efb546b59fac32554f53e6144bf5e31 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json @@ -0,0 +1,44 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "7.0", + "7.2", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0+PTX" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "8h63pMZd2X1U1F0BSuLBIrjwJswYa4PadBgVeuQh3TQ=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..f87a373bd97d5ccf9200cfc6f8f7debfce8609b6 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUNS8KbUm2cJXtb02rFC9nJVxtICkwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjQ1WhcNMjYwOTA0MTQzMjQ1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAENFPz3Y5iCCmTuig61eJS2m7AkAYXJpKf756ZHxXjCTcmDJKMeeZR6afXw98gxbVtZgWw9kCEJ8oYAdKpnkK8MaOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUjTIH8truz6DFiEuJSaiB+xiA4RgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzM8PkAAAQDAEYwRAIgPlqgkirqsnGD526jzoDauqThlzdqP2y1fOSsFpQiAt8CIFAWtI2Jcp7I8tZ9nZnlg8FPVKvmJGsQVIy+TuwQhdfsMAoGCCqGSM49BAMDA2cAMGQCMHpyjroXuk7Itte7Vle0VonOEJuixoLD9ScFZyTGCmq1wJTpkpQau8agFW6Ofi39PwIwU1YQY01F8MfQp6eSn5tfNWgpNzOtEj5MUImq3KVmFrfpZ9rka6A8pqCAF4PDpiyA"}, "tlogEntries":[{"logIndex":"2712275010", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531765", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCoEQXv6CZ0h7sbt8KJ18ENZJjSTuP5/0hXAG/oSkQFxAIhAMkJPs0/BpNa3zMKUrZakWOQ8HtZBu0kifcflHIrhd2J"}, "inclusionProof":{"logIndex":"2590370748", "rootHash":"WOJLb40JriEpuLbrlz/LQ2ON2KZ63diIlkyYJTO/V00=", "treeSize":"2590371118", "hashes":["hoCwHV5wS2eX/GEAEYnRA32fARQgHdQI0NgY9gpsxck=", "DzhieRVOg5oDP9nufVqjN/xs+TCiaUmdMD53txGNuj8=", "xNQW5rvIe6HsBxQ584Tt0yO28Z5Kg/XU5pgVhCYEi2k=", "jt8hmsedaygU1zcccj2AkWYwOYhPRhZvMonuKb19UiQ=", "Fz/nPuKIZRFOvxwHgpV6RSBMZZNC7OEEuIQHk0LVOYI=", "4FGqcl05Fiv3t/BwnWcFl7hjF5mj8fbRiMOszVejPRA=", "N48MgevKmVgqqIUhg510nDYN6OX9j8KSH/voGD6kX8w=", "kD8wHY8/mWQf9NSx1KlBvrv/pCjwnpcGizYj75j9/1Y=", "V7TZwsfTzq2H/ncdMBBxWAa/YpFiAOmK4N7fGo+7YH0=", "A7DmwgbRSrc2LD36Ce5iqpi/zjfceiCQMFKHmZVjcEA=", "ux91q8ob/83kyNNg3zkFrSzwdug06vebJ9S2BIdw0z0=", "an7xDmqrJ5PN89Dsx6EqYZmD0qEyAoXB9zuf+Z0A9kc=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590371118\nWOJLb40JriEpuLbrlz/LQ2ON2KZ63diIlkyYJTO/V00=\n\n— rekor.sigstore.dev wNI9ajBFAiAd8xsy+pEONo4M6B960TxtTwa4GSAkmyQa9EbRbu0fxQIhAJkyFwVHEtEpNvo5T1NG/bQEcvAINjFrndZSq4yTVaYF\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5NTZkZDlkYzdiZmMyNmRiYWMwYjVjMDQxMzUxODRkMmY0ZjE0MzA0OTJmZjA4M2JlYjU2ODBjMzcyYWU0Y2VhIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNySVRQZEFBME5IeU05MTlxZXhzWnNrQ2xtZ3JXaVhUYjlMMWRCSFVubHZ3SWdjZkVQMTBVTG4zRTQyNmxaZCtKdG85ODAxLzh2UUpId3FwVlprSFo4aHZrPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlRsTTRTMkpWYlRKalNsaDBZakF5Y2taRE9XNUtWbmgwU1VOcmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BSTVZkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxFeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZPUmxCNk0xazFhVU5EYlZSMWFXYzJNV1ZLVXpKdE4wRnJRVmxZU25CTFpqYzFObG9LU0hoWWFrTlVZMjFFU2t0TlpXVmFValpoWmxoM09UaG5lR0pXZEZwblYzYzVhME5GU2podldVRmtTM0J1YTBzNFRXRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZxVkVsSUNqaDBjblY2TmtSR2FVVjFTbE5oYVVJcmVHbEJORkpuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazA0VUd0QlFVRlJSQXBCUlZsM1VrRkpaMUJzY1dkcmFYSnhjMjVIUkRVeU5tcDZiMFJoZFhGVWFHeDZaSEZRTW5reFprOVRjMFp3VVdsQmREaERTVVpCVjNSSk1rcGpjRGRKQ2poMFdqbHVXbTVzWnpoR1VGWkxkbTFLUjNOUlZrbDVLMVIxZDFGb1pHWnpUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVaHdlV3B5YjFnS2RXczNTWFIwWlRkV2JHVXdWbTl1VDBWS2RXbDRiMHhFT1ZOalJscDVWRWREYlhFeGQwcFVjR3R3VVdGMU9HRm5SbGMyVDJacE16bFFkMGwzVlRGWlVRcFpNREZHT0UxbVVYQTJaVk51TlhSbVRsZG5jRTU2VDNSRmFqVk5WVWx0Y1ROTFZtMUdjbVp3V2pseWEyRTJRVGh3Y1VOQlJqUlFSSEJwZVVFS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgQfFgdVC2uA283k4wG2/MOiPRUiQ6piqgf0m0omrRB7kCFQCAuYhrVCVT57rCY2E3/EpcX8xauxgPMjAyNjA5MDQxNDIyNDVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MjI0NVowLwYJKoZIhvcNAQkEMSIEIHnA9HVgYJMUrqjDANy9gy5N+bDW/DCFCLhj3ajj1IWUMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCZbxfe1smZvPl76BDtLzO/UD9RYjrf8cAVVOsykfT62iuAFyy2lkAgEpxJL6KHmEYCMQDUp28X/vEUZRrAYEbx/vGyTudTjf+ak7XqFSxHXVJjPCTv4y6eaBx6BnEBNrSVVWo="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"lW3Z3Hv8JtusC1wEE1GE0vTxQwSS/wg761aAw3KuTOo="}, "signature":"MEUCIQCrITPdAA0NHyM919qexsZskClmgrWiXTb9L1dBHUnlvwIgcfEP10ULn3E426lZd+Jto9801/8vQJHwqpVZkHZ8hvk="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-x86_64-linux/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..f3d2b0a9a383b8e8fc26947d29e3f280614a2959 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ccc666dc1cf38c3114cf46a41c893a5516d0fc6d68b5c144356e76df05ff08e +size 74854720 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_ops.py b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch213-cxx11-cu130-x86_64-linux/cpp_functions.py b/build/torch213-cxx11-cu130-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch213-cxx11-cu130-x86_64-linux/layers.py b/build/torch213-cxx11-cu130-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..7905b3c23ef730aecb52ac12cda2d526dbd24978 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "TMxmbcHPOMMRTPRqQciTpVFtD8bWi1wUQ1bnbfBf8I4=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..124b7dbecfbba5e84fd5da652c69f5c389729e98 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtCgAwIBAgIUFng7IaiS6DcAEK1EHv1ld6FprIkwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjUxWhcNMjYwOTA0MTQzMjUxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEAfNwk9j+mov18tY+YJuYa8YNjGFXAj/Yc8dHrll/ZDcB6e0kgdwF3zfe0MYcj76deRrylfU5Af2TUG889jRmBKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUR+HegA6+22p2w55NouwHyFNn2uUwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzNCFsAAAQDAEYwRAIgDas3h2e3EU+dnTZxnSc3qF7CNSjZmrM3ekKFXZMmwwUCICkOylD3KsWasx5YiryT0EkKYin+78nQmL1kcMxexuT6MAoGCCqGSM49BAMDA2kAMGYCMQCG8zrKsEVY3UPMmKhczJPZfG8GpBTMo+DyuePxX27nZqRpWaOBCPx6ul2nfFaLAGACMQD6VG2IxStxI1roTxRKn7vK+LLd63GvB8NL/ExA80zVVQQvOFCuHI7tgL7QTFACz10="}, "tlogEntries":[{"logIndex":"2712276137", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531771", "inclusionPromise":{"signedEntryTimestamp":"MEUCIGadSeKmzAGHrjPxl8bvHa6we6oyrEZyQg9bt7TPHvAdAiEAgDnHb3btVNkHWMv9S6mW12K+5xU/mmD/Sm8TbAH1s1w="}, "inclusionProof":{"logIndex":"2590371875", "rootHash":"FD4RtxjaXvqzyhL8fK45M9POAwfrWoPNxZnrcPZzoHI=", "treeSize":"2590372130", "hashes":["+uFFkIpsyGgj6F1kqukEDjyAcqnU+6jOmz4DH/YfXvw=", "4R2actF4DeT013caMxcAIU5myT0vTkQNdKz8kdp8V7k=", "KZAa5Mt5ZNK/ZqE2UhfcILjRCACpIC6wwoXLqd/FcN4=", "mCKuzAGn4Em4TVoP0iLMrQ9UC5oof5FZxrfxdIeYwXI=", "rh2WhSLJ92WU4txzthkYBdIpxYN76cWiL0ng7EACAqs=", "Hns2juWMCyRqiF6cisfU31JLW8b3o3ph2U2ZqDezMbA=", "ER8h92bO0Kq6UpkuLrCoiifQ3SLokTc/N9trlL9RFTg=", "dWUoVFc56lw215SZWOa2tDKeAeerJBqHpb3xWAaaAuE=", "yU27m1URS+cvYLdBse4fR8NrFpYa5oddLjgt+GIbP6o=", "t8nRtM4piqY9M2vLuqc30bnlaHH5a1Mc5+3BX+VcR48=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590372130\nFD4RtxjaXvqzyhL8fK45M9POAwfrWoPNxZnrcPZzoHI=\n\n— rekor.sigstore.dev wNI9ajBFAiATHYwVsYNgFh7kY2hKuMlICTcx6ZkaXVpY6cibhYsOvgIhAN1QAO4VXnHZsCKOwhfYnFuJ55el40A6pUa1Fmhlf6X3\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkYmUzYWY0NzEyMTM2OWJiMDNmNzIwMDEwOTVkNDI5ZTc0MDljNjVhZTFiNjAwYTJkNGI2MjhhMjZhZWMyMjRlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUR0b0o0MXllZFRJL3AwZ2lWM3dJMHRrZXB2NzZtREhGWXJta3oyanFwWEJ3SWhBSVNEclM0UkplSC9Wb25telhkWTRLeUlqMk5sT25Ra0lKc1BkY0JyZnB2SCIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SRFowRjNTVUpCWjBsVlJtNW5OMGxoYVZNMlJHTkJSVXN4UlVoMk1XeGtOa1p3Y2tscmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BWZUZkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZCWms1M2F6bHFLMjF2ZGpFNGRGa3JXVXAxV1dFNFdVNXFSMFpZUVdvdldXTTRaRWdLY214c0wxcEVZMEkyWlRCcloyUjNSak42Wm1Vd1RWbGphamMyWkdWU2NubHNabFUxUVdZeVZGVkhPRGc1YWxKdFFrdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZTSzBobENtZEJOaXN5TW5BeWR6VTFUbTkxZDBoNVJrNXVNblZWZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazVEUm5OQlFVRlJSQXBCUlZsM1VrRkpaMFJoY3pOb01tVXpSVlVyWkc1VVduaHVVMk16Y1VZM1EwNVRhbHB0Y2swelpXdExSbGhhVFcxM2QxVkRTVU5yVDNsc1JETkxjMWRoQ25ONE5WbHBjbmxVTUVWclMxbHBiaXMzT0c1UmJVd3hhMk5OZUdWNGRWUTJUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlhMEZOUjFsRFRWRkRSemg2Y2tzS2MwVldXVE5WVUUxdFMyaGpla3BRV21aSE9FZHdRbFJOYnl0RWVYVmxVSGhZTWpkdVduRlNjRmRoVDBKRFVIZzJkV3d5Ym1aR1lVeEJSMEZEVFZGRU5ncFdSekpKZUZOMGVFa3hjbTlVZUZKTGJqZDJTeXRNVEdRMk0wZDJRamhPVEM5RmVFRTRNSHBXVmxGUmRrOUdRM1ZJU1RkMFowdzNVVlJHUVVONk1UQTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgLaJh1aWg7ofZx3Toj1yuSM6AIW/dEdxW1nnEwXyp1jUCFQC2ic5BNY2RkoZ0ll5krqye7+B38BgPMjAyNjA5MDQxNDIyNTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MjI1MVowLwYJKoZIhvcNAQkEMSIEIP+JnfDVRzYf87aVz2P0LMzIpSAlBdNHgkpzo8Gf2yuQMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMG1KfJ7CeRECYJoMm1nwzJxPmFAFP0nH7CNFnWZrtbF29lrVDS7BAmA51yBFPUkm2AIwFm920GCh78DpLL1tqnP/7d8GmD6vGPqtY1d7AvL+zRXjuBlvf9sy9JFALTUwYD+X"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"2+OvRxITabsD9yABCV1CnnQJxlrhtgCi1LYoomrsIk4="}, "signature":"MEYCIQDtoJ41yedTI/p0giV3wI0tkepv76mDHFYrmkz2jqpXBwIhAISDrS4RJeH/VonmzXdY4KyIj2NlOnQkIJsPdcBrfpvH"}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-x86_64-linux/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..98a6c6247cdc97af90c430e0b3cc7c7762eecb8d --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c960e2218e9c78baad1940fca084f582fe5e1b8efc6aa375adb8cdee64195051 +size 75006024 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_ops.py b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch213-cxx11-cu132-x86_64-linux/cpp_functions.py b/build/torch213-cxx11-cu132-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch213-cxx11-cu132-x86_64-linux/layers.py b/build/torch213-cxx11-cu132-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..cb3801b8568b101af7d0ecd51c74ea08a29a14b6 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "yWDiIY6ceLqtGUD8oIT1gv5eG478aqN1rbjN7mQZUFE=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..435cf29d91f5c952ea5fac10a57dcf7612a4e639 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUMNzhX+sttr0EzZbGBv2OVR/u24MwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjU0WhcNMjYwOTA0MTQzMjU0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEgLDbN+t8HKxqSCOeyuEBOSxQTi14hCfJHh19PdB+16+EewNrG6T17Io6WJ1iZHcd/syW4+RQd+ZWuH/S7F3fjaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU2ZrgpwSytSmP/dgBdjMzsFY2CLQwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzNFSgAAAQDAEcwRQIgSn2vXY/k2h4zij01aCh62eq8TZRHysXF/sZyMjkYUuMCIQCTU/CUi+ia5Z1vPe1cfaivnEIVFi0y5lNDxWOT3UM5gjAKBggqhkjOPQQDAwNnADBkAjACvo/wfvt+ZAENlR8Nt8XAK+kewjb6q9IBR50/GxDP2r/wuvJkzo88tTgd/2KMcDoCMAYNQynvnu+lOLDN1RimCVXK75/Z407yL3xmrzNYE+MJfNNR2ticAF+ZSMb36R0p6g=="}, "tlogEntries":[{"logIndex":"2712276860", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531775", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDaBfIujK3VAxqwxy9CkheoUR25ysRVfWXPT8XChPDhbwIgSQiWIVR8jDn1Ds+o1V3BOZfRTaBdqEuW5PhP6lz0L6A="}, "inclusionProof":{"logIndex":"2590372598", "rootHash":"3zBBzKYl/gvEFwbi1B/z3d1UjloEwLob0eoei78rdn4=", "treeSize":"2590373015", "hashes":["n/9x7tyglNeFETZzTp6rY0oO0kHNhE+L4EH3BenuRSk=", "pnQ20MCI/n8QskCvw/uD9My8FXCRL4Yn4rgMaS3rLl8=", "yFDbS8U775HFh4dd0e9F47VCJyzFg594N/pMYfwGN3I=", "7eDEBCh3E5OeSY3+TF7U0oSUYuiQ6jwA4CbDhbKmb+4=", "hZNG19EtQwgp95zH4FchSYzPAJKQ1dv4hpFW65uvvnw=", "iQD/9FISAFkSGpzq7gidoXdJeC/OpVac7jdsJ1UiyyQ=", "eivSjj7ro5EoRhYTxfsxtHZxjI7rdGJ50drXuuvI38A=", "+2KqRqe/Mg8wq37GaCVTyFtPkf4snXNx204HZ38AHKA=", "bXcePXVPVlcEbEJ8/OVx0b5fajcxzgZM0w+VBkrrm08=", "YK6duE+sfXoZkn8Q1c5ZHamOoon8m2K4S1HBsLhWHAQ=", "xnHQZYKaSGd9Toi7TJoxtEEV77mbPJ9v49jSm9hSWAo=", "t8nRtM4piqY9M2vLuqc30bnlaHH5a1Mc5+3BX+VcR48=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590373015\n3zBBzKYl/gvEFwbi1B/z3d1UjloEwLob0eoei78rdn4=\n\n— rekor.sigstore.dev wNI9ajBGAiEAljScYS7ar4L/oMs3mitUcomcdSqmNezoNykPLY/4QRkCIQCz2LS4MF7z2YesCuRzlusPBT1v0ZbK+w2WOdB3QbKWRw==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkNGVkNDg4MDc1ZTU2MGYyNmVhMjI4MWI5YTk1MWQ0MWZkNTMyMTE1ZDdhZWI2YzExODBkY2UwNTkyNmQ5Y2U4In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURGeFVGQWNMREdpMnQwQjd2UndPdGZUSDlObUJHdWRNUDJiUFNyblN4T2VnSWhBSVpnR3BRVVZyY05jVlM4MC9qT2cydnJrdW1ubjZ5dnNzRXk5akdUS0RwNiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlRVNTZhRmdyYzNSMGNqQkZlbHBpUjBKMk1rOVdVaTkxTWpSTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BWTUZkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxVd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZuVEVSaVRpdDBPRWhMZUhGVFEwOWxlWFZGUWs5VGVGRlVhVEUwYUVObVNraG9NVGtLVUdSQ0t6RTJLMFZsZDA1eVJ6WlVNVGRKYnpaWFNqRnBXa2hqWkM5emVWYzBLMUpSWkN0YVYzVklMMU0zUmpObWFtRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV5V25KbkNuQjNVM2wwVTIxUUwyUm5RbVJxVFhwelJsa3lRMHhSZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazVHVTJkQlFVRlJSQXBCUldOM1VsRkpaMU51TW5aWVdTOXJNbWcwZW1scU1ERmhRMmcyTW1WeE9GUmFVa2g1YzFoR0wzTmFlVTFxYTFsVmRVMURTVkZEVkZVdlExVnBLMmxoQ2pWYU1YWlFaVEZqWm1GcGRtNUZTVlpHYVRCNU5XeE9SSGhYVDFRelZVMDFaMnBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJRM1p2TDNjS1puWjBLMXBCUlU1c1VqaE9kRGhZUVVzcmEyVjNhbUkyY1RsSlFsSTFNQzlIZUVSUU1uSXZkM1YyU210NmJ6ZzRkRlJuWkM4eVMwMWpSRzlEVFVGWlRncFJlVzUyYm5VcmJFOU1SRTR4VW1sdFExWllTemMxTDFvME1EZDVURE40YlhKNlRsbEZLMDFLWms1T1VqSjBhV05CUml0YVUwMWlNelpTTUhBMlp6MDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgXErvz094FXSKTPzqjTAM5/ueBZNY+oYOpaekCtINbW8CFFytH/5qa4Cgzq/P27UhFKkac8hKGA8yMDI2MDkwNDE0MjI1NFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQyMjU0WjAvBgkqhkiG9w0BCQQxIgQgNZXJK7sgCjnza/kjAiNzz4h7Mi42LvvkKOQMxj61MnowgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxANaLfSZanoFoFuKapQgMiKpVaIhnrFOhqwi7bij9alwArl0FE0wSDoO5K1v1gFYzzAIwWwzo5kSKM3htipODQvYCQ8TSa7HQpG7AyURQ5uskHCCDy1XzG80jHRTd9sZMJLob"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"1O1IgHXlYPJuoigbmpUdQf1TIRXXrrbBGA3OBZJtnOg="}, "signature":"MEYCIQDFxUFAcLDGi2t0B7vRwOtfTH9NmBGudMP2bPSrnSxOegIhAIZgGpQUVrcNcVS80/jOg2vrkumnn6yvssEy9jGTKDp6"}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..c2fc857f6f9be2934bcbebe85ce81fd00f545127 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c767cfeadd09f5f4c32b220212e7bcb4a30fc4439eca808784aceae0887d510 +size 78828744 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch214-cxx11-cu126-x86_64-linux/cpp_functions.py b/build/torch214-cxx11-cu126-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch214-cxx11-cu126-x86_64-linux/layers.py b/build/torch214-cxx11-cu126-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/metadata.json b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..9c697ca92f8eadab7c673dd4ef278c857c8a095d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json @@ -0,0 +1,44 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "7.0", + "7.2", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0+PTX" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "nHZ8/q3Qn19MMrIgIS57y0ow/EQ57KgIeErOrgiH1RA=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..7403b04d44d049b9e1fbdb2ba99583ffabb7ed80 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUX3nCIwizfqd420GeKURtRYZWZZQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMjU4WhcNMjYwOTA0MTQzMjU4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEFDTAScgFcdk4w02sVXw+hPFYXV1nyfTSrcUa99595J2yv60p5z6yMZsFvJ26hplV6lco8iA8SauYfAR0/lz59KOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUJa0BAkAZjdQLsIu+wC056CslEecwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzNI+oAAAQDAEcwRQIhALnb/VhrM01mNrWLbPENdB+KfI9p5ZqZrdFEBGT5GdkbAiAnLcKyQRDfekfVbnDTaWB+IJ4kxEdhwnBdjLIYqrFLVjAKBggqhkjOPQQDAwNoADBlAjANZKOoU9YS1lNCvOV4nsnus9HGoDzjX5y1quw6i3Rq6qkgmkJuQi8m5zp0Ks7heUMCMQC5Pj6D2u/m3nn+za5Dr7ofmVivVtUqplhouIyTMlxcm1Qo7RTkGpUJGvkB3XTGRgE="}, "tlogEntries":[{"logIndex":"2712278093", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531781", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCCDQcj91Tx6neog6tgVxHnDRzy9mUQfvLEoietoPrhHgIgVeJ6H2dOj7K2IeOGJa5BvRD5V3+e15LIQbX1EiiXTAk="}, "inclusionProof":{"logIndex":"2590373831", "rootHash":"mS6d2w9Yexafn0f8D8k1bdmfYZvCa61SBQ26J2zga8E=", "treeSize":"2590374174", "hashes":["cQf8Y/GeNI9Q/bqIYuvb71rIESmna4cp6bbD5paQOdk=", "nLGRU51CWANMreNFumVez9S9RxVgoKaQrluFSVcgR1M=", "q40AaV6gPAL3of/tZOPrcHGwQrhAS8Iswv3y2Nz0mos=", "sPs1VtwovLMXOIE60ROd1Ybj6EDEw/Fvhofm7WZarVc=", "ZihQqROa+CXiMT1fdKzMNKFB2ADaTgZCsCINuhdiAjE=", "KPoiPkxZie0a9iieWNN38t1Mjb+b6V0xPnwcAQjnoz0=", "Am17RejWPbkeMvU4DjLjmZhXu810VaF/mujXJgfHkZE=", "pywrNoqFd6256oeu2Wi63Q7eBpkgbl1JFfKlB2rHXN0=", "RFqobRBrcOXxeURSxmi6P4EoTVNmcVGzn9pMRHCAYeg=", "tOxPxNOAs8vBN4lKmAkaU0EhUIXa1F14VNzCbSQOZV8=", "pQLGdzKh06tR0Po68Y2+yI2+u1azZf2cgF74eD/o6Ls=", "umDboCcZnwAfz9I2MEymJ+/7uzkoaoX+gkS47OpKmr4=", "t8nRtM4piqY9M2vLuqc30bnlaHH5a1Mc5+3BX+VcR48=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590374174\nmS6d2w9Yexafn0f8D8k1bdmfYZvCa61SBQ26J2zga8E=\n\n— rekor.sigstore.dev wNI9ajBFAiA69iNUY6wbV+vKS0/zDmfj2uLAtrTUVYXl3lVN5QucZwIhALf3gE54e9NxkF4mGzJAG3VhE2r3n/ye2AYx8pf+fWIn\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIxZjIxMmRiMDI0ZDE4ZjNhY2Q5YTYyNDZkMGFlM2MzNTk5YjQ4NTQxZDk0NmM3ZmFhODY0N2ZhMTYzODMzMjBjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUMwYVpHbXR0THB1V0E3Wm11Vi9ZZ202KzFiQXozR25qbFFQVW9KYXMxTWZ3SWhBUDlwaUlOb2MzNnBkd2s2R2JWRVZMTmpOSHNZU1RMejVsdHVLc1NUQWwzciIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVldETnVRMGwzYVhwbWNXUTBNakJIWlV0VlVuUlNXVnBYV2xwUmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUV3BWTkZkb1kwNU5hbGwzVDFSQk1FMVVVWHBOYWxVMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZHUkZSQlUyTm5SbU5rYXpSM01ESnpWbGgzSzJoUVJsbFlWakZ1ZVdaVVUzSmpWV0VLT1RrMU9UVktNbmwyTmpCd05YbzJlVTFhYzBaMlNqSTJhSEJzVmpac1kyODRhVUU0VTJGMVdXWkJVakF2YkhvMU9VdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZLWVRCQ0NrRnJRVnBxWkZGTWMwbDFLM2RETURVMlEzTnNSV1ZqZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazVKSzI5QlFVRlJSQXBCUldOM1VsRkphRUZNYm1JdlZtaHlUVEF4YlU1eVYweGlVRVZPWkVJclMyWkpPWEExV25GYWNtUkdSVUpIVkRWSFpHdGlRV2xCYmt4alMzbFJVa1JtQ21WclpsWmlia1JVWVZkQ0swbEtOR3Q0UldSb2QyNUNaR3BNU1ZseGNrWk1WbXBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJUbHBMVDI4S1ZUbFpVekZzVGtOMlQxWTBibk51ZFhNNVNFZHZSSHBxV0RWNU1YRjFkelpwTTFKeE5uRnJaMjFyU25WUmFUaHROWHB3TUV0ek4yaGxWVTFEVFZGRE5RcFFhalpFTW5VdmJUTnViaXQ2WVRWRWNqZHZabTFXYVhaV2RGVnhjR3hvYjNWSmVWUk5iSGhqYlRGUmJ6ZFNWR3RIY0ZWS1IzWnJRak5ZVkVkU1owVTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgFCdZHOIsKiNn0qt/6y9sz+L4pQyChmJTDGObuPFWiQgCFEaGJIbI2qX59C/kC1Kjf5nARzE+GA8yMDI2MDkwNDE0MjI1OFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQyMjU4WjAvBgkqhkiG9w0BCQQxIgQg3IoiTa61wfPjYep/itRK1tj+/lzVbewwfbGRBY+H/h0wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAM5v8wUui733nMJZ/lhrmhScjj2d+cc7RGFietK8S23SsTOWt6RS1BDGOmPgC1EFvQIxAIQked/gkYW7mmakQBGMCAYYOrFRFMoBaSpZu/RxRfB5VDst5Pa1Hs61ymcRTjtmTA=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"HyEtsCTRjzrNmmJG0K48NZm0hUHZRsf6qGR/oWODMgw="}, "signature":"MEYCIQC0aZGmttLpuWA7ZmuV/Ygm6+1bAz3GnjlQPUoJas1MfwIhAP9piINoc36pdwk6GbVEVLNjNHsYSTLz5ltuKsSTAl3r"}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..341b8a4a5adddac4ef261e1e909e10573f26004f --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:36851b6ada5ba33a0d11fe196f5a8ee1cfb7a1ebb95b900f4f5ea78d86331dcc +size 74938800 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch214-cxx11-cu130-x86_64-linux/cpp_functions.py b/build/torch214-cxx11-cu130-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch214-cxx11-cu130-x86_64-linux/layers.py b/build/torch214-cxx11-cu130-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/metadata.json b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..f4b108226ff29227d5772dfd0a4fd165d94eb0a5 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "NoUbatpbozoNEf4Zb1qO4c+3oeu5W5APT16njYYzHcw=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..8195b9bcf7e3bf564a85662d872c618dbaf38ee5 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUByHUr7ld121rShRqX3/j/1eVtlMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMzA2WhcNMjYwOTA0MTQzMzA2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE1pLopbliRjdlLyUBemBVTNhIXbFoAM6REz0f9Kjwv9w8tQjHZ2MG7TzGObgGn35O2sR3p1CzGGJqGSBbH9JRkKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUC90N/XQ0hJv9GoP1rWlajlca9PMwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzNQ48AAAQDAEYwRAIgIKQYcei/or08f6toZbYWPhdBVq23BSvWbGKoLFFQBKMCIBX3EWzljawROJ8PqC1/oVDsUZ06TDRzWprYFUVrYnLAMAoGCCqGSM49BAMDA2cAMGQCMB/FOgkp4z9Sa674w9dQVrxt3f4sbtQMhLLdi0eWAiFG9AVwUFp0JXIgH7p8N77plAIwHmhOD77THSvvbirkzGMdz79whljZ+HbVXuSUloxtIuV9/FznFFdJ6RAGiWSczP/4"}, "tlogEntries":[{"logIndex":"2712279113", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531786", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCgscutCijakRfKn1ktW3zefv3Pe4OY0u34uL78l9s3LgIhAM3SjnISUN4uSTTdg5/x/ktR/LrBSm0OqDaXFjb2GhtN"}, "inclusionProof":{"logIndex":"2590374851", "rootHash":"Y4hhdV8boWtaY4KEUcecokF+XGnqbXM1fFSE3SLNOhA=", "treeSize":"2590375215", "hashes":["sFmXkes9ycMxNf8D7wx/ARnWeHSfDxOQmmitv7KbNVg=", "Qgr5Rbe1Clg5eKaZxk84EqYCT0yXxOU4qiPRkuYosDY=", "ONXdCpFxQqmdMgCGq9ato5eDYChWCjuF1npStSj6RmU=", "+FpMfHjvTS7XevNbYb82an5xvkyeKB20gcMDUspLLK0=", "4sAl+ExkDozr1tptGXpphWnyzR4ykBmJe6gWcfgfCD0=", "RSB8ZQE01jvF8kgGkeWe8uCddCMaQhvJYQRh82fz1Kc=", "geB3F9IuClf8gbfwjkH26jQMA1eQ3hUkwpT9TZJXHr0=", "zjefqTwaXGHXjXt1fT34b9FtzFeV4uVXzNTRVCa2uc8=", "7SPnPV/HpBw1gEVnJF9hPaB4VHNVMs0JBnPD3t+zHYc=", "r5iGlBwsppIQ6/qYFjbNNIoDfXlSpgr7ICN2CTZVLkA=", "7ZQDlUKfpYyElQO9sT0UXuHhy9+hTAD/AszLEceiIpw=", "W93RUf9D/+sJn+qU9XdapIH75v3rh+OVEdOHM6PMPkA=", "t8nRtM4piqY9M2vLuqc30bnlaHH5a1Mc5+3BX+VcR48=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590375215\nY4hhdV8boWtaY4KEUcecokF+XGnqbXM1fFSE3SLNOhA=\n\n— rekor.sigstore.dev wNI9ajBFAiEAsCL82lWE3MkDaH4tEf8VOBgMEsSvAV01OGAG/BtI5ooCIBVZQt5RL1KjR/88eIYXBeH09DiFkpdBRf3r9i5u+a9k\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4Y2JkN2JkOTcxODQ5YmUzZDEwYTNlNjhiYTdkNjA0ZTQ4N2YwMGQxMWM2OGQzMTgxZThjZGFlMzYwYTk3MDQ3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRHdSMU5wZ0dxdFkvTjZqR05nZnQzYWlvTWFaY0FxMkl6T3hlZHV4dy9zQUFpRUF6NmsxNmhYdC9kN1IvK3k2L1djTzVCUnRLNkpEUTRrYVN6c25qcFg5dGZNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlFubElWWEkzYkdReE1qRnlVMmhTY1Znekwyb3ZNV1ZXZEd4TmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUWHBCTWxkb1kwNU5hbGwzVDFSQk1FMVVVWHBOZWtFeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV4Y0V4dmNHSnNhVkpxWkd4TWVWVkNaVzFDVmxST2FFbFlZa1p2UVUwMlVrVjZNR1lLT1V0cWQzWTVkemgwVVdwSVdqSk5SemRVZWtkUFltZEhiak0xVHpKelVqTndNVU42UjBkS2NVZFRRbUpJT1VwU2EwdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZET1RCT0NpOVlVVEJvU25ZNVIyOVFNWEpYYkdGcWJHTmhPVkJOZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazVSTkRoQlFVRlJSQXBCUlZsM1VrRkpaMGxMVVZsalpXa3ZiM0l3T0dZMmRHOWFZbGxYVUdoa1FsWnhNak5DVTNaWFlrZExiMHhHUmxGQ1MwMURTVUpZTTBWWGVteHFZWGRTQ2s5S09GQnhRekV2YjFaRWMxVmFNRFpVUkZKNlYzQnlXVVpWVm5KWmJreEJUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVSXZSazluYTNBS05IbzVVMkUyTnpSM09XUlJWbko0ZERObU5ITmlkRkZOYUV4TVpHa3daVmRCYVVaSE9VRldkMVZHY0RCS1dFbG5TRGR3T0U0M04zQnNRVWwzU0cxb1R3cEVOemRVU0ZOMmRtSnBjbXQ2UjAxa2VqYzVkMmhzYWxvclNHSldXSFZUVld4dmVIUkpkVlk1TDBaNmJrWkdaRW8yVWtGSGFWZFRZM3BRTHpRS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgYE0fzGxfoGRd9HBhrVDhU/RtqJPRl6PMEL2+dp4H02wCFDEwfalEFKLeP00FLhHKqzwRaA0bGA8yMDI2MDkwNDE0MjMwNlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQyMzA2WjAvBgkqhkiG9w0BCQQxIgQgp+yyWwuKWgMERNqZ8TbvvKvk0derktGWWSw/zJcdpPgwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwNZsi+7HEV3x8HujKrcTviC9GTy7wJttPdBTGIYXz3LzBETbcfvRE1S1vNbGUkzsdAjEAtaVeStpK5bcC+dF6pIB05sPD3RImyfoZJSseSn4XJLJTo+xHcDIyM6N/URz1suu4"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"jL172XGEm+PRCj5oun1gTkh/ANEcaNMYHoza42CpcEc="}, "signature":"MEUCIDwR1NpgGqtY/N6jGNgft3aioMaZcAq2IzOxeduxw/sAAiEAz6k16hXt/d7R/+y6/WcO5BRtK6JDQ4kaSzsnjpX9tfM="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/__init__.py @@ -0,0 +1,12 @@ +from . import layers +from .causal_conv1d_interface import causal_conv1d_fn, causal_conv1d_update +from .causal_conv1d_varlen import causal_conv1d_varlen_states + +__all__ = [ + # wrappers + "layers", + # originals + "causal_conv1d_fn", + "causal_conv1d_update", + "causal_conv1d_varlen_states", +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..baccb8325cb04392a45a3221cc16d3e817600816 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f1e22ef77a7ad2e6bd99d882905328b8d34a7b83e095e4c01f71a977fde1bdd8 +size 74982024 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _causal_conv1d_cuda_056e28b +ops = torch.ops._causal_conv1d_cuda_056e28b + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_causal_conv1d_cuda_056e28b::{op_name}" diff --git a/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_interface.py @@ -0,0 +1,242 @@ +# Copyright (c) 2024, Tri Dao. + +import torch +import torch.nn.functional as F + +from .cpp_functions import causal_conv1d_fwd_function, causal_conv1d_bwd_function, causal_conv1d_update_function + + +class CausalConv1dFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(2) != 1 and x.stride(1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + if seq_idx is not None: + assert ( + initial_states is None + ), "initial_states must be None if seq_idx is not None" + assert ( + not return_final_states + ), "If seq_idx is not None, we don't return final_states_out" + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + if initial_states is not None and ( + initial_states.stride(2) != 1 and initial_states.stride(1) != 1 + ): + initial_states = initial_states.contiguous() + if return_final_states: + assert ( + x.stride(1) == 1 + ), "Only channel-last layout support returning final_states_out" + if final_states_out is not None: + assert ( + final_states_out.stride(2) == 1 or final_states_out.stride(1) == 1 + ) + else: + batch, dim, seqlen = x.shape + width = weight.shape[1] + final_states_out = torch.empty( + batch, width - 1, dim, device=x.device, dtype=x.dtype + ).transpose(1, 2) + else: + final_states_out = None + ctx.activation = activation in ["silu", "swish"] + out = causal_conv1d_fwd_function( + x, weight, bias, seq_idx, initial_states, final_states_out, ctx.activation + ) + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = return_final_states + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out if not return_final_states else (out, final_states_out) + + @staticmethod + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + None, + ctx.return_dinitial_states, + ctx.activation, + ) + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + dinitial_states if initial_states is not None else None, + None, + None, + None, + ) + + +def causal_conv1d_fn( + x, + weight, + bias=None, + seq_idx=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, dim, seqlen) + """ + return CausalConv1dFn.apply( + x, + weight, + bias, + seq_idx, + initial_states, + return_final_states, + final_states_out, + activation, + ) + + +def causal_conv1d_ref( + x, + weight, + bias=None, + initial_states=None, + return_final_states=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if return_final_states: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not return_final_states else (out, final_states_out) + + +def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None, conv_state_indices=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + activation = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + out = causal_conv1d_update_function( + x, conv_state, weight, bias, activation, cache_seqlens, conv_state_indices + ) + if unsqueeze: + out = out.squeeze(-1) + return out + + +def causal_conv1d_update_ref(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/causal_conv1d_varlen.py @@ -0,0 +1,86 @@ +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +@triton.jit +def _causal_conv1d_varlen_states( + X, + CU_SEQLENS, + STATES, + state_len, + dim, + stride_x_seqlen, stride_x_dim, + stride_states_batch, stride_states_seqlen, stride_states_dim, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr +): + batch_idx = tl.program_id(2) + STATES += batch_idx * stride_states_batch + end_idx = tl.load(CU_SEQLENS + batch_idx + 1) + start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len) + rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] * stride_x_dim, + mask=(rows[:, None] >= start_idx) & (cols[None, :] < dim), + other=0) + rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M) + tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[None, :] * stride_states_dim, + x, + mask=(rows_states[:, None] >= 0) & (cols[None, :] < dim)) + + +def causal_conv1d_varlen_states(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.empty(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + BLOCK_M = min(triton.next_power_of_2(state_len), 16) + BLOCK_N = min(triton.next_power_of_2(dim), 256) + grid = (triton.cdiv(dim, BLOCK_N), triton.cdiv(state_len, BLOCK_M), batch) + with getattr(torch, x.device.type).device(x.device.index): + _causal_conv1d_varlen_states[grid]( + x, + cu_seqlens, + states, + state_len, + dim, + x.stride(0), x.stride(1), + states.stride(0), states.stride(2), states.stride(1), + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N + ) + return states + + +def causal_conv1d_varlen_states_ref(x: Tensor, cu_seqlens: Tensor, state_len: int) -> Tensor: + """ + Forward pass only, does not support backward pass. + Parameters: + x: (total_tokens, dim) + cu_seqlens: (batch + 1), must already be sorted. The cumulative sum of the sequence lengths, starting from 0. + state_len: int. For each cu_seqlens, how many elements from x should be copied to the state. + If some of those elements belong to a different sequence, the value of the states will be zero. + Return: + states: (batch, dim, state_len) + """ + _, dim = x.shape + batch = cu_seqlens.shape[0] - 1 + cu_seqlens = cu_seqlens.contiguous() + states = torch.zeros(batch, state_len, dim, dtype=x.dtype, device=x.device).transpose(1, 2) + for i in range(batch): + end_idx = cu_seqlens[i + 1] + start_idx = torch.maximum(cu_seqlens[i], end_idx - state_len) + states[i, :, -(end_idx - start_idx):] = x[start_idx:end_idx].T + return states diff --git a/build/torch214-cxx11-cu132-x86_64-linux/cpp_functions.py b/build/torch214-cxx11-cu132-x86_64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/cpp_functions.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, Tri Dao. + +import torch + +from ._ops import ops + +def causal_conv1d_fwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + final_states_out: torch.Tensor | None, + silu_activation: bool, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + seq_idx=seq_idx, + initial_states=initial_states, + out=out, + final_states_out=final_states_out, + silu_activation=silu_activation, + ) + return out + + +def causal_conv1d_bwd_function( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + dout: torch.Tensor, + seq_idx: torch.Tensor | None, + initial_states: torch.Tensor | None, + dfinal_states: torch.Tensor | None, + dx: torch.Tensor | None, + return_dinitial_states: torch.Tensor, + silu_activation: bool, +) -> tuple[torch.Tensor | None]: + batch_size, dim = x.size()[:2] + width = weight.size(-1) + + if dx is None: + dx = torch.empty_like(x) + dweight = torch.zeros_like(weight, dtype=torch.float32) + dbias = None + if bias is not None: + dbias = torch.zeros_like(bias, dtype=torch.float32) + dinitial_states = None + if return_dinitial_states: + dinitial_states = torch.empty(batch_size, width - 1, dim, device=x.device, dtype=x.dtype).transpose(1, 2) + + ops.causal_conv1d_bwd( + x=x, + weight=weight, + bias=bias, + dout=dout, + seq_idx=seq_idx, + initial_states=initial_states, + dfinal_states=dfinal_states, + dx=dx, + dweight=dweight, + dbias=dbias, + dinitial_states=dinitial_states, + silu_activation=silu_activation, + ) + + dweight = dweight.type_as(weight) + if dbias is not None: + dbias = dbias.type_as(bias) + return dx, dweight, dbias, dinitial_states + + +def causal_conv1d_update_function( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + cache_seqlens: torch.Tensor | None, + conv_state_indices: torch.Tensor | None, +) -> torch.Tensor: + out = torch.empty_like(x) + ops.causal_conv1d_update( + x=x, + conv_state=conv_state, + weight=weight, + bias=bias, + out=out, + silu_activation=silu_activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + ) + return out diff --git a/build/torch214-cxx11-cu132-x86_64-linux/layers.py b/build/torch214-cxx11-cu132-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/layers.py @@ -0,0 +1,50 @@ +import torch +import torch.nn as nn + +from .causal_conv1d_interface import causal_conv1d_fn as kernel_causal_conv1d_fn +from .causal_conv1d_interface import causal_conv1d_update as kernel_causal_conv1d_update + + +class causal_conv1d_fn(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + **kwargs, + ): + # For varlen + seq_idx = kwargs.pop("seq_idx", None) + + return kernel_causal_conv1d_fn( + x=hidden_states, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + ) + + +class causal_conv1d_update(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + ): + return kernel_causal_conv1d_update( + x=hidden_states, + conv_state=conv_state, + weight=weight, + bias=bias, + activation=activation, + ) + + +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/metadata.json b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..53a8fb3984d168c37e251daeaa661413d5eb90da --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,47 @@ +{ + "name": "causal-conv1d", + "id": "_causal_conv1d_cuda_056e28b", + "version": 2, + "license": "BSD-3-Clause", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "ybfee7cD5pV/RtZZSki1skLZM+COZVoqTa+IlL+zY7E=", + "_causal_conv1d_cuda_056e28b.abi3.so": "8eIu93p60ua9mdiCkFMouNNKe4PgleTAH3Gpd/3hvdg=", + "_ops.py": "in6wiEQzq3wMOxiS5jap8cHfMQFJoUmFV73UTM6S8S4=", + "causal_conv1d_interface.py": "Mm0fBbamT4dihp+Z53RzIFM/dn5Y8nLu1vJ4W4W53jQ=", + "causal_conv1d_varlen.py": "j6tdAEf/x0aBoNx3mEA9qo4xAxmH6fuUUdVnq7tuGgM=", + "cpp_functions.py": "gUG0/NiMaAmAHLZQ4TnnD5LyKJp9G/BlINOFXXcnlU8=", + "layers.py": "RH7iOTcjqtn3buFO7fGHfC/Pgu47eaX97ANzfwLD6C8=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "081dcd473f8eca27226eef1f9c47c45b86c756b5", + "dirty": false + }, + "kernel": { + "commit": "056e28b85220c032b4dc1bb3dadebddd56dc228e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..3a46b16360491dbd130cb61138802dc2137de481 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUXSDkrQy3GZvyLXKZLtVMDIwZ0iAwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQyMzEwWhcNMjYwOTA0MTQzMzEwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEn31PBs5G5cP+QErxzPd0JDwXi4j0t9xbWQnlf0T4NrpELqjrFXp27XL/tBle6YgQKcJVMYzs4jraiBAy4ajgfKOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUTOvi7VzNgR1cXP25479ooGAiuHkwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzNUtIAAAQDAEgwRgIhAMygmI+8qzhejA+Xp7HCpkrJ+pu0N5lvH7xUu6CWlHOrAiEAoWTqe+lnAMYcXNFV+LukE4R8PwUpQoq1xODMesIx9rswCgYIKoZIzj0EAwMDaAAwZQIxANPkz+4o98MYrpFDiY+xMApXfUaJ7KKPEp8OBQ65rZtJjGEkSqLOL2/eBOKfWcgI7QIwJDjiv9hpEJqzc3wwaR3C3M6uEcLRSNN8JXlV9KOI5B8v88LGs106OcW2Llnq7IUc"}, "tlogEntries":[{"logIndex":"2712280149", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1788531792", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQC00BKjWCFeQ/vaS3lyMi1CMHHgGDbWFeiDWP9g7TNMCwIgLTe+jBVEPnQzujCnvqOy2Ae1p+C34LFgqLDwnfJ9ciY="}, "inclusionProof":{"logIndex":"2590375887", "rootHash":"E1wUV3ViFvRCUrnHEi9Jtx7mIwb4QEbGJDhhrpS5QeM=", "treeSize":"2590376027", "hashes":["ADL1xAj0dBwlRUYxLmFx6gfjuC54TPOedhn+IA5+yc8=", "I1ARUpIFe04CiCPfJdqmAb8VOzvjJ8pxHVLF7Y0FUu0=", "SvlZdoz2fyu/ZKTHoLMLAWqgnWoEwCzyVM7+jyVGKEA=", "u7rbJO6r+9BBzcdTxVJ3ANcx36qRNsFAeijIQKAMKAE=", "4mDHvJXaS0A5NHsE6DmjxoGyLpc9Wvi6tao8SMzvJLE=", "tbH3syeoVPO1D8NheZ6hR9MN/U2AIa39/Bi0oiBescM=", "jq8NOCglBk6NTBMHfAV4zBpslcR6foRL7yUKvxM6wZE=", "VNPdxKGG7ZZq9QV0nS/2qcn575QJbKWSbVNLo9Pqf0I=", "0GOSVpZ6KGFEmrcovBBUC1+51uTj09Ri+bq1qQLNeqw=", "L/YAJOZsdls37VNJgRupiQZ+3/baWTZm72IV3HXVcYA=", "kT0CBuImcJWrDmmSFNL3zx+n9cqRy2aCHWAOGMcdTDQ=", "W93RUf9D/+sJn+qU9XdapIH75v3rh+OVEdOHM6PMPkA=", "t8nRtM4piqY9M2vLuqc30bnlaHH5a1Mc5+3BX+VcR48=", "dpE4JeS09wbniSRpzbHmp3WEjZfaXFuaxQ/Jz4h8Qdo=", "5BBv7rWK6NFgtdHnmpxaOG4IgQ2e23dxOnb7iVqIJfs=", "CVXlwy6p/0g8OGYc+N5iYlD4qo/oIV4lgJstV8/jX9Y=", "JLnTjrlsTfbD8CvvizQuzVwJPXlC98tLzdnXZ1HqaTI=", "9yqRlModEw/R4m6jPiFNn9ya4s8ZGiXYyREc5jvu3XM=", "wgA+dB9j8L++Xwb+XT+v0evEwWqMLyd201Sc69PSLzg=", "N+foNlyVQOzU7HZsPPSKIQU/1+E+tmJj33u2r6dwmBs=", "I1FtczpUgFViy1Ae7CAfwglj2nSpp5nZrcsYauARxMs=", "fnPrvt7afEuTlIhaMcrGUdu9+heAmk4jOJ3+Br8/B8M=", "eXSLrLqKSxRTixSP/k4YimriB2lufao6trqTrsxxlyg=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2590376027\nE1wUV3ViFvRCUrnHEi9Jtx7mIwb4QEbGJDhhrpS5QeM=\n\n— rekor.sigstore.dev wNI9ajBFAiBqHtAIDT9dQZ9RLikg7KQfzsmFkqXD3aF4QBHB+2JkBwIhAN/HUjXplbThtjCiV/cYZW6r8IbR9RbUFDpjFfbTndCg\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3MTA4MTVkZDJlNjc0NmRlNTk0ZTJmM2FjYzc1YzFjMDM4MDIwNWIyODIzOTRjNjc5YzhjODI1YjEwMWMzZDdmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRTU2akVzbHBzeC95VnRka3lsVEVjWVkrNVg5bWI2K1puWkZDTjlESkFZSEFpRUEzSTQrVi8wT0JsVXcxRXJyR3JZakNvaXViQkdPSXQxK1o5UTZPV01uejV3PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVldGTkVhM0pSZVROSFduWjVURmhMV2t4MFZrMUVTWGRhTUdsQmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjVUWHBGZDFkb1kwNU5hbGwzVDFSQk1FMVVVWHBOZWtWM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1TXpGUVFuTTFSelZqVUN0UlJYSjRlbEJrTUVwRWQxaHBOR293ZERsNFlsZFJibXdLWmpCVU5FNXljRVZNY1dweVJsaHdNamRZVEM5MFFteGxObGxuVVV0alNsWk5XWHB6TkdweVlXbENRWGswWVdwblprdFBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZVVDNacENqZFdlazVuVWpGaldGQXlOVFEzT1c5dlIwRnBkVWhyZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhlazVWZEVsQlFVRlJSQXBCUldkM1VtZEphRUZOZVdkdFNTczRjWHBvWldwQksxaHdOMGhEY0d0eVNpdHdkVEJPTld4MlNEZDRWWFUyUTFkc1NFOXlRV2xGUVc5WFZIRmxLMnh1Q2tGTldXTllUa1pXSzB4MWEwVTBVamhRZDFWd1VXOXhNWGhQUkUxbGMwbDRPWEp6ZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkplRUZPVUdzS2VpczBiems0VFZseWNFWkVhVmtyZUUxQmNGaG1WV0ZLTjB0TFVFVndPRTlDVVRZMWNscDBTbXBIUld0VGNVeFBUREl2WlVKUFMyWlhZMmRKTjFGSmR3cEtSR3BwZGpsb2NFVktjWHBqTTNkM1lWSXpRek5OTm5WRlkweFNVMDVPT0VwWWJGWTVTMDlKTlVJNGRqZzRURWR6TVRBMlQyTlhNa3hzYm5FM1NWVmpDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg5ivKfrE+67aRw+u8iNl3dckYg7xQU+gbGOP0WGyE7NkCFB2bvxcb5QkMozp9HTEeq/nETXkPGA8yMDI2MDkwNDE0MjMxMFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQyMzEwWjAvBgkqhkiG9w0BCQQxIgQgxal5Q24xbb/DbPGvcf5oJujCGnnLMXV4w89gLk68ihEwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKVC8mdcCUgDUfkUomwOO/2WrRHAcSVKpUfmQDxP6KQAj6KqmaxeLSiEUkTv9p46SwIxALiX+VfuPjw4XLuDqNVMJmA7K6w0nksfcyyAT5+JqCR9diDhnrW7jIqj+9uz3SUovw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"cQgV3S5nRt5ZTi86zHXBwDgCBbKCOUxnnIyCWxAcPX8="}, "signature":"MEUCIE56jEslpsx/yVtdkylTEcYY+5X9mb6+ZnZFCN9DJAYHAiEA3I4+V/0OBlUw1ErrGrYjCoiubBGOIt1+Z9Q6OWMnz5w="}} \ No newline at end of file