The Decision Engine Behind Fused Attention: How Transformer Engine Orchestrates cuDNN on Blackwell

A modern Transformer attention call can look deceptively simple:

output = attention(q, k, v)

On an NVIDIA GPU, however, that call is rarely a single implementation decision. The result depends on the tensor data type, Q/K/V layout, query and key/value lengths, head dimensions, mask and bias features, dropout, training versus inference, CUDA Graph capture, GPU architecture, and the installed cuDNN release. Reproducibility can add one more constraint. NVIDIA Transformer Engine (TE) exists in part to make that combinatorial problem manageable.

This article follows the path from a framework-level attention request to cuDNN’s fused scaled dot-product attention (SDPA) backend. It explains what cuDNN’s Graph API actually is, how TE’s backend selector encodes a versioned compatibility policy, why fused attention is fast, what changes on Blackwell-class GPUs, and how deterministic training changes the set of eligible kernels. The source-level discussion is pinned to the TransformerEngine main-branch revision e40c5000b825a538c3f1f01e5edffb7f51b1924f; conditions may evolve in later releases. [1]

Central idea: TE is not only a collection of kernels. It is an integration and policy layer that turns an attention request into a support-constrained cuDNN Graph problem, selects a compatible backend, and leaves room for framework-level fallback when the fused route is not eligible.

Attention optimization is a dispatch problem

For each attention call, TE must answer a practical question: which implementation is valid for this exact request? That is a stricter question than “which implementation is fast?” A kernel can be excellent for one combination of precision, layout, and sequence length but invalid for another.

The selector in transformer_engine/common/fused_attn/fused_attn.cpp receives the parameters that matter to this decision: training mode; Q/K/V types and layout; bias, mask, and softmax types; dropout; numbers of attention and GQA groups; query and KV maximum sequence lengths; QK and V head dimensions; sliding-window bounds; a request for maximum logits; CUDA Graph capture; and a determinism request. [1]

This table is the most useful mental model for reading TE source: it is a compatibility matrix expressed as code, not a simple “if Blackwell, choose cuDNN” switch.

What cuDNN Graph API means — and what it does not

cuDNN’s Graph API is a declarative programming model introduced in cuDNN 8.0. Instead of asking the programmer to call a fixed function for every individual operation, it lets the programmer describe a data-flow graph of tensor operations. The graph is a mathematical specification; it is decoupled from the engine that will ultimately execute it. cuDNN can then select from one or more engine configurations that implement the graph. [2]

For an SDPA-style computation, the logical expression is:

[
O = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d}} + \text{bias/mask}\right)V.
]

In training, the graph may also include dropout state and softmax statistics needed by the backward pass. At a logical level, the execution chain is:

Q, K, V
  → QKᵀ
  → scale / bias / mask
  → softmax
  → optional dropout
  → PV
  → O (+ statistics for training)

The Graph API allows a client such as TE to describe that computation and attach semantic attributes, such as the attention scale, mask configuration, dropout inputs, and deterministic backward request. cuDNN’s attention interface exposes an SDPA backward attribute named set_deterministic_algorithm(bool) in C++, with the corresponding Python argument use_deterministic_algorithm; the documented default is False. [3]

The typical Graph API lifecycle is as follows.

cuDNN documents three heuristic styles: Mode A prioritizes low CPU latency and broad coverage, Mode B aims for more generally accurate ranking at higher CPU cost, and fallback heuristics look for functional rather than performance-optimal alternatives. Applications can additionally auto-tune candidate configurations for a specific device and problem shape. [2]

Graph API is not CUDA Graph

The similar names hide a meaningful distinction. cuDNN Graph API describes the mathematical operation graph and lets cuDNN choose an execution engine. CUDA Graph captures and replays a sequence of GPU launches to reduce CPU launch overhead. They can be used in the same application, but they solve different scheduling problems. TE’s selector includes a CUDA Graph capture flag because some cuDNN-version and shape combinations have special capture constraints; that flag does not mean the two graph systems are identical. [1]

Why fusion matters

A naïve attention implementation might materialize the score matrix, launch separate kernels for scaling, masking, softmax, dropout, and the final matrix multiplication, and repeatedly write intermediate tensors to global memory. Fused SDPA targets the opposite behavior: keep tiles on chip where possible, avoid materializing the full (S_q \times S_{kv}) score matrix, and combine compatible stages into a small number of tightly scheduled kernels.

This is where the Graph API’s separation between what is computed and how it is scheduled becomes valuable. TE supplies the required attention semantics; cuDNN may select a specialized attention engine, a precompiled implementation, or a runtime-fusion route when a graph pattern is supported. The engine’s exact implementation is deliberately opaque to the application. [2] [3]

A careful note on kernel pipelining

It is tempting to infer every detail of cuDNN SDPA from public FlashAttention papers and blogs. That would be too strong: cuDNN kernel source and its exact scheduling choices are not public. However, FlashAttention-3 is a useful public architectural analogue for understanding the hardware pressure that fused attention must manage.

FlashAttention-3 describes two complementary forms of overlap on Hopper: inter-warpgroup pipelining, which overlaps work assigned to different warp groups, and intra-warpgroup overlap using asynchronous memory movement and asynchronous matrix-multiply machinery. The article discusses TMA (Tensor Memory Accelerator) and WGMMA (Warpgroup Matrix Multiply-Accumulate) as hardware mechanisms that can overlap tile movement with computation. [4]

Illustrative attention-tile pipeline (architectural model, not cuDNN source)

TMA:     load tile n+1 ───────────────────────────────► shared memory
WGMMA:                compute QKᵀ / PV on tile n ─────►
Other work:           normalize / coordinate next tile ►

The important engineering lesson is not that every cuDNN kernel follows this exact diagram. It is that attention performance depends on keeping data movement, matrix multiplication, and normalization work from serializing unnecessarily. A backend selector must therefore recognize not just a GPU generation but also the detailed feature set that a particular engine supports.

Reading TE’s selector: FP8 first, then FP16/BF16

The selector starts with an FP8 branch. At the pinned revision, it requires FP8 Q and KV data types, SM 90 or later, no bias, and a set of versioned restrictions on layout, head dimensions, mask type, softmax type, ragged offsets, and return_max_logit. [1]

One current source comment is especially useful for comparing Hopper and Blackwell eligibility in the cuDNN 9.7 route:

// sm90: fwd d<=256, bwd d=128 only
// sm100: fwd d<=128, bwd d<=128

In the corresponding condition, Blackwell-class SM 100 uses QK and V head dimensions no larger than 128 for that route. The code also contains later version-specific cases, including a cuDNN 9.21 condition that allows QK dimensions up to 192 with V up to 128 under its stated constraints. These are source-level eligibility rules, not universal promises for all Blackwell attention configurations. [1]

If FP8 is not eligible, TE evaluates FP16/BF16 support. The central boolean is commonly named flag_arb; if every relevant condition passes, TE selects NVTE_F16_arbitrary_seqlen, the arbitrary-sequence-length fused attention backend. The conditions cover the following dimensions:

This is why upgrading cuDNN sometimes changes attention behavior without any model-code change. A release can unlock a new layout, a head-dimension case, a mask combination, or a fixed bug; TE’s source explicitly tracks those differences.

Blackwell: feature availability is a versioned contract

The TE source uses sm_arch_ >= 100 for Blackwell-class paths. Its current F16/BF16 branch illustrates why the right question is not “does Blackwell support fused attention?” but rather “does this shape and feature set qualify on this Blackwell system?”

At the pinned revision, examples include:

The same source carries a protective condition for Bottom-Right Causal Mask (BRCM) cross-attention on SM 100 with older cuDNN releases. The important operational point is that these gates are not arbitrary conservatism: they encode known support and correctness boundaries. [1]

Determinism: repeatability changes the eligible set

In ML systems, a deterministic algorithm is usually understood as one that produces the same bit pattern for the same inputs under a fixed supported execution environment. That qualification matters. Different GPU architectures, math modes, library releases, and distributed communication topologies can still alter the numerical path.

Why can GPU backward passes be non-deterministic? A common cause is a concurrent reduction. Multiple thread blocks may contribute gradient values to the same output location. A high-throughput implementation can rely on atomic accumulation, and the order in which independent contributions arrive is not fixed. Because floating-point addition is not associative, a changed accumulation order can change the final rounded result. cuDNN’s reproducibility documentation lists several operation/algorithm combinations that are not guaranteed reproducible because they use atomic operations. [5]

The conceptual trade-off is straightforward:

The exact deterministic SDPA reduction strategy used by cuDNN is not disclosed in the public documentation. It should therefore not be represented as a specific two-stage reduction, loop kernel, or fixed warp schedule. The public contract is the API attribute; the exact kernel is an implementation detail. [3]

TE’s current Blackwell gate for deterministic fused attention

TE’s selector turns this public contract into a concrete eligibility check. For the Blackwell arbitrary-sequence-length FP16/BF16 path, the pinned source comments state:

// pre-9.18.1: fwd: deterministic; bwd: non-deterministic
// 9.18.1+: fwd: deterministic; bwd: non-deterministic/deterministic

The following condition allows deterministic training only when all three constraints hold:

is_training && deterministic &&
cudnn_runtime_version >= 91801 &&
dropout == 0.0 &&
bias_type == NVTE_Bias_Type::NVTE_NO_BIAS

If that condition is not satisfied, TE does not mark the arbitrary-sequence-length fused cuDNN path as eligible. What happens next depends on the caller and installed alternatives: a framework can use another compatible backend or a non-fused implementation, while a direct low-level request may report that no fused backend applies. It is therefore more precise to say that TE removes this fused candidate than to promise a particular fallback in every environment. [1]

The selector also explicitly rejects deterministic fused attention on SM 120 at this revision. This is a reminder that SM >= 100 is not a single homogeneous feature level; precise architecture and release checks matter. [1]

A practical workflow for users of TE and cuDNN

High-performance attention problems are easier to debug when they are treated as a compatibility problem first and a performance problem second.

1. Record the environment. Log GPU compute capability, driver and CUDA versions, cuDNN version, Transformer Engine version/commit, and framework version. The selector is version-aware, so these are not incidental details.
2. Record the request. Capture dtype, Q/K/V layout, Q/KV lengths, head dimensions, GQA grouping, mask, bias, dropout, training mode, and CUDA Graph capture status.
3. Check eligibility before benchmarking. Confirm whether the requested fused path is supported. A failed fused-attention selection often reflects a valid guardrail rather than a performance regression.
4. Separate correctness mode from throughput mode. Enable deterministic algorithms for gradient checking, unit tests, regression triage, and scientific reproducibility. Benchmark non-deterministic training separately if production throughput is the goal.
5. Benchmark on the real deployment configuration. Heuristics estimate ranking, but the meaningful result is measured latency and throughput for the actual shapes, driver, GPU, and memory conditions.

Closing perspective

The deepest lesson in TE’s fused-attention implementation is architectural rather than syntactic. The high-level attention API stays stable while the backend policy evolves with new GPU generations, new cuDNN releases, expanded mask/layout support, FP8 and Blackwell-specific capabilities, and reproducibility requirements.

cuDNN Graph API provides the right division of responsibility: TE states the computation and constraints; cuDNN maps them to an execution engine. TE’s backend selector adds the systems policy that turns a large support matrix into predictable behavior. On Blackwell, that policy is increasingly important: new precision and shape capabilities coexist with version gates and deterministic-training restrictions that must be handled explicitly.

For users, the practical payoff is clear. Treat fused attention as an eligibility-aware subsystem, not a single opaque kernel. When the environment and request satisfy the contract, TE and cuDNN can expose a highly optimized route. When they do not, the correct response is to inspect the constraints, choose a compatible mode, and benchmark the resulting path—rather than assuming that a newer GPU alone guarantees every optimization.

References

1. NVIDIA, *TransformerEngine `fused_attn.cpp`*, main revision `e40c5000b825a538c3f1f01e5edffb7f51b1924f`. Source

2. NVIDIA, *cuDNN Frontend: Graphs*. Source

3. NVIDIA, *cuDNN Frontend: Attention*. Source

4. Tri Dao, *FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision*. Source

5. NVIDIA, *cuDNN Backend: Reproducibility (Determinism)*. Source

6. NVIDIA, *Transformer Engine Attention Example and Backend Guide*. Source