Causal depthwise conv1d in CUDA, with a PyTorch interface
Project Links
Meta
Author: Tri Dao
Requires Python: >=3.9
Classifiers
Programming Language
- Python :: 3
License
- OSI Approved :: BSD License
Operating System
- Unix
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.
-
Locate your ROCm installation directory. This is typically found at
/opt/rocm/, but may vary depending on your installation. -
Apply the Patch. Run with
sudoin case you encounter permission issues.patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch
1.6.2.post1
May 09, 2026
1.6.1
Mar 10, 2026
1.6.0
Jan 12, 2026
1.5.3.post1
Oct 10, 2025
1.5.2
Jul 18, 2025
1.5.1
Jul 18, 2025
1.5.0.post8
Dec 06, 2024
1.5.0.post5
Dec 06, 2024
1.4.0
Jun 29, 2024
1.3.0.post1
Jun 18, 2024
1.2.2.post1
May 26, 2024
1.2.1
May 26, 2024
1.2.0.post2
Mar 02, 2024
1.2.0.post1
Mar 02, 2024
1.1.3.post1
Feb 01, 2024
1.1.2.post1
Feb 01, 2024
1.1.1
Dec 25, 2023
1.1.0
Dec 10, 2023
1.0.2
Dec 10, 2023
1.0.1
Dec 08, 2023
1.0.0
Dec 03, 2023
