diff --git a/build/torch212-cxx11-cu126-aarch64-linux/__init__.py b/build/torch212-cxx11-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..89c50591f954428710daeac53bae219cfe8e76d5 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a438e8fbe2a8bf179fccaa5738cbc3d3696d2207b6a2ba5ef161ebcebac50ac +size 78911952 diff --git a/build/torch212-cxx11-cu126-aarch64-linux/_ops.py b/build/torch212-cxx11-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch212-cxx11-cu126-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/layers.py b/build/torch212-cxx11-cu126-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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-aarch64-linux/metadata.json b/build/torch212-cxx11-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b963522b39814387ff2e90fa89e4fc3580ca3667 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-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": "ekOOj74qi/F5/Mqlc4y8PTaW0iB7aiul7xYevOusUKw=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..33105ff0d7458903b4b8c87d637c0f28845ac549 --- /dev/null +++ b/build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUf+RWxlCUhmQ+Szg0CRh0AGloyK8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzIwWhcNMjYwOTA0MTQyNzIwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEZJQTz23TCiPXmxyyZlysngiWVrOBqZDcdp6Of9b+qz7qylCA28EVRGfPDZB4U7N15kgbIjJiO0cWtbeg+n63/aOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUsjGLyb+sWt4kj+XWKRgpdPUoxcYwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzH+4wAAAQDAEgwRgIhAJBiF3l2I8nvvQcCK4f6tBpX4FF23/D1eZmHFHct9xdEAiEA3LaSWab+5EGu5PwO+MWM7Stop/8i6MIQQSQ0y4AJp4owCgYIKoZIzj0EAwMDaAAwZQIwHndYsbaHT5xa/6dshYwdU/1Ihs+sBls2I+FRxbhSXcxt/f53eIjQVEfb7GjHm8grAjEAtY8Rp9JXzgvaoH/H430fhqQf9fLL1avxebvvDEtKtE5oyEJT0L1Ro8AOxzVVEarC"},"tlogEntries":[{"logIndex":"2712206770","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531441","inclusionPromise":{"signedEntryTimestamp":"MEQCH3OrB5T9uqPjwrYaBorlK+d5/j96QiFXEsjIE0gP7zMCIQClO0Oh1UiYb5EO/kmpT7CZL1TjchF9hUY5VFBpimsV7Q=="},"inclusionProof":{"logIndex":"2590302508","rootHash":"dNSC2zhHBEloks7ZHoUnRaofBpbG28rCK7d4aT3A7c8=","treeSize":"2590302688","hashes":["I+CVSQOtmuKccwgZuKBONdCVEcailVdxVUnrrZ6IhWw=","VQzWAftcU6GhyrLQVMfvC7buOtTZWvl5gtpGEA+qEBE=","byh0tau9mN9C3+5s/mFds8SAh8dHpOhCXvthk/GLpVU=","VMkr1GEdvxuJrZxzPuuHj0VvDQCoH8ieLdGa/jbmDpA=","e6q38MDDetc1Kd/w6bGjtjSum4Hf8ejgX+Ep+byFFnc=","yUi4MwNjCxY8uQzo394W4vlfU+5//Bu7Q60vPlk5Dpg=","3xXk9Q8vetv0MSH2nLmACbznUDXHsaTz87tfHkjKQ80=","jp48I9+TwWbqE0BGWvyF+u5dAWLebmpOi+fltZI9hC0=","4kbZ3uoeMKtWhuQVYghcVMlVvFdvSwqMkTnkZfzyo3E=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590302688\ndNSC2zhHBEloks7ZHoUnRaofBpbG28rCK7d4aT3A7c8=\n\n— rekor.sigstore.dev wNI9ajBEAiACBnn4LHo7HzddvTKk09GvBfTpDKGF9yBCjia0TqxDcgIgT/VMFFL6VUtuiuMEh/z/u6DOnVg1d7jIM5ecfUh/azs=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzYTY4ZDMxMjRlMmUwNmVkNTI2Nzc4NTRhOGEyZTUxMTI0YTE3MmU0YjMwYjdhOWQzYjFiNDU5ZDg2MWFkZjlmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNUa2FGVFlGdW0rWE4wUlZ3aWwrNlFaOW91N0lMYVN6d0UzTVFoK2U5L1pBSWdkc1ZSZU5SM1FjM1FKSllJZlVmVFFsdmcxMUNPc2xic3dqeUV1Mk4wVE5JPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVlppdFNWM2hzUTFWb2JWRXJVM3BuTUVOU2FEQkJSMnh2ZVVzNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBKZDFkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWtsM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZhU2xGVWVqSXpWRU5wVUZodGVIbDVXbXg1YzI1bmFWZFdjazlDY1ZwRVkyUndOazhLWmpsaUszRjZOM0Y1YkVOQk1qaEZWbEpIWmxCRVdrSTBWVGRPTVRWcloySkpha3BwVHpCalYzUmlaV2NyYmpZekwyRlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ6YWtkTUNubGlLM05YZERScmFpdFlWMHRTWjNCa1VGVnZlR05aZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2dyTkhkQlFVRlJSQXBCUldkM1VtZEphRUZLUW1sR00yd3lTVGh1ZG5aUlkwTkxOR1kyZEVKd1dEUkdSakl6TDBReFpWcHRTRVpJWTNRNWVHUkZRV2xGUVROTVlWTlhZV0lyQ2pWRlIzVTFVSGRQSzAxWFRUZFRkRzl3THpocE5rMUpVVkZUVVRCNU5FRktjRFJ2ZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkpkMGh1WkZrS2MySmhTRlExZUdFdk5tUnphRmwzWkZVdk1VbG9jeXR6UW14ek1ra3JSbEo0WW1oVFdHTjRkQzltTlRObFNXcFJWa1ZtWWpkSGFraHRPR2R5UVdwRlFRcDBXVGhTY0RsS1dIcG5kbUZ2U0M5SU5ETXdabWh4VVdZNVpreE1NV0YyZUdWaWRuWkVSWFJMZEVVMWIzbEZTbFF3VERGU2J6aEJUM2g2VmxaRllYSkRDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgZbaI/qtRQuxvB1GM4N258VxnlBTY6jZP9s1RugtaL1wCFBsFWBBa8RWjYttnrXgPJzYPtwXZGA8yMDI2MDkwNDE0MTcyMFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQxNzIwWjAvBgkqhkiG9w0BCQQxIgQgg0/opwKjAfUOl2e9kAymhhauZlwyMvIzg2V9/95nzN4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAJ1bj0YMMqMFT1/U5jUVYAjRUXQp4baUDq9of2++1A26fBqgv/JNDVVVGq/j04Ce+gIwc5sTU92/EL9WbtUnj0L70336Ta2N+ims/0JR3/ig7/4NxWT/qHV+8nbeBYhx7p+F"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"OmjTEk4uBu1SZ3hUqKLlESShcuSzC3qdOxtFnYYa358="},"signature":"MEUCIQCTkaFTYFum+XN0RVwil+6QZ9ou7ILaSzwE3MQh+e9/ZAIgdsVReNR3Qc3QJJYIfUfTQlvg11COslbswjyEu2N0TNI="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-aarch64-linux/__init__.py b/build/torch212-cxx11-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6c48c986a1d8e471989576c0899be677407ecd2d --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49e5ababd9692f41e928027eecd1897d156a16187fc778e2b10a3036468157ad +size 74934456 diff --git a/build/torch212-cxx11-cu130-aarch64-linux/_ops.py b/build/torch212-cxx11-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch212-cxx11-cu130-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/layers.py b/build/torch212-cxx11-cu130-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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-aarch64-linux/metadata.json b/build/torch212-cxx11-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e11dd407d324eb13ec5e02d0c4802d2d2688330f --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-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": "SeWrq9lpL0HpKAJ+7NGJfRVqFhh/x3jisQowNkaBV60=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..b25ab575757cc2dea7674d0d3b46180df58cfb02 --- /dev/null +++ b/build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUGgkpT387TNaEbCPHkP6ocsdul34wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzIyWhcNMjYwOTA0MTQyNzIyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEMAdRZcjHgOQqmTbGZ+tjINm7sA1s+DFeXMbxW82mPc4627QjXc9uQGfF+9JRYAkQc6sbA4Zje61kcwG3Pr3/3KOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUZDx33npEdmZAaFcd5ZCidnw/EeIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIA/UAAAQDAEcwRQIgbr7ru/sijDi3JUZ/IIEam3imrzH+QlBTB2gPwP/Vd7ECIQCx7w28xnOMWUYqTglUdeMSZ5JxOeSdatmS87C1CrdJFjAKBggqhkjOPQQDAwNoADBlAjEAh7KL+RTp5AE7xVcWZKItFuHr9H330drX6bUC22dYWGVAOOcB13KFgVwWn2rJdbHGAjBpwB4cHrZYg9j8bITJf5SrvkUz8PpZqtKLYeoaMqZvOI9TG3Iot22s33/jMpofGes="},"tlogEntries":[{"logIndex":"2712207293","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531443","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCvLZ9HNSrcl+79asYzS48IqyLTXO2tTg04KG8EgnKB+wIhANhz+A9zhnqNiKRwl8zEGy31IqSp2vf9++To3/t/Ibz9"},"inclusionProof":{"logIndex":"2590303031","rootHash":"WfmdI6CxfC8d543stVXCU4To3gKfMtJjE+BNFbx9qUU=","treeSize":"2590303500","hashes":["2cxFvrGkhX5zREL0kn26ODNvYj3IZZCPT9G9aD+suW0=","rGZ5EUOASjIshNK5g6sFQmwujYe9PYUNx+Cd21hkTug=","s4ewz/bVA19TPG2PTUvVv8XyyOw6vNXC+PyorYetMLg=","eYH/hy0lG0mYPgsh7OYY628MVLQpLqKczI+qui51vTo=","54dKAMqxTV45UZREE3X5VZZ5ok9xmSBwzO+bmXZLGU0=","Gfrrd3wiiL/m2p6MgXLWcGz/0aMRwannidIRFan6Zb8=","pNtypuTy5shS+Q/WhqCkH06v+2jNOm6odfO3NTTNJpc=","2EzrN0aKNSZB0v52y2meIimGsze5uJI/RSnLUqHJC3A=","7me4LqbSAx/c0mXz4Zw8nO8yEjDLlw2g5gWRc3fY3Cs=","g6+15tz0IvMIiVVN/ZXRO4KFNFolmaE27hj4XhxyJ1s=","LiITXVFFflQQ5udIm3jdVx/NCeh6sQJK9NkWd+1RTKY=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590303500\nWfmdI6CxfC8d543stVXCU4To3gKfMtJjE+BNFbx9qUU=\n\n— rekor.sigstore.dev wNI9ajBGAiEA7PfiQhQ5ShWfR13X9al1dEWBdMfj2kpGej2WQtpPbvMCIQD8/JtqKrOPa3KEXqS+LHp0Cs7yEVwx8Em07OpfThqvpA==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3ZWQ4NmM4NDBmYmMxNGRlZGExYjlhZGMzYjU2ZGM0M2ZmMzFhZmEzMWQ4NWEyZmUxYWRiZGMwYzFjMGExMjUyIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURrUGZ1bm9MbEwyZFUwWTczdWdYTFFMeXNEcitNTzB5bDRkei9oY2dDT1RBSWhBSjJsN2ZpUGcvSE1wTldkaFhFK3JzbFlVZWIxN0duN1JGMXJpZHlqKzFLRSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlIyZHJjRlF6T0RkVVRtRkZZa05RU0d0UU5tOWpjMlIxYkRNMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBKZVZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWtsNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZOUVdSU1dtTnFTR2RQVVhGdFZHSkhXaXQwYWtsT2JUZHpRVEZ6SzBSR1pWaE5ZbmdLVnpneWJWQmpORFl5TjFGcVdHTTVkVkZIWmtZck9VcFNXVUZyVVdNMmMySkJORnBxWlRZeGEyTjNSek5RY2pNdk0wdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZhUkhnekNqTnVjRVZrYlZwQllVWmpaRFZhUTJsa2JuY3ZSV1ZKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xCTDFWQlFVRlJSQXBCUldOM1VsRkpaMkp5TjNKMUwzTnBha1JwTTBwVldpOUpTVVZoYlROcGJYSjZTQ3RSYkVKVVFqSm5VSGRRTDFaa04wVkRTVkZEZURkM01qaDRiazlOQ2xkVldYRlVaMnhWWkdWTlUxbzFTbmhQWlZOa1lYUnRVemczUXpGRGNtUktSbXBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRV2czUzB3S0sxSlVjRFZCUlRkNFZtTlhXa3RKZEVaMVNISTVTRE16TUdSeVdEWmlWVU15TW1SWlYwZFdRVTlQWTBJeE0wdEdaMVozVjI0eWNrcGtZa2hIUVdwQ2NBcDNRalJqU0hKYVdXYzVhamhpU1ZSS1pqVlRjblpyVlhvNFVIQmFjWFJMVEZsbGIyRk5jVnAyVDBrNVZFY3pTVzkwTWpKek16TXZhazF3YjJaSFpYTTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgcyvSoqPI9zF5dbD+ig3TZfW/GppdoJ0rUSLCWO5xrzACFQDe6+O5l5/BFEmYKVAT9Z9NchbVKhgPMjAyNjA5MDQxNDE3MjJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MTcyMlowLwYJKoZIhvcNAQkEMSIEIGtkh8QvWhkt6wBoy9Vqx3AicJfUBlYbyrUYEJD7l3plMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMDC6KRMQkWGRZEj5Nh3fVjDpogS6PnT3saO0+efomZMzTl+T5NSdpzNjGdCoX95skwIxAOnJXu6R2B1+nX773sh0aDijaQIQi6lsC7MfI8p61XrsT9CYPAZ/udtou5IGf6bX4w=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"fthshA+8FN7aG5rcO1bcQ/8xr6MdhaL+GtvcDBwKElI="},"signature":"MEYCIQDkPfunoLlL2dU0Y73ugXLQLysDr+MO0yl4dz/hcgCOTAIhAJ2l7fiPg/HMpNWdhXE+rslYUeb17Gn7RF1ridyj+1KE"}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-aarch64-linux/__init__.py b/build/torch212-cxx11-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch212-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..e56b833fb66e2144e59614d21b4b34de8031d6aa --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:012229d08b1266fcd8a9ed38c7a50661a1ced324a906df2d7959639c54623aa1 +size 75002488 diff --git a/build/torch212-cxx11-cu132-aarch64-linux/_ops.py b/build/torch212-cxx11-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch212-cxx11-cu132-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/layers.py b/build/torch212-cxx11-cu132-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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-aarch64-linux/metadata.json b/build/torch212-cxx11-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c155a580b92cdbc8cdc176c8c1ef177e7bc22ddb --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-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": "ASIp0IsSZvzYqe04x6UGYaHO0ySpBt8teVljnFRiOqE=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..75a22ff947f55988e68fdd8ae6acd3c3a4f63672 --- /dev/null +++ b/build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtKgAwIBAgIUE11Yy+ls4GvbxTdWZpelUx3pOsUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzI1WhcNMjYwOTA0MTQyNzI1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEos1zKlalJf+qWXNxwCCQ3l3hISTLzv9z2JnijhSrBDK0112VKo9LzW9zCfJ5sBRVNXAwGSyYn4IuBTXSkIPc/KOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU3uUqPVaWx4VhKgMmEFla7c/6Le4wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIDycAAAQDAEgwRgIhAP8vqZGkNaGnRIlGm001aFQakZnqCkLWvXcKa3mUvT0VAiEAzpDz+kFK/O6JbiClUdJ/7R9xZUBNSro0/S3+9i7kbJAwCgYIKoZIzj0EAwMDZwAwZAIwFg3piE1u7sseou8AYoSX0lyNwjqG2tVvZ3ORoesJiGuIphZs0HtRJH6VoHzAd3mtAjBRp8ouHhP5TJbwowE102cxAIREBcLdXm3KxSHW6qSByNUkDaJZjSF5Hzj7BSh0H5s="},"tlogEntries":[{"logIndex":"2712207943","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531445","inclusionPromise":{"signedEntryTimestamp":"MEUCIGjRq6DLUFLe7ri4ykDwm/iyiNKHTVkG2LWuzZ9+pfRLAiEArbtiQYSuPAxNnbcE0SJ7LUvaqQ2+dWajZhl7ZJcCbOQ="},"inclusionProof":{"logIndex":"2590303681","rootHash":"42aTA6sdGh1OHznEAybnbgyDXBJTcMqRB18u9vTyrs8=","treeSize":"2590304182","hashes":["E94iD6zRcB7SEYeWwIB7KeTmpU2PxVFRyDH1b95Eht0=","si0hTWeBvTS7sNBwBYeJaN09D50AAz2Z18GvzP0pIPE=","Jh0boxaEq8QqEw1W9BWsjgUe5ix/wLPJtD9coWAoG4g=","4Ffg7NpSwd2Xz/3X2+xifKmZYOsaObTFXuHtvmmuBdw=","q8cRbDXdKTGaMQub+k/AQ4QaNQqELpaWAlS/yX6ovqo=","3UB/k0PDCeTfEybl9iefMnlb8qiXTKYnhj7PoCyD53M=","IHmh9RWSp3jsMdeBPMl+gJnZiTFoCRxnshs6S6ZywHE=","4n0Eo6WnyJCDEqFpK46WRQE1wQKg/3AjcYsnwZQYip0=","MpPwp638fmb2Q2ZLj6zGkCJwTVI7pLz000sbmluXGg4=","T9gvtU8aznESDhqcF7J3TmgYgwpxETlztDF8ytqtcU0=","lEwXj6GVlLVdEynO9nlxd0jbX2m6mEtQjOdrv/C8Gxs=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590304182\n42aTA6sdGh1OHznEAybnbgyDXBJTcMqRB18u9vTyrs8=\n\n— rekor.sigstore.dev wNI9ajBFAiEAo63FaOTQOMQ/zR2FYz32qOMiJq+S9of2/BekRwENsOUCIE/wlmEyCJIaRGBm1Z/UrFJX+FI3id/PMxC6wiSni4e+\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmOWFiYTBlODc5NThjYTgxNDlmNDc2NDQ4ZjQ3ZjRkNTFmYTg3YTIxZGY2ZjYxYjc3MmU1ZDI3NjJjZDBhMTkyIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJR29xN3lPc1B4NTlOS1lVdTZmZUpEemdCcEdDMkE2Y2tyN3daOHpIV3Yxd0FpQSt2NVNSWUFrdUMrMDNKUFZubTFZamRhRi83d21JZ0pLdW9KSm1sblhJbGc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5STFowRjNTVUpCWjBsVlJURXhXWGtyYkhNMFIzWmllRlJrVjFwd1pXeFZlRE53VDNOVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBKTVZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWtreFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ2Y3pGNlMyeGhiRXBtSzNGWFdFNTRkME5EVVROc00yaEpVMVJNZW5ZNWVqSktibWtLYW1oVGNrSkVTekF4TVRKV1MyODVUSHBYT1hwRFprbzFjMEpTVms1WVFYZEhVM2xaYmpSSmRVSlVXRk5yU1ZCakwwdFBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV6ZFZWeENsQldZVmQ0TkZab1MyZE5iVVZHYkdFM1l5ODJUR1UwZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xFZVdOQlFVRlJSQXBCUldkM1VtZEphRUZRT0haeFdrZHJUbUZIYmxKSmJFZHRNREF4WVVaUllXdGFibkZEYTB4WGRsaGpTMkV6YlZWMlZEQldRV2xGUVhwd1JIb3JhMFpMQ2k5UE5rcGlhVU5zVldSS0x6ZFNPWGhhVlVKT1UzSnZNQzlUTXlzNWFUZHJZa3BCZDBObldVbExiMXBKZW1vd1JVRjNUVVJhZDBGM1drRkpkMFpuTTNBS2FVVXhkVGR6YzJWdmRUaEJXVzlUV0RCc2VVNTNhbkZITW5SV2Rsb3pUMUp2WlhOS2FVZDFTWEJvV25Nd1NIUlNTa2cyVm05SWVrRmtNMjEwUVdwQ1VncHdPRzkxU0doUU5WUktZbmR2ZDBVeE1ESmplRUZKVWtWQ1kweGtXRzB6UzNoVFNGYzJjVk5DZVU1VmEwUmhTbHBxVTBZMVNIcHFOMEpUYURCSU5YTTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgx11+2BUb5an2kg3ZkmIken4V8n+PHgicgoF+6/kLn8MCFDzydz6ysqQqk1bR4/FjQvJZfPRKGA8yMDI2MDkwNDE0MTcyNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQxNzI1WjAvBgkqhkiG9w0BCQQxIgQgG5PJSk9HUtrVpi73syVGSJmMLFJiXL0LhrATICpADgswgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwYsUcdGXlsofK5VZOw7Jq46InIO9qERSh5WD6pg2qQ/jrLODT67Ye74g/RRGiCiZUAjEA7gHiUtNalSQKyuOuBxUJwvh9LFxzljf5ESiI3QU6fxr50zoZcaQrwF3puHDsK1ZJ"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"+aug6HlYyoFJ9HZEj0f01R+oeiHfb2G3cuXSdizQoZI="},"signature":"MEQCIGoq7yOsPx59NKYUu6feJDzgBpGC2A6ckr7wZ8zHWv1wAiA+v5SRYAkuC+03JPVnm1YjdaF/7wmIgJKuoJJmlnXIlg=="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-aarch64-linux/__init__.py b/build/torch213-cxx11-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d47221bee4213502cdb794ad1afa82426b7daa82 --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e166fbb8d695ba179fe3d197a4daf5da55e8d8ae398a28340d28cb80ddb29a2e +size 78911912 diff --git a/build/torch213-cxx11-cu126-aarch64-linux/_ops.py b/build/torch213-cxx11-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch213-cxx11-cu126-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/layers.py b/build/torch213-cxx11-cu126-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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-aarch64-linux/metadata.json b/build/torch213-cxx11-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..10bcc226394bc55c3585dcc817463cbcf19d74c4 --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-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": "4Wb7uNaVuhef49GXpNr12lXo2K45iig0DSjLgN2ymi4=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..4df7bdf8cb849392d075679c79e524dbfb9344ac --- /dev/null +++ b/build/torch213-cxx11-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUYpMV7O/BP0Z4rM86oJuHThhtOPMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzI5WhcNMjYwOTA0MTQyNzI5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEtd8DSwDgT/q1uAOyFNq6HbWgETeCAMKSgs7V5jGiTJevEV6Bxxx2wMqbTvN/K4OLQqTHUEHCApj3lToKireaPKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUA4ju9yrjdKqA6I/lrGzC20Y3Ey0wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIHiMAAAQDAEcwRQIhAIJx4rQrXJaqm90bvLvBZX1+FdxmWZ5D5LPxqZF1zFjMAiAfLb2aMh4c80jEyO94m2VlikVa4zffWavpx9lPaN0aejAKBggqhkjOPQQDAwNoADBlAjEAmngOuXX+KpzTxCM90SP8GvQyVTNnI+TmN84VD5XTLCAMOLlrtbwdHsrRWEtEcpg1AjAx3SVgNTRo96mdvYNMhN6O2dSF3pNYTyZzD+94eGRK0eehqcS6dq7BLWKHE3PRyeI="},"tlogEntries":[{"logIndex":"2712209126","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531451","inclusionPromise":{"signedEntryTimestamp":"MEYCIQD5DJ5liR1Q9OmFxwOzEUFrHuIk8+eEuJb6LMKjBuOuOQIhAOfPGdxOw+gz5hDOQLGFKkVLFBIb/FX02WeQDeDIaFZV"},"inclusionProof":{"logIndex":"2590304864","rootHash":"E7HGGeT4aeWLU2Te7Fx/tfzmYi/JL8QhfUV3ApKxBok=","treeSize":"2590304907","hashes":["BwPMMPYyLAYc4+tnv+WCUxHP7hHfp36cb2UHdIlMzjA=","M/Aghb36880tBfbaisSSzrwyZMx1IzFBdZ7UvYHCH8k=","lWznuXEYmQ/+6JuR5wZRtZYsJK28bCrp8yYpFdnPohg=","HL9+kMb843VTY9B8hAb7tTUydDzY1KordoGlv4hBIcA=","LXdgt0UINWyxd0GGPNBKAsyAp+NVZzS1FREAdREMrKA=","e2hyvQrpy7PFMLCqN+D9OGTG686qdnFaCJR68lgm2Ok=","CDT0BN7tUbNcQSKcpfgpqgDNdKTyZvrsdrxg5N008Ls=","Tw7XDM9uiRXrf2kJFUdcOr1cEvTDrUL/YxEzflG6/d8=","HeOLbwKuDOT5caHN88H+iFCJfQDVT2lRnVzhrysbU9s=","Ynsf1tbb9300W+zDms6HJywTVRd4XY1FxuqzcgA2qmc=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590304907\nE7HGGeT4aeWLU2Te7Fx/tfzmYi/JL8QhfUV3ApKxBok=\n\n— rekor.sigstore.dev wNI9ajBGAiEAo54YYv519wo4gdz19zduwQUpjO7FDlYrKcp+DZE7aNYCIQD6f6XGdRUx2gmSr6SKhmAV/69mMAjxMT4eH9flx8TPHA==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3NDJmZDBjOTUyNTFjODNmZTNkODhjMTZhOTNiZTE5YTBmMGY2YjIzOGVkNDBlNDc1ODdhNGFkZDAzM2E5ZTczIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRzJpU1pPWUFZeHJhWng2SnByNXdBNU14Q2tCdzZzcm5YeC9ybUpMQ2wrMkFpRUFuUHF3YVd5RFd1NnZuSTZTZi9RaWVnY1l4eERpYlVPYmMwOFVYeWd6SmxnPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVldYQk5WamRQTDBKUU1GbzBjazA0Tm05S2RVaFVhR2gwVDFCTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBKTlZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWtrMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYwWkRoRVUzZEVaMVF2Y1RGMVFVOTVSazV4TmtoaVYyZEZWR1ZEUVUxTFUyZHpOMVlLTldwSGFWUktaWFpGVmpaQ2VIaDRNbmROY1dKVWRrNHZTelJQVEZGeFZFaFZSVWhEUVhCcU0yeFViMHRwY21WaFVFdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZCTkdwMUNqbDVjbXBrUzNGQk5ra3ZiSEpIZWtNeU1Ga3pSWGt3ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xJYVUxQlFVRlJSQXBCUldOM1VsRkphRUZKU25nMGNsRnlXRXBoY1cwNU1HSjJUSFpDV2xneEswWmtlRzFYV2pWRU5VeFFlSEZhUmpGNlJtcE5RV2xCWmt4aU1tRk5hRFJqQ2pnd2FrVjVUemswYlRKV2JHbHJWbUUwZW1abVYyRjJjSGc1YkZCaFRqQmhaV3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRVzF1WjA4S2RWaFlLMHR3ZWxSNFEwMDVNRk5RT0VkMlVYbFdWRTV1U1N0VWJVNDRORlpFTlZoVVRFTkJUVTlNYkhKMFluZGtTSE55VWxkRmRFVmpjR2N4UVdwQmVBb3pVMVpuVGxSU2J6azJiV1IyV1U1TmFFNDJUekprVTBZemNFNVpWSGxhZWtRck9UUmxSMUpMTUdWbGFIRmpVelprY1RkQ1RGZExTRVV6VUZKNVpVazlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg9MbyzODSVI1+mksO2aeeDDt9sIzjvOVh1G6yveT9ARkCFD4+qIV4rD4nbewFJAQ2ACSKk7cKGA8yMDI2MDkwNDE0MTcyOVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQxNzI5WjAvBgkqhkiG9w0BCQQxIgQgvWisheq77VWz6KciYnw9KZeKaYEjfe+rgOsznqQfyk0wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAN+KHU2tmJlCnvhy76dzZQIiwyJCqF0g4Bp/uduwfuojEyEiI7Pr6OwqI85Ck0K2ogIwVvdluktsA9iyAtreVgdieMtgj9V2LMtOMe2iaApgijicJNCG9vUX3oaPWTsuMP//"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"dC/QyVJRyD/j2IwWqTvhmg8PayOO1A5HWHpK3QM6nnM="},"signature":"MEUCIG2iSZOYAYxraZx6Jpr5wA5MxCkBw6srnXx/rmJLCl+2AiEAnPqwaWyDWu6vnI6Sf/QiegcYxxDibUObc08UXygzJlg="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-aarch64-linux/__init__.py b/build/torch213-cxx11-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..18a5c7d39fdd9e98a4a205301a09b4bbbe6dd152 --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce373b8049bc86d3309df58ee8d6d23f5a79bf991a8150ff66dffa1c2b40a995 +size 74934224 diff --git a/build/torch213-cxx11-cu130-aarch64-linux/_ops.py b/build/torch213-cxx11-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch213-cxx11-cu130-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/layers.py b/build/torch213-cxx11-cu130-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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-aarch64-linux/metadata.json b/build/torch213-cxx11-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..4a3f3661e7bd160806b020e14443cc2443f232e6 --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-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": "zjc7gEm8htMwnfWO6NbSP1p5v5kagVD/Zt/6HCtAqZU=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..6b827b1b01514e27ef5679a1fef827211f6a4f7b --- /dev/null +++ b/build/torch213-cxx11-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUQp4elL/n1oL9cROBz/D/atQYC5gwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzMyWhcNMjYwOTA0MTQyNzMyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEIt2cmRhMqz64ycaO4h3gILa+Z+0dPT7KJ8qyT4lSOuVYBLeQBMqEBOPsgZMSmiAr6Hs3+YtfTmFNJm7pjHzA2qOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUMA35A7DcJXHm6RVbr1ko7Epg6OwwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIK+IAAAQDAEcwRQIhAPnwZ7W6CHe9gMb7EHaEafC1MRm35jE2ebGxgeryowEdAiASdtlxKs/62fOcKh/WQHMm892vQPhFiwJu/lS8iW7pDzAKBggqhkjOPQQDAwNoADBlAjBlvRXFE3MPsi3TSbA5wB/lfOOkk8Qk2o/7EyOPNe3axj7m9NlA7W9KtN2bgvK3tUYCMQCnFrOSocucrPWmtCRT/H4xYBPY3ztTuzGaWHfL8uff8iUB1tzY5ZBQLjktRRdhWvg="},"tlogEntries":[{"logIndex":"2712210243","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531455","inclusionPromise":{"signedEntryTimestamp":"MEQCIFgfugGpPezwpsrim07ywoV77dwDmNz4fHNLCYJDpetsAiBt9aLPpLTMU1mM9fyFUs2IwtStmoDxttzOX/11NyR2BQ=="},"inclusionProof":{"logIndex":"2590305981","rootHash":"QeH0Hu+gnSc1orhRbA0NGre4Rm4M3Bf7fTkIhSl8fI0=","treeSize":"2590306306","hashes":["e5RgmPDCGyvO57/4WjmpeaIKb0pbQGesMhNpXOO86Pc=","x1O8W/5AKzBdHTOfn4YgIKmTTtPfK8B91XQqR2sMmlI=","Mpt7/AVW+v3CQrGmyiYW1TFlrDU1Z3KjQK9cAFJqXrg=","fRD5/zktUsVofDngcLl7zW3vU5bNFc+xpfOWHrAKUn8=","VBvVHSqOLtgAuNiiXMCdVRlgVXzCwSBYyyXdWWrqmRI=","V+shyTkOAnhceHAvgUbLxUBoMa+0GxC8O2bHOexuB84=","5DaZWMbs2p7ZWGB6ZwVwl3i9kOYGQKmysee/WlwPGTE=","A4F7dKS2cP53btGWtIH8Tv+xrHtsletofyEtPzcF/k8=","7DoQ3ggxKtzBEfzB/Q7Bl2k8hZGsOZkunhG1uLobSpE=","46f6poNmQe6XU7iE73fENBdfxybdQIWnxw5JZFpmHVM=","rh0X0kaj+rt1oOdlpoMW2u3Ef9e5khn4K2m4JvfYyJQ=","Ynsf1tbb9300W+zDms6HJywTVRd4XY1FxuqzcgA2qmc=","JxEJ83kMMSvspIYj2SF9rag7dIVF24Xvxk5o3+DY7GU=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590306306\nQeH0Hu+gnSc1orhRbA0NGre4Rm4M3Bf7fTkIhSl8fI0=\n\n— rekor.sigstore.dev wNI9ajBGAiEAimD4WAU/GpvSAAPC9iM5XKmV9CiZS5yAfb2ZUjkbqqMCIQD0ukzvz8aTadS1Tw9zaA9QtxwU0jKI/oQLSbaMFRXxKg==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3ZjJjMTI2YTgxMzAzNGU0OGIzOWRhNTVjNWU5NTk1ZTMyZTY1YTk3YTc0NDg0ZmE4ZGU4NGM5ZDY1YmVjMWJkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUQrNzIzUGJJM0E5Sk1MZ2Q1eXl4TFJsZlJVSjRwNXJrRUZNK0htVUs5elZnSWhBSnA3SjduaEtWNy9NeDlINk5MSlV6aS82enN5QVM1NVQ0TzNsOG4zT2RnZyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlVYQTBaV3hNTDI0eGIwdzVZMUpQUW5vdlJDOWhkRkZaUXpWbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBOZVZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWsxNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZKZERKamJWSm9UWEY2TmpSNVkyRlBOR2d6WjBsTVlTdGFLekJrVUZRM1MwbzRjWGtLVkRSc1UwOTFWbGxDVEdWUlFrMXhSVUpQVUhObldrMVRiV2xCY2paSWN6TXJXWFJtVkcxR1RrcHROM0JxU0hwQk1uRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZOUVRNMUNrRTNSR05LV0VodE5sSldZbkl4YTI4M1JYQm5OazkzZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xMSzBsQlFVRlJSQXBCUldOM1VsRkphRUZRYm5kYU4xYzJRMGhsT1dkTllqZEZTR0ZGWVdaRE1VMVNiVE0xYWtVeVpXSkhlR2RsY25sdmQwVmtRV2xCVTJSMGJIaExjeTgyQ2pKbVQyTkxhQzlYVVVoTmJUZzVNblpSVUdoR2FYZEtkUzlzVXpocFZ6ZHdSSHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENiSFpTV0VZS1JUTk5VSE5wTTFSVFlrRTFkMEl2YkdaUFQydHJPRkZyTW04dk4wVjVUMUJPWlROaGVHbzNiVGxPYkVFM1Z6bExkRTR5WW1kMlN6TjBWVmxEVFZGRGJncEdjazlUYjJOMVkzSlFWMjEwUTFKVUwwZzBlRmxDVUZremVuUlVkWHBIWVZkSVprdzRkV1ptT0dsVlFqRjBlbGsxV2tKUlRHcHJkRkpTWkdoWGRtYzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg5u2w9yKM7b9wXozgKEcrtwPD1jZo/8PeihjGd76l5J0CFFNJK+lnlqiZyqCUA/j0L2UC8XxSGA8yMDI2MDkwNDE0MTczMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTA0MTQxNzMyWjAvBgkqhkiG9w0BCQQxIgQgTdiZ35nRdKUuhlSdVYAOgR63LDH4jJpLUJX+ise4hv4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwBlBJKtiIuN7rOlMqRdZLJp0fM+ajElPMZEM0L4TiT15Rb+2nCe1t9AlcskI+Rn9YAjAoRWKLrISIZK58MXcqJj0TcxfSWYwz+abJmQnoeml8t+VhU6KS4WuUbl8bxFMogKI="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"fywSaoEwNOSLOdpVxelZXjLmWpenRIT6jehMnWW+wb0="},"signature":"MEYCIQD+723PbI3A9JMLgd5yyxLRlfRUJ4p5rkEFM+HmUK9zVgIhAJp7J7nhKV7/Mx9H6NLJUzi/6zsyAS55T4O3l8n3Odgg"}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-aarch64-linux/__init__.py b/build/torch213-cxx11-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch213-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..9a397cb0dffd9d23e1f0e023676f97d11ee0d623 --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:79bcf114bcc754bc73b331505be2a3655b2e8348bfbef3b5883ac06719b9de83 +size 75002256 diff --git a/build/torch213-cxx11-cu132-aarch64-linux/_ops.py b/build/torch213-cxx11-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch213-cxx11-cu132-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/layers.py b/build/torch213-cxx11-cu132-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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-aarch64-linux/metadata.json b/build/torch213-cxx11-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..3578cc24ce6f7fa843dc6bd4babf04e4dae33ca6 --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-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": "ebzxFLzHVLxzszFQW+KjZVsug0i/vvO1iDrAZxm53oM=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..8a29a1c44a5809b5b438f927935b37e46dbeda3f --- /dev/null +++ b/build/torch213-cxx11-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtGgAwIBAgIUDe0T3ztOqlpCKG0/TVF1Z/RnRO0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzM4WhcNMjYwOTA0MTQyNzM4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEphwwYFui60pOhf7aiWqvrPD8Doj/yjata4Nu8f0Pa5iw6hrFQxLeWWw6x5hLdm7Of2DuwwbIly2KsNvk9WtnGKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU263yEb+vpw6qmorUx6KUI8iuPoUwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIQcwAAAQDAEcwRQIgRIFozTKTZqvE56qHvW45qfWxzGTc9wLYwpb/5FQcYB0CIQCewu44cn16Nn3sRt+87PGu28vFOEs0sPJ3e56fRezglDAKBggqhkjOPQQDAwNpADBmAjEA9GC6gaccujI8lvWCB9fZuOAxTasDR8one7AuIIvpQIjQszdrQCtEJjaVBjf9FMsoAjEAyFnV0uZz8qX0Mak7rP7xK84USTq5tji7mdP1xVoIr0EypWx+SAN89MsR6W3aohJ1"},"tlogEntries":[{"logIndex":"2712211216","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531460","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCrzGIz63Op+BsAqbxKH4ZQsDvQJh1OoFInqacywE5KOgIhAKawiZDX50yGexwV0F96WTk/I1JJIVnWuBHnjEue9jQj"},"inclusionProof":{"logIndex":"2590306954","rootHash":"nH9lLvFChjylGyRkU1MF93xFyB+B8Ufft2PRbE8r1aA=","treeSize":"2590307026","hashes":["YX3fX8y7Tm4Cy4VUTGGdTTwo2g284TBz0SXQm9MVSZo=","+OqyzqTKaRs3gqgp0mufxxKDFpjS3VUljAUUaUJ7uBY=","L2pW0E3OifumNv2D3jXGMkevj5l+u3Qrba71I83fyDY=","DOaeDrh4pU6JnzPibwGV7YINzUbUIwDN3gORHNk/Q5k=","6BNZ9pwxF/97VGf5x4YsC5oTHxkyKbNfxEfKhMzt4Ww=","Ch6MkNKBjIkWPiUZ6K8ifr4YwuIshhj/QbH7SUUcObI=","M5D2fqyxdTgEgKK47mGsX5GXaXeNIulyYXuf2l5m3fA=","zf+3Ehr9UuKbcP4X8ALmOVIoQRukpN3rEoXLaRGquso=","yQ93L/Fd5OLAMs8PAr/g0rIiI2CUgq2QhPUkNzx/Cjc=","bZskCqT+sD+OR/yUokKG9vukKBzl/3CkQ1GhGyrF3k0=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590307026\nnH9lLvFChjylGyRkU1MF93xFyB+B8Ufft2PRbE8r1aA=\n\n— rekor.sigstore.dev wNI9ajBGAiEApexUV1mM6qUL6R/ecff7NWucf07Up/DLPLgPU9MRTUgCIQDAD76BOVSLBztBRzBX3eDqgxBvcxEphFEcN5wzCWe+6w==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiZTMwMTc3NzY3NzBiZjZiMDNlYjkwN2M4MGY2N2U2YzhhMTk0ZTc2OTcyODc1NTUzNTMzY2Y4MDFiMDhhNjRlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURlTkJzQnJTUEdHNEd3UEpkYU45YWpVNkpnZElMeGprbVJKWThXRkpYVm1RSWhBTHNRUGlNeERiMGZvcTJUeFpZNU04MGtWc0dqSVppbXF2WmNPTGc2aGZIYSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5SSFowRjNTVUpCWjBsVlJHVXdWRE42ZEU5eGJIQkRTMGN3TDFSV1JqRmFMMUp1VWs4d2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBOTkZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWswMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ3YUhkM1dVWjFhVFl3Y0U5b1pqZGhhVmR4ZG5KUVJEaEViMm92ZVdwaGRHRTBUblVLT0dZd1VHRTFhWGMyYUhKR1VYaE1aVmRYZHpaNE5XaE1aRzAzVDJZeVJIVjNkMkpKYkhreVMzTk9kbXM1VjNSdVIwdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV5TmpONUNrVmlLM1p3ZHpaeGJXOXlWWGcyUzFWSk9HbDFVRzlWZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xSWTNkQlFVRlJSQXBCUldOM1VsRkpaMUpKUm05NlZFdFVXbkYyUlRVMmNVaDJWelExY1daWGVIcEhWR001ZDB4WmQzQmlMelZHVVdOWlFqQkRTVkZEWlhkMU5EUmpiakUyQ2s1dU0zTlNkQ3M0TjFCSGRUSTRka1pQUlhNd2MxQktNMlUxTm1aU1pYcG5iRVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbkJCUkVKdFFXcEZRVGxIUXpZS1oyRmpZM1ZxU1Roc2RsZERRamxtV25WUFFYaFVZWE5FVWpodmJtVTNRWFZKU1had1VVbHFVWE42WkhKUlEzUkZTbXBoVmtKcVpqbEdUWE52UVdwRlFRcDVSbTVXTUhWYWVqaHhXREJOWVdzM2NsQTNlRXM0TkZWVFZIRTFkR3BwTjIxa1VERjRWbTlKY2pCRmVYQlhlQ3RUUVU0NE9VMXpValpYTTJGdmFFb3hDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgX07kX9zDFkPAAqGbjaNVerNECjDiDnwMY9alTTKELbMCFQDr27b5n7bM3a5X8Uxs3vVGD8XTfhgPMjAyNjA5MDQxNDE3MzhaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MTczOFowLwYJKoZIhvcNAQkEMSIEIExu7kmh1hYb9jgdKUMrGVC3IgOXn46L+if3TwOflxxpMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHj7Daw7IhbTfMz0p95BpwMgVlb/RCVO6YcsGv08azZaGKW0T7fLdktHMV5MWQN+hwIxANXYTkcm1ZZE/ioZZLwF/O/QD1IrpX9ln150AgR9e/i4xasD81JkgQhzlTxVVlq1Zg=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"vjAXd2dwv2sD65B8gPZ+bIoZTnaXKHVVNTPPgBsIpk4="},"signature":"MEYCIQDeNBsBrSPGG4GwPJdaN9ajU6JgdILxjkmRJY8WFJXVmQIhALsQPiMxDb0foq2TxZY5M80kVsGjIZimqvZcOLg6hfHa"}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-aarch64-linux/__init__.py b/build/torch214-cxx11-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..2adc952391a749b6b497a311f33d3edc6362371a --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:62011129918d50f3a1258de28008ffbeea5d64da11045f98ba49b498d5f110d5 +size 78916688 diff --git a/build/torch214-cxx11-cu126-aarch64-linux/_ops.py b/build/torch214-cxx11-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch214-cxx11-cu126-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/layers.py b/build/torch214-cxx11-cu126-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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-aarch64-linux/metadata.json b/build/torch214-cxx11-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..73e4dc9d53b287ab381f8a7fc7d92599eb5c235b --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-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": "YgERKZGNUPOhJY3igAj/vupdZNoRBF+Yukm0mNXxENU=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..3ed0d6293574a85f41b48f8dc2544903b5b58c30 --- /dev/null +++ b/build/torch214-cxx11-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUUTWj7PAfLwHOLWsNMKNiI6bID/4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzQxWhcNMjYwOTA0MTQyNzQxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE7cGN80YO96NYh8huNP+JqQNQ118DUk/AsM3Vse3vTfIY7tlLvrDRhxHx1CNgZCjfy/G7d095JTPWbpVEetn43KOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUA5ts1ZX0cb8k9ok2PwLEtRoETmgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzITQ4AAAQDAEYwRAIgFYv5SdwBynD2HT9ZG61ut1Sfdt8mf4lN54Hnv1L4CzgCIBhrCZd9RyLnHn6+LsOh3d38I6yIHfFhKN8NfU5A+BY3MAoGCCqGSM49BAMDA2cAMGQCMGkeMEOIJGfMkzC/ntvvFI6/JrOOpkkVgJy9xZvqXfJdeZ1wlkGiVgwoc2GUIrC8+QIwWWXecj3ilgsD9LRKwPsI/kbjT5YBy+yDHIBUJdpbnhrK3/ZkLZooHdCoR/G7GDAL"},"tlogEntries":[{"logIndex":"2712211887","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531463","inclusionPromise":{"signedEntryTimestamp":"MEUCIQD913UBgBnNOu4npQrUquQiWbXwSMOOyOiqtwwkWdiqPwIgO/GgHTekha4uAVKpC8XEm10ZTzb/vAhxsDjYyuAlBBQ="},"inclusionProof":{"logIndex":"2590307625","rootHash":"iPy4Xx2bRoLI6aoGlgJvOwaJjqbhcVIOibFuBhLQw6I=","treeSize":"2590308000","hashes":["i8pzSbtIO4jJWAmwRFfERCBTQ2OaKY03wx8XNbxEMoc=","igYerqB0VJqbdOgpmCwAdbHPLE57paz5McPaw2fxqwk=","AcJ7o9z+Lif2ux96q1d6abpPDz5WRgnBOYDMeXlqSTE=","idBauuo5Gy3jKnflFw8KW42PTjSsVUShzggf6fOSk6Y=","xGv+2rMAB+JWw+TgIYgWfb/ULjSZ8PrXpqhQ6loPCoI=","Sc86HapCxXuchQXOB+rllvnYKA+WvBZh6R1Wn9J3DbE=","C3kpiX3tB9V5UTfdELykWoteJmn7hcEd7uw3RfkG0Ek=","UMtavoKHF0nY5aIZbrUd+HMxvJuRr7QFc2bVeC2Rec8=","vO4PmCjXeGtzFe+NY1hxKSwnDoTlU9Mp77kRLNrUFeE=","h7ojtjSf9rQXc39QHd7oe3sp2hNvmvx1Gl4pJcK5zQE=","h8U/Mw868cCHBNTL9fzD/L5wiigaCY+uB7tEWHkzU/0=","bZskCqT+sD+OR/yUokKG9vukKBzl/3CkQ1GhGyrF3k0=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590308000\niPy4Xx2bRoLI6aoGlgJvOwaJjqbhcVIOibFuBhLQw6I=\n\n— rekor.sigstore.dev wNI9ajBFAiEA8Oif6RzCoaoqmv0AW30eBWrmIrU25osuG+oCGAWh/2sCIGfos04XjkxmjRanD/vP5Xk1V1xsg/hAIvdTnb+r4IL1\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4NTNmY2QwNDYwNTgwYjg0NWRhOTIyZDNkYWJhYjIzODgyNDE4MTBkNTFkOWM0Y2JhMDBhMTg5ZGJkZWJiZTMwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUMrUEdFVnVyS3VpR3hCanhpcXpsb0Y2RkFGbXlCamUvMVJTZURCOWNvREVnSWdIL1JOMzU1VDVLWHlHQmNBZWNwN0pQZk9pQXc3Y0hidVhvMnVTZ2kybHRFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlZWUlhhamRRUVdaTWQwaFBURmR6VGsxTFRtbEpObUpKUkM4MGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBSZUZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWxGNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUzWTBkT09EQlpUemsyVGxsb09HaDFUbEFyU25GUlRsRXhNVGhFVldzdlFYTk5NMVlLYzJVemRsUm1TVmszZEd4TWRuSkVVbWg0U0hneFEwNW5Xa05xWm5rdlJ6ZGtNRGsxU2xSUVYySndWa1ZsZEc0ME0wdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZCTlhSekNqRmFXREJqWWpock9XOXJNbEIzVEVWMFVtOUZWRzFuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xVVVRSQlFVRlJSQXBCUlZsM1VrRkpaMFpaZGpWVFpIZENlVzVFTWtoVU9WcEhOakYxZERGVFptUjBPRzFtTkd4T05UUklibll4VERSRGVtZERTVUpvY2tOYVpEbFNlVXh1Q2todU5pdE1jMDlvTTJRek9FazJlVWxJWmtab1MwNDRUbVpWTlVFclFsa3pUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVZHJaVTFGVDBrS1NrZG1UV3Q2UXk5dWRIWjJSa2syTDBweVQwOXdhMnRXWjBwNU9YaGFkbkZZWmtwa1pWb3hkMnhyUjJsV1ozZHZZekpIVlVseVF6Z3JVVWwzVjFkWVpRcGphak5wYkdkelJEbE1Va3QzVUhOSkwydGlhbFExV1VKNUszbEVTRWxDVlVwa2NHSnVhSEpMTXk5YWEweGFiMjlJWkVOdlVpOUhOMGRFUVV3S0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgmCHN/ElVJvlD9W/Et48/S2SJK3rYhj/WwFmYy5xsAkYCFQDURVkjFuozCMc7OBD2yEO7QkVgVBgPMjAyNjA5MDQxNDE3NDFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MTc0MVowLwYJKoZIhvcNAQkEMSIEIHv1Q06+rJNbY+9QMoJs9otswawYA/WlDElN0Pf6Ts93MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDjJNC3FmnbJwPSJncUBRGgleaaFPsLWHh0odKoPqWNojsB/Ol4K5l71qAxD2BVLL4CMQC73Z+dQK/wJNtjtTc6IO4ROJ5s/WVd9KvCDKMtI5YOXjm6cqYtlwrsSagVPVJp7LY="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"hT/NBGBYC4RdqSLT2rqyOIJBgQ1R2cTLoAoYnb3rvjA="},"signature":"MEUCIQC+PGEVurKuiGxBjxiqzloF6FAFmyBje/1RSeDB9coDEgIgH/RN355T5KXyGBcAecp7JPfOiAw7cHbuXo2uSgi2ltE="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-aarch64-linux/__init__.py b/build/torch214-cxx11-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..623f8835757ff50d4c3dd2abda0664051a582837 --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0b54705535cf62a1e9b6634ec3293b2f1f46524a9e63253d4d90d018f6b2375 +size 74869264 diff --git a/build/torch214-cxx11-cu130-aarch64-linux/_ops.py b/build/torch214-cxx11-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch214-cxx11-cu130-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/layers.py b/build/torch214-cxx11-cu130-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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-aarch64-linux/metadata.json b/build/torch214-cxx11-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..af052a81033869feb47eff3488909f8d7849d251 --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-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": "oLVHBVNc9ioem2Y07DKTsvH0ZSSp5jJT1NkNAY9rI3U=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..f05edbaccc921b01f44bc3013b6d8b56b292dd2d --- /dev/null +++ b/build/torch214-cxx11-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHTTCCBtKgAwIBAgIUZ1/7f0zpLavQM9+qIq1P0LCOBV8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzQ3WhcNMjYwOTA0MTQyNzQ3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAESBVzzxQiKPwefTXknksUrM57jAgvoH/H97ltGf2lHEf7wplF4E/3Cr+jXzSxTcjyuz2d9yENPGEoSjX2CDJuOqOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUp/0oqtO9nOIzODKtCS+5I74/prAwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIY18AAAQDAEgwRgIhAO27RVtrlPC8NE9NqOuGQX6e7SB9OLZBu+xw/ydqzenUAiEAzwr5Jjm3GiYERRXwok2y9t6VOHc5/xNT4nS1Hi7m11MwCgYIKoZIzj0EAwMDaQAwZgIxAKqYcPmYxCJZOo/rO2b3adl08j0P/5GWZjKBvPRxJYPt6uWcCZE7Om5x7ebzCZf6bwIxAImSR5+BhxNd5DIj1nDqdWHdf7FOegevojuntCrwhBBRyBzP2+3teoIHbKMTlQAZ+Q=="},"tlogEntries":[{"logIndex":"2712213170","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531469","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCDf19EIzC7BhQqppR91DiluEyNT357cM0MtvMQgGEyEAIhAIIYqPWpubh98+r6mtnIYytMFubIwu0f8g3m7KzB0oST"},"inclusionProof":{"logIndex":"2590308908","rootHash":"CGaXKNzDSaKb7JPCHhQxRsCwCR17qJVX5quXcRoj5l4=","treeSize":"2590308970","hashes":["N/G6pD4MePVhmwKpVGJ7haFZ/LEjA4Ah8xSstcXZ49o=","ZJByqzHcek2qFeezZj7wWVJl9tB0xkqp6j77W87Gf4s=","RFdWOfDH2mieXH+Mln81wyBqWJR31Cwj3Sk7euPohts=","/GiwLqGbQppDfVKM74bYAn1iYHSkuUSnwKJDE2DGWw8=","Oz1PzoZSE7AmV6AWlH937/xS6WHbSKdeFNHQZwPi21k=","m9W/h8ytzobcMoAHXSqLGXoj3eIvhfx/dKfdjmfjr1w=","kV3xGBvsFhX2fOT1uBiPKR+LUBEI8xlugdo3sLuesAE=","W4sGQUcL2i3YRDDpvurMXD52kkVPqhWQzXWP+MDEXAM=","HH9t0vqxOQmp65YCmpAX1S8uUf/m+J1f2tUJki7zbCQ=","bZskCqT+sD+OR/yUokKG9vukKBzl/3CkQ1GhGyrF3k0=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590308970\nCGaXKNzDSaKb7JPCHhQxRsCwCR17qJVX5quXcRoj5l4=\n\n— rekor.sigstore.dev wNI9ajBGAiEAkeqzpizM+CyPf7+2zWdKR12pkpzYiryQ/PMeeapuQiACIQDY0mS4OTTZEwD9aBDjXSkXEO8iYrELNRUwkr3AbI9sIQ==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0MDFjMDZkZTI1OWJkYjVmNTc3MTFkNjEwMjI1ZDRhN2FmMWM4MDExYzgyMjc4OTk1ZWYyMGM1N2JiYWIyMGM3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRkg4a1FCZVFGWDlMUm1NNlBPaXkzMXRSTjVGNWZlZ2VlbGw3czhtN0R5T0FpQU1UUS9mUnhDTEFCZTVISUtrRWllMi9aa25iUUNuUThtaVBuclRaYVFJYnc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVWRU5EUW5STFowRjNTVUpCWjBsVldqRXZOMll3ZW5CTVlYWlJUVGtyY1VseE1WQXdURU5QUWxZNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBSTTFkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWxFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZUUWxaNmVuaFJhVXRRZDJWbVZGaHJibXR6VlhKTk5UZHFRV2QyYjBndlNEazNiSFFLUjJZeWJFaEZaamQzY0d4R05FVXZNME55SzJwWWVsTjRWR05xZVhWNk1tUTVlVVZPVUVkRmIxTnFXREpEUkVwMVQzRlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ3THpCdkNuRjBUemx1VDBsNlQwUkxkRU5US3pWSk56UXZjSEpCZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xaTVRoQlFVRlJSQXBCUldkM1VtZEphRUZQTWpkU1ZuUnliRkJET0U1Rk9VNXhUM1ZIVVZnMlpUZFRRamxQVEZwQ2RTdDRkeTk1WkhGNlpXNVZRV2xGUVhwM2NqVkthbTB6Q2tkcFdVVlNVbGgzYjJzeWVUbDBObFpQU0dNMUwzaE9WRFJ1VXpGSWFUZHRNVEZOZDBObldVbExiMXBKZW1vd1JVRjNUVVJoVVVGM1dtZEplRUZMY1ZrS1kxQnRXWGhEU2xwUGJ5OXlUekppTTJGa2JEQTRhakJRTHpWSFYxcHFTMEoyVUZKNFNsbFFkRFoxVjJORFdrVTNUMjAxZURkbFlucERXbVkyWW5kSmVBcEJTVzFUVWpVclFtaDRUbVExUkVscU1XNUVjV1JYU0dSbU4wWlBaV2RsZG05cWRXNTBRM0ozYUVKQ1VubENlbEF5S3pOMFpXOUpTR0pMVFZSc1VVRmFDaXRSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgefr3igDdcMjpSsIoO2e33L0JKQF7jW13l02sLkIgcIECFQDr4JACsomIh5bojwtWm1gEpHF9SxgPMjAyNjA5MDQxNDE3NDdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MTc0N1owLwYJKoZIhvcNAQkEMSIEIAwxPnrTi0Myn5a0Ce6PGgCu1jSmf0Pj6dOTfm8srCbpMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMAp0pefL4mNyAEr+KWEmMCCG2ySKB46E2VXzDvABz3OCQevAJTDA4U9xUteKXMA51AIwCaNkTYH6pytYZ0i5znmk1a+b4M2vq3RyjNVKgto1YPvMED0D93A4rxpq65MvW39+"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"QBwG3iWb219XcR1hAiXUp68cgBHIIniZXvIMV7urIMc="},"signature":"MEQCIFH8kQBeQFX9LRmM6POiy31tRN5F5fegeell7s8m7DyOAiAMTQ/fRxCLABe5HIKkEie2/ZknbQCnQ8miPnrTZaQIbw=="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-aarch64-linux/__init__.py b/build/torch214-cxx11-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a80bd24a974027db3cef7988137e7d6e39ec29 --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so b/build/torch214-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..dc296ef87a27acf7228d134e97f3382f2ae6e295 --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9060f88613b646edb6c5bfde4b18caeef648987efefec6050385662af798d6d1 +size 75002824 diff --git a/build/torch214-cxx11-cu132-aarch64-linux/_ops.py b/build/torch214-cxx11-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..470f80eca502956c063147fd602a9be409790b3d --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_interface.py b/build/torch214-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..d46d56c415e91e16a475a7261e01658b2259d377 --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/causal_conv1d_varlen.py b/build/torch214-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..e319c19f4a19727497e8aa0859ee42b57c47eb53 --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/cpp_functions.py b/build/torch214-cxx11-cu132-aarch64-linux/cpp_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1ddb9f83ceb4e9f72754fe39340738d47b6aea1b --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/layers.py b/build/torch214-cxx11-cu132-aarch64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf7af48fbd31dee1845c9e1dd1f591f1e5f95ef --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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-aarch64-linux/metadata.json b/build/torch214-cxx11-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..47f5275436dd8ed7c865c1080e4a969a8f95c1ae --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-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": "kGD4hhO2Ru22xb/eSxjK7vZImH7+/sYFA4VmKveY1tE=", + "_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-aarch64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..7ff1f003f5d6251e210e613b4ed3e65445413ce7 --- /dev/null +++ b/build/torch214-cxx11-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUaGnXcpeZVXfQ+f6WzqIrq+7PQEYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTA0MTQxNzUxWhcNMjYwOTA0MTQyNzUxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEZilJ6cNLTYb2/Y3VZ3qU0UwHxgA1Z5TigbFnu9Itz7vpXEDlc4IESy4XGw3AYh6dL/s45D63ana5hWO9QjWgtaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUj5KWJ+nUbwFRNi/8C54NQB2X0YIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMDU2ZTI4Yjg1MjIwYzAzMmI0ZGMxYmIzZGFkZWJkZGQ1NmRjMjI4ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDA1NmUyOGI4NTIyMGMwMzJiNGRjMWJiM2RhZGViZGRkNTZkYzIyOGUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzM4ODA5ODMwMzcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoGzIcx4AAAQDAEcwRQIgKYZNLZs2KXp7qBb0qZE3l5aIJg6zXfczo1jFOrEiwvICIQDRDoHs03yvJtM9co9wWJNUfoFWuQP6Wat7aftYO/50ADAKBggqhkjOPQQDAwNoADBlAjAqvBJBrsbhA44sOSGgatq/S2M0kVDu1S8yB/Wkk5Q/F6MXIOpSiZ9FG3YbBJdtA4gCMQDljlx7ysuTYizumxSxzMNSCWiL8CgpdCH0Q6Tmwt3ljdPzrJcxb0nXycYvphG1opQ="},"tlogEntries":[{"logIndex":"2712213894","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1788531472","inclusionPromise":{"signedEntryTimestamp":"MEUCIQDI/0voKBpaiH/RJ3H30dLSeSaXbXkikfnaWOOZ6O1s9wIgJK62OQoJ8VJRdkY5SgqR7h6SxD3ZXANWYRTqBM7QQj4="},"inclusionProof":{"logIndex":"2590309632","rootHash":"b5hkN3iMdOrkWzel7I2Z0llMXXYQiNUjVGVVDv0G5zg=","treeSize":"2590309762","hashes":["8Y4WzRo6S6DE3wVDArU7gi5O53U1/svP4zvC0aK3XAw=","v8JUavlVl1I6Hx797Mvjn1WZ0c2ZGqVRriRbvpX+ySQ=","2DSVPVC7/wqpmhw03Ogjf9JNCEyiGc7PGW32jcaVx0E=","B2fsVknsuMmosObGhe638AtcZCGBw8vaxzHI7dRCNnw=","nqO+7WahH76UF/DJQVhLrJY9gOcNbAea5QOlsSf+BQU=","x+ObHxrB+3nGfOiT6DvOUvp+cdknnocOZZV0k7Z/cNk=","t9i8tbaDkzOo1W+w2EpkAVjWopIUtG4ui/mlrmgLR6Q=","TaRt3BCRWI+AvTRf+DLMaBjvumTD/23xd7ARmZA1sVI=","dibzsaaX/WX4I+KvB9JPRVPva4LdDM995vBf2Ma0QTw=","itZx1puagQ5wXsd3ocJVvhdVDRV0E04lXVMbO8A+i18=","HH9t0vqxOQmp65YCmpAX1S8uUf/m+J1f2tUJki7zbCQ=","bZskCqT+sD+OR/yUokKG9vukKBzl/3CkQ1GhGyrF3k0=","r+69k0fOjzxcXnOwIlD352Kuyjs22Maer1/AFrceTYo=","ohrY7Q07dQsFxHshSCrWurSq9MSwWE8Ea+QdavE4zrM=","8Ne5O/slxyxV3WMOHVD3Kjoe/OUxlahj83eaPxftJ6U=","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\n2590309762\nb5hkN3iMdOrkWzel7I2Z0llMXXYQiNUjVGVVDv0G5zg=\n\n— rekor.sigstore.dev wNI9ajBFAiB3ZK7vaTkg9XJa8E1ud7LdniaQBfD6wGNKEwmGiRTlGQIhAMoxySLNhC2JqQczgIRV3IYNxUuU/r/Fv93TzxocLjlY\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjZTA5NTcwOGQ5MDFjMmU2Y2ZmODc1OTY0OWRlNDZlZTM5OTJlZWI0NWViNmUwYzYxZjI1MWM5ZWZjZWE3NzQ1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSHozRy9iSjVrelRkamN5elRQSWdvdUpnbTc0Y2dKZHZYU2VwUWVwSnFvMUFpRUEvRGZPYmhkS0NGekg1RHJEeUgySXA4S0t2MjVxVFR3bVR4bURESkxYejdvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVllVZHVXR053WlZwV1dHWlJLMlkyVjNweFNYSnhLemRRVVVWWmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUVRCTlZGRjRUbnBWZUZkb1kwNU5hbGwzVDFSQk1FMVVVWGxPZWxWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZhYVd4S05tTk9URlJaWWpJdldUTldXak54VlRCVmQwaDRaMEV4V2pWVWFXZGlSbTRLZFRsSmRIbzNkbkJZUlVSc1l6UkpSVk41TkZoSGR6TkJXV2cyWkV3dmN6UTFSRFl6WVc1aE5XaFhUemxSYWxkbmRHRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZxTlV0WENrb3JibFZpZDBaU1Rta3ZPRU0xTkU1UlFqSllNRmxKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SRlV5V2xSSk5GbHFaekZOYWtsM1dYcEJlazF0U1RCYVIwMTRDbGx0U1hwYVIwWnJXbGRLYTFwSFVURk9iVkpxVFdwSk5GcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGhaYlVsNldrZEdhMXBYU21zS1drZFJNVTV0VW1wTmFrazBXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFFVlRKYVZFazBXV3BuTVUxcVNYZFpla0Y2VFcxSk1GcEhUWGdLV1cxSmVscEhSbXRhVjBwcldrZFJNVTV0VW1wTmFrazBXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkJNVTV0VlhrS1QwZEpORTVVU1hsTlIwMTNUWHBLYVU1SFVtcE5WMHBwVFRKU2FGcEhWbWxhUjFKclRsUmFhMWw2U1hsUFIxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWswMFQwUkJOVTlFVFhkTmVtTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUhla2xqZURSQlFVRlJSQXBCUldOM1VsRkpaMHRaV2s1TVduTXlTMWh3TjNGQ1lqQnhXa1V6YkRWaFNVcG5ObnBZWm1ONmJ6RnFSazl5UldsM2RrbERTVkZFVWtSdlNITXdNM2wyQ2twMFRUbGpiemwzVjBwT1ZXWnZSbGQxVVZBMlYyRjBOMkZtZEZsUEx6VXdRVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJjWFpDU2tJS2NuTmlhRUUwTkhOUFUwZG5ZWFJ4TDFNeVRUQnJWa1IxTVZNNGVVSXZWMnRyTlZFdlJqWk5XRWxQY0ZOcFdqbEdSek5aWWtKS1pIUkJOR2REVFZGRWJBcHFiSGczZVhOMVZGbHBlblZ0ZUZONGVrMU9VME5YYVV3NFEyZHdaRU5JTUZFMlZHMTNkRE5zYW1SUWVuSktZM2hpTUc1WWVXTlpkbkJvUnpGdmNGRTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgEyvQkKazdH6GlvAw0rgP452FDxSKt84xho0rCrfWDh4CFQDqDnw5TfJBCMS0GXWOStKInTq7QBgPMjAyNjA5MDQxNDE3NTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkwNDE0MTc1MVowLwYJKoZIhvcNAQkEMSIEIFQ+ngENpLxvAFylbkUXdPi+QRp+2WGP01OCxTUgX9rJMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCur89yQ1XPCCePR95L3E6Ju3G+9vq1JJ9gO0t0ZNTIf1Coa+LZ0fgUEduBQ0v9DCwCMFrE4i9/LrD/XGop4HtCoLDEu7Ch1e4SmGDayOFiaGkDm5XGx63yD6TixxM3c3lAtA=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"zglXCNkBwubP+HWWSd5G7jmS7rRetuDGHyUcnvzqd0U="},"signature":"MEUCIHz3G/bJ5kzTdjcyzTPIgouJgm74cgJdvXSepQepJqo1AiEA/DfObhdKCFzH5DrDyH2Ip8KKv25qTTwmTxmDDJLXz7o="}} \ No newline at end of file