Implementation:NVIDIA TransformerEngine Fused Attn FP8

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

Overview

Header declaring the forward and backward functions for fused attention using FP8 (E4M3/E5M2) precision, enabling reduced memory bandwidth and higher throughput on Hopper and newer GPUs.

Description

fused_attn_fp8.h declares the highest-performance attention path in TransformerEngine:

• fused_attn_fp8_fwd: Forward pass accepting separate Q, K, V tensors with FP8 quantization. Includes an input_output_S scaling state tensor for FP8 softmax statistics. Supports GQA groups.
• fused_attn_fp8_bwd: Backward pass with separate gradient outputs for dQ, dK, dV, using FP8 intermediate values (M, ZInv, S) for memory-efficient backward computation.

Both functions are guarded by CUDNN_VERSION >= 8900 and leverage FP8 tensor cores on Hopper GPUs for maximum throughput.

Usage

Used by the fused attention dispatch layer when the data type is FP8 (E4M3 or E5M2) on Hopper or newer GPUs.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/fused_attn/fused_attn_fp8.h

Lines

1--40

Signature

namespace transformer_engine {
#if (CUDNN_VERSION >= 8900)

void fused_attn_fp8_fwd(
    size_t batch, size_t num_attn_heads, size_t num_gqa_groups,
    size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim,
    bool is_training, float attn_scale, float p_dropout,
    NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
    NVTE_Mask_Type mask_type,
    const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V,
    Tensor *input_output_S, Tensor *output_O,
    NVTETensorPack *Aux_CTX_Tensors,
    const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv,
    const Tensor *rng_state, Tensor *workspace,
    cudaStream_t stream, cudnnHandle_t handle);

void fused_attn_fp8_bwd(...);

#endif
}  // namespace transformer_engine

Import

#include "fused_attn/fused_attn_fp8.h"

I/O Contract

Inputs

Name	Type	Required	Description
input_Q	const Tensor*	Yes	FP8 query tensor
input_K	const Tensor*	Yes	FP8 key tensor
input_V	const Tensor*	Yes	FP8 value tensor
input_output_S	Tensor*	Yes	FP8 scaling state for softmax
batch	size_t	Yes	Batch size

Outputs

Name	Type	Description
output_O	Tensor*	FP8 attention output
Aux_CTX_Tensors	NVTETensorPack*	Auxiliary context for backward

Usage Examples

// Called internally by the fused_attn dispatch layer for FP8
fused_attn_fp8_fwd(batch, num_attn_heads, num_gqa_groups,
    max_seqlen_q, max_seqlen_kv, head_dim,
    is_training, attn_scale, dropout,
    qkv_layout, bias_type, mask_type,
    Q_fp8, K_fp8, V_fp8, S, O, &aux,
    cu_seqlens_q, cu_seqlens_kv,
    rng_state, workspace, stream, handle);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements