Metadata-Version: 2.2
Name: subquadratic-ops-torch-cu13
Version: 0.1.1
Summary: subquadratic-ops-torch-cu13 - GPU Accelerated Torch Extensions for Subquadratic Operations
Author-Email: Alireza Moradzadeh <amoradzadeh@nvidia.com>
License: Other/Proprietary License (NVIDIA Proprietary Software)
         
         # Software License Agreement
         
         LICENSE AGREEMENT FOR NVIDIA MATH LIBRARIES SOFTWARE DEVELOPMENT KITS
         
         This license agreement(“Agreement”) is a legal agreement between you and NVIDIA Corporation (“NVIDIA”) and governs your use of the NVIDIA math libraries software development kit as available at NVIDIA’s discretion (each, a “SDK”).
         
         Each SDK has its own set of software and materials, but here is a description of the types of items that may be included in a SDK: source code, header files, APIs, data sets and assets (examples include images, textures, models, scenes, videos, native API input/output files), binary software, sample code, libraries, utility programs, programming code and documentation.
         
         This Agreement can be accepted only by an adult of legal age of majority in the country in which the SDK is used.
         
         If you are entering into this Agreement on behalf of a company or other legal entity, you represent that you have the legal authority to bind the entity to this Agreement, in which case "you" will mean the entity you represent.
         
         If you don't have the required age or authority to accept this Agreement, or if you don't accept all the terms and conditions of this Agreement, do not download, install or use the SDK.
         
         You agree to use the SDK only for purposes that are permitted by (a) this Agreement, and (b) any applicable law, regulation or generally accepted practices or guidelines in the relevant jurisdictions.
         
             License.
         
         1.1 Grant
         
         Subject to the terms of this Agreement, NVIDIA hereby grants you a non-exclusive, non-transferable license, without the right to sublicense (except as expressly provided in this Agreement) to:
         
             Install and use the SDK, and
         
             Distribute the binary files, files identified as samples, and headers as incorporated into a software application that meets the distribution requirements indicated in this Agreement.
         
         1.2 Distribution Requirements
         
         These are the distribution requirements for you to exercise the distribution grant:
         
             Your application must have material additional functionality, beyond the included portions of the SDK.
         
             The distributable portions of the SDK shall only be accessed by your application.
         
             The following notice shall be included in modifications and derivative works of sample source code distributed: "This software contains source code provided by NVIDIA Corporation."
         
             Unless a developer tool is identified in this Agreement as distributable, it is delivered for your internal use only.
         
             The terms under which you distribute your application must be consistent with the terms of this Agreement, including (without limitation) terms relating to the license grant and license restrictions and protection of NVIDIA's intellectual property rights. Additionally, you agree that you will protect the privacy, security and legal rights of your application users.
         
             You agree to notify NVIDIA in writing of any known or suspected distribution or use of the SDK not in compliance with the requirements of this Agreement, and to enforce the terms of your agreements with respect to distributed SDK.
         
         1.3 Authorized Users
         
         You may allow employees and contractors of your entity or of your subsidiary(ies) to access and use the SDK from your secure network to perform work on your behalf.
         
         If you are an academic institution you may allow users enrolled or employed by the academic institution to access and use the SDK from your secure network.
         
         You are responsible for the compliance with the terms of this Agreement by your authorized users. If you become aware that your authorized users didn't follow the terms of this Agreement, you agree to take reasonable steps to resolve the non-compliance and prevent new occurrences.
         
         1.4 Pre-Release SDK
         
         The SDK versions identified as alpha, beta, preview or otherwise as pre-release, may not be fully functional, may contain errors or design flaws, and may have reduced or different security, privacy, accessibility, availability, and reliability standards relative to commercial versions of NVIDIA software and materials. Use of a pre-release SDK may result in unexpected results, loss of data, project delays or other unpredictable damage or loss. You may use a pre-release SDK at your own risk, understanding that pre-release SDKs are not intended for use in production or business-critical systems. NVIDIA may choose not to make available a commercial version of any pre-release SDK. NVIDIA may also choose to abandon development and terminate the availability of a pre-release SDK at any time without liability. 1.5 Updates
         
         NVIDIA may, at its option, make available patches, workarounds or other updates to this SDK. Unless the updates are provided with their separate governing terms, they are deemed part of the SDK licensed to you as provided in this Agreement.
         
         You agree that the form and content of the SDK that NVIDIA provides may change without prior notice to you. While NVIDIA generally maintains compatibility between versions, NVIDIA may in some cases make changes that introduce incompatibilities in future versions of the SDK.
         
         1.6 Components Under Other Licenses.
         
         The SDK may include NVIDIA or third-party components with separate legal notices or terms as may be described in proprietary notices accompanying the SDK, such as components governed by open source software licenses. If and to the extent there is a conflict between the terms in this license and the license terms associated with a component, the license terms associated with the components control only to the extent necessary to resolve the conflict.
         
         1.7 Reservation of Rights
         
         NVIDIA reserves all rights, title and interest in and to the SDK not expressly granted to you under this Agreement.
         
             Limitations.
         
         The following license limitations apply to your use of the SDK:
         
         2.1 The SDK is licensed for you to develop applications only for use in systems with NVIDIA GPUs.
         
         2.2 You may not reverse engineer, decompile or disassemble, or remove copyright or other proprietary notices from any portion of the SDK or copies of the SDK.
         
         2.3 Except as expressly provided in this Agreement, you may not copy, sell, rent, sublicense, transfer, distribute, modify, or create derivative works of any portion of the SDK.
         
         2.4 Unless you have an agreement with NVIDIA for this purpose, you may not indicate that an application created with the SDK is sponsored or endorsed by NVIDIA.
         
         2.5 You may not bypass, disable, or circumvent any encryption, security, digital rights management or authentication mechanism in the SDK.
         
         2.6 You may not use the SDK in any manner that would cause it to become subject to an open source software license. As examples, licenses that require as a condition of use, modification, and/or distribution that the SDK be (i) disclosed or distributed in source code form; (ii) licensed for the purpose of making derivative works; or (iii) redistributable at no charge.
         
         2.7 You acknowledge that the SDK as delivered is not tested or certified by NVIDIA for use in connection with the design, construction, maintenance, and/or operation of any system where the use or failure of such system could result in a situation that threatens the safety of human life or results in catastrophic damages (each, a "Critical Application"). Examples of Critical Applications include use in avionics, navigation, autonomous vehicle applications, ai solutions for automotive products, military, medical, life support or other life critical applications. NVIDIA shall not be liable to you or any third party, in whole or in part, for any claims or damages arising from such uses. You are solely responsible for ensuring that any product or service developed with the SDK as a whole includes sufficient features to comply with all applicable legal and regulatory standards and requirements.
         
         2.8 You agree to defend, indemnify and hold harmless NVIDIA and its affiliates, and their respective employees, contractors, agents, officers and directors, from and against any and all claims, damages, obligations, losses, liabilities, costs or debt, fines, restitutions and expenses (including but not limited to attorney's fees and costs incident to establishing the right of indemnification) arising out of or related to products or services that use the SDK in or for Critical Applications, and for use of the SDK outside of the scope of this Agreement, or not in compliance with its terms.
         
             Ownership.
         
         3.1 NVIDIA or its licensors hold all rights, title and interest in and to the SDK and its modifications, including their respective intellectual property rights. This SDK may include software and materials from NVIDIA's licensors, and these licensors are intended third party beneficiaries that may enforce this Agreement with respect to their intellectual property rights.
         
         3.2 You may, but don't have to, provide to NVIDIA suggestions, feature requests or other feedback regarding the SDK, including possible enhancements or modifications to the SDK. For any feedback that you voluntarily provide, you hereby grant NVIDIA and its affiliates a perpetual, non-exclusive, worldwide, irrevocable license to use, reproduce, modify, license, sublicense (through multiple tiers of sublicensees), and distribute (through multiple tiers of distributors) it without the payment of any royalties or fees to you. NVIDIA will use feedback at its choice.
         
             No Warranties.
         
         THE SDK IS PROVIDED BY NVIDIA "AS IS" AND "WITH ALL FAULTS." TO THE MAXIMUM EXTENT PERMITTED BY LAW, NVIDIA AND ITS AFFILIATES EXPRESSLY DISCLAIM ALL WARRANTIES OF ANY KIND OR NATURE, WHETHER EXPRESS, IMPLIED OR STATUTORY, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE, NON-INFRINGEMENT, OR THE ABSENCE OF ANY DEFECTS THEREIN, WHETHER LATENT OR PATENT. NO WARRANTY IS MADE ON THE BASIS OF TRADE USAGE, COURSE OF DEALING OR COURSE OF TRADE.
         
             Limitations of Liability.
         
         TO THE MAXIMUM EXTENT PERMITTED BY LAW, NVIDIA AND ITS AFFILIATES SHALL NOT BE LIABLE FOR ANY SPECIAL, INCIDENTAL, PUNITIVE OR CONSEQUENTIAL DAMAGES, OR ANY LOST PROFITS, LOSS OF USE, LOSS OF DATA OR LOSS OF GOODWILL, OR THE COSTS OF PROCURING SUBSTITUTE PRODUCTS, ARISING OUT OF OR IN CONNECTION WITH THIS AGREEMENT OR THE USE OR PERFORMANCE OF THE SDK, WHETHER SUCH LIABILITY ARISES FROM ANY CLAIM BASED UPON BREACH OF CONTRACT, BREACH OF WARRANTY, TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY OR ANY OTHER CAUSE OF ACTION OR THEORY OF LIABILITY. IN NO EVENT WILL NVIDIA'S AND ITS AFFILIATES TOTAL CUMULATIVE LIABILITY UNDER OR ARISING OUT OF THIS AGREEMENT EXCEED US$10.00. THE NATURE OF THE LIABILITY OR THE NUMBER OF CLAIMS OR SUITS SHALL NOT ENLARGE OR EXTEND THIS LIMIT.
         
         These exclusions and limitations of liability shall apply regardless if NVIDIA or its affiliates have been advised of the possibility of such damages, and regardless of whether a remedy fails its essential purpose. These exclusions and limitations of liability form an essential basis of the bargain between the parties, and, absent any of these exclusions or limitations of liability, the provisions of this Agreement, including, without limitation, the economic terms, would be substantially different.
         
             Termination.
         
         6.1 This Agreement will continue to apply until terminated by either you or NVIDIA as described below.
         
         6.2 If you want to terminate this Agreement, you may do so by stopping to use the SDK.
         
         6.3 NVIDIA may, at any time, terminate this Agreement if: (i) you fail to comply with any term of this Agreement and the non-compliance is not fixed within thirty (30) days following notice from NVIDIA (or immediately if you violate NVIDIA's intellectual property rights); (ii) you commence or participate in any legal proceeding against NVIDIA with respect to the SDK; or (iii) NVIDIA decides to no longer provide the SDK in a country or, in NVIDIA's sole discretion, the continued use of it is no longer commercially viable.
         
         6.4 Upon any termination of this Agreement, you agree to promptly discontinue use of the SDK and destroy all copies in your possession or control. Your prior distributions in accordance with this Agreement are not affected by the termination of this Agreement. Upon written request, you will certify in writing that you have complied with your commitments under this section. Upon any termination of this Agreement all provisions survive except for the licenses granted to you.
         
             General.
         
         If you wish to assign this Agreement or your rights and obligations, including by merger, consolidation, dissolution or operation of law, contact NVIDIA to ask for permission. Any attempted assignment not approved by NVIDIA in writing shall be void and of no effect. NVIDIA may assign, delegate or transfer this Agreement and its rights and obligations, and if to a non-affiliate you will be notified.
         
         You agree to cooperate with NVIDIA and provide reasonably requested information to verify your compliance with this Agreement.
         
         This Agreement will be governed in all respects by the laws of the United States and of the State of Delaware as those laws are applied to contracts entered into and performed entirely within Delaware by Delaware residents, without regard to the conflicts of laws principles. The United Nations Convention on Contracts for the International Sale of Goods is specifically disclaimed. You agree to all terms of this Agreement in the English language.
         
         The state or federal courts residing in Santa Clara County, California shall have exclusive jurisdiction over any dispute or claim arising out of this Agreement. Notwithstanding this, you agree that NVIDIA shall still be allowed to apply for injunctive remedies or an equivalent type of urgent legal relief in any jurisdiction.
         
         If any court of competent jurisdiction determines that any provision of this Agreement is illegal, invalid or unenforceable, such provision will be construed as limited to the extent necessary to be consistent with and fully enforceable under the law and the remaining provisions will remain in full force and effect. Unless otherwise specified, remedies are cumulative.
         
         Each party acknowledges and agrees that the other is an independent contractor in the performance of this Agreement.
         
         The SDK has been developed entirely at private expense and is "commercial items" consisting of "commercial computer software" and "commercial computer software documentation" provided with RESTRICTED RIGHTS. Use, duplication or disclosure by the U.S. Government or a U.S. Government subcontractor is subject to the restrictions in this Agreement pursuant to DFARS 227.7202-3(a) or as set forth in subparagraphs (b)(1) and (2) of the Commercial Computer Software - Restricted Rights clause at FAR 52.227-19, as applicable. Contractor/manufacturer is NVIDIA, 2788 San Tomas Expressway, Santa Clara, CA 95051.
         
         The SDK is subject to United States export laws and regulations. You agree that you will not ship, transfer or export the SDK into any country, or use the SDK in any manner, prohibited by the United States Bureau of Industry and Security or economic sanctions regulations administered by the U.S. Department of Treasury's Office of Foreign Assets Control (OFAC), or any applicable export laws, restrictions or regulations. These laws include restrictions on destinations, end users and end use. By accepting this Agreement, you confirm that you are not a resident or citizen of any country currently embargoed by the U.S. and that you are not otherwise prohibited from receiving the SDK.
         
         Any notice delivered by NVIDIA to you under this Agreement will be delivered via mail, email or fax. You agree that any notices that NVIDIA sends you electronically will satisfy any legal communication requirements. Please direct your legal notices or other correspondence to NVIDIA Corporation, 2788 San Tomas Expressway, Santa Clara, California 95051, United States of America, Attention: Legal Department.
         
         This Agreement constitutes the entire agreement of the parties with respect to the subject matter of this Agreement and supersedes all prior negotiations or documentation exchanged between the parties relating to this subject matter. Any additional and/or conflicting terms on documents issued by you are null, void, and invalid. Any amendment or waiver under this Agreement shall be in writing and signed by representatives of both parties.
         
         If the distribution terms in this Agreement are not suitable for your organization, or for any questions regarding this Agreement, please contact NVIDIA at nvidia-compute-license-questions@nvidia.com.
         
         (v. February 10, 2022)
         
         ---
         
         **Note:** This package also includes third-party components with their respective licenses and attributions. Please see the accompanying `Third_party_attr.txt` file for complete third-party license information.
         
Project-URL: Homepage, https://github.com/nvidia-digital-bio/subquadraticOps-docs
Project-URL: Documentation, https://nvidia-digital-bio.github.io/subquadraticOps-docs/index.html
Project-URL: Repository, https://github.com/nvidia-digital-bio/subquadraticOps-docs
Requires-Dist: scikit-build-core>=0.10
Requires-Dist: nanobind>=1.5.0
Requires-Dist: warp-lang>=1.8.0
Requires-Dist: nvidia-ml-py
Description-Content-Type: text/markdown

# subquadratic_ops_torch

## Introduction

`subquadratic_ops_torch` provides CUDA kernels for subquadratic operations e.g. long and short convolution.
It contains pytorch bindings to optimized kernels.

## Installation

Please install using `pip install subquadratic-ops-torch-cu[12,13]`

## Documentation

For detailed usage information of the kernels, please refer to the doc-strings
in their respective function.

## Usage

You can import the library from python:

```python
import subquadratic_ops_torch as subq
```

Kernels are primarily exposed as function call underlying `torch.ops`, which also provide a lower-level
interface as `torch.library` operators.
This allows you to export models using this operations using `torch.export`,
and running inference on them using TensorRT.


## Support and Feedback

Please contact the developers for any issues you might encounter.


## Requirements

- CUDA-compatible NVIDIA GPU (Ampere+)
- CUDA Toolkit 12.0 or higher
- Python 3.11-13

## Modules

### B2B CausalConv1d

#### Operation

Back-to-back causal conv1d for the striped hyena 2 architecture used in Evo2 model.
The operation is performed in a causal manner, meaning each position only attends to previous positions in the sequence.
In code terms,
```python
in_dim = 8192

width_proj = 8
width_mixer = 128

dtype = torch.float32

class Conv1DModel(nn.Module):
    def __init__(self, in_dim, width, dtype, skip_bias=False):
        super(Conv1DModel, self).__init__()

        self.conv = nn.Conv1d(
            in_dim,
            in_dim,
            width,
            padding=width - 1,
            groups=in_dim,
            bias=False,
            dtype=dtype,
            device="cuda:0",
        )
        self.width = width
        self.weight = self.conv.weight.reshape(-1, width)
        if skip_bias:
            self.skip_bias = nn.Parameter(torch.zeros(in_dim, dtype=dtype, device="cuda:0").reshape(1, -1, 1))
        else:
            self.skip_bias = None

    def forward(self, x):
        seqlen = x.shape[-1]
        out = self.conv(x)
        return out[..., :seqlen]

def model(x, conv1d_proj, conv1d_mixer):
    xv = conv1d_proj(x)
    z = xv[:,1::3, :] * xv[:, 2::3, :]
    y = conv1d_mixer(z) + conv1d_mixer.skip_bias * z
    return y * xv[:, ::3, :]
x = torch.randn(batch_size, 3*in_dim, seq_dim)
conv1d_proj = Conv1DModel(3*in_dim, width_proj)
conv1d_mixer = Conv1DModel(in_dim, width_mixer, True)

y = model(x, conv1d_proj, conv1d_mixer)
```
is equivalent to,
```python
weight_proj = torch.randn(3*in_dim, width_proj).to(dtype)
weight_mixer = torch.randn(in_dim, width_mixer).to(dtype)
skip_bias = torch.randn(in_dim).to(dtype)

b2b_causal_conv1d(x, weight_proj, weight_mixer, skip_bias)

```

#### Supported Kernel Sizes for B2B Causal Conv1d

| Kernel Type | Supported Sizes |
|------------|-----------------|
| Projection | 2, 3, 4, 8, 16, 32 |
| Mixer | 2, 3, 4, 5, 6, 7, 8, 16, 32, 64, 128, 256 |

#### CausalConv1d

Causal conv1d : the convolution operation is performed in a causal manner, meaning each position only attends to previous positions in the sequence.
In code terms,
```python
in_dim = 8192
width = 8
dtype = torch.float32
model = nn.Conv1d(
            in_dim,
            in_dim,
            width,
            padding=width - 1,
            groups=in_dim,
            bias=False,
            dtype=dtype,
            device="cuda:0",
        )
y = model(x)
```
is equivalent to,
```python
weight = torch.randn((in_dim, width))
causal_conv1d(x, weight)
```

#### Supported Kernel Sizes for Causal Conv1d

| Kernel Type | Supported Sizes |
|------------|-----------------|
| CausalConv1d | 2, 3, 4, 5, 6, 7, 8, 16, 32, 64, 128, 256 |


### FFT CausalConv1d


FFT Causal Conv1d : the convolution operation is performed in a causal manner, meaning each position only attends to previous positions in the sequence. It uses real FFT and IFFT instead of direct summation for convolution.
In code terms,
```python
in_dim = 8192
width = 1024
dtype = torch.float32
weight = torch.randn(1, in_dim, width)
def model(x, w):
    fft_size = x.shape[-1] * 2
    xf = torch.fft.rfft(x, n=fft_size, dim=-1)
    wf = torch.fft.rfft(w, n=fft_size, dim=-1)
    return torch.fft.irfft(xf*wf, n=fft_size, dim=-1)[..., x.shape[-1]]

y = model(x, weight)

```
is equivalent to,
```python
weight = torch.randn((in_dim, width))
y = fft_causal_conv1d(x, weight)

```
