Releases · NVIDIA/cuEquivariance

Release list

v0.11.1 Latest

Latest

jomitchellnv released this 07 Aug 07:27

v0.11.1

eaa4e82

0.11.1

Patch release focused on PyTorch compatibility. No API or behavior changes — the pure-Python frontend is unchanged from 0.11.0. This is a coordinated version bump across the frontend and the compiled cuequivariance-ops-* wheels.

Changed

[ops-torch] Widened PyTorch version support. The cuequivariance-ops-torch-cu12 / -cu13 wheels now import and run across PyTorch 2.11, 2.12, and 2.13. Previously the compiled extension was ABI-locked to the exact PyTorch minor it was built against, so any other version failed with an undefined-symbol error at import.

Bug Fixes

[ops-torch] torch>=2.11 is now a real runtime dependency. It was previously declared only in an optional test extra, so pip did not enforce it — a user on PyTorch < 2.11 could install the wheel and then hit a cryptic import failure. pip now blocks the incompatible install up front.

Assets 2

v0.11.0

jomitchellnv released this 04 Aug 22:40

v0.11.0

ad21650

Added

[Torch] cuet.triangle_attention accepts kv_lengths, an int32 per-row key/value length tensor (shape [B, N, 1, 1, 1]) for right-padded sequences.
Passing kv_lengths selects the SM100f (B200 / B300) or sm120f (RTX 6000 Pro Blackwell) length fast path when available; on other GPUs it is converted to the equivalent dense prefix mask. kv_lengths and mask are mutually exclusive (passing both raises). Added cuet.mask_to_kv_lengths to convert and validate prefix masks.
[Torch/JAX] Native sm120f Triangle Attention kernel for RTX 6000 Pro Blackwell, replacing the generic kernel used on that SKU in 0.10.
[Torch] cuet.attention_pair_bias exposes the optional generalized projection parameters b_proj_k, b_proj_v, w_ln_q/b_ln_q, w_ln_k/b_ln_k, and b_proj_g as keyword-only arguments (default None), preserving the strict contract on single_repr; w_proj_z may be omitted when is_cached_z_proj=True. (#293)
[JAX] Added a dictionary-based interface for organizing and processing features by representation type. (#271)

Bug Fixes

[Torch] Triangle Attention no longer silently mis-handles dense masks on Blackwell. In 0.10, passing any mask= on Blackwell (B200 / B300) took the SM100f fast path. SM100f is a length/prefix kernel — it attends [0, kv_length) per row — so it collapsed any mask into a single per-row length. A full or right-padded (prefix) mask was fast and correct, but an arbitrary (non-prefix) mask was fast and silently incorrect: interior gaps were dropped, with no warning and no error. In 0.11, SM100f rejects any dense mask= (including an all-ones mask) and routes it to the correct fallback kernel. Arbitrary masks are now always correct.
[Torch] Fixed TorchScript and Pytorch 2.2 compatibility issues in cuet.ChannelWiseTensorProduct and cuet.SegmentedPolynomial. (#277)

Breaking Changes

[Torch] cuet.attention_pair_bias now takes the single representation single_repr plus the node-LayerNorm (w_ln_a/b_ln_a) and Q/K/V/gate projection weights instead of pre-projected q/k/v and a separate s, and performs the projections internally. The pre-projected q/k/v and s are removed, and the op is now self-attention only (square N × N) — rectangular U × V (U != V) cross-attention is no longer supported. (#290)
[Torch] On Blackwell (B200 / B300), a dense mask= passed to cuet.triangle_attention (including an all-ones mask) no longer takes the SM100f fast path — it routes to the correct fallback, which is slower than the fast path. Build kv_lengths once upstream, then pass kv_lengths=cuet.mask_to_kv_lengths(mask) for the fast path. A one-time warning points to this whenever a mask would have reached SM100f.
[Torch] PyTorch 2.11.X is now required.

Notes

Which masking choice accelerates what (Torch, B200/B300):
- Fully-valid sequences → pass no mask and no kv_lengths. This is the fastest path, and it routes exactly as it did in 0.10 (unaffected by the mask change).
- Right-padded / prefix sequences → pass kv_lengths. The kernel skips the padded keys entirely, while native PyTorch computes them and masks afterwards — so cuEquivariance's advantage over native PyTorch is largest here.
- Arbitrary (non-prefix) masks → pass mask=. Correct in 0.11, but runs the fallback.
Build kv_lengths once, upstream. cuet.mask_to_kv_lengths does a one-time host sync (prefix-shape validation; it raises on a non-prefix mask), so compute it in the collator / data-processing stage. The op itself is host-sync-free with kv_lengths: the only range-check .item() is opt-in (CUEQ_TRIATTN_VALIDATE_KV_LENGTHS) and skipped under torch.compile; the clamp is device-side.
[JAX] The mask → kv_lengths changes are Torch-only. The JAX path auto-selects the SM100f fast path from its mask and is unaffected.
Recommended usage. For best performance, run cuEquivariance ops under torch.compile(model, mode="reduce-overhead") (CUDA graphs), which removes per-call launch/dispatch overhead — most impactful at small shapes and low batch, where the kernels are launch-bound. Warm up first: the first few iterations (≈3) go to tracing and CUDA-graph capture and are slower; the speedups appear only once the graph is set up. CUDA graphs assume static shapes and re-capture when shapes change.
Hardware support. Triangle Attention: SM100f fast path on B200 / B300 (via kv_lengths); native sm120f on RTX 6000 Pro Blackwell; generic kernel on H100 / H200 / A100 / L40. APB inference: H100 / H200, B200, B300, RTX 6000 Pro Blackwell, A100 (not L40). APB training: H100 / H200, B200, B300 (not A100, RTX 6000 Pro Blackwell, or L40).
Data type support. Triangle Attention supports BF16, FP16, and FP32/TF32 — the FP32/TF32 kernel requires hidden_dim <= 32 and hidden_dim % 4 == 0; larger head dimensions must use BF16 or FP16. The specialized sm100f and sm120f kernels are BF16/FP16 only. APB's optimized paths for inference and training use BF16 and, for select routes, FP16.
APB shape support. The optimized attention_pair_bias path supports pair width 128 and 256, head dimension 8-64, and num heads 4-16.
Triangle Multiplication remains unchanged in 0.11.
Known issue. A minor per-call overhead affects the Triangle Attention kv_lengths (padded-sequence) inference path, most visible at short sequences. It does not affect correctness, training, or the no-mask path. A fix is in progress.

Masked and Right-Padded Sequences on B200 / B300

On B200 / B300, supplying kv_lengths selects the SM100f length fast path, and the advantage over native PyTorch grows when sequences are masked or right-padded. Native PyTorch materializes the dense mask and computes the padded keys before masking them; cuEquivariance passes a compact per-row length and skips them in the kernel entirely.

Known issues

A minor per-call overhead affects the Triangle Attention kv_lengths (padded-sequence) inference path, most visible at short sequences. It does not affect correctness, training, or the no-mask path. A fix is in progress.
CUDA 12.8 wheels do not support Blackwell's SM100F fast path. In the next release we will upgrade our CUDA 12.8 wheels to use CUDA 12.9 which will support TriAttn’s SM100F fast path on Blackwell.
TriMul training is slower than native PyTorch (with torch.compile) at batch sizes > 1 due to a slowdown in the backwards pass – we will work to resolve this in our next release.
There is an issue using torch version 2.11+ which we will fix in a minor update (v0.11.1) coming soon

Documentation

[Torch/JAX] Expose generalized projection params in attention_pair_bias wrapper (#293)
[Torch] Triangle Attention kv_lengths documentation (#287)
[JAX] Added an example showing how to implement MACE with Equinox. (#281)

Assets 2

v0.10.0

phiandark released this 22 Apr 01:20

v0.10.0

5196d21

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

Contributors

mariogeiger, phiandark, and 3 other contributors

Assets 2

v0.9.1 (patch release)

mariogeiger released this 16 Mar 17:25

v0.9.1

b137a8c

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.

Assets 2

v0.9.0

mariogeiger released this 17 Feb 22:56

v0.9.0

2b9cb54