doc-string updates+ rel-notes (#220)

* docstring updates for APB

* corrected doctest skip

* skip doctest for apb caching

* triattn docstring update for balckwell kernel

* release notes updated

* update torch/jax for triatn rel notes

* version update

Author: hsadasivan

CHANGELOG.md

@@ -1,5 +1,17 @@
 ## Latest Changes
 
+## 0.8.0 (2025-12-05)
+
+### Added
+- [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
+
 ## 0.7.0 (2025-10-13)
 
 ### Added

VERSION

@@ -1 +1 @@
-0.8.0rc4
+0.8.0

cuequivariance_torch/cuequivariance_torch/primitives/triangle.py

@@ -80,6 +80,7 @@ def triangle_attention(
         (1) Context is saved for backward pass. You don't need to save it manually.
         (2) Kernel precision (fp32, bf16, fp16) is based on input dtypes. For tf32, set it from torch global scope
         (3) Triangle attention kernel supports: 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.
+        (4) Blackwell-optimized kernels (for compute capabilities 10.0 and 10.3) provide superior performance especially for long sequences and higher head dimensions. These kernels require the sequence length N to be a multiple of 8 for the forward pass; pad the sequence if necessary. Currently, this feature is supported only for cu13 builds.
 
     Example:
         >>> import torch
@@ -333,8 +334,12 @@ def attention_pair_bias(
         - The proj_z output is experimental to prevent breakage when caching
           of pair bias tensor is enabled in the next release.
         - Tested for bf16, fp16, fp32 and tf32. torch.set_float32_matmul_precision maybe used to toggle between fp32/tf32.
+        - Currently, the kernel provides superior performance only when DH (head dimension) is a multiple of 32.
+          For non-multiples of 32, we also recommend using graph compilation techniques like torch.compile, in addition.
 
     Examples:
+        Basic usage without caching:
+
         >>> import torch
         >>> from cuequivariance_torch import attention_pair_bias
         >>> if torch.cuda.is_available():  # doctest: +SKIP
@@ -365,7 +370,7 @@ def attention_pair_bias(
         ...     b_ln_z = torch.randn(z_dim,
         ...                     device=device, dtype=torch.bfloat16)
         ...     # Perform operation
-        ...     output, proj_z = attention_pair_bias(
+        ...     output = attention_pair_bias(
         ...         s=s,
         ...         q=q,
         ...         k=k,
@@ -378,9 +383,36 @@ def attention_pair_bias(
         ...         w_proj_o=w_proj_o,
         ...         w_ln_z=w_ln_z,
         ...         b_ln_z=b_ln_z,
+        ...         return_z_proj=False,
         ...     )
         ...     print(output.shape)  # torch.Size([1, 32, 64])
         torch.Size([1, 32, 64])
+
+        Example with caching (recommended for inference when z doesn't change):
+
+        >>> # Check cache and determine if z is already projected
+        >>> if model_cache is not None and "proj_z" in model_cache:  # doctest: +SKIP
+        ...     z = model_cache["proj_z"]
+        ...     is_cached_z = True
+        ... else:
+        ...     is_cached_z = False
+        >>>
+        >>> # Call attention_pair_bias
+        >>> o, proj_z = attention_pair_bias(  # doctest: +SKIP
+        ...     s=s, q=q, k=k, v=v, z=z, mask=mask,
+        ...     num_heads=num_heads,
+        ...     w_proj_z=w_proj_z if not is_cached_z else None,
+        ...     w_proj_g=w_proj_g,
+        ...     w_proj_o=w_proj_o,
+        ...     w_ln_z=w_ln_z if not is_cached_z else None,
+        ...     b_ln_z=b_ln_z if not is_cached_z else None,
+        ...     return_z_proj=True,
+        ...     is_cached_z_proj=is_cached_z,
+        ... )
+        >>>
+        >>> # Cache proj_z for next call
+        >>> if model_cache is not None and "proj_z" not in model_cache:  # doctest: +SKIP
+        ...     model_cache["proj_z"] = proj_z
     """
 
     try: