Repository navigation
Switch to Python binding for cuDNN frontend - #3693
Open
vcherepanov-nv wants to merge 58 commits into
Open
vcherepanov-nv wants to merge 58 commits into
vcherepanov-nv wants to merge 58 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
 [Medium impact] The latest changes appear safe to merge; the two previous concerns remain non-blocking. FindingsSummaryMoves PyTorch and JAX cuDNN attention to shared Python graph-building helpers.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart LR
PT[PyTorch attention] --> Shared[Shared Python graph helpers]
JAX[JAX attention] --> Shared
Shared --> PTExec[PyTorch frontend execution]
Shared --> Plans[Serialized cuDNN plans]
Plans --> FFI[JAX native FFI executor]
PTExec --> GPU[GPU attention]
FFI --> GPU
Reviews (34) · Last reviewed commit: "Fix attention test backend availability ..." · 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>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Validate model-precision backward graphs when NVTE_FP8_DPA_BWD=0, retaining the original QKV layout and nominal dtype. Normalize inference dropout before preparing runtime graph inputs. Add regression coverage for FP8 backward selection and inference output against zero-dropout execution and FP32 SDPA. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Map paged K/V dimensions to BHSD consistently with their strides. Add FP16/BF16 regression coverage for all six paged layouts against FP32 SDPA. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Keep both JAX API documentation additions and preserve ragged-attention and configurable-device imports. Retain the shared cuDNN graph helpers and port configurable device checks and device-index normalization into the shared PyTorch adapter. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Return No_Backend instead of raising from PyTorch and JAX backend selection when the cuDNN frontend package is missing, its version does not match the C++ build, or it rejects the graph with a frontend error. Report a clear error when building TE/JAX without nvidia-cudnn-frontend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Resolve NVTE_FUSED_ATTN_CACHE_DEBUG and the rank filter once, on first use, instead of re-reading the environment on every cache event in the attention hot path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Restore the original handling in PyTorch and JAX backend selection: only a configuration rejected by cuDNN falls back to another backend, while a missing frontend, a frontend version mismatch, or a device failure is raised. Falling back hid broken installs behind a silent switch to unfused attention and broke test_fused_attn_query_propagates_setup_errors. This reverts the attention changes from cad3188 and keeps its build_tools/jax.py error message. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Retain supported graphs and serialized bindings without the arbitrary 256-entry eviction limit, matching the existing PyTorch cache policy and avoiding unnecessary rebuilds. Validation: focused JAX attention checks (12 passed, 1 skipped), pre-commit, changed-file pylint and license checks passed. FP8-backward dropout is unsupported by cuDNN on this Blackwell configuration. Full JAX lint reports existing tuple-unpacking warnings in unchanged files; the full license check fails only on pre-existing untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
| return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" | ||
|
|
||
|
|
||
| @lru_cache(maxsize=None) |
Contributor
There was a problem hiding this comment.
Changing both caches to maxsize=None keeps every supported cuDNN graph and serialized plan alive for the process lifetime. Dense requests use their actual sequence lengths as cache keys, so long-lived JAX workers handling changing shapes keep accumulating entries, even after discarding their compiled functions. There is no cleanup hook for these caches.
This is a non-blocking memory-growth concern. Keep a cache limit or release entries when their compiled functions are discarded.
Check hardware and recipe exclusions before constructing descriptors and reuse the shared MXFP8 support policy. Validate backward independently of training mode, including model-precision fallback and trainable bias or sink parameters with frozen Q/K/V. Keep dropout tied to training mode and order graph build before execute. Expose the existing common sequence-length conversion kernel through a utility C API and private PyTorch binding. Convert both cumulative-length arrays and zero-fill batch padding in one launch, preserving direct cumulative inputs where supported. Validation: rebuilt the editable PyTorch/common installation; 520 focused PyTorch tests and 32 common C++ tests passed on GB200 with CUDA 13.4, cuDNN 9.27, and cuDNN Frontend 1.31. Both framework lint suites, pre-commit, and changed-file license checks passed. Coverage includes independent numerical references, backward selection, nondefault CUDA streams, and CUDA graph replay. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Merge validation, RNG sizing, and max-logit reduction into the cuDNN Python forward and backward entry points. Move shared attention enums, layout mappings, and FP8 metadata constants to pytorch.constants and update callers. Preserve package-level exports, function signatures, and guarded graph execution. Validation: 479 focused PyTorch tests passed on GB200, including packed FP8 attention, max-logit reduction, CUDA graph replay, and torch.compile backend selection. Fresh import and signature checks, PyTorch/common Python lint, pre-commit, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Consolidate attention tests, retire the custom FP8 harness, and retain compact metadata coverage. Use the TE unfused reference and move paged-layout cases into the KV-cache suite. Document C attention API deprecation, clarify max-logit naming, and avoid redundant device guards while retaining safe cross-device stream binding. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Remove the backward token-bound and fake-statistics shape tests requested in review, along with the unused SimpleNamespace import. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Filter unserializable SDPA plans before support checking for ordinary and score-modified attention. Pin queries to eligible plans and retain compilation fallback across backend candidates. Reuse the common cuDNN handle in FFI initialization and execution, group attention options before CP defaults, and remove unused bias metadata. Add regressions for excluded-plan rejection and compilation fallback. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Keep the existing backend query and numeric results while moving its FusedAttnBackend enum into Python. Update JAX, Flax and test consumers, remove the unused pybind enum, and retain descriptive backend logging. Validation: rebuilt the JAX extension; 57 focused tests passed. Pre-commit, focused pylint, C++ lint and changed-file licenses passed. The full lint launcher reports no diagnostics beyond the existing ones. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Honor NVTE_DPA_FP8CS_O_in_F16 for current-scaled attention. Include the resolved O precision in graph requests, describe and quantize FP8 O with its inverse scale, and retain model-precision forward outputs. Ignore the option for MXFP8 and model-precision backward. Document the precision option and the absence of delayed-scaling JAX attention. Extend existing numerical, context-parallel and cache tests. Validation: 132 focused tests passed, one skipped, including eight two-GPU context-parallel cases. Existing numerical tolerances retained. Pre-commit, focused pylint and changed-file license checks passed. Full JAX lint reports only the existing diagnostics. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Request heuristic mode A without FALLBACK, matching the common fused attention path on main. Update the existing score-modified attention guard expectation. Validation: 158 focused PyTorch tests passed, including Flex Attention, cache reuse, ragged numerics, and delayed/current/MXFP8 attention. Black, whitespace, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Keep F16, tensor-scaled FP8, and MXFP8 dispatch in separate sections, with forward and backward calls together. Clarify that fp8 describes the effective precision already resolved by the frontend. Rename the shared preparation module to mxfp8_utils and update both frontend imports and the existing JAX tests. Validation: 182 focused PyTorch tests passed, covering all three FP8 recipes and both backward precision settings. Black, syntax, whitespace, and changed-file license checks passed. JAX runtime validation was unavailable because JAX and its extension are not installed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Carry actual bias strides into support queries so strided bias inputs reuse queried graphs. Allocate contiguous dBias to satisfy cuDNN for permuted bias inputs. Add graph identity, compilation count, and independent numerical checks for FP16/BF16 bias layouts. Extend optional-output descriptor coverage. Validation: 251 focused PyTorch tests passed, with 12 existing skips. Repository pre-commit, Python lint, syntax, whitespace, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Remove the explicit architecture/version support matrix and the quantization/scale-byte test requested in review. Retain the remaining scale-layout, numerical, and integration coverage. Validation: 31 focused MXFP8 tests passed; pre-commit and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
| assert actual_grad.dtype == dtype | ||
| np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-3) | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
Removing test_quantization_and_scale_bytes drops the independent check of the scale bytes passed to cuDNN. The remaining shape test checks sizes, not byte placement. The numerical tests also no longer exercise head widths 96 and 160.
This is non-blocking, but mistakes in rowwise and columnwise scale packing become harder to catch. Keep the byte comparison, or replace it with an equivalent test.
Knowledge Base Used: Verification and CI matrix
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Remove the private descriptor-bucketing test requested in review. Retain real FFI numerical coverage for ragged padding and gradients. Validation: six ragged-padding cases passed; pre-commit and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Move F16/CS context-parallel and MXFP8 sharding/communication tests into the existing distributed attention suite. L0's selection now excludes those cases and the L1 launcher collects them, with all numerical statements and parameter matrices preserved. Remove the legacy common-API routing and redundant Flax MXFP8 mode forwarding tests. Reuse the existing numerical helpers without adding another launcher or test infrastructure. Validation: 47 distributed and 92 single-GPU cases passed. Confirmed L0 excludes the relocated cases and L1 collects all 47. Pre-commit and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Use a frontend-owned common cache for ordinary and score-modified attention. Share synchronized graph lookup, compile-once state, and optional serialization while retaining callback identity, unsafe-callback bypass, and engine eligibility rules. Consolidate PyTorch score-modified key, build, and lookup helpers around AttentionGraphSpec and a shared live entry type. Preserve JAX serialized results and request-specific binding memoization. Validation: 257 focused PyTorch checks passed with 12 existing skips. All 51 Flex Attention tests passed again after the final fixture edits. Repository pre-commit, both Python lint launchers, syntax, whitespace, and changed-file license checks passed. JAX runtime validation is unavailable in this container. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Reuse the existing FP32 attention reference for CS and MXFP8, replace broad descriptor and initialization matrices with focused contracts, and cover realistic scaled MHA, GQA, and cross-attention cases. Keep Flax-specific sink and prequantized boundary checks, and merge context-parallel storage coverage into a curated distributed matrix. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Keep the exact JAX cuDNN Frontend runtime pin while adopting the minimum versions from main: Frontend 1.31.0 and JAX 0.5.3. Align wheel dependencies and installation documentation. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Fused attention backward with dBias does not support deterministic execution. Skip the bias graph cache test when deterministic algorithms are required, while retaining its nondeterministic coverage. Validation: all 12 deterministic cases skipped and all 12 nondeterministic cases passed on GB200. Pre-commit, PyTorch Python lint, and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Align THD FP8 MHA token counts to 16 for weight-gradient GEMMs and reserve enough space in the final sequence for alignment. Limit batch-bucketing identity assertions to ragged graphs. Skip the deterministic fused eval/backward numerical test below SM90. Validation on GB200: 243 focused tests passed with 50 expected skips. Checked 96 graph-metadata cases with simulated SM80/90/100/120 capabilities and token alignment/bounds for 384 configurations. Pre-commit, PyTorch Python lint, and changed-file license checks passed. Actual H100/A100 execution remains for CI. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Use graph_attributes for forward and backward FFI calls instead of repeating the nine serialized-graph attributes. Validation: pre-commit, JAX Python lint, syntax, and changed-file license checks passed. Original-versus-new FFI argument parity passed for inference forward, training forward, and backward. JAX numerical tests were not run because JAX is not installed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Share per-framework diagnostic collectors between ordinary SDPA and Flex Attention. Report cache events and graph-build timings, plus PyTorch executions, while retaining JAX tracing-only reporting and correct uncached-callback behavior. Document the shared summaries and add coverage for debug levels 0, 1, and 2 with cacheable and uncacheable callbacks. Validation: 70 PyTorch tests passed on GB200 and 21 isolated JAX unit cases passed with framework stubs and mocked cuDNN. Both framework Python lints, pre-commit, syntax, and changed-file license checks passed. Native JAX execution was not tested because JAX and its extension are not installed. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> # Conflicts: # tests/pytorch/attention/test_attention.py
Contributor
Comments Outside DiffThese findings could not be posted inline.
|
8 of 13 tasks
Match backward support queries and autograd state to the test execution mode. Check the custom FP8 test's model-precision reference separately from its fused FP8 path to avoid depending on leaked emulation settings. Validate all 44 previously extra-skipped cases and the shared 5,259-case selection. Outcomes now match main: 2,734 passed and 2,525 skipped, with no failures. 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
Migrate PyTorch and JAX cuDNN attention to the cuDNN Frontend Python API,
with shared graph-building helpers and framework-specific execution.
The public C attention APIs remain available for compatibility but are now
deprecated. The migrated PyTorch and JAX attention paths no longer call them.
Type of change
Changes
and cache diagnostics across PyTorch and JAX.
precompiled cuDNN plans.
context parallelism, and add ALiBi support.
regression coverage for dropout replay and backward precision.
cuDNN cannot currently serialize those plans for execution through the
native executor.
Checklist: