0.11.0-rc2

What's new

Triangle Attention: kv_lengths fast path
Added a kv_lengths argument to triangle_attention, plus a mask_to_kv_lengths
helper. Providing per-row key/value lengths routes variable-length (right-padded)
inputs to the optimized fused attention kernel on Blackwell (SM100f), instead of the
dense-mask path.
kv_lengths and mask are mutually exclusive.
Dense boolean masks are still supported; when a dense mask forces execution off the
fast path, a warning is now emitted so callers can switch to kv_lengths for best
performance.
Attention Pair Bias
New fused attention-with-pair-bias primitive with an optimized Blackwell (SM100f)
implementation.
Kernel / performance improvements
Continued improvements to the Blackwell (SM100 / SM103) fused-attention kernels.
Notes
This is a release candidate intended for validation ahead of 0.11.0.

Release candidate for 0.11.0

Release candidate for 0.11.0

What's new

Triangle Attention: kv_lengths fast path

• Added a kv_lengths argument to triangle_attention, plus a mask_to_kv_lengths
  helper. Providing per-row key/value lengths routes variable-length (right-padded)
  inputs to the optimized fused attention kernel on Blackwell (SM100f), instead of the
  dense-mask path.
• kv_lengths and mask are mutually exclusive.
• Dense boolean masks are still supported; when a dense mask forces execution off the
  fast path, a warning is now emitted so callers can switch to kv_lengths for best
  performance.

Attention Pair Bias

• New fused attention-with-pair-bias primitive with an optimized Blackwell (SM100f)
  implementation.

Kernel / performance improvements

• Continued improvements to the Blackwell (SM100 / SM103) fused-attention kernels.

Notes

• This is a release candidate intended for validation ahead of 0.11.0.

v0.10.0

Added

