[DRAFT] Switch to Python binding for cuDNN frontend

build_tools/jax.py

@@ -6,6 +6,7 @@
 
 import os
 import warnings
+from importlib.metadata import version as distribution_version
 from pathlib import Path
 from packaging import version
 
@@ -27,7 +28,12 @@
 
 def install_requirements() -> List[str]:
     """Install dependencies for TE/JAX extensions."""
-    return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.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.
+    frontend_version = distribution_version("nvidia-cudnn-frontend")
+    if version.parse(frontend_version) < version.parse("1.29.0"):
+        raise RuntimeError("Building TE/JAX requires nvidia-cudnn-frontend>=1.29.0.")
+    return ["jax", "flax>=0.7.1", f"nvidia-cudnn-frontend=={frontend_version}", "triton"]
 
 
 def test_requirements() -> List[str]:

build_tools/wheel_utils/build_wheels.sh

@@ -24,7 +24,7 @@ git checkout $TARGET_BRANCH
 git submodule update --init --recursive
 
 # Install deps
-/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' 'apache-tvm-ffi>=0.1.12'
+/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.29.0' 'apache-tvm-ffi>=0.1.12'
 
 if $BUILD_METAPACKAGE ; then
         cd /TransformerEngine

docs/api/jax.rst

@@ -59,3 +59,9 @@ Modules
   :members: __call__
 
 .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules
+
+Attention functions
+-------------------
+
+.. autoapifunction:: transformer_engine.jax.attention.fused_attn
+.. autoapifunction:: transformer_engine.jax.attention.is_fused_attn_kernel_available

docs/envvars.rst

@@ -238,7 +238,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 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 plan compilation, and each execution. JAX reports Python graph-cache and build events during tracing; execution events apply to the PyTorch and common C++ paths. 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
 
@@ -280,7 +280,7 @@ 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: Enable FP8 in the backward pass of DotProductAttention in PyTorch and JAX. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. In JAX this applies when attention is configured with ``fp8=True``.
 
 .. envvar:: NVTE_DPA_FP8CS_O_in_F16
 

docs/examples/jax/attention.rst

@@ -16,6 +16,42 @@ 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.
 
+cuDNN execution and FP8
+----------------------
+
+JAX fused attention constructs and selects cuDNN execution plans with the
+``nvidia-cudnn-frontend`` Python package. The compiled JAX FFI executor caches
+those plans and its own cuDNN handles per device. Ordinary attention and
+``score_mod`` share this executor; neither calls TE common's attention APIs.
+The Python frontend and the headers used to build the JAX extension must have
+the same version, since execution plans are serialized between the two.
+
+Pass ``fp8=True`` to ``transformer_engine.jax.attention.fused_attn`` or to the
+Flax ``DotProductAttention``, ``MultiHeadAttention`` or ``TransformerLayer``
+module to enable cuDNN FP8 attention. Inputs remain FP16 or BF16. Query, key
+and value are quantized to E4M3 using their current tensor maximum; backward
+quantizes the output gradient to E5M2. Outputs and input gradients retain the
+input dtype, and softmax statistics use FP32. Packed layouts use one scale per
+packed tensor. This option is independent of the GEMM quantization recipe and
+does not maintain delayed-scaling history.
+
+``NVTE_FP8_DPA_BWD=0`` selects FP16/BF16 backward while keeping FP8 forward.
+The backward pass uses dequantized forward operands, so both passes operate on
+the same quantized query, key and value. The default, ``1``, uses FP8 for both.
+
+FP8 availability depends on cuDNN, GPU architecture, layout and the selected
+features. Unsupported combinations produce a diagnostic rather than silently
+using a different precision. For example, cuDNN 9.26 rejects FP8 backward with
+dropout on Blackwell; use ``NVTE_FP8_DPA_BWD=0`` for that combination.
+Tensor-scaled FP8 is supported; MXFP8 attention,
+paged KV-cache APIs and quantized module inputs/outputs require further JAX API
+and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only.
+
+``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

@@ -16,6 +16,30 @@ Prerequisites
 3. |driver link|_ supporting CUDA 12.1 or later.
 4. `cuDNN 9.12 <https://developer.nvidia.com/cudnn>`__ or later.
 
+The PyTorch and JAX cuDNN attention backends use shared graph builders through
+the cuDNN frontend Python API. Both installations require
+``nvidia-cudnn-frontend`` version 1.29.0 or later. JAX builds pin their runtime
+dependency to the exact frontend version that supplies the build headers.
+Graph descriptions are cached by tensor metadata
+and target device or architecture. Support queries select supported plans without
+compiling their kernels. PyTorch compiles on first execution and uses its CUDA
+generator for dropout, including CUDA graph capture and replay. JAX compiles
+during tracing when workspace sizes are needed, then serializes plans for its
+native FFI executor. The JAX extension and the frontend Python package must use
+matching cuDNN frontend versions; rebuild the extension after upgrading that
+package.
+
+Automatic ``score_mod`` graph caching includes the callback's code and snapshots
+of builtin scalar and container defaults, including keyword-only defaults.
+Callbacks with tensor or opaque stateful defaults bypass automatic caching.
+Stateful callbacks can provide ``score_mod_graph_cache_key`` to explicitly
+describe their graph structure; pass runtime data through ``score_mod_tensors``
+and ``score_mod_bprop_tensors``.
+
+On SM12x GPUs, including SM120 and SM121, packed attention uses dense
+per-sequence softmax statistics internally. Graph construction and context
+parallelism use the same statistics layout.
+
 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.

tests/jax/test_build_tools.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Dependency metadata for builds that deserialize cuDNN frontend graphs."""
+
+from email.parser import Parser
+from io import StringIO
+
+import pytest
+from packaging.requirements import Requirement
+from setuptools import Distribution
+
+from build_tools import jax as jax_build
+
+
+@pytest.mark.parametrize("frontend_version", ("1.29.0", "1.30.0"))
+def test_jax_metadata_pins_build_frontend(monkeypatch, frontend_version):
+    """Built metadata must reject runtime serializers other than the build version."""
+
+    def distribution_version(package):
+        assert package == "nvidia-cudnn-frontend"
+        return frontend_version
+
+    monkeypatch.setattr(jax_build, "distribution_version", distribution_version)
+    dist = Distribution(
+        {
+            "name": "transformer_engine_jax",
+            "version": "0.0.0",
+            "install_requires": jax_build.install_requirements(),
+        }
+    )
+    metadata = StringIO()
+    dist.metadata.write_pkg_file(metadata)
+    requirements = [
+        Requirement(value)
+        for value in Parser().parsestr(metadata.getvalue()).get_all("Requires-Dist")
+    ]
+    frontend = next(req for req in requirements if req.name == "nvidia-cudnn-frontend")
+    assert frontend.specifier.contains(frontend_version)
+    for other_version in ("1.28.0", "1.29.0", "1.29.1", "1.30.0", "1.30.1", "1.31.0"):
+        assert frontend.specifier.contains(other_version) == (other_version == frontend_version)
+
+
+def test_jax_build_requires_minimum_frontend(monkeypatch):
+    """Pinning the installed version must preserve the minimum build requirement."""
+    monkeypatch.setattr(jax_build, "distribution_version", lambda _package: "1.28.0")
+    with pytest.raises(RuntimeError, match="nvidia-cudnn-frontend>=1.29.0"):
+        jax_build.install_requirements()