Switch to Python binding for cuDNN frontend by vcherepanov-nv · Pull Request #3693 · NVIDIA/TransformerEngine

vcherepanov-nv · 2026-10-16T00:30:14Z

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

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Migrate PyTorch and JAX cuDNN attention to Python binding of cuDNN FE API API, sharing graph construction, plan caching, diagnostics, and flex attention support.
Close JAX feature gaps: add FP8/MXFP8 attention, ALiBi.
Fix existing callback-cache collisions, SM12x handling, JAX dropout replay, and backward-precision support queries.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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>

vcherepanov-nv · 2026-10-19T21:23:25Z

@greptile full review

greptile-apps · 2026-10-19T21:34:34Z

Confidence Score: 5/5

[High impact] The PR appears safe to merge; no new actionable issue was established.

Summary

Moves PyTorch and JAX cuDNN attention to shared Python graph builders while keeping the public common APIs available.

PyTorch attention builds and runs cuDNN graphs through Python.
JAX attention runs graphs built by the cuDNN Python frontend.
JAX attention adds FP8, MXFP8, and ALiBi options.
Flex attention uses the shared cuDNN graph-building path.

The previous and current commits have different merge bases with the target commit, so this review used the full PR diff. GPU tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  PT[PyTorch attention] --> Shared[Shared Python graph builders]
  JAX[JAX attention] --> Shared
  Shared --> PTPlan[PyTorch cached plans]
  Shared --> Serialized[Serialized JAX plans]
  PTPlan --> Python[cuDNN Python execution]
  Serialized --> FFI[JAX native FFI executor]
  Python --> GPU[GPU]
  FFI --> GPU

_{Reviews (7) · Last reviewed commit: "Merge branch 'main' into py-all" · Reviewed by Greptile}

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>

vcherepanov-nv · 2026-10-19T23:20:57Z

@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>

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>

cyanguwa · 2026-11-07T03:46:33Z

+#
+# See LICENSE for license information.
+
+"""Framework-independent MXFP8 attention operand preparation rules."""


Would this more of a mxfp8_utils.py file, rather than mxfp8 attention?

cyanguwa · 2026-11-07T03:48:35Z

+    quantization groups 32 sequence positions. The two representations must be
+    quantized independently from the higher-precision operand.
+    """
+    return ((True, backward), (True, backward), (backward, not backward))


Would this be similar to get_fused_attn_spec? Could that be reused here in some way?

cyanguwa · 2026-11-07T03:51:37Z

+    raise TypeError("Opaque callback defaults require an explicit structural cache key.")
+
+
+def callback_cache_key(callback, is_tensor, uncacheable):


How special is score_mod at this point, compared to the regular SDPA workflow? Can the caching mechanism be merged, so there's only one cache, and one implementation of caching?

cyanguwa · 2026-11-07T03:53:26Z

+        backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend
+        message = ""
+        if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS:
+            message = "Fused attention does not support pre-scale bias."


This is probably already covered in the common get_fused_attn_backend?

cyanguwa · 2026-11-07T03:53:39Z

+        if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS:
+            message = "Fused attention does not support pre-scale bias."
+        elif self.q_dtype not in (jnp.float16, jnp.bfloat16):
+            message = "Fused attention requires FP16/BF16 inputs."


What about the pre-quantized case?

cyanguwa · 2026-11-07T03:55:33Z

+                cp_striped_window_size=None,
+                stripe_size=None,
+                return_max_logit=self.return_max_logit,
+                fp8=self.fp8,


Maybe reorder the args a bit for aesthetics.

cyanguwa · 2026-11-07T03:56:07Z

+                backend = (
+                    NVTE_Fused_Attn_Backend.NVTE_FP8
+                    if self.fp8
+                    else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen


Maybe remove NVTE_Fused_Attn_Backend and the Python equivalent one?

cyanguwa · 2026-11-07T03:57:10Z

            bias_batch = bias_heads = 0
        else:
            *bias_batch_shape, bias_heads, _, _ = bias_aval.shape
            bias_batch = reduce(operator.mul, bias_batch_shape)


If we don't need bias_batch here, we can probably remove this if/else. It was needed by the C API FusedAttnConfig before this PR.

cyanguwa · 2026-11-07T03:57:35Z

-            *bias_batch_shape, bias_heads, _, _ = bias_aval.shape
-            bias_batch = reduce(operator.mul, bias_batch_shape)
-
-        window_size_left, window_size_right = config.effective_window_size


You may need to keep this line for CP use cases.

cyanguwa · 2026-11-07T04:00:27Z

+
+
+@lru_cache(maxsize=256)
+def _supported_graph(spec, capabilities, frontend_version, backend_version):


We don't seem to have a cache size cap on the PyTorch side?

cyanguwa · 2026-11-07T04:01:54Z

+            tensor(name, (graph_batch + 1, 1, 1, 1), (1, 1, 1, 1), "INT64")
+    stats_dim = (graph_batch, h, graph_sq, 1)
+    stats_stride = (h * graph_sq, 1, h, 1) if ragged_stats else (h * graph_sq, graph_sq, 1, 1)
+    stats_offset = "offset_stats" if ragged_stats else ""


What about ragged and not ragged_stats (sm12x) cases?

cyanguwa · 2026-11-07T04:03:38Z

+def finalize_graph(cudnn, graph):
+    """Validate, select and compile a supported execution plan."""
+    check_graph_support(cudnn, graph)
+    deselect_frost_plans(graph)


This probably should be called in check_support, rather than here? Otherwise check_support would return "true, there's support", but here it can't execute?

cyanguwa · 2026-11-07T04:04:17Z

+def deselect_frost_plans(graph):
+    """Temporarily exclude FROST SDPA plans from the JAX C++ serialization path."""
+    # cuDNN cannot currently serialize FROST SDPA plans for C++ execution.
+    # Remove this workaround once cuDNN fixes its serialization path.


Please mention this in the PR description, together with a few other major points implemented in this PR. Thanks!

cyanguwa · 2026-11-07T04:05:25Z

    workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8)
    dq, dk, dv, _ = ffi.ffi_call(
-        "te_fused_attn_score_mod_backward_ffi",
+        "te_cudnn_graph_ffi",


Can flex attn implementation be merged onto the main SDPA path, and this file be deleted? Same on the PyTorch side.

cyanguwa · 2026-11-07T04:08:23Z

+  cudnnHandle_t value = nullptr;
+  Handle() { CheckCudnn(cudnnCreate(&value)); }
+  ~Handle() { cudnnDestroy(value); }
+};


Handle and GetHandle are already implemented in common.

cyanguwa · 2026-11-07T04:08:59Z

+  alignas(16) std::array<uint8_t, 16> value{};
+};
+
+const std::vector<Scalar> &GetScalars(ffi::Span<const uint8_t> values) {


GetScalars uses another global lock, another global cache. Can it use the same one as the main graph?

cyanguwa · 2026-11-07T04:10:54Z

            non-differentiable auxiliary inputs.
        score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or
            Python/NumPy scalars made available to `score_mod_bprop`.
+        fp8 (bool | str): True enables tensor-scaled cuDNN FP8 SDPA; "mxfp8"


We probably should mention DS, CS and MXFP8, rather than just tensor-scaled SDPA. Some of the the arch-support details can be omitted here, such as sm100 or later.

CS also has an option of using O in FP8 or F16 in bwd - please check if the new Python API supports that, because we have a test about it in test_atteniton.py.

cyanguwa · 2026-11-07T04:15:13Z

+                params["q"].shape[0] if self.q_format == "thd" else 0,
+                params["k"].shape[0] if self.kv_format == "thd" else 0,
+            )
+        self.b, self.h, self.hg, self.dq, self.dv, self.tq, self.tk = dimensions


We should bucket b as well?

cyanguwa · 2026-11-07T04:17:02Z

+        )
+        _event(self, "EXECUTE")
+
+    def build(self):


Reorder build and execute? build() first, then execute()?

cyanguwa · 2026-11-07T04:18:05Z

+            "FP8 fused attention only supports FP8DelayedScaling, FP8CurrentScaling, or MXFP8"
+            " recipes!"
+        )
+    if attn.mxfp8 and version < 92100:


Should we call mxfp8_support_error here?

cyanguwa · 2026-11-07T04:18:22Z

+        return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}."
+    try:
+        fwd = _query_graph(config, False)
+        reason = _support_reason(


Move this before _query_graph?

cyanguwa · 2026-11-07T04:19:03Z

+            return FusedAttnBackend.No_Backend, reason
+        fwd.entry()
+        if config.is_training:
+            _query_graph(config, True).entry()


Sometimes we need to check if bwd is supported even when is_training=False. main added a separate flag for it. Could you please ensure parity here?

cyanguwa · 2026-11-07T04:19:26Z

+        return FusedAttnBackend.No_Backend, str(exc)
+
+
+def fused_attn_fwd(params):


Should we retire cpp_extensions/fused_attn.py all together?

cyanguwa · 2026-11-07T04:19:59Z

+    # Import lazily: the frontend's quantized tensor classes also import this module.
+    from . import cudnn_fused_attn
+
+    out, aux_ctx_tensors, max_tensor = cudnn_fused_attn.fused_attn_fwd(


Maybe call it max_s_tensor?

cyanguwa · 2026-11-07T04:20:13Z

+        # thd (newer cuDNN runtimes, non-sm120): out [tq, h, d],    Stats [tq, h, 1],    Max [tq, h, 1]
        # thd (older cuDNN runtimes or sm120):   output_tensors: out [tq, h, d],    Stats [b, h, sq, 1], Max [b, h, sq, 1]
        # bshd:                                  output_tensors: out [b, sq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1]
        # sbhd:                                  output_tensors: out [sq, b, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1]


Maybe we can remove the other output_tensors: as well.

cyanguwa · 2026-11-07T04:22:11Z

+    from . import cudnn_fused_attn
+
+    output_tensors = cudnn_fused_attn.fused_attn_bwd(
+        {


Is this order of the args by design or can we straighten them to the FwdParams order, or group them a bit more logically? (For aesthetics/readability only)

cyanguwa · 2026-11-07T04:25:15Z

+        seed, q_cu, kv_cu, q_offsets, kv_offsets = args[5:10]
+        # JAX PRNG keys are explicit: replaying a key replays the dropout mask.
+        rng = jnp.broadcast_to(
+            jnp.concatenate((seed.reshape(-1)[:2], jnp.zeros(2, jnp.uint32))), result_avals[3].shape


The offset shouldn't be hardcoded to 0 I think.

cyanguwa · 2026-11-07T04:33:05Z

@@ -558,6 +568,7 @@ def get_attention_backend(
        if not logger.hasHandlers():


Should we cache the result of get_attention_backend so next time we don't spend so much CPU time checking the support matrix again?

cyanguwa · 2026-11-07T04:39:15Z

+def compile_graph(cudnn, graph, run_stage=None):
+    """Compile the selected plan and return its nonempty workspace size."""
+    run = run_stage or (lambda _name, function: function())
+    run("build_plans", lambda: graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE))


Would this compile again, if the plan is already compiled? Should we add a "once_only" flag like on main?

cyanguwa · 2026-11-07T04:41:23Z

+    attn.scalar("attn_scale", torch.tensor([p["attn_scale"]], dtype=torch.float32), by_value=True)
+    if "padding" in p["attn_mask_type"]:
+        for name, cu in (("seq_q", p["cu_seqlens_q"]), ("seq_kv", p["cu_seqlens_kv"])):
+            attn.add(name, cu[1:] - cu[:-1], (attn.b, 1, 1, 1), (1, 1, 1, 1))


This might be causing too many tiny kernels. Maybe switch to cu_seqlens_to_actual_seqlens?

vcherepanov-nv added 8 commits October 11, 2026 21:09

Merge branches 'py-jax' and 'py-pyt' into py-all

0c3ef6b

greptile-apps Bot reviewed Oct 19, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py Resolved

Comment thread transformer_engine/jax/cpp_extensions/cudnn_attention.py Resolved

Comment thread transformer_engine/jax/csrc/extensions/cudnn_graph.cpp Resolved

vcherepanov-nv and others added 17 commits October 20, 2026 02:15

Merge branch 'main' into py-all

4a70c52

Skip deterministic cuDNN execution tests below SM90

9e0f6e7

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>

Require precompiled cuDNN plans during JAX graph deserialization

6bb6482

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Temporarily disable FROST plans in JAX cuDNN attention

710e419

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Fix SBHD paged-cache dimensions in PyTorch attention

5ab71e0

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>

cyanguwa reviewed Nov 7, 2026

View reviewed changes

		raise TypeError("Opaque callback defaults require an explicit structural cache key.")


		def callback_cache_key(callback, is_tensor, uncacheable):



		@lru_cache(maxsize=256)
		def _supported_graph(spec, capabilities, frontend_version, backend_version):

		return FusedAttnBackend.No_Backend, str(exc)


		def fused_attn_fwd(params):

		@@ -558,6 +568,7 @@ def get_attention_backend(
		if not logger.hasHandlers():

Conversation

vcherepanov-nv commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

vcherepanov-nv commented Oct 19, 2026

Uh oh!

greptile-apps Bot commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vcherepanov-nv commented Oct 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Nov 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Nov 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa Nov 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

vcherepanov-nv commented Oct 16, 2026 •

edited

Loading

greptile-apps Bot commented Oct 19, 2026 •

edited

Loading

cyanguwa Nov 7, 2026 •

edited

Loading

cyanguwa Nov 7, 2026 •

edited

Loading

cyanguwa Nov 7, 2026 •

edited

Loading