Switch to Python binding for cuDNN frontend

build_tools/jax.py

@@ -6,7 +6,9 @@
 
 import os
 import warnings
+from importlib.metadata import PackageNotFoundError, version as distribution_version
 from pathlib import Path
+from packaging import version
 
 import setuptools
 
@@ -26,7 +28,17 @@
 
 def install_requirements() -> List[str]:
     """Install dependencies for TE/JAX extensions."""
-    return ["jax>=0.5.3", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.31.0", "triton"]
+    # Python serializes graphs for C++ built with headers from this distribution.
+    # Keep the runtime serializer at the exact version used by this build.
+    try:
+        frontend_version = distribution_version("nvidia-cudnn-frontend")
+    except PackageNotFoundError as exc:
+        raise RuntimeError(
+            "Building TE/JAX requires nvidia-cudnn-frontend>=1.31.0 to be installed."
+        ) from exc
+    if version.parse(frontend_version) < version.parse("1.31.0"):
+        raise RuntimeError("Building TE/JAX requires nvidia-cudnn-frontend>=1.31.0.")
+    return ["jax>=0.5.3", "flax>=0.7.1", f"nvidia-cudnn-frontend=={frontend_version}", "triton"]
 
 
 def test_requirements() -> List[str]:

docs/api/c/fused_attn.rst

@@ -6,4 +6,15 @@
 fused_attn.h
 ============
 
+.. deprecated:: 2.22
+
+   The C fused-attention execution and support-query APIs, including their v2
+   variants, configuration and parameter accessors, and C++ parameter wrappers,
+   will be removed in the next major release. Use
+   :class:`transformer_engine.pytorch.DotProductAttention` or
+   :func:`transformer_engine.jax.attention.fused_attn` instead.
+
+   Layout enums and helpers, RNG utilities, KV-cache operations, and
+   context-parallel tensor utilities are not deprecated by this change.
+
 .. doxygenfile:: fused_attn.h

docs/api/jax.rst

@@ -60,7 +60,6 @@ Modules
 
 .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules
 
-
 Mixture of Experts
 ------------------
 Router, routing kernels, and grouped dense for MoE layers. See

docs/envvars.rst

@@ -240,7 +240,7 @@ backend-selection overview.
 
    :Type: ``int`` (0, 1 or 2), optionally followed by ``:<ranks>``
    :Default: ``0``
-   :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of the cache counters and the mean time of each cuDNN build stage. ``2`` additionally traces every event as it happens: each graph built, each graph cuDNN accepts and the cache keeps, each lookup and whether it hit or missed, each first execution that compiles kernels, and each execution. When the launcher exports a rank, only rank 0 logs; append ``:<ranks>`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3.
+   :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints a summary of cache counters and build timings at process exit. ``2`` also prints individual events. Ordinary SDPA and Flex Attention contribute to the same per-framework summary; Flex Attention uses the ``f16`` forward/backward counters. Uncacheable callbacks report graph construction and compilation without cache hits, misses, or insertions. At both levels, JAX diagnostics describe graph construction and caching during tracing, rather than each execution of a compiled function. When the launcher exports a rank, only rank 0 logs; append ``:<ranks>`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3.
 
 .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO
 
@@ -282,13 +282,17 @@ FP8 Configuration
 
    :Type: ``int`` (0 or 1)
    :Default: ``1``
-   :Description: Enable FP8 in the backward pass of DotProductAttention. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward.
+   :Description: When FP8 attention is enabled in PyTorch or JAX, control whether the backward pass also uses FP8. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward.
 
 .. envvar:: NVTE_DPA_FP8CS_O_in_F16
 
    :Type: ``int`` (0 or 1)
    :Default: ``1``
-   :Description: For Float8CurrentScaling in DotProductAttention, use FP16/BF16 for the output tensor in the backward pass. ``1`` = use F16/BF16 output in backward, ``0`` = use FP8 output in backward.
+   :Description: For current-scaled FP8 attention in PyTorch or JAX, control the precision
+      of the forward output tensor consumed by the backward pass. ``1`` = use FP16/BF16
+      output in backward, ``0`` = quantize it to FP8. This option does not affect MXFP8
+      attention or FP16/BF16 backward selected with ``NVTE_FP8_DPA_BWD=0``. In JAX,
+      public forward outputs are unchanged; set this option before tracing/compilation.
 
 .. envvar:: NVTE_DPA_FP8_RECIPE
 

docs/examples/jax/attention.rst

@@ -16,6 +16,24 @@ FlexAttention-style customization. The API also supports different Q/K and V
 head dimensions, as used by the attention operation after the projection
 stages in DeepSeek-style MLA.
 
+FP8 attention
+--------------
+
+Set ``fp8=True`` on the Flax ``DotProductAttention``, ``MultiHeadAttention``
+or ``TransformerLayer`` module to enable FP8 attention with current tensor
+scaling. Set ``fp8="mxfp8"`` for block-scaled MXFP8 attention. These options
+are independent of the GEMM quantization recipe. Inputs, outputs and input
+gradients retain their FP16/BF16 dtype.
+
+:envvar:`NVTE_FP8_DPA_BWD` controls whether the backward pass also uses FP8.
+Availability depends on the GPU, installed cuDNN version, and attention
+configuration.
+
+``AttnBiasType.ALIBI`` (``attn_bias_type="alibi"`` in Flax) enables cuDNN's
+ALiBi mask without an explicit bias tensor. Explicit dropout keys have JAX's
+usual replay semantics: using the same key repeats the mask, including when
+recomputing a forward pass for its backward pass.
+
 For long contexts, selected BSHD and THD configurations support context
 parallelism with Ring or AllGather collectives. Exact fused-kernel availability
 depends on the input shape, dtype, GPU architecture, and feature combination;

docs/installation.rst

@@ -17,6 +17,11 @@ Prerequisites
 4. `cuDNN 9.12 <https://developer.nvidia.com/cudnn>`__ or later.
 5. `cuDNN Frontend 1.31.0 <https://pypi.org/project/nvidia-cudnn-frontend/>`__ or later.
 
+PyTorch and JAX cuDNN attention require ``nvidia-cudnn-frontend`` version
+1.31.0 or later. JAX builds pin the runtime dependency to the frontend version
+used for the build headers. The JAX extension and frontend Python package must
+use matching versions; rebuild the extension after upgrading that package.
+
 If the CUDA Toolkit headers are not available at runtime in a standard
 installation path, e.g. within `CUDA_HOME`, set
 `NVTE_CUDA_INCLUDE_PATH` in the environment.

docs/release_notes.md

@@ -6,7 +6,8 @@
 
 ## Fixed Issues
 
-- 
+- PyTorch fused-attention backend selection checks backward support when gradients are needed,
+  including in evaluation mode.
 
 ## Breaking Changes in This Release
 

tests/cpp/operator/test_splits_to_offsets.cu

@@ -13,10 +13,59 @@
 #include <gtest/gtest.h>
 
 #include <transformer_engine/transformer_engine.h>
+#include <transformer_engine/utils.h>
 #include "../test_common.h"
 
 class SplitsToOffsetsTestSuite : public ::testing::TestWithParam<std::tuple<size_t, int64_t>> {};
 
+TEST(SequenceLengthsTest, ConvertAndPad) {
+  using namespace transformer_engine;
+  for (const size_t batch : {0, 3, 129, 513}) {
+    for (const size_t padding : {0, 7}) {
+      const size_t capacity = batch + padding;
+      SCOPED_TRACE(::testing::Message() << "batch=" << batch << ", capacity=" << capacity);
+      test::Tensor cu_q("cu_q", std::vector<size_t>{batch + 1}, DType::kInt32);
+      test::Tensor cu_kv("cu_kv", std::vector<size_t>{batch + 1}, DType::kInt32);
+      test::Tensor seq_q("seq_q", std::vector<size_t>{capacity}, DType::kInt32);
+      test::Tensor seq_kv("seq_kv", std::vector<size_t>{capacity}, DType::kInt32);
+      auto *q = cu_q.rowwise_cpu_dptr<int32_t>();
+      auto *kv = cu_kv.rowwise_cpu_dptr<int32_t>();
+      q[0] = kv[0] = 0;
+      for (size_t i = 0; i < batch; ++i) {
+        q[i + 1] = q[i] + i % 7;
+        kv[i + 1] = kv[i] + i % 11;
+      }
+      cu_q.from_cpu();
+      cu_kv.from_cpu();
+      nvte_cu_seqlens_to_actual_seqlens(cu_q.data(), cu_kv.data(), seq_q.data(), seq_kv.data(),
+                                      nullptr);
+      seq_q.to_cpu();
+      seq_kv.to_cpu();
+      for (size_t i = 0; i < capacity; ++i) {
+        EXPECT_EQ(seq_q.rowwise_cpu_dptr<int32_t>()[i], i < batch ? i % 7 : 0);
+        EXPECT_EQ(seq_kv.rowwise_cpu_dptr<int32_t>()[i], i < batch ? i % 11 : 0);
+      }
+    }
+  }
+}
+
+TEST(SequenceLengthsTest, RejectInvalidMetadata) {
+  using namespace transformer_engine;
+  test::Tensor cu("cu", std::vector<size_t>{3}, DType::kInt32);
+  test::Tensor lengths("lengths", std::vector<size_t>{2}, DType::kInt32);
+  test::Tensor small("small", std::vector<size_t>{1}, DType::kInt32);
+  test::Tensor wrong_type("wrong_type", std::vector<size_t>{2}, DType::kInt64);
+  test::Tensor wrong_rank("wrong_rank", std::vector<size_t>{1, 2}, DType::kInt32);
+  EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), small.data(), lengths.data(),
+                                                  lengths.data(), nullptr));
+  EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), small.data(),
+                                                  small.data(), nullptr));
+  EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), wrong_type.data(),
+                                                  wrong_type.data(), nullptr));
+  EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), wrong_rank.data(),
+                                                  wrong_rank.data(), nullptr));
+}
+
 TEST_P(SplitsToOffsetsTestSuite, TestSplitsToOffsets) {
   const size_t num_tensors = std::get<0>(GetParam());
   const int64_t logical_last_dim = std::get<1>(GetParam());