Repository navigation
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>
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>
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>
Explicitly enable FP8 backward before the custom MHA support query so it matches the direct FP8 backward call and skips unsupported H100 configurations. Restore the parameterized tests' FP8 environment settings with monkeypatch to avoid test-order dependencies. Validation: pre-commit and L0_license pass. Focused pytest collection is blocked by the installed native extension missing DType.kFloat8E8M0. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Retain immutable aligned scalar buffers across asynchronous cuDNN execution and CUDA graph replay to prevent corrupted FP8 GQA key and value gradients. Extend numerical coverage to different attention scales sharing compiled plans. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Preserve the shared cuDNN graph lifecycle while applying main's score_mod backend protections to PyTorch and JAX. Exclude Python SDPA engines before planning and reject them if selected, retaining JAX's serialization restrictions. Validation: all 45 PyTorch Flex Attention tests and three isolated JAX finalizer regression cases passed. Pre-commit and the PyTorch and JAX Python lint launchers passed. Full JAX runtime validation was unavailable because JAX is not installed. The license launcher reports failures only in the pre-existing, untracked nccl_ep directory. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Keep the JAX attention overview focused on supported module options and remove version-specific support tables, workarounds, and implementation details. Retain installation requirements, clarify cache diagnostics and backward precision, move callback cache-key guidance to API docstrings, and remove internal JAX functions from the public API reference. Validation: pre-commit, both framework Python lint suites, changed-file license checks, and docstring-only AST comparisons passed. The full Sphinx build could not proceed without generated Doxygen XML in this container. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Separate logical batch sizes from graph capacities in the PyTorch cuDNN attention path. Bucket eligible ragged backward graphs and converted-length forward graphs, padding sequence metadata and terminal offsets. Use cumulative lengths and ragged-offset multipliers for eligible dropout-free forward graphs, matching the common path's version gates. Extend cache-reuse and numerical coverage across batch buckets, including packed padding, mixed layouts, FP8 backward precision, and dropout checked against an independent FP32 reference. Validation: 394 selected PyTorch cases passed and 281 were skipped. Pre-commit, both framework Python lint suites, and edited-file license checks passed. Validation used GB200 with PyTorch 2.15.0a0, CUDA 13.4, cuDNN 9.27, and cuDNN Frontend 1.31. JAX runtime validation was unavailable because JAX is not installed. 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>
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>
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
…ough V scales A columnwise MXFP8 V scale is shared by 32 consecutive tokens, so under ordinary causal MXFP8 attention a later token in a query's own 32-token block can change how earlier, visible V values are rounded: outputs and gradients of token i depend on tokens after i (NVIDIA#3248). Add MXFP8BlockScaling(fp8_dpa_causal_own_block=True), defaulting to NVTE_FP8_DPA_CAUSAL_OWN_BLOCK. cuDNN runs MXFP8 attention over completed 32-token blocks (sdpa_mxfp8 causal_exclude_own_block=32), a Triton kernel adds each query's own block on a per-token-scaled copy of V and merges the partials through their log-sum-exps, and the backward runs in FP16/BF16 on Q/K decoded from their MXFP8 bytes and V decoded from the per-token copy. Backend selection queries the graphs that run (the own-block forward with a constant scale and the FP16/BF16 backward), so a cuDNN Frontend without the option is rejected at selection. No other backend serves the option: unfused FP8 emulation and context parallelism are refused, and the no-backend error names it. Adds tests, docs and a benchmark. Signed-off-by: Romit <romitjain@sarvam.ai> Co-authored-by: Tejesh <tejesh@sarvam.ai> Co-authored-by: Pushkar <pushkar@sarvam.ai>
github-actions
Bot
added
the
community-contribution
5 tasks done
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
A columnwise MXFP8 V scale is shared by 32 consecutive tokens. Under an ordinary causal mask, a later token in the query's own 32-token block can raise that shared scale and change how earlier, visible V values are rounded, so the output (and gradients) for token
idepend on tokens afteri. This PR addsMXFP8BlockScaling(fp8_dpa=True, fp8_dpa_causal_own_block=True), which keeps P·V in MXFP8 for every completed 32-token block and removes the leak:j < 32 * floor(i / 32)): cuDNN MXFP8 SDPA with the newsdpa_mxfp8(causal_exclude_own_block=32)op attribute (Add sdpa_mxfp8(causal_exclude_own_block=32) for causal MXFP8 attention without future-token leakage cudnn-frontend#2563). Every V scale it reads covers visible tokens only. This is nearly all of the attention FLOPs.Builds on #3693 (cuDNN Python attention graphs); this branch is based on its head
c689b371. Until #3693 merges, this PR's diff also shows its commits; the change proposed here is the last commit: review it on its own.Fixes #3248
Type of change
Changes
MXFP8BlockScaling.fp8_dpa_causal_own_block(default fromNVTE_FP8_DPA_CAUSAL_OWN_BLOCK, off by default), documented in the recipe anddocs/envvars.rst.transformer_engine/common/triton/mxfp8_causal_attention.pyand its PyTorch wrapper: own-block forward, token-normalized V, backward operand decode.FusedAttentionforward/backward: runs the cuDNN completed-block graph, the Triton own-block pass and the high-precision backward when the option is on.causal_exclude_own_blockis rejected at selection with a clear message.benchmarks/attention/benchmark_mxfp8_causal_own_block.py: BF16 vs MXFP8 vs MXFP8 own-block, first-call, forward, forward+backward, kernel breakdown and error vs BF16.tests/pytorch/attention/test_attention.py: support rules, selection-time graph query (with a fake old frontend), end-to-end numerics and causality on BSHD/SBHD, MHA/GQA.Requirements
SM100/SM103, a cuDNN Frontend whose
sdpa_mxfp8acceptscausal_exclude_own_block, separate BSHD or SBHD Q/K/V with D = 128, equal Q and KV lengths, causal mask, no dropout, bias, softmax offset or context parallelism.Performance
B200,
benchmarks/attention/benchmark_mxfp8_causal_own_block.py --batch 1 --heads 16 --kv-heads {4,16} --seqlens 8192 16384 32768, BSHD, causal, median of 50 after 10 warm-up, CUDA events around the module call (includes quantization and Python launch). Forward, ms (ratios: own-block / BF16, own-block / MXFP8):Forward + backward, ms:
The MXFP8 forward has a fixed cost of about 1 ms at batch 1 (quantization and launch: both MXFP8 paths take the same time at 8K and 16K), so at 8K-16K both MXFP8 paths are slower than BF16 in the forward. At 32K the own-block forward is 9-16% faster than BF16 and 6-10% slower than ordinary (leaky) MXFP8; forward + backward is 3-5% faster than BF16 from 16K. At 8K, own-block forward + backward beats ordinary MXFP8 because its backward is BF16, not FP8. The cuDNN completed-block kernel itself runs within 0-1.6% of plain causal MXFP8 at every shape measured (see NVIDIA/cudnn-frontend#2563).
Checklist:
Testing
B200 (cc 10.0), driver 580.178.04, NGC PyTorch 26.08 (PyTorch 2.14.0a0, CUDA 13.4, cuDNN 9.25.0), cuDNN Frontend from NVIDIA/cudnn-frontend#2563:
python -m pytest tests/pytorch/attention/test_attention.py -k "own_block or fused_attn_backend_message or query_backward_precision"93 passed.
python -m pytest tests/pytorch/attention/test_attention.py -k "test_dpa_fp8_vs_f16 or test_mha_fp8_vs_f16 or test_dpa_fp8_extra_state or fused_attn_query or fused_attn_python"2423 passed, 2407 skipped, 2 failed. The two failures are
test_dpa_fp8_extra_state[dtype0-large]and[dtype1-large], which fail identically on the base commit: TE refuses to load the test's pickledLayerNormLinearextra state. They are unrelated to this PR.pylint 3.3.1 (
--rcfile=pylintrc, theqa/L0_pytorch_lintversion) on the changed library files: 10.00/10. Black 24.4.2 (line length 100) leaves every changed file unchanged.Merge order: needs a cuDNN Frontend release that contains
sdpa_mxfp8(causal_exclude_own_block=...). With an older frontend the option is rejected at backend selection with a clear error, and nothing else changes.