Skip to content

[Draft] Port cuDNN frontend attention to Python API #3508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vcherepanov-nv wants to merge 42 commits into from
Closed

[Draft] Port cuDNN frontend attention to Python API #3508

Show file tree
Hide file tree
Changes from 21 commits
Commits
Show all changes
42 commits
Select commit Hold shift + click to select a range
1e4052e
Port JAX fused attention to cuDNN Python frontend
vcherepanov-nv Aug 26, 2026
df64163
Test shared cuDNN frontend version check
vcherepanov-nv Aug 26, 2026
e42c3ac
Port PyTorch cuDNN attention to Python frontend
vcherepanov-nv Aug 26, 2026
35d7158
Fix JAX ragged cuDNN attention graphs
vcherepanov-nv Aug 26, 2026
04188f2
Fix F16 attention inference auxiliaries
vcherepanov-nv Sep 2, 2026
31c043d
Key F16 attention graphs by logical batch
vcherepanov-nv Sep 2, 2026
87ef877
Use logical graph batch for dense attention tensors
vcherepanov-nv Sep 2, 2026
fd1aa12
Skip ragged KV token hint on SM120
vcherepanov-nv Sep 2, 2026
6b61996
Mask THD padding before max-logit reduction
vcherepanov-nv Sep 2, 2026
8424cc4
Zero THD outputs on the slow fill path
vcherepanov-nv Sep 2, 2026
95c2b5a
Key FP8 backward graphs by encoding
vcherepanov-nv Sep 2, 2026
31be282
Key paged attention graphs by table layout
vcherepanov-nv Sep 2, 2026
92ed239
Fix JAX cuDNN graph build compatibility
vcherepanov-nv Sep 2, 2026
c141994
Merge branch 'main' into attn-pyt-pygraph
vcherepanov-nv Sep 10, 2026
618d2d6
Merge branch 'main' into attn-jax-pygraph
vcherepanov-nv Sep 10, 2026
24af929
Merge branch 'attn-jax-pygraph' into attn-pygraph
vcherepanov-nv Sep 10, 2026
bfe6b7d
Remove obsolete common fused attention code
vcherepanov-nv Sep 10, 2026
3a4be7d
Factor shared cuDNN attention helpers
vcherepanov-nv Sep 10, 2026
b7b2cac
[JAX] Add FP8 cuDNN attention parity
vcherepanov-nv Sep 11, 2026
d6f38aa
[JAX] Fix attention regression failures
vcherepanov-nv Sep 11, 2026
13ae9af
[JAX] Fix MXFP8 attention scale strides
vcherepanov-nv Sep 11, 2026
3e0bdca
Merge branch 'main' into attn-pygraph-parity
vcherepanov-nv Sep 11, 2026
028c583
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 11, 2026
614a123
Merge branch 'main' into attn-pygraph-parity
cyanguwa Sep 11, 2026
478ce7c
Remove obsolete fused attention backend override
vcherepanov-nv Sep 11, 2026
150d5ed
Restore fused attention rejection diagnostics
vcherepanov-nv Sep 11, 2026
4e5efd1
Restore fused attention graph cache diagnostics
vcherepanov-nv Sep 11, 2026
984ddbd
Restore PyTorch FP8 THD attention support
vcherepanov-nv Sep 11, 2026
4c6ecc5
Reject affected cuDNN versions for JAX MXFP8 attention
vcherepanov-nv Sep 11, 2026
138dc2b
Gate JAX FP8 attention by recipe support
vcherepanov-nv Sep 11, 2026
51471fa
Match JAX current-scaling attention quantizers
vcherepanov-nv Sep 11, 2026
2d7909c
Merge branch 'main' into attn-pygraph-parity
vcherepanov-nv Sep 11, 2026
13cec24
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 11, 2026
a5248dc
Fix FlexAttention execution mock signature
vcherepanov-nv Sep 12, 2026
a497b13
Fix FP8 attention checkpoint test opt-in
vcherepanov-nv Sep 12, 2026
848cd07
Avoid retaining cuDNN graph workspaces
vcherepanov-nv Sep 12, 2026
a12ab7d
Fix paged attention table descriptors
vcherepanov-nv Sep 12, 2026
d4e3915
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 12, 2026
8e9b0ac
Fix cuDNN graph execution device context
vcherepanov-nv Sep 14, 2026
1969670
Fix attention Python lint findings
vcherepanov-nv Sep 14, 2026
c6f373e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 14, 2026
033c701
Fix JAX attention lint warnings
vcherepanov-nv Sep 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,7 @@ System Requirements
* Compiler: GCC 9+ or Clang 10+ with C++17 support
* Python: 3.12 recommended

* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+, nvidia-cudnn-frontend 1.25.0+
* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+, nvidia-cudnn-frontend 1.27.0+

* **Notes:** FP8 features require Compute Capability 8.9+ (Ada/Hopper/Blackwell)

Expand Down
13 changes: 5 additions & 8 deletions build_tools/build_ext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,8 +213,9 @@ def run(self) -> None:
shutil.rmtree(nccl_ep_dir)

def build_extensions(self):
# For core lib + JAX install, fix build_ext from pybind11.setup_helpers
# to handle CUDA files correctly.
# For JAX builds, fix build_ext from pybind11.setup_helpers to handle
# CUDA files correctly. This is also required by the standalone JAX
# wheel, where framework_extension_only is true.
if "pytorch" not in get_frameworks():
# Ensure at least an empty list of flags for 'cxx' and 'nvcc' when
# extra_compile_args is a dict.
Expand All @@ -226,8 +227,7 @@ def build_extensions(self):

