Implementation:NVIDIA TransformerEngine Fused Attn Dispatch

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, Optimization
Last Updated	2026-02-07 14:00 GMT

Overview

Main dispatch layer for TransformerEngine's fused multi-head attention, routing attention forward and backward passes to the appropriate backend implementation based on sequence length, data type, and QKV layout configuration.

Description

fused_attn.cpp serves as the central entry point for all fused attention operations in TransformerEngine. It abstracts away backend selection and tensor layout complexity from higher-level framework bindings.

Key capabilities:

• QKV layout mapping: nvte_get_qkv_layout_group maps the extensive NVTE_QKV_Layout enum (25+ variants including SBHD, BSHD, THD, and paged-KV formats) to layout groups and formats.
• Backend dispatch: Routes to one of three backends:
  • Max-512 seqlen: Optimized cuDNN path for short sequences (fused_attn_max_512_fwd/bwd)
  • Arbitrary seqlen: Flash Attention via cuDNN for long sequences (fused_attn_arbitrary_seqlen_fwd/bwd)
  • FP8 attention: FP8 precision path (fused_attn_fp8_fwd/bwd)
• Tensor unpacking: Helper functions (make_tensor_view, calculate_qkv_stride, calculate_qkv_unpacked_shape) to unpack packed QKV/KV tensors by computing strides and creating tensor views with pointer offsets.

Usage

This dispatch layer is called by framework bindings (PyTorch, JAX) when invoking fused attention. Users typically interact with it through high-level APIs like te.DotProductAttention.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/fused_attn/fused_attn.cpp

Lines

1--1351

Signature

NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout);
NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout);

void nvte_fused_attn_fwd(
    NVTETensor Q, NVTETensor K, NVTETensor V, NVTETensor Bias,
    NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors,
    NVTETensor cu_seqlens_q, NVTETensor cu_seqlens_kv,
    NVTETensor rng_state, ...);

void nvte_fused_attn_bwd(
    NVTETensor Q, NVTETensor K, NVTETensor V,
    NVTETensor O, NVTETensor dO, NVTETensor S,
    NVTETensor dQ, NVTETensor dK, NVTETensor dV, ...);

Import

#include "transformer_engine/fused_attn.h"

I/O Contract

Inputs

Name	Type	Required	Description
Q	NVTETensor	Yes	Query tensor
K	NVTETensor	Yes	Key tensor
V	NVTETensor	Yes	Value tensor
qkv_layout	NVTE_QKV_Layout	Yes	Memory layout of QKV tensors
attn_scale	float	Yes	Attention scaling factor
p_dropout	float	Yes	Dropout probability

Outputs

Name	Type	Description
O	NVTETensor	Attention output tensor
Aux_CTX_Tensors	NVTETensorPack*	Auxiliary tensors saved for backward pass

Usage Examples

#include "transformer_engine/fused_attn.h"

// The dispatch layer is called via the C API
nvte_fused_attn_fwd(Q, K, V, Bias, S, O, &aux_ctx,
                    cu_seqlens_q, cu_seqlens_kv, rng_state,
                    max_seqlen_q, max_seqlen_kv,
                    is_training, attn_scale, p_dropout,
                    qkv_layout, bias_type, mask_type,
                    window_size_left, window_size_right,
                    workspace, stream);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements