[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,10 @@ 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.quantize_attention_qkv
+.. 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,80 @@ 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 THD attention requires cuDNN 9.23 or later. FP8 THD backward with off-by-one
+or learnable softmax requires cuDNN 9.26 or later; select
+``NVTE_FP8_DPA_BWD=0`` to use FP16/BF16 backward instead. For FP16/BF16 THD
+learnable-softmax backward with packed statistics, cuDNN versions before 9.26
+remain eligible for value-head dimensions 64, 128, and 256. Other value-head
+dimensions require cuDNN 9.26 or later. Execution-plan support still depends on
+the complete attention configuration.
+
+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.
+Pass ``fp8="mxfp8"`` to select MXFP8 attention on SM100 or later. MXFP8 uses
+E4M3 operands with E8M0 block scales along the attention contraction axes in
+both passes. Inputs, outputs and gradients remain FP16/BF16. Sequence lengths
+must be divisible by 128 and head dimensions by 32. This mode currently supports
+separate, QKV-packed and KV-packed BSHD inputs. THD, SM12x, and cuDNN
+9.23.0/9.23.1 are excluded. It requires
+cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operands
+for MXFP8 backward, matching the PyTorch MXFP8 path.
+Attention bias, ALiBi, dropout and ``return_max_logit`` are excluded from this mode.
+Exact head-dimension availability is checked by cuDNN. Frontend 1.29.0 has an
+incorrect MXFP8 backward scale stride for Q/K heads wider than 128; these
+configurations require ``NVTE_FP8_DPA_BWD=0``. This check protects both frontends.
+
+To reuse quantizations, call ``quantize_attention_qkv((q, k, v))`` and pass its
+result to ``fused_attn(..., fp8="mxfp8", quantized_qkv=quantized)``. Each operand
+is a JAX ``ScaledTensor2x`` with independent rowwise/columnwise E4M3 data and
+unswizzled E8M0 scales in BHSD order. Equivalent operands from a producer may
+be supplied directly. The original BSHD arrays supply the gradient targets and
+the original FP16/BF16 operands for ``NVTE_FP8_DPA_BWD=0``. These auxiliary
+buffers are non-differentiable. Scale reordering happens after partitioning,
+so data and tensor parallel sharding preserve per-head scale boundaries.
+Ring and AllGather context parallelism communicate FP16/BF16 operands and
+quantize each local attention step. Prequantized inputs are decoded before CP
+communication, as in the PyTorch MXFP8 implementation. Separate and packed
+BSHD storage, GQA, and causal dual-chunk load balancing are supported. Every
+local attention step must meet the 128-token alignment requirement; 256 tokens
+per CP rank accommodates the half-sequence steps used by causal load balancing
+and AllGather. The existing JAX CP restrictions on sliding windows, padding,
+dropout and non-vanilla softmax still apply.
+
+Paged KV-cache APIs, quantized module outputs, and automatic quantized
+projection-to-attention integration 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,54 @@
+# 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
+from pathlib import Path
+
+import pytest
+from packaging.requirements import Requirement
+from setuptools import Distribution
+
+# CI runs outside the checkout root and tests the installed TE package.
+# Expose the source-only build tools just for this import.
+with pytest.MonkeyPatch.context() as _build_import_path:
+    _build_import_path.syspath_prepend(str(Path(__file__).resolve().parents[2]))
+    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()