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).

Kernel Factory TriMul kernels

On CUDA 13 wheels, triangle_multiplicative_update runs prebuilt Kernel Factory (KF) TriMul kernels for BF16 calls on the exact shapes each kernel was validated on, on H100 SXM (80GB HBM3), H200, GH200 480GB, B200 and RTX PRO 6000 Blackwell (Server and Workstation Edition). Every other call keeps the existing implementation. The kernels ship as prepackaged cubins and need nothing extra; their outputs match the existing implementation within its numerical tolerance, not bitwise. (Development builds may register a kernel in its CuTe DSL JIT form; that needs nvidia-cutlass-dsl[cu13]==4.7.0 and apache-tvm-ffi==0.1.14.post0, which are not dependencies of the published package. Without them, dispatch keeps the existing implementation.)

A call reaches KF only when the pair AND the module's parameters are BF16, or when FP32 inputs are passed with precision="bfloat16" (the op rounds them first). The automatic mixed precision default, a BF16 pair with FP32 parameters (with or without BF16 autocast), keeps the existing implementation in 0.13 (an FP32-input KF variant is BIO-1329).

In training, a second backward through one forward (retain_graph=True) recomputes the forward and is exact in eager mode, but transiently needs about 1.9x the first backward's memory. Replayed from a CUDA graph, such a second backward gives NaN gradients by design (loud, never silently wrong; the first backward of every replay is exact), where the existing implementation returns correct gradients. Double backward (create_graph=True) raises.

Some KF kernels synchronise their thread blocks with software grid barriers and assume exclusive use of the GPU: under MPS sharing without SM partitions they can hang.

The environment variable CUEQ_TRIMUL_KF_CAMPAIGN selects the mode. It is read once, when cuequivariance_ops_torch is imported:

Any other value warns and uses the default. Values are case-insensitive.

KF under torch.compile

KF covers the allow-listed shapes, and under torch.compile each one gets its own compiled graph of the calling function. Dynamo keeps at most torch._dynamo.config.recompile_limit graphs per function (8 by default), so with the default torch.compile (automatic dynamic shapes) up to recompile_limit - 2 distinct listed shapes per process get KF once shapes start to vary; further new shapes use the cuEq 0.12 route on one shared dynamic graph (logged once). A shape that already has a KF graph keeps it, results are unaffected, and the calling function never exceeds the limit because of KF when its batch size is fixed. For KF on many shapes, compile with dynamic=False and raise torch._dynamo.config.recompile_limit above the number of shapes. Under DDP or multi-node training each rank applies the limit on its own, so ranks that see shapes in different orders may run KF on different shapes: correctness is unaffected, and a small speed imbalance between ranks is possible.

KF TriMul kernels are used for training only through the packed op (triangle_multiplicative_update); the split-parameter entry (triangle_multiplicative_update_split_parameters) does not use KF for training in 0.13 (it uses the SM90 training owner where that applies, else the PyTorch formula), even where the training shape query answers True or prepare_..._training reports "owner": "kf".

Exported models use the cuEq 0.12 route in 0.13: torch.export (strict or non-strict) and ONNX export (and so TensorRT) never record a KF kernel, and neither does a trace with symbolic sizes outside torch.compile (e.g. make_fx or FakeTensorMode).

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.