Kernels:

kernels-community
/

causal-conv1d

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Sep 4

Commit

512a4a9

verified ·

1 Parent(s): 2999c83

Uploaded using `kernel-builder`.

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

build/torch212-cxx11-cu126-aarch64-linux/__init__.py +12 -0
build/torch212-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch212-cxx11-cu126-aarch64-linux/_ops.py +9 -0
build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py +86 -0
build/torch212-cxx11-cu126-aarch64-linux/cpp_functions.py +96 -0
build/torch212-cxx11-cu126-aarch64-linux/layers.py +50 -0
build/torch212-cxx11-cu126-aarch64-linux/metadata.json +44 -0
build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore +1 -0
build/torch212-cxx11-cu130-aarch64-linux/__init__.py +12 -0
build/torch212-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch212-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py +86 -0
build/torch212-cxx11-cu130-aarch64-linux/cpp_functions.py +96 -0
build/torch212-cxx11-cu130-aarch64-linux/layers.py +50 -0
build/torch212-cxx11-cu130-aarch64-linux/metadata.json +47 -0
build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore +1 -0
build/torch212-cxx11-cu132-aarch64-linux/__init__.py +12 -0
build/torch212-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch212-cxx11-cu132-aarch64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py +86 -0
build/torch212-cxx11-cu132-aarch64-linux/cpp_functions.py +96 -0
build/torch212-cxx11-cu132-aarch64-linux/layers.py +50 -0
build/torch212-cxx11-cu132-aarch64-linux/metadata.json +47 -0
build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore +1 -0
build/torch213-cxx11-cu126-aarch64-linux/__init__.py +12 -0
build/torch213-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch213-cxx11-cu126-aarch64-linux/_ops.py +9 -0
build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py +86 -0
build/torch213-cxx11-cu126-aarch64-linux/cpp_functions.py +96 -0
build/torch213-cxx11-cu126-aarch64-linux/layers.py +50 -0
build/torch213-cxx11-cu126-aarch64-linux/metadata.json +44 -0
build/torch213-cxx11-cu126-aarch64-linux/metadata.json.sigstore +1 -0
build/torch213-cxx11-cu130-aarch64-linux/__init__.py +12 -0
build/torch213-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py +86 -0
build/torch213-cxx11-cu130-aarch64-linux/cpp_functions.py +96 -0
build/torch213-cxx11-cu130-aarch64-linux/layers.py +50 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +47 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json.sigstore +1 -0
build/torch213-cxx11-cu132-aarch64-linux/__init__.py +12 -0
build/torch213-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so +3 -0
build/torch213-cxx11-cu132-aarch64-linux/_ops.py +9 -0
build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py +242 -0
build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py +86 -0

build/torch212-cxx11-cu126-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:7a438e8fbe2a8bf179fccaa5738cbc3d3696d2207b6a2ba5ef161ebcebac50ac
+size 78911952

build/torch212-cxx11-cu126-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch212-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu126-aarch64-linux/cpp_functions.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu126-aarch64-linux/layers.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu126-aarch64-linux/metadata.json ADDED Viewed

	@@ -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
+    }
+  }
+}

build/torch212-cxx11-cu126-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -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="}}

build/torch212-cxx11-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:49e5ababd9692f41e928027eecd1897d156a16187fc778e2b10a3036468157ad
+size 74934456

build/torch212-cxx11-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch212-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu130-aarch64-linux/cpp_functions.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu130-aarch64-linux/layers.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -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
+    }
+  }
+}

build/torch212-cxx11-cu130-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -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"}}

build/torch212-cxx11-cu132-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:012229d08b1266fcd8a9ed38c7a50661a1ced324a906df2d7959639c54623aa1
+size 75002488

build/torch212-cxx11-cu132-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch212-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu132-aarch64-linux/cpp_functions.py ADDED Viewed

	@@ -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

build/torch212-cxx11-cu132-aarch64-linux/layers.py ADDED Viewed

	@@ -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",
+]

build/torch212-cxx11-cu132-aarch64-linux/metadata.json ADDED Viewed

	@@ -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
+    }
+  }
+}

build/torch212-cxx11-cu132-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -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=="}}

build/torch213-cxx11-cu126-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch213-cxx11-cu126-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e166fbb8d695ba179fe3d197a4daf5da55e8d8ae398a28340d28cb80ddb29a2e
+size 78911912

build/torch213-cxx11-cu126-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch213-cxx11-cu126-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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

build/torch213-cxx11-cu126-aarch64-linux/cpp_functions.py ADDED Viewed

	@@ -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

build/torch213-cxx11-cu126-aarch64-linux/layers.py ADDED Viewed

	@@ -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",
+]

build/torch213-cxx11-cu126-aarch64-linux/metadata.json ADDED Viewed

	@@ -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
+    }
+  }
+}

build/torch213-cxx11-cu126-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -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="}}

build/torch213-cxx11-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch213-cxx11-cu130-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:ce373b8049bc86d3309df58ee8d6d23f5a79bf991a8150ff66dffa1c2b40a995
+size 74934224

build/torch213-cxx11-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch213-cxx11-cu130-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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

build/torch213-cxx11-cu130-aarch64-linux/cpp_functions.py ADDED Viewed

	@@ -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

build/torch213-cxx11-cu130-aarch64-linux/layers.py ADDED Viewed

	@@ -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",
+]

build/torch213-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -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
+    }
+  }
+}

build/torch213-cxx11-cu130-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -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"}}

build/torch213-cxx11-cu132-aarch64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch213-cxx11-cu132-aarch64-linux/_causal_conv1d_cuda_056e28b.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:79bcf114bcc754bc73b331505be2a3655b2e8348bfbef3b5883ac06719b9de83
+size 75002256

build/torch213-cxx11-cu132-aarch64-linux/_ops.py ADDED Viewed

	@@ -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}"

build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_interface.py ADDED Viewed

	@@ -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)

build/torch213-cxx11-cu132-aarch64-linux/causal_conv1d_varlen.py ADDED Viewed

	@@ -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