• Python 3.14 support finalized, including a fix for stale tuple hashes in SegmentedTensorProduct after in-place operand mutation, and updated CI matrix (#272)
• [Torch/JAX] cuet.triangle_attention/cuex.triangle_attention: new faster sm100f (CC 10.0/10.3) forward kernel for hidden_dim ≤ 256, bwd hidden_dim ≤ 128; bias is cast to q/k/v dtype (instead of always float32) under sm100f; non-contiguous input tensors are handled internally — no manual contiguity assertion is required as long as shape requirements are met; updated docstrings. Only available on cu13 builds (#260)
• [JAX] MACE flax.nnx example restructured to use nnx.split + @jax.jit on (graphdef, state) instead of @nnx.jit on the module, removing the Python-side nnx graph traversal overhead from each training/inference step (#261)
• [JAX] NVTX markers added to the MACE examples to make step boundaries visible in nsys profiles (#266)

Bug fix

• [Torch] SegmentedPolynomial checkpoint portability: GPU-saved models now load correctly on CPU. Implemented via __reduce__ on SegmentedPolynomialFromUniform1dJit, SegmentedPolynomialFusedTP, SegmentedPolynomialIndexedLinear, and SegmentedPolynomial, plus graceful fallback when specific cuequivariance_ops_torch extensions (e.g. uniform_1d) are unavailable (#270)
• [Torch] Replaced deprecated is_fx_tracing with is_fx_symbolic_tracing (#270)
• [JAX] Restrict PTX 88 to sm_121 for CUDA 12.9+, avoiding breakage on other architectures (addresses the known issue noted in the 0.9.0 release) (#250)
• [Torch/JAX] cuet.attention_pair_bias/cuex.attention_pair_bias: fixed incorrect results when the hidden dimension is not a multiple of 32; the previous torch fallback for these cases is removed as the kernel now handles them correctly

Notes

• [Torch] The CUEQ_TORCH_COMPILE environment variable (experimental) enables torch.compile for cuet.triangle_attention; useful for non-contiguous tensor inputs on Ampere/Hopper architectures

Documentation

• Fixed tutorial format issues (#274)

What's Changed

• Fix eq / lt on segmented polynomial types for Python and JAX compatibility by @mariogeiger in #258
• [merge after release] Restrict PTX 88 to sm_121 for CUDA 12.9+ by @hsadasivan in #250
• doc string and api update for triattn by @hsadasivan in #260
• api: remove dim_order from triangle_attention by @hsadasivan in #262
• Removed dim_order from triAttn by @phiandark in #263
• nnx.split by @mariogeiger in #261
• add nvtx marker by @paulz-nv in #266
• Fixing some torch segmented_polynomial support by @phiandark in #270
• Add Python 3.14 support and fix CI setup by @mariogeiger in #272
• Fix tutorials doc format issues by @LiamZhang100 in #274
• Add skill.md files by @mariogeiger in #269
• Release 0.10.0 by @mariogeiger in #275

New Contributors

• @paulz-nv made their first contribution in #266

Full Changelog: v0.9.0...v0.10.0

v0.9.1 (patch release)

Bug fix

• [Torch/JAX] Fixed a rare overflow in the uniform 1d kernel when indices are large, by casting index arithmetic from 32-bit to 64-bit
• [Torch/JAX] Disabled parallel compilation of uniform 1d kernels by default (CUEQUIVARIANCE_OPS_PARALLEL_COMPILE now defaults to 0). Parallel compilation caused issues in multi-GPU setups. It will be re-enabled by default in a future release once the underlying bugs are resolved.

v0.9.0

0.9.0 (2026-02-17)

Added

• GB10 (DGX Spark) support
• Support for Python 3.13 and 3.14
• [JAX] Support for Triton 3.6.0
• [JAX] flax.nnx MACE example
• [Torch/JAX] Deterministic indexing mode for uniform 1d kernels
• [Torch/JAX] Parallel JIT compilation for uniform_1d kernels with per-kernel caching, significantly reducing compilation time. New optional environment variable CUEQUIVARIANCE_OPS_NVRTC_CACHE_DIR allows setting a directory for caching compiled kernels.
• Documentation: new tutorials for JAX and PyTorch segmented polynomials

Bug fix

• [JAX] Fixed Triton tuning issue for triangular multiplicative update
• [JAX] Compatibility with JAX 0.8.2: fixed FFI interface and dtype casting issues when x64 mode is not enabled
• [JAX] Improved triangle attention error messages
• [Torch/JAX] Fixed yx_rotation descriptor
• [Torch] TensorRT QDP plugin workaround

Breaking Changes

• [Torch/JAX] The environment variable CUEQUIVARIANCE_OPS_USE_JIT no longer exists. JIT compilation is now the default behavior for uniform_1d kernels (already since few releases).
• [Torch/JAX] Renamed filter_drop_unsued_operands to filter_drop_unused_operands (typo fix)
• [Torch/JAX] Removed nvfatbin optional dependency
• [Torch] Removed deprecated primitive classes: TensorProduct, EquivariantTensorProduct, SymmetricTensorProduct, and IWeightedSymmetricTensorProduct. Use cuet.SegmentedPolynomial with method='uniform_1d' instead, or the high-level APIs (cuet.ChannelWiseTensorProduct, cuet.FullyConnectedTensorProduct, cuet.SymmetricContraction). Attempting to import these classes will raise an ImportError with migration instructions.
• [Torch] Removed deprecated low-level wrapper classes: TensorProductUniform1d, TensorProductUniform4x1d, TensorProductUniform3x1dIndexed, TensorProductUniform4x1dIndexed, and SymmetricTensorContraction from cuequivariance_ops_torch. Use torch.ops.cuequivariance.uniform_1d or cuet.SegmentedPolynomial instead.

Notes

• [JAX] DGX Spark/GB10 (sm_121) with CUDA 12.9: This release uses PTX 87, which works correctly for most architectures but is not compatible with DGX Spark/GB10 on CUDA 12.9. To enable DGX Spark/GB10 support with CUDA 12.9, refer to #250 for a simple frontend integration tweak that restricts PTX 88 to sm_121 only. This fix will be merged after the 0.9.0 release.

v0.8.1

0.8.1 (2026-01-09)

Bug fix

• [Torch] Fixed torch.compile compatibility for non-contiguous tensors in backward pass for cuet.triangle_attention, cuet.triangle_multiplicative_update, and cuet.attention_pair_bias. This resolves stride mismatch errors in TorchInductor when compiling models using these operations (#223)

v0.8.0

Added

• Support for CUDA 13 on ARM
• [Torch/JAX] Blackwell-optimized BF16/FP16 forward and backward kernels for cuet.triangle_attention (runs on compute capabilities 10.0 and 10.3). These kernels provide superior performance especially for long sequences and higher head dimensions. This is only supported on cu13 builds as of this release

Bug fix

• [Torch/JAX] Fixed index overflow and out of bound issues leading to illegal memory access in cuet.triangle_attention

Notes

• [Torch/JAX] Blackwell-optimized kernels require the sequence length N to be a multiple of 8 for the forward pass; pad the sequence if necessary
• [Torch/JAX] Blackwell-optimized kernels are currently supported only for CUDA 13 builds

What's Changed

• 16bit math datatype for uniform_1d by @mariogeiger in #188
• [JAX] code examples by @mariogeiger in #156
• [docs] add mace perf figure in docs by @mariogeiger in #192
• Adding examples to the documentation by @phiandark in #191
• Small SegmentedPolynomial improvements by @phiandark in #210
• fix by @mariogeiger in #212
• Added wrappers for ONNX/TRT initialization by @borisfom in #214
• [JAX] triatt 16bit fix by @mariogeiger in #213
• [Torch] Improve error messages for uniform 1d by @mariogeiger in #217
• triatt jax bwd dtype issue by @mariogeiger in #219
• doc-string updates+ rel-notes by @hsadasivan in #220

Full Changelog: v0.7.0...v0.8.0

v0.7.0

Added

• Support for CUDA 13
• Support for Python 3.13
• [Torch] MACE example in the documentation
• [JAX] MACE and NequIP examples in the folder cuequivariance_jax/examples
• [Torch/JAX] The Segmented Polynomial operation's math_dtype argument now accepts method-specific string values, with each method supporting different options
• [Torch] ONNX export and TensorRT runtime plugin support for triangle attention, triangle multiplication and attention with pairwise bias
• [Torch] Support for caching the "bias tensor z" calculated from the proj_z linear layer in the Attention with pairwise bias kernel (as implemented in the Boltz code)

Bug fix

• Correct the documentation CUEQ_TRITON_TUNING_MODE -> CUEQ_TRITON_TUNING
• [JAX] Make triangle_multiplicative_update use the same tuning cache files as its PyTorch counterpart

What's Changed

• [JAX] expand math_dtype argument by @mariogeiger in #153
• [JAX] add (partial) jax.vmap support for triangle attention and triangle multiplicative update by @mariogeiger in #155
• [JAX] Use nondiff_argnums instead of nondiff_argnames by @mariogeiger in #158
• [JAX] trimul batch dims and int64 support by @mariogeiger in #157
• Merge release to main by @mariogeiger in #169
• apb w/ caching api changes by @hsadasivan in #174
• [JAX] fix benchmarking function by @mariogeiger in #161
• Enabling richer string math_dtype for SegmentedPolynomial by @phiandark in #172
• [JAX] make trimul cache name/keys similar to their torch counterpart by @mariogeiger in #173
• api defaults change for apb with caching by @hsadasivan in #179
• setup copy-pr-bot by @mariogeiger in #180
• update ruff by @mariogeiger in #181
• GHA tests by @mariogeiger in #182
• Update tolerance parameters in spherical harmonics test by @mariogeiger in #185
• update aot flag name in docstring by @hsadasivan in #186
• [JAX] add sharding notes in docstrings by @mariogeiger in #197
• update changelog and version by @mariogeiger in #196

Full Changelog: v0.6.0...v0.7.0

v0.6.1

Latest Changes

0.6.1 (2025-09-04)

Added

• [Torch/JAX] Support for variable leading batch dimensions in triangle multiplicative update
• [Torch/JAX] Triangle attention kernel support for additional input configs: all hidden_dim<=32 and divisible by 4 for tf32/fp32, and for all hidden_dim<=128 and divisible by 8 for bf16/fp16. In the rare instance that the kernel does not support an input config, fallback to torch is enabled instead of erroring out.
• [Torch/JAX] Tuned config for RTX PRO 6000 GPUs for triangle multiplicative update.
• [JAX] vmap support for triangle multiplicative update and triangle attention
• [Torch] Improved error reporting on import failure with traceback information for stacktrace

Bug fix

• [Torch/JAX] Fixed illegal memory access issue stemming from int32 indexing for longer sequences in triangle multiplicative update and attention with pair bias.
• [JAX] Moved to using nondiff_argnums instead of nondiff_argnames to be compatible with older JAX versions

v0.6.0

Added

• [Torch] New feature: Added cuet.attention_pair_bias (support for caching the pair bias tensor & further kernel acceleration coming up soon. There maybe API related changes for this in the next release)
• [Torch/JAX] Added method argument to cuet.SegmentedPolynomial/cuex.segmented_polynomial to give users control over which backend solution is used (naive, uniform_1d, fused_tp, indexed_linear).
• [Torch/JAX] Added torch fallback option based on sequence lengths for triangle kernels and attention pair bias. The user may control this by setting env vars: CUEQ_TRIMUL_FALLBACK_THRESHOLD, CUEQ_TRIATTN_FALLBACK_THRESHOLD, CUEQ_ATTENTION_PAIR_BIAS_FALLBACK_THRESHOLD. The corresponding APIs default to torch fallback for seq_lens < 100 for optimal performances.
• [Torch/JAX] Added support for the optional projection and gating biases in the input and output of cuex.triangle_multiplicative_update
• [JAX] Added JAX bindings for triangle operations with cuex.triangle_attention and cuex.triangle_multiplicative_update

Bug fix

• [Torch] Added cueuivariance_ops_torch.init_triton_cache() for users to initialize triton cache before calling torch compiled triangular multiplicative update. If not used, Torch compile would break if directly applied on cuex.triangle_multiplicative_update.
• [Torch/JAX] Fixed the illegal memory access error for long sequences in triangle attention. This increases the usable limits on sequence lengths.

Breaking Changes

• Dropped support for CUDA 11. Only CUDA 12 is now supported (cuequivariance-ops-torch-cu12, cuequivariance-ops-jax-cu12).
• [Torch/JAX] Simplified precision arg of triangular multiplicative update to just two: None (defaults to triton language dot's default for non-32b input and for 32b input, tf32/tf32x3 based on 1/0 value set in torch.backends.cuda.matmul.allow_tf32) and IEEE-754.
• [Torch/JAX] We have moved away from the default round-towards-zero (RZ) implementation to round-nearest (RN) for better tf32 accuracy in cuex.triangle_multiplicative_update. In rare circumstances, this may cause minor differences in results observed.

Known Issues

• [JAX] The function cuex.triangle_multiplicative_update requires triton<=3.3.1. We are waiting for an update of the package jax-triton.
• [PyTorch] The function cuet.triangle_multiplicative_update requires triton>=3.4.0 on Blackwell GPUs.
• As a consequence of the two last point, cuex.triangle_multiplicative_update can't run on Blackwell GPUs.
• cuet.attention_pair_bias does not support caching of projected pairwise tensor. We are working on adding support for this.

What's Changed

• Add methods to slice operands of descriptors by @mariogeiger in #127
• [Torch] skip many tests by @mariogeiger in #129
• [JAX] triangle attention by @mariogeiger in #122
• [JAX] triangle multiplication by @mariogeiger in #126
• [JAX] Rename a variable by @mariogeiger in #134
• [JAX] method API by @mariogeiger in #132
• [Torch+JAX] add/improve measure clock ticks by @mariogeiger in #138
• [JAX] add bias to tri-mul by @mariogeiger in #139
• Improve CI stuff by @mariogeiger in #140
• Enhance benchmarking in measure_clock_ticks function by @mariogeiger in #141
• [PyTorch] SegmentedPolynomial "method" argument by @phiandark in #130
• Improvements to the PyTorch SegmentedPolynomial by @phiandark in #142
• Update trimul inputs by @phiandark in #145
• Adding attention_pair_bias to frontend by @phiandark in #146
• Update benchmarking scripts by @mariogeiger in #147
• [WIP] Update changelog and documentation by @mariogeiger in #148
• Update CHANGELOG.md by @mariogeiger in #151

New Contributors

• @phiandark made their first contribution in #130

Full Changelog: v0.5.1...v0.6.0