Repository navigation
[DRAFT] Switch to Python binding for cuDNN frontend - #3693
Draft
vcherepanov-nv wants to merge 18 commits into
Draft
vcherepanov-nv wants to merge 18 commits into
vcherepanov-nv wants to merge 18 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>
Collaborator
Author
|
@greptile full review |
Contributor
 [High impact] The latest changes appear safe to merge and address all three previous findings. SummaryThis PR moves PyTorch and JAX attention graph construction to shared cuDNN Python helpers. JAX runs serialized plans through its native executor; PyTorch runs plans through the Python frontend.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart LR
PT[PyTorch attention] --> Shared[Shared Python graph builders]
JAX[JAX attention] --> Shared
Shared --> Plans[cuDNN execution plans]
Plans --> PTExec[PyTorch Python execution]
Plans --> Serialized[Serialized JAX plans]
Serialized --> FFI[JAX native executor]
PTExec --> GPU[GPU]
FFI --> GPU
Reviews (2) · Last reviewed commit: "Fix cuDNN attention support guards and o..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Restore specialized FP16/BF16 THD learnable-softmax kernels before cuDNN 9.26 and reject all affected FP8 THD sink-token backward paths in JAX. Preserve the FP16/BF16 backward fallback. Skip fills for fully written JAX outputs while retaining initialization for ragged buffers and unwritten optional results. Add support-policy, buffer-binding, and numerical regression coverage, and document the version restrictions. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Collaborator
Author
|
@greptile full review |
Temporarily expose the checkout root while importing build_tools so the new tests collect from the CI working directory. Restore the import path immediately to preserve installed-package resolution. Validation: all three tests pass with direct pytest and from the JAX CI working directory; import-path restoration and git diff --check pass. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Match the existing backend policy when the Python execution test forces cuDNN. Skip deterministic training below SM90 before disabling the fallback backends. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Centralize operand directions, scale padding and reordering, alignment checks, and known backend exclusions. Adapt the existing PyTorch preparation without changing tensor ownership or GPU execution. Validation: 14 JAX metadata tests and changed-file pre-commit checks pass. PyTorch runtime validation is unavailable in this JAX environment. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Expose fp8="mxfp8" through fused_attn and Flax modules. Prepare independent contraction-axis quantizations and per-head reordered E8M0 scales using existing common quantization kernels, then execute shared MXFP8 graphs through the existing JAX FFI executor. Validation on four GB200 GPUs with cuDNN 9.26 / Frontend 1.29: 30 focused MXFP8 tests and 15 existing attention regressions pass. Changed-file pre-commit and focused JAX pylint pass. Full license checking encounters existing failures including untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Support packed QKV/KV storage, GQA masks, sliding windows, sink softmax, unequal head dimensions, and deterministic execution. Keep original model-precision dO layout and use the common fused quantizer for columnwise operands to avoid constant-zero NaNs in the pure-JAX fallback. Validate 60 focused cases across these capabilities, FP16/BF16 backward, and checkpointing on GB200 with cuDNN 9.26. Add explicit SM12x and dropout exclusions. Changed-file pre-commit and focused JAX pylint pass. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Accept independent rowwise/columnwise ScaledTensor2x operands alongside model-precision gradient targets. Consume supplied FP8 buffers directly in the shared cuDNN graph, reorder scales after partitioning, and preserve original operands for model-precision backward. Add quantize_attention_qkv, API documentation, and nine passing GB200 tests for packed layouts, both backward modes, supplied values, validation, and two-GPU data/tensor sharding. Focused JAX pylint and changed-file pre-commit pass. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Decode supplied MXFP8 operands before model-precision Ring/AllGather communication and quantize each local attention step. Support packed storage, causal load balancing, prequantized Flax BSHD/SBHD calls, checkpointing, and vmapped training. Preserve scale boundaries for replicated, data-parallel, and tensor-parallel execution. Reject the confirmed cuDNN Frontend 1.29.0 backward scale-stride defect for Q/K heads wider than 128 in the shared graph builder. Model-precision backward remains available via NVTE_FP8_DPA_BWD=0. Validation on four GB200 GPUs (SM100), cuDNN 9.26 / Frontend 1.29.0: full focused suite passes 109 tests, followed by 15 final checks including the added replicated-mesh case (110 distinct MXFP8 cases total). Existing attention, score-modification, and build-tools regressions: 58 passed, one skipped. JAX Python lint, changed-file pre-commit, and changed-file license checks pass. Tree-wide license checking encounters existing failures; PyTorch runtime validation is unavailable in this environment. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Query fused attention backend support before running FP8 tests instead of assuming Hopper supports FP16/BF16 outputs. Preserve cuDNN rejection reasons and match context-parallel queries to their per-device shapes and layouts. 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
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: