Subquadratic Ops Torch Documentation#

subquadratic_ops_torch provides high-performance CUDA kernels for core operations in popular subquadratic attention architectures, such as Hyena, Mamba, etc. These operations include long and short depthwise convolutions, in multiple dimensions. The library contains PyTorch bindings to these optimized kernels which can be used to accelerate models that rely on these operations.

Kernels are primarily exposed as function calls underlying torch.ops, which also provide a lower-level interface as torch.library operators.

Installation#

Please install using pip install subquadratic-ops-torch-cu13

Usage#

You can import the library from python:

import subquadratic_ops_torch as subq

Requirements#

  • CUDA-compatible NVIDIA GPU (Ampere+)

  • CUDA 13 (the wheel is CUDA-version-specific: subquadratic-ops-torch-cu13 depends on nvidia-cudnn-cu13, so a CUDA 12 environment fails at runtime, not at install time)

  • Python 3.11–3.14

  • b2b_causal_conv1d only: cuDNN 9.24 or newer plus nvidia-cudnn-frontend 1.27.0 or newer. It is served by the cuDNN frontend rather than by a kernel in this wheel, does not support float64, and raises rather than falling back if the frontend is unavailable. See the v0.3.0 breaking changes in Release Notes.

Support and Feedback#

Please contact the developers for any issues you might encounter.

  • Alireza Moradzadeh, amoradzadeh at nvidia.com

  • Saee Paliwal, saeep at nvidia.com