# Define new _compile method that redirects to NVCC for .cu and .cuh files.
original_compile_fn = self.compiler._compile
if not framework_extension_only:
self.compiler.src_extensions += [".cu", ".cuh"]
self.compiler.src_extensions += [".cu", ".cuh"]

def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None:
# Copy before we make any modifications.
Expand All @@ -236,10 +236,7 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None:
try:
original_compiler = self.compiler.compiler_so

if (
os.path.splitext(src)[1] in [".cu", ".cuh"]
and not framework_extension_only
):
if os.path.splitext(src)[1] in [".cu", ".cuh"]:
nvcc_bin = nvcc_path()
if nvcc_bin is None:
raise RuntimeError(
Expand Down
21 changes: 15 additions & 6 deletions build_tools/jax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,27 +5,35 @@
"""JAX related extensions."""

import os
from importlib.metadata import version as get_package_version
from pathlib import Path
from packaging import version
from typing import List

import setuptools
from packaging import version

from .utils import (
get_cuda_include_dirs,
all_files_in_dir,
cudnn_frontend_include_path,
debug_build_enabled,
setup_mpi_flags,
get_cuda_include_dirs,
nccl_ep_enabled,
nccl_include_path,
nccl_lib_path,
nccl_ep_enabled,
setup_mpi_flags,
)
from typing import List


def install_requirements() -> List[str]:
"""Install dependencies for TE/JAX extensions."""
return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"]
# Serialized cuDNN graphs use a version-specific wire format, so the Python
# frontend used at runtime must match the headers used to build the extension.

@jberchtold-nvidia jberchtold-nvidia Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! Can you also add a TODO in the comment here to remind us to upgrade to the Python-side JAX API for the pygraph API once available so we can remove this strictness. But I agree, this is a great check to have now where we do have that strict serialization version match requirement

frontend_version = get_package_version("nvidia-cudnn-frontend")
return [
"jax",
"flax>=0.7.1",
f"nvidia-cudnn-frontend=={frontend_version}",
Comment on lines +31 to +35

@greptile-apps greptile-apps Bot Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Metadata Requires Preinstalled Dependency

This computes install_requires by querying the locally installed cuDNN frontend. In a documented --no-build-isolation installation or a direct setup.py build where that dependency is not already installed, metadata generation raises PackageNotFoundError before dependencies can be resolved. In a combined PyTorch/JAX build, an installed frontend older than 1.28 also produces an exact JAX pin that conflicts with PyTorch's >=1.28.0 requirement, making the package impossible to install.

Knowledge Base Used: Build, extensions, and packaging

]


def test_requirements() -> List[str]:
Expand Down Expand Up @@ -89,6 +97,7 @@ def setup_jax_extension(
csrc_source_files = Path(csrc_source_files)
extensions_dir = csrc_source_files / "extensions"
sources = all_files_in_dir(extensions_dir, name_extension="cpp")
sources += all_files_in_dir(extensions_dir, name_extension="cu")

@jberchtold-nvidia jberchtold-nvidia Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we move the single new .cu file into TE common instead? That way building the framework extension only requires a C++ compiler, not nvcc. I believe Przemek and Kirthi worked to split the build like this a year or so ago to ensure this split. Can you check with them on the latest guidance here if this is still the case? And maybe also loop in Frederick. Thanks!


# Header files
include_dirs = get_cuda_include_dirs()
Expand Down
2 changes: 2 additions & 0 deletions build_tools/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@ def install_requirements() -> List[str]:
"packaging",
"pydantic",
"nvdlfw-inspect",
# PyTorch cuDNN attention is built and executed with the Python graph
# API; FP8/MXFP8 graph capture requires Frontend 1.28 or newer.
"nvidia-cudnn-frontend>=1.28.0",
]

Expand Down
2 changes: 1 addition & 1 deletion build_tools/wheel_utils/build_wheels.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,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'
/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.27.0'

if $BUILD_METAPACKAGE ; then
cd /TransformerEngine
Expand Down
4 changes: 3 additions & 1 deletion qa/L0_jax_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@ pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to instal
: ${XML_LOG_DIR:=/logs}
mkdir -p "$XML_LOG_DIR"

python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
python3 -m pytest --import-mode=importlib -v --junitxml=$XML_LOG_DIR/pytest_test_common_attention_helpers.xml $TE_PATH/tests/test_common_attention_helpers.py || test_fail "test_common_attention_helpers.py"
python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py --ignore=$TE_PATH/tests/jax/test_fp8_fused_attn.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fp8_fused_attn.xml $TE_PATH/tests/jax/test_fp8_fused_attn.py || test_fail "tests/jax/test_fp8_fused_attn.py"

@jberchtold-nvidia jberchtold-nvidia Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does test_fp8_fused_attn need to be called out specifically here? Shouldn't it be picked automatically via the pattern matching of -k 'not distributed' similarly to test_custom_call_compute.py?

I wonder if we're also double testing score mod below this line? I may have missed that in earlier review

python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_score_mod.xml $TE_PATH/tests/jax/test_fused_attn_score_mod.py || test_fail "tests/jax/test_fused_attn_score_mod.py"
NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py"

Expand Down
1 change: 1 addition & 0 deletions qa/L0_pytorch_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export NVTE_FLASH_ATTN_V4=0

pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"

python3 -m pytest --import-mode=importlib --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_common_attention_helpers.xml $TE_PATH/tests/test_common_attention_helpers.py || test_fail "test_common_attention_helpers.py"
NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py"
Expand Down
Loading