cuequivariance-ops-torch

Introduction

cuequivariance_ops_torch provides CUDA kernels for the cuEquivariance project's PyTorch components. As such, it contains pytorch bindings to optimized kernels that cuEquivariance's operations map down to. In general, we advice that you access those kernels through cuEquivariance, but you may also find them useful on their own.

Installation

Please install using either pip install cuequivariance-ops-torch-cu12 or pip install cuequivariance-ops-torch-cu13 (depending on the CUDA toolkit you wish to use).

Documentation

For detailed usage information of the kernels, please refer to the doc-strings in their respective modules. For higher-level documentations, refer to cuEquivariance.

Pairformer inference primitives

The Torch package also exposes lower-level primitives for integrating Pairformer inference graphs:

• attention_pair_bias_mask_from_normalized projects an already-normalized pair representation and applies an arbitrary Boolean [B,U,V] mask. The implementation is shape-generic; BF16 D=256, H=16, and square sequence lengths 384 and 512 are the performance-validated configurations. FP16 and BF16 inputs use native low-precision tensor-core products with FP32 accumulation; FP32 inputs retain the higher-accuracy TF32x3 path.
• pair_update_finalize_masked_attention_bias finalizes a pair residual and, on the shapes admitted for the device's exact compute capability, also returns the next layer's already-masked BF16 attention bias, which attention_pair_bias_masked_cache(..., is_cached_z_proj=True) consumes in place of attention_pair_bias. The canonical shape bounds and measurement scope are in _apb_route_policy.py: _pair_update_finalize_admits for exact capabilities (9, 0), (10, 0) and (12, 0), and _sm100_apb_block_kf_admits for the embedded block. The consumer computes any cache it is handed. On (10, 0) and (10, 3) both halves are taken over by the embedded Kernel Factory APB block (apb_kf_sm100_torch: six CUBINs per exact SM, sm_100a and sm_103a from one source, launched through the CUDA driver, no CuTeDSL at runtime) inside its admitted hull for D_z=384, H=12, DH=64, when the producer has no pair-projection bias and the consumer receives the packed QKV projection with bias, both q/k norm weights and the gate/output biases.
• pairformer_combined_triangle_multiplication combines incoming and outgoing triangle multiplication after caller-provided non-affine LayerNorm with epsilon 1e-5. Its fused direction join applies the affine middle LayerNorm with epsilon 0.03. This operation intentionally fails closed outside BF16 inference with B=1, D=256, sequence length 384 or 512, and an SM100 GPU. The caller guarantees that the second normalized carrier is the separately materialized spatial transpose of the first; the hot path validates metadata and distinct storage but does not compare tensor contents. It returns the base update before any caller-owned peri-LayerNorm, output mask, dropout, or residual addition; model-level dropout must be disabled for inference. Use pack_pairformer_grouped_projection to convert grouped projection weights to the expected payload/gate layout.

Both operations are registered as torch.library custom operators with fake implementations for torch.compile tracing. See their Python docstrings for the complete input contracts.

Calling convention for declined cells

Call triangle_multiplicative_update unconditionally. Under torch.compile a BF16 cell no kernel owner admits (any architecture, hidden width 256 with a pair width in 256..512 -- the measured band; inference or training) runs the package's compiler-visible reference formula, which is written the way a model's own module is (one gated projection per operand under autograd, one fused projection GEMM without it, contiguous contraction operands) and compiles to the same kernels or better, so a declined call costs no more than the caller's native path (measured on L40S, A100 and H200: inference 1.10-1.50x, training 0.98-1.01x of RFD4's compiled module; the generic Triton route the op used to take there measured 0.58-0.88x). An eager no-grad call no owner admits runs the fused generic route from 2**15 pair rows (B * S * S) and the reference formula below, where the generic route's dispatch cost is not repaid. The shape predicate (triangle_multiplicative_update_inference_shape_is_supported) exists for callers that must decide before the tensors exist; it is not needed to avoid a penalty.

Usage

You can import the library from python:

import cuequivariance_ops_torch

Kernels are primarily exposed as torch.nn.Module, but also provide a lower-level interface as torch.library operators. Generally, the module is responsible for proper input transformation and initialization, and the operator execute the kernel. 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 cuEquivariance developers for any issues you might encounter.