[Draft] Port cuDNN frontend attention to Python API by vcherepanov-nv · Pull Request #3508 · NVIDIA/TransformerEngine

vcherepanov-nv · 2026-09-11T03:05:57Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Construct and plan standard fused-attention forward and backward graphs through the cuDNN frontend Python API. Share graph serialization and binding infrastructure with flex attention, including packed and ragged tensor bindings. Replace the JAX attention-specific TE common bridge with a graph-agnostic XLA FFI executor and a JAX-local capture-safe RNG state helper. Keep TE common unchanged for other frameworks. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Update the experimental flex-attention version test to exercise the cuDNN frontend compatibility check in the shared JAX graph module instead of retaining aliases for the old private helpers. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Build and execute FusedAttention graphs through the cuDNN Frontend Python API for FP16, BF16, FP8, and MXFP8 configurations. Share graph runtime utilities with FlexAttention and mirror backend selection in Python. Remove the legacy PyTorch fused-attention bindings while retaining TE common for other frameworks. Add a graph-safe Philox reservation helper and require cuDNN Frontend 1.27.0. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Restore cuDNN-supported ragged graph bucketing and provide external element offsets from JAX for forward and backward execution. Use XLA FFI buffer byte sizes directly so uint32 RNG buffers do not require a TE dtype conversion. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Always generate and return softmax statistics and the Philox state for F16 forward graphs so context-parallel inference can combine partition outputs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Include the effective bucketed THD graph batch in forward and backward cache keys to prevent reuse across incompatible ragged sequence counts. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Describe dense peers of ragged QKV tensors with the bucketed graph batch while preserving their physical storage strides. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Avoid passing the packed KV token bound to cuDNN THD backward on SM120, where token bucketing is disabled and the graph sequence extent is per-sequence. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Build a device-side validity mask from padded query offsets so unwritten inter-sequence and tail gaps cannot affect ragged max-logit results. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Always initialize ragged F16 outputs and gradient storage because cuDNN leaves inter-sequence and tail padding unwritten, including when fast zero fill is disabled. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Include the FP8 dtypes of attention inputs, dO, and gradient outputs so delayed-scaling recipes with identical uint8 buffers cannot share incompatible graphs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Record both page-table tensor descriptors in F16 forward cache keys so differing maximum-page dimensions and strides build distinct graphs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Enable NVCC compilation for standalone JAX wheels, preserve legacy THD graph dimensions before cuDNN 9.6, and pin the runtime frontend package to the build-time version. Add regression coverage for the THD bucketing boundary. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

# Conflicts: # build_tools/pytorch.py # pyproject.toml # transformer_engine/pytorch/csrc/extensions/attention.cpp

# Conflicts: # build_tools/jax.py # transformer_engine/jax/csrc/extensions/attention.cpp

PyTorch and JAX now construct and execute cuDNN attention graphs in their framework-specific Python frontends, leaving the common graph implementations unused. Remove those implementations and their public entry points while retaining shared layout and support helpers. Move the graph-safe seed and offset extraction helper into the generic common utilities because PyTorch dropout and stochastic rounding still use it. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Centralize the FP16/BF16 support policy, graph lifecycle, ragged bucketing, mask normalization, and score-mod cache keys used by the JAX and PyTorch frontends. Preserve framework-specific capability differences explicitly and add regression coverage. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Implement graph-backed JAX dot-product attention for DelayedScaling, Float8CurrentScaling, and MXFP8BlockScaling recipes. Quantize Q/K/V and dO internally while preserving FP16/BF16 module inputs, outputs, and gradients, and propagate delayed-scaling state through the custom VJP. Factor FP8 backend selection, layout parsing, mask normalization, graph operation dispatch, tensor strides, MXFP8 padding, and scale swizzling into shared helpers. Migrate the PyTorch cuDNN attention path to these helpers without intentionally changing its supported configurations. Align the JAX frontend with PyTorch by exposing cuDNN ALiBi attention, supporting explicit bottom-right diagonal alignment, and consolidating previously framework-specific compatibility policy flags. Add common policy and graph-dispatch tests, JAX numerical forward/backward coverage for all three FP8 recipes across FP16 and BF16 boundaries, MXFP8 scale-layout coverage, and ALiBi and bottom-right parity tests. Wire the common suite into both L0 framework jobs and the JAX FP8 suite into L0 JAX CI. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Update the score-mod fused-attention test double for the newly propagated bottom-right diagonal argument so plumbing tests continue to model the public call signature. Provide dtype, dimensions, and strides for unused MXFP8 amax tensors in forward and backward cuDNN graphs. This satisfies frontend tensor validation while keeping those tensors excluded from graph outputs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

JAX transposes compact MXFP8 scale-inverse buffers to contiguous BHSD before applying the cuDNN F8_128x4 swizzle. Describe those buffers with matching BHSD strides instead of the original BSHD data layout, preventing invalid GPU memory accesses during MXFP8 attention. Add a focused regression test for the graph scale dimensions, strides, and reordering contract. The complete JAX FP8 attention test file passes on an SM100 CUDA device. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

greptile-apps · 2026-09-11T03:11:30Z

Confidence Score: 4/5

The PR is not yet safe to merge because JAX package metadata generation still requires nvidia-cudnn-frontend to be installed before dependency resolution.

Findings

Metadata Requires Preinstalled Dependency ▶

Summary

Centralizes FP16 and FP8 cuDNN attention capability checks and rejection reasons.
Adds Python cuDNN graph construction, serialization, caching, and cache diagnostics for PyTorch and JAX.
Adds PyTorch tanh logit softcapping across supported dense and context-parallel paths.
Expands JAX FP8 scaling-mode support and guards known unsafe cuDNN configurations.
Updates native bindings, build inputs, documentation, and framework test coverage.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    App[PyTorch or JAX attention request] --> Policy[Shared capability policy]
    Policy -->|Unsupported| Fallback[Fallback or diagnostic reason]
    Policy -->|Supported| Builder[Python cuDNN graph builder]
    Builder --> Cache[Framework graph cache]
    Cache --> Executor[Native cuDNN graph executor]
    Executor --> Output[Attention output and backward state]
    Builder -. events and timings .-> Debug[Cache diagnostics]
    Cache -. hits and misses .-> Debug

_{Reviews (2) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..."}

greptile-apps · 2026-09-11T03:11:49Z

+    frontend_version = get_package_version("nvidia-cudnn-frontend")
+    return [
+        "jax",
+        "flax>=0.7.1",
+        f"nvidia-cudnn-frontend=={frontend_version}",


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

# Conflicts: # tests/jax/test_fused_attn.py # tests/jax/test_fused_attn_score_mod.py # tests/pytorch/test_torch_compile.py # transformer_engine/common/fused_attn/fused_attn.cpp # transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu # transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h # transformer_engine/common/fused_attn/fused_attn_fp8.cu # transformer_engine/common/fused_attn/fused_attn_fp8.h # transformer_engine/common/fused_attn/utils.cu # transformer_engine/common/fused_attn/utils.h # transformer_engine/common/include/transformer_engine/fused_attn.h # transformer_engine/jax/cpp_extensions/attention.py # transformer_engine/jax/csrc/extensions.h # transformer_engine/jax/csrc/extensions/attention.cpp # transformer_engine/jax/csrc/extensions/pybind.cpp # transformer_engine/jax/flax/transformer.py # transformer_engine/pytorch/attention/dot_product_attention/utils.py # transformer_engine/pytorch/csrc/extensions.h # transformer_engine/pytorch/csrc/extensions/attention.cpp # transformer_engine/pytorch/csrc/extensions/pybind.cpp

for more information, see https://pre-commit.ci

cyanguwa · 2026-09-11T15:47:53Z

+    return cudnn.pygraph(**kwargs)
+
+
+def build_cudnn_graph(cudnn, graph, *, description: str) -> int:


I don't know if we should move this function to transformer_engine/common/attention, because this might be only specific to SDPA graphs? I'm not sure if other graphs have the same member methods. Thanks.

cyanguwa · 2026-09-11T15:49:10Z

I think GitHub didn't resolve the "update branch" correctly. Some changes from main are not showing up here. Sorry - could you please fix the merge? Thanks.

cyanguwa · 2026-09-11T15:51:25Z

+
+
+@dataclass(frozen=True)
+class FusedAttentionConfig:


Both #2964 and #3480 are going in the direction of packing up config parameters/input tensors into structs or dataclasses. Could you please compare with those two PRs and see if we should make any adjustments here (to the member list of this dataclass or any others)? Thanks.

cyanguwa · 2026-09-11T15:52:10Z

Do we need a separate file for F16? I only see fp8.py here.

cyanguwa · 2026-09-11T15:54:02Z

+_UID_DV = 21
+_UID_DBIAS = 22
+_UID_DSINK = 23
+_UID_ATTN_SCALE = 24


We can package up these UIDs into a dataclass of IntEnum (F16_UIDs, FP8_UIDs, etc) and share them between PyTorch and Jax.

cyanguwa · 2026-09-11T15:58:52Z

+        tuple((tuple(aval.shape), dtype_name(aval.dtype)) for aval in avals),
+        get_cudnn_version(),
+        _device_arch(),
+    )


Please take a look at #2964's graph_cache.h and graph_cache_debug.h files and see if we can apply any of the techniques there, for example, make_cache_key(). Thanks.

cyanguwa · 2026-09-11T16:00:41Z

This file is 2700 lines long. I think we probably should split it up to 3 files: F16, FP8, and MXFP8. It'd help with readability and maintainability, and cuDNN FE has 3 graph builders too (graph.sdpa, sdpa_fp8, sdpa_mxfp8). Thanks.

cyanguwa

Thanks for the PR! I left a few comments even though it's still in draft mode. Will review it again once it's stabilized more. Thanks.

Also, a very rough branch for this migration using my AI: https://github.com/cyanguwa/TransformerEngine/tree/fused_attn_py_migration. It's incomplete and untested. But there might be some ideas to borrow. Just for your reference. Thanks.

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> # Conflicts: # transformer_engine/common/fused_attn/utils.cu

for more information, see https://pre-commit.ci

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

for more information, see https://pre-commit.ci

vcherepanov-nv added 21 commits August 26, 2026 21:20

Fix F16 attention inference auxiliaries

04188f2

Always generate and return softmax statistics and the Philox state for F16 forward graphs so context-parallel inference can combine partition outputs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Key F16 attention graphs by logical batch

31c043d

Include the effective bucketed THD graph batch in forward and backward cache keys to prevent reuse across incompatible ragged sequence counts. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Use logical graph batch for dense attention tensors

87ef877

Describe dense peers of ragged QKV tensors with the bucketed graph batch while preserving their physical storage strides. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Skip ragged KV token hint on SM120

fd1aa12

Avoid passing the packed KV token bound to cuDNN THD backward on SM120, where token bucketing is disabled and the graph sequence extent is per-sequence. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Mask THD padding before max-logit reduction

6b61996

Build a device-side validity mask from padded query offsets so unwritten inter-sequence and tail gaps cannot affect ragged max-logit results. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Zero THD outputs on the slow fill path

8424cc4

Always initialize ragged F16 outputs and gradient storage because cuDNN leaves inter-sequence and tail padding unwritten, including when fast zero fill is disabled. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Key FP8 backward graphs by encoding

95c2b5a

Include the FP8 dtypes of attention inputs, dO, and gradient outputs so delayed-scaling recipes with identical uint8 buffers cannot share incompatible graphs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Key paged attention graphs by table layout

31be282

Record both page-table tensor descriptors in F16 forward cache keys so differing maximum-page dimensions and strides build distinct graphs. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Merge branch 'main' into attn-pyt-pygraph

c141994

# Conflicts: # build_tools/pytorch.py # pyproject.toml # transformer_engine/pytorch/csrc/extensions/attention.cpp

Merge branch 'main' into attn-jax-pygraph

618d2d6

# Conflicts: # build_tools/jax.py # transformer_engine/jax/csrc/extensions/attention.cpp

Merge branch 'attn-jax-pygraph' into attn-pygraph

24af929

vcherepanov-nv requested review from KshitijLakhani, cyanguwa, jberchtold-nvidia, ksivaman and ptrendx as code owners September 11, 2026 03:05

vcherepanov-nv marked this pull request as draft September 11, 2026 03:08

greptile-apps Bot reviewed Sep 11, 2026

View reviewed changes

pre-commit-ci Bot and others added 2 commits September 11, 2026 03:51

[pre-commit.ci] auto fixes from pre-commit.com hooks

028c583

for more information, see https://pre-commit.ci

Merge branch 'main' into attn-pygraph-parity

614a123

cyanguwa reviewed Sep 11, 2026

View reviewed changes

cyanguwa self-assigned this Sep 11, 2026

vcherepanov-nv and others added 14 commits September 11, 2026 20:43

Remove obsolete fused attention backend override

478ce7c

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Restore fused attention rejection diagnostics

150d5ed

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Restore fused attention graph cache diagnostics

4e5efd1

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Restore PyTorch FP8 THD attention support

984ddbd

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Reject affected cuDNN versions for JAX MXFP8 attention

4c6ecc5

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Gate JAX FP8 attention by recipe support

138dc2b

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Match JAX current-scaling attention quantizers

51471fa

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Merge branch 'main' into attn-pygraph-parity

2d7909c

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> # Conflicts: # transformer_engine/common/fused_attn/utils.cu

[pre-commit.ci] auto fixes from pre-commit.com hooks

13cec24

for more information, see https://pre-commit.ci

Fix FlexAttention execution mock signature

a5248dc

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Fix FP8 attention checkpoint test opt-in

a497b13

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Avoid retaining cuDNN graph workspaces

848cd07

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Fix paged attention table descriptors

a12ab7d

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

d4e3915

for more information, see https://pre-commit.ci

Navigation Menu

Uh oh!

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

[Draft] Port cuDNN frontend attention to Python API#3508
vcherepanov-nv wants to merge 38 commits into
NVIDIA:mainfrom
vcherepanov-nv:attn-pygraph-parity

vcherepanov-nv commented Sep 11, 2026

Uh oh!

greptile-apps Bot commented Sep 11, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Sep 11, 2026

Uh oh!

cyanguwa Sep 11, 2026

Uh oh!

cyanguwa Sep 11, 2026

Uh oh!

cyanguwa Sep 11, 2026 •

edited

Loading

Uh oh!

cyanguwa Sep 11, 2026

Uh oh!

cyanguwa Sep 11, 2026 •

edited

Loading

Uh oh!

cyanguwa Sep 11, 2026 •

edited

Loading

Uh oh!

cyanguwa Sep 11, 2026 •

edited

Loading

Uh oh!

cyanguwa left a comment •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return cudnn.pygraph(**kwargs)


		def build_cudnn_graph(cudnn, graph, *, description: str) -> int:



		@dataclass(frozen=True)
		class FusedAttentionConfig:

Uh oh!

Conversation

vcherepanov-nv commented Sep 11, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 4/5

Findings

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

greptile-apps Bot commented Sep 11, 2026 •

edited

Loading

cyanguwa Sep 11, 2026 •

edited

Loading

cyanguwa Sep 11, 2026 •

edited

Loading

cyanguwa Sep 11, 2026 •

edited

Loading

cyanguwa Sep 11, 2026 •

edited

Loading

cyanguwa left a comment •

edited

Loading