Repository navigation
[DRAFT] Switch to Python binding for cuDNN frontend - #3693
Draft
vcherepanov-nv wants to merge 8 commits into
Draft
vcherepanov-nv wants to merge 8 commits into
vcherepanov-nv wants to merge 8 commits into
Conversation
Build and select cuDNN attention plans through the Python frontend for ordinary and score-modified attention. Share graph serialization and a JAX-owned C++ FFI executor with per-device plan and handle caches and CUDA command-buffer support. Remove JAX's calls to TE common attention APIs while leaving TE common and PyTorch unchanged. Expose current tensor-scaled FP8 forward and backward through fused_attn and Flax attention modules, retaining FP16/BF16 inputs, outputs, and input gradients. Add ALiBi and preserve packed and ragged layouts, GQA, masks, sliding windows, dropout replay, attention sinks, max logits, and Ring and AllGather context parallelism. Preserve neutral statistics for empty ragged ring blocks and differentiation with is_training=False. Document FP8 availability, scaling, remaining feature gaps, and the NVTE_FP8_DPA_BWD switch. Add independent numerical references and focused regression coverage for the new precision and execution paths. Validate with the editable JAX build, the requested attention suites (1,241 passed, 3,901 skipped), and additional FP8, Flax, deterministic, two-GPU context-parallel, and CUDA command-buffer checks. JAX L0 lint, pre-commit, and changed-file license checks pass. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Move cuDNN attention support queries, graph construction, plan caching, and forward/backward execution into the PyTorch Python frontend. Cover FP16/BF16, FP8 and MXFP8 while preserving layout, masking, softmax, ragged and paged KV, quantization, and context-parallel behavior. Share graph lifecycle helpers with experimental flex attention. Remove the obsolete PyTorch bindings to common attention execution and retain a capture-safe C++ Philox helper for PyTorch generator state. Keep the common and JAX attention implementations intact. Add numerical-reference, execution-path isolation, capability-exclusion, and CUDA generator replay regressions, and document the Python frontend dependency and graph caching behavior. Validation on GB200 with cuDNN 9.27 and frontend 1.30: - Editable PyTorch build succeeded. - Requested attention/flex suites: 2,869 passed, 2,425 skipped. - Final focused numerical coverage: 629 passed, 560 skipped. - Deterministic attention: 2,231 passed, 3,038 skipped with NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 and NVTE_FLASH_ATTN_V4=0. - Compiled attention: 52 passed; CUDA graph and two-/four-GPU context-parallel coverage passed. - Default/custom generator replay passed; FP8 dropout replay was skipped because cuDNN does not support the tested Blackwell configuration. - Changed-code lint, formatting, and whitespace checks passed. Whole-repository checks still report an existing Python import-order violation in fused_mla_q_uproj.py and missing license headers in the untracked nccl_ep directory. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Share tensor descriptors, SDPA construction, graph lifecycle stages, diagnostics, and score-modification cache utilities between PyTorch and JAX. Keep allocation, quantization, RNG, and execution in framework adapters. Separate support checks from compilation, fix JAX anonymous callback cache keys, preserve empty ragged capacities, and align frontend requirements and documentation. Validation: editable PyTorch build succeeded. Attention and Flex suites reported 2,877 passed and 2,426 skipped; both checkpoint-loading failures passed on retry with NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1. Pre-commit, Python lint, license checks, and JAX descriptor checks passed. JAX build and runtime validation remain for a separate session. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Restore bucketed total-token bounds for FP16/BF16 packed backward graphs and forward Q and KV bounds independently to preserve mixed layouts. Extend graph reuse coverage to backward and compare packed and mixed layout outputs and gradients against FP32 SDPA. Update the torch.compile backend-selection regression to patch the Python support probe after removal of the C++ binding. Pin each JAX build's runtime cuDNN Frontend dependency to the distribution supplying its build headers while retaining the minimum supported version. Document the exact-version requirement and test generated dependency metadata for frontend versions 1.29.0 and 1.30.0. Validation: 382 compile tests passed, 24 skipped, and 1 xpassed; 34 focused attention tests and 3 dependency-metadata tests passed. Pre-commit, JAX Python lint, and changed-file license checks passed. JAX runtime validation was unavailable because JAX is not installed. Full PyTorch lint and repository license checks report existing issues in unrelated files and untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Address the remaining issues from PR NVIDIA#3289 in the shared cuDNN graph paths. Key plain callbacks on their code and typed snapshots of positional and keyword defaults so callbacks with different graph structure cannot reuse the same cached plan. Disable implicit caching for tensor, opaque, or recursive defaults while retaining the explicit structural-key protocol. Apply dense per-sequence statistics and dimensions throughout the SM12x family in common C++, PyTorch, and JAX. Share the Python layout predicate with context-parallel LSE handling and extend the related backend gates. Add forward/backward cache-lookup regressions for code and default changes, SM121/SM129 layout and eligibility coverage, an SM130 control, and matching JAX configuration tests. Document the caching and statistics conventions. Validation: - Rebuild the editable PyTorch installation with its cached configuration. - Compile suite: 382 passed, 24 skipped, 1 xpassed. - Focused attention regressions: 93 passed. - Common C API: four FP16/BF16 dense/packed support queries passed. - Two-GB200 THD context-parallel forward/backward regression passed. - Pre-commit, C++ lint, JAX Python lint, and changed-file license checks passed. JAX and physical SM121 runtime validation were unavailable. PyTorch Python lint still reports the pre-existing import-order issue in fused_mla_q_uproj.py. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Use the shared architecture predicate for JAX Max tensors and PyTorch fake saved statistics. Add regression coverage for SM121 and SM129 alongside neighboring architecture families. Validation: 19 focused JAX tests, pre-commit, and both Python lint suites passed. PyTorch tests require the unavailable torch package. The license check reports existing untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Retain both the SM12x fused-attention restriction tests from py-all and the FA4 D=256 sequence-metadata eligibility tests from main. Validation: 19 focused JAX attention tests, pre-commit on overlapping files, and PyTorch Python lint passed. PyTorch runtime tests could not collect because torch is unavailable. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
This branch has not been deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Description
This PR changes TE to use Python binding for cuDNN backend.
The public nvte_* functions remain in TE/Common, but are no longer used by FW-specific code
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: