Oven logo

Oven

causal-conv1d1.5.0.post8

Published

Causal depthwise conv1d in CUDA, with a PyTorch interface

pip install causal-conv1d

Package Downloads

Weekly DownloadsMonthly Downloads

Project URLs

Requires Python

>=3.9

Dependencies

    Causal depthwise conv1d in CUDA with a PyTorch interface

    Features:

    • Support fp32, fp16, bf16.
    • Kernel size 2, 3, 4.

    How to use

    from causal_conv1d import causal_conv1d_fn
    
    def causal_conv1d_fn(x, weight, bias=None, activation=None):
        """
        x: (batch, dim, seqlen)
        weight: (dim, width)
        bias: (dim,)
        activation: either None or "silu" or "swish"
    
        out: (batch, dim, seqlen)
        """
    

    Equivalent to:

    import torch.nn.functional as F
    
    F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)[..., :seqlen]
    

    Additional Prerequisites for AMD cards

    Patching ROCm

    If you are on ROCm 6.0, run the following steps to avoid errors during compilation. This is not required for ROCm 6.1 onwards.

    1. Locate your ROCm installation directory. This is typically found at /opt/rocm/, but may vary depending on your installation.

    2. Apply the Patch. Run with sudo in case you encounter permission issues.

       patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch