Overview
This issue tracks the integration of the cuDNN Python Scaled Dot Product Attention (SDPA) API into Transformer Engine (TE), PyTorch, and JAX. The goal is a unified, high-performance interface that leverages multiple backends while maintaining flexibility for framework-level customizations and supporting both eager and compiled execution modes.
Architectural Strategy: The Intercept Layer
The core of this work is a common intercept layer shared across Torch and JAX ops. This layer deviates from the standard "graph-first" policy due to the rigid nature of SDPA operations.
Heuristics & Backend Dispatch
A small heuristics layer beneath the framework-specific ops routes requests between:
- Custom CuteDSL backends — high-performance kernels for specific shapes or features
- cuDNN pygraph — the standard cuDNN Frontend graph implementation
- Direct integrated backends — direct calls to optimized kernels (e.g., OSS SDPA)
Core Design Principles
| Principle |
Description |
| Decoupled fwd/bwd |
call_fwd_fn and call_bwd_fn are separate, enabling mix-and-match backends (e.g., cuDNN fprop + FlashAttention bprop) |
| Framework isolation |
JAX path must not depend on import torch |
| FFI path |
Low-overhead Foreign Function Interface for the cuDNN backend to minimize CPU overhead |
| Context management |
Optimized context workspaces (hyperparameters, intermediate states) can be passed fwd→bwd without backend restrictions |
Implementation Scope
Autograd Registration
Autograd registration and context passing are orthogonal. The intercept layer manages execution; framework-specific wrappers (or users) register cuDNN implementations with the respective autograd systems.
Execution Mode Compatibility
torch.eager — direct execution via the Python API intercept layertorch.compile — integration with executors (similar to Lightning Thunder) for automatic backend registration during graph capture
Current Status
cu_seqlens / ragged offsets — now supported in the Python API, exclusive to the "Unified" path- Architecture coverage — testing focused on SM80 and SM90 for both fprop and bprop; SM100+ support coming later
- D256 + THD performance gap — ongoing work to close the performance gap for this configuration
Target Timeline
| Component |
Target |
| SDPA Python API |
End of June 2026 |
| OSS SDPA Performance |
Q2 2026 |
| Next-gen GPU backend integration |
Post upcoming hardware release |
| JAX Integration |
Ongoing |
| Transformer Engine testing |
Continuous |
Next Steps
Overview
This issue tracks the integration of the cuDNN Python Scaled Dot Product Attention (SDPA) API into Transformer Engine (TE), PyTorch, and JAX. The goal is a unified, high-performance interface that leverages multiple backends while maintaining flexibility for framework-level customizations and supporting both eager and compiled execution modes.
Architectural Strategy: The Intercept Layer
The core of this work is a common intercept layer shared across Torch and JAX ops. This layer deviates from the standard "graph-first" policy due to the rigid nature of SDPA operations.
Heuristics & Backend Dispatch
A small heuristics layer beneath the framework-specific ops routes requests between:
Core Design Principles
call_fwd_fnandcall_bwd_fnare separate, enabling mix-and-match backends (e.g., cuDNN fprop + FlashAttention bprop)import torchImplementation Scope
Autograd Registration
Autograd registration and context passing are orthogonal. The intercept layer manages execution; framework-specific wrappers (or users) register cuDNN implementations with the respective autograd systems.
Execution Mode Compatibility
torch.eager— direct execution via the Python API intercept layertorch.compile— integration with executors (similar to Lightning Thunder) for automatic backend registration during graph captureCurrent Status
cu_seqlens/ ragged offsets — now supported in the Python API, exclusive to the "Unified" pathTarget Timeline
Next Steps