Skip to main content

Causal depthwise conv1d in CUDA, with a PyTorch interface

Project description

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 
    

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

causal_conv1d-1.5.3.post1.tar.gz (24.2 kB view details)

Uploaded Source

File details

Details for the file causal_conv1d-1.5.3.post1.tar.gz.

File metadata

  • Download URL: causal_conv1d-1.5.3.post1.tar.gz
  • Upload date:
  • Size: 24.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.18

File hashes

Hashes for causal_conv1d-1.5.3.post1.tar.gz
Algorithm Hash digest
SHA256 aba1b717484472d0b2f2e40520a1c03f35fe5155555bd753d1c324afc56ba468
MD5 31b9e83cf5b08342a1e75e435d143433
BLAKE2b-256 64cb104778c728dc3d5ea3bf65a484e3a4cdbe894bdaa2586320e2f61d007b8c

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page