jp6/cu122/: causal-conv1d-1.4.0 metadata and description

Homepage Simple index

Causal depthwise conv1d in CUDA, with a PyTorch interface

author Tri Dao
author_email tri@tridao.me
classifiers
  • Programming Language :: Python :: 3
  • License :: OSI Approved :: BSD License
  • Operating System :: Unix
description_content_type text/markdown
requires_dist
  • torch
  • packaging
  • ninja
requires_python >=3.8

Because this project isn't in the mirror_whitelist, no releases from root/pypi are included.

File Tox results History
causal_conv1d-1.4.0-cp310-cp310-linux_aarch64.whl
Size
16 MB
Type
Python Wheel
Python
3.10

Causal depthwise conv1d in CUDA with a PyTorch interface

Features:

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