v0.12.0

0.12.0

Added

• [Performance] Improvements to triangle multiplicative update under torch.compile (inference/training, geometric mean improvement across all supported shapes in v0.11.1): 1.2x/1.1x on A100, 1.2x/1.1x on H100/H200, 1.3x/1.15x on B200 / B300, on RTX Pro 6000 Blackwell Server Edition.
• [Torch] Expanded cuet.triangle_multiplicative_update kernel coverage for pair width 384 and hidden width 256: inference and training on H100 / H200, and inference on B200 / B300 and RTX PRO 6000 Blackwell. The H100 / H200 inference path also accepts sequence lengths that are not multiples of 128.
• [Torch] cuet.attention_pair_bias accepts norm_kind="layer_norm" (default) or "rms_norm". RMSNorm requires every normalization bias to be None. (#299)
• [Torch] Expanded Attention Pair Bias training kernel coverage on H100 / H200 for RMSNorm with single width 768, pair width 384, 12 heads, and head dimension 64 under BF16 autocast and torch.compile.
• [ops-torch] Added an RMSNorm pair-update / Attention Pair Bias inference handoff on H100 / H200, B200 / B300, and RTX PRO 6000 Blackwell for supported shapes, including pair width 384 with 12 heads and head dimension 64. pair_update_finalize_masked_attention_bias returns the updated pair representation and an optional masked attention-bias cache; pass that cache to attention_pair_bias_masked_cache with is_cached_z_proj=True.
• [ops-torch] Added module-based Attention Pair Bias entry points, packed Q/K/V projection parameters, and attention_pair_bias_shape_is_supported. Module-based entry points preserve normalization for non-affine LayerNorm and RMSNorm modules.
• [ops-torch] Added split-parameter Triangle Multiplicative Update entry points, preparation helpers, and shape-support and inference-route queries for callers that manage separate projection parameters or select kernels before allocating tensors.
• [Torch/JAX] Expanded Triangle Attention layout support on RTX PRO 6000 Blackwell (sm120f) to include non-contiguous Q/K/V layouts with a contiguous head dimension and 16-byte-aligned outer strides. Updated the head-dimension-64 forward kernel.
• [Torch] Triangle Attention with kv_lengths on B200 / B300 now handles zero-length rows through CUDA postprocessing in eager inference. Under torch.compile, row zeroing remains visible to the compiler.
• [Torch/JAX] Added Vera Rubin (sm_107) to the cuequivariance-ops-*-cu13 wheel builds, using the existing package names. cuEquivariance is currently functional on Vera Rubin for the purpose of ecosystem integration, and performance improvements will be released in the near future.
  • Note: Users must bring their own Rubin compatible container and/or PyTorch version and install the wheel inside of it.

Bug Fixes

• [Torch/JAX] Fixed a warp-reduction bug in fused tensor products that could silently produce incorrect results in segmented polynomials, tensor products, and indexed linear layers. The fix applies across GPU architectures.
• [Torch] Fixed the Triangle Multiplicative Update training performance regression at batch sizes greater than 1 for the covered BF16 shapes.
• [Torch] cuet.triangle_multiplicative_update now applies eps to both LayerNorms in the reference implementation. Fixed an assertion failure in the generic path when the pair and hidden widths differ.
• [Torch] Fixed graph breaks caused by dispatch warnings under torch.compile.
• [Torch] Fixed Triangle Attention backward handling of sliced inputs and K/V tensors with different layouts on B200 / B300. Misaligned input views are copied to aligned storage before the SM100f / SM120f kernels run.
• [JAX] Fixed the custom_vjp output-structure mismatch in cuex.triangle_attention that caused jax.grad and jax.value_and_grad to fail. Added coverage for differentiation through checkpointed scans. (#295, #300)

Breaking Changes

• [Torch] Triangle Attention output on SM100f / SM120f can now preserve the query tensor's non-contiguous layout. Call .contiguous() before operations that require contiguous storage. Auxiliary lse and max outputs remain contiguous.
• [Torch] cuet.attention_pair_bias raises ValueError when w_proj_z is None unless is_cached_z_proj=True. (#299)
• [JAX] cuequivariance-jax now requires JAX >= 0.9.1, matching the XLA FFI API used by the kernel bindings.

Notes

• CUDA 12 wheel builds now use CUDA 12.9.1 (previously CUDA 12.8). The cu12 builds now include the SM100f and SM120f Triangle Attention kernels.
• CUDA 13 wheel builds now use CUDA 13.4.1 (previously CUDA 13.0), which now supports Rubin.
• [Torch] Under torch.compile, Triangle Multiplicative Update calls outside the specialized kernels' supported shapes use a compiler-visible reference formula for hidden width 256 and pair widths 256–512, with BF16 inputs or FP32 inputs under BF16 autocast at default precision. Other shapes and precision modes retain their existing dispatch rules.
• [Torch] Shape support also depends on dtype, layout, gradient mode, and the kernels available in the installed build. H100 / H200 TriMul training kernels still require sequence lengths divisible by 128. The APB RMSNorm inference handoff is exposed through cuequivariance_ops_torch; setting norm_kind="rms_norm" on cuet.attention_pair_bias alone does not select it. If the finalizer returns no cache, use the ordinary attention call.