Implementation:NVIDIA TransformerEngine Fused Attn F16 Arbitrary Seqlen

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 with FP16/BF16 data types and arbitrary (greater than 512) sequence lengths, implemented via the cuDNN Flash Attention backend.

Description

fused_attn_f16_arbitrary_seqlen.h declares two functions guarded by CUDNN_VERSION >= 8900:

• fused_attn_arbitrary_seqlen_fwd: Forward pass supporting full feature set including GQA groups, sliding window attention, paged KV cache, dropout, variable-length sequences, and all QKV layouts.
• fused_attn_arbitrary_seqlen_bwd: Backward pass with deterministic mode option and gradient computation for Q, K, V, bias, and softmax offsets.

This is the most general fused attention path in TransformerEngine, supporting long sequences with the complete set of cuDNN Flash Attention features.

Usage

Used by the fused attention dispatch layer when the sequence length exceeds 512 and the data type is FP16 or BF16.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h

Lines

1--52

Signature

namespace transformer_engine {
#if (CUDNN_VERSION >= 8900)

void fused_attn_arbitrary_seqlen_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_qk, size_t head_dim_v,
    size_t num_tokens_q, size_t num_tokens_kv,
    size_t num_pages_k, size_t num_pages_v,
    size_t page_size_k, size_t page_size_v,
    bool is_training, bool return_max_logit, float attn_scale, float p_dropout,
    NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
    NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type,
    int64_t window_size_left, int64_t window_size_right,
    bool bottom_right_diagonal,
    const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V,
    const Tensor *input_Bias, const Tensor *input_SoftmaxOffset,
    Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors,
    const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv,
    const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded,
    const Tensor *page_table_k, const Tensor *page_table_v,
    const Tensor *rng_state, Tensor *workspace,
    cudaStream_t stream, cudnnHandle_t handle);

void fused_attn_arbitrary_seqlen_bwd(...);

#endif
}  // namespace transformer_engine

Import

#include "fused_attn/fused_attn_f16_arbitrary_seqlen.h"

I/O Contract

Inputs

Name	Type	Required	Description
input_Q	const Tensor*	Yes	Query tensor
input_K	const Tensor*	Yes	Key tensor
input_V	const Tensor*	Yes	Value tensor
batch	size_t	Yes	Batch size
num_attn_heads	size_t	Yes	Number of attention heads
num_gqa_groups	size_t	Yes	Number of GQA groups
handle	cudnnHandle_t	Yes	cuDNN handle

Outputs

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

Usage Examples

// Called internally by the fused_attn dispatch layer
// Not typically called directly by user code
fused_attn_arbitrary_seqlen_fwd(batch, heads, gqa_groups,
    max_seqlen_q, max_seqlen_kv, head_dim, head_dim_v,
    num_tokens_q, num_tokens_kv, 0, 0, 0, 0, 0, 0,
    is_training, false, attn_scale, dropout,
    qkv_layout, bias_type, mask_type, softmax_type,
    window_left, window_right, false,
    Q, K, V, Bias, nullptr, O, &aux,
    cu_seqlens_q, cu_seqlens_kv, nullptr, nullptr,
    nullptr, nullptr, rng_state, workspace, stream, handle);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements