-
Notifications
You must be signed in to change notification settings - Fork 824
[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.
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
Draft
vcherepanov-nv
wants to merge
42
commits into
NVIDIA:main
Choose a base branch
from
vcherepanov-nv:attn-pygraph-parity
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from all 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 df64163
Test shared cuDNN frontend version check
vcherepanov-nv e42c3ac
Port PyTorch cuDNN attention to Python frontend
vcherepanov-nv 35d7158
Fix JAX ragged cuDNN attention graphs
vcherepanov-nv 04188f2
Fix F16 attention inference auxiliaries
vcherepanov-nv 31c043d
Key F16 attention graphs by logical batch
vcherepanov-nv 87ef877
Use logical graph batch for dense attention tensors
vcherepanov-nv fd1aa12
Skip ragged KV token hint on SM120
vcherepanov-nv 6b61996
Mask THD padding before max-logit reduction
vcherepanov-nv 8424cc4
Zero THD outputs on the slow fill path
vcherepanov-nv 95c2b5a
Key FP8 backward graphs by encoding
vcherepanov-nv 31be282
Key paged attention graphs by table layout
vcherepanov-nv 92ed239
Fix JAX cuDNN graph build compatibility
vcherepanov-nv c141994
Merge branch 'main' into attn-pyt-pygraph
vcherepanov-nv 618d2d6
Merge branch 'main' into attn-jax-pygraph
vcherepanov-nv 24af929
Merge branch 'attn-jax-pygraph' into attn-pygraph
vcherepanov-nv bfe6b7d
Remove obsolete common fused attention code
vcherepanov-nv 3a4be7d
Factor shared cuDNN attention helpers
vcherepanov-nv b7b2cac
[JAX] Add FP8 cuDNN attention parity
vcherepanov-nv d6f38aa
[JAX] Fix attention regression failures
vcherepanov-nv 13ae9af
[JAX] Fix MXFP8 attention scale strides
vcherepanov-nv 3e0bdca
Merge branch 'main' into attn-pygraph-parity
vcherepanov-nv 028c583
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 614a123
Merge branch 'main' into attn-pygraph-parity
cyanguwa 478ce7c
Remove obsolete fused attention backend override
vcherepanov-nv 150d5ed
Restore fused attention rejection diagnostics
vcherepanov-nv 4e5efd1
Restore fused attention graph cache diagnostics
vcherepanov-nv 984ddbd
Restore PyTorch FP8 THD attention support
vcherepanov-nv 4c6ecc5
Reject affected cuDNN versions for JAX MXFP8 attention
vcherepanov-nv 138dc2b
Gate JAX FP8 attention by recipe support
vcherepanov-nv 51471fa
Match JAX current-scaling attention quantizers
vcherepanov-nv 2d7909c
Merge branch 'main' into attn-pygraph-parity
vcherepanov-nv 13cec24
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] a5248dc
Fix FlexAttention execution mock signature
vcherepanov-nv a497b13
Fix FP8 attention checkpoint test opt-in
vcherepanov-nv 848cd07
Avoid retaining cuDNN graph workspaces
vcherepanov-nv a12ab7d
Fix paged attention table descriptors
vcherepanov-nv d4e3915
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 8e9b0ac
Fix cuDNN graph execution device context
vcherepanov-nv 1969670
Fix attention Python lint findings
vcherepanov-nv c6f373e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 033c701
Fix JAX attention lint warnings
vcherepanov-nv File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This computes
install_requiresby querying the locally installed cuDNN frontend. In a documented--no-build-isolationinstallation or a directsetup.pybuild where that dependency is not already installed, metadata generation raisesPackageNotFoundErrorbefore 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.0requirement, making the package impossible to install.Knowledge Base Used: Build, extensions, and packaging