fused_attn.h#

Enums and functions for fused attention.

Typedefs

typedef void *NVTEFusedAttnConfig#

Opaque fused-attention configuration handle.

typedef void *NVTEFusedAttnFwdParams#

Opaque fused-attention forward-parameter handle.

typedef void *NVTEFusedAttnBwdParams#

Opaque fused-attention backward-parameter handle.

Enums

enum NVTE_QKV_Layout#

Memory layouts of QKV tensors. S, B, H, D, and T stand for sequence length, batch size, number of heads, head size, and the total number of tokens in a batch, i.e. t = sum(s_i) for i = 0...b-1. SBHD and BSHD-based layouts are used when sequences in a batch are of equal length or padded to the same length, and THD-based layouts are used when sequences have different lengths in a batch. Paged_KV-based layouts are used for paged attention.

Values:

enumerator NVTE_SB3HD#

SB3HD layout

enumerator NVTE_SBH3D#

SBH3D layout

enumerator NVTE_SBHD_SB2HD#

SBHD_SB2HD layout

enumerator NVTE_SBHD_SBH2D#

SBHD_SBH2D layout

enumerator NVTE_SBHD_SBHD_SBHD#

SBHD_SBHD_SBHD layout

enumerator NVTE_BS3HD#

BS3HD layout

enumerator NVTE_BSH3D#

BSH3D layout

enumerator NVTE_BSHD_BS2HD#

BSHD_BS2HD layout

enumerator NVTE_BSHD_BSH2D#

BSHD_BSH2D layout

enumerator NVTE_BSHD_BSHD_BSHD#

BSHD_BSHD_BSHD layout

enumerator NVTE_T3HD#

T3HD layout

enumerator NVTE_TH3D#

TH3D layout

enumerator NVTE_THD_T2HD#

THD_T2HD layout

enumerator NVTE_THD_TH2D#

THD_TH2D layout

enumerator NVTE_THD_THD_THD#

THD_THD_THD layout

enumerator NVTE_SBHD_BSHD_BSHD#

SBHD_BSHD_BSHD layout

enumerator NVTE_BSHD_SBHD_SBHD#

BSHD_SBHD_SBHD layout

enumerator NVTE_THD_BSHD_BSHD#

THD_BSHD_BSHD layout

enumerator NVTE_THD_SBHD_SBHD#

THD_SBHD_SBHD layout

enumerator NVTE_Paged_KV_BSHD_BSHD_BSHD#

Paged_KV_BSHD_BSHD_BSHD layout

enumerator NVTE_Paged_KV_BSHD_SBHD_SBHD#

Paged_KV_BSHD_SBHD_SBHD layout

enumerator NVTE_Paged_KV_SBHD_BSHD_BSHD#

Paged_KV_SBHD_BSHD_BSHD layout

enumerator NVTE_Paged_KV_SBHD_SBHD_SBHD#

Paged_KV_SBHD_SBHD_SBHD layout

enumerator NVTE_Paged_KV_THD_BSHD_BSHD#

Paged_KV_THD_BSHD_BSHD layout

enumerator NVTE_Paged_KV_THD_SBHD_SBHD#

Paged_KV_THD_SBHD_SBHD layout

enumerator NVTE_BHSD_BHSD_BHSD#

BHSD_BHSD_BHSD layout

enumerator NVTE_QKV_Layout_NOT_SET#

Not set

enum NVTE_QKV_Layout_Group#

QKV layout groups.

Values:

enumerator NVTE_3HD#

3HD QKV layouts, i.e. BS3HD, SB3HD, T3HD

enumerator NVTE_H3D#

H3D QKV layouts, i.e. BSH3D, SBH3D, TH3D

enumerator NVTE_HD_2HD#

HD_2HD QKV layouts, i.e. BSHD_BS2HD, SBHD_SB2HD, THD_T2HD

enumerator NVTE_HD_H2D#

HD_H2D QKV layouts, i.e. BSHD_BSH2D, SBHD_SBH2D, THD_TH2D

enumerator NVTE_HD_HD_HD#

HD_HD_HD QKV layouts, i.e. BSHD_BSHD_BSHD, SBHD_SBHD_SBHD, THD_THD_THD

enumerator NVTE_Paged_KV_HD_HD_HD#

Paged_KV_HD_HD_HD QKV layouts, e.g. Paged_KV_BSHD_BSHD_BSHD, Paged_KV_THD_SBHD_SBHD

enumerator NVTE_SD_SD_SD#

SD_SD_SD QKV layouts, e.g. BHSD_BHSD_BHSD

enum NVTE_QKV_Format#

QKV formats.

Values:

enumerator NVTE_SBHD#

SBHD QKV format, i.e. SB3HD, SBH3D, SBHD_SB2HD, SBHD_SBH2D, SBHD_SBHD_SBHD, Paged_KV_SBHD_SBHD_SBHD

enumerator NVTE_BSHD#

BSHD QKV format, i.e. BS3HD, BSH3D, BSHD_BS2HD, BSHD_BSH2D, BSHD_BSHD_BSHD, Paged_KV_BSHD_BSHD_BSHD

enumerator NVTE_THD#

THD QKV format, i.e. T3HD, TH3D, THD_T2HD, THD_TH2D, THD_THD_THD

enumerator NVTE_BSHD_2SBHD#

BSHD format for Q and SBHD format for KV, i.e. BSHD_SBHD_SBHD, Paged_KV_BSHD_SBHD_SBHD

enumerator NVTE_SBHD_2BSHD#

SBHD format for Q and BSHD format for KV, i.e. SBHD_BSHD_BSHD, Paged_KV_SBHD_BSHD_BSHD

enumerator NVTE_THD_2BSHD#

THD format for Q and BSHD format for KV, i.e. THD_BSHD_BSHD, Paged_KV_THD_BSHD_BSHD

enumerator NVTE_THD_2SBHD#

THD format for Q and SBHD format for KV, i.e. THD_SBHD_SBHD, Paged_KV_THD_SBHD_SBHD

enumerator NVTE_BHSD#

BHSD QKV format, e.g. BHSD_BHSD_BHSD

enumerator NVTE_QKV_Format_NOT_SET#

Not set

enum NVTE_Bias_Type#

Bias types.

Values:

enumerator NVTE_NO_BIAS#

No bias

enumerator NVTE_PRE_SCALE_BIAS#

Bias before scale

enumerator NVTE_POST_SCALE_BIAS#

Bias after scale

enumerator NVTE_ALIBI#

ALiBi

enum NVTE_Mask_Type#

Attention mask types.

Values:

enumerator NVTE_NO_MASK#

No masking

enumerator NVTE_PADDING_MASK#

Padding attention mask

enumerator NVTE_CAUSAL_MASK#

Causal attention mask (aligned to the top left corner)

enumerator NVTE_PADDING_CAUSAL_MASK#

Padding and causal attention mask (aligned to the top left corner)

enumerator NVTE_CAUSAL_BOTTOM_RIGHT_MASK#

Causal attention mask (aligned to the bottom right corner)

enumerator NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK#

Padding and causal attention mask (aligned to the bottom right corner)

enum NVTE_Softmax_Type#

Attention softmax types as described in Efficient Streaming Language Models with Attention Sinks (https://arxiv.org/pdf/2309.17453v3). For a given attention score S = Q*K^T, different softmax types perform different operations on S, NVTE_VANILLA_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), NVTE_OFF_BY_ONE_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and NVTE_LEARNABLE_SOFTMAX: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), where alpha is a learnable parameter of shape [H].

Values:

enumerator NVTE_VANILLA_SOFTMAX#

Vanilla softmax

enumerator NVTE_OFF_BY_ONE_SOFTMAX#

Off-by-one softmax

enumerator NVTE_LEARNABLE_SOFTMAX#

Learnable softmax

enum NVTE_Fused_Attn_Backend#

Fused attention backends.

Values:

enumerator NVTE_No_Backend#

No supported backend

enumerator NVTE_F16_arbitrary_seqlen#

cuDNN-based FP16/BF16 fused attention for any sequence length

enumerator NVTE_FP8#

cuDNN-based FP8 fused attention

enum NVTEFusedAttnConfigAttribute#

Attributes for NVTEFusedAttnConfig.

This enum is used to index the FusedAttnConfig struct. The order of its fields matches that of the declaration fields and the attr_sizes array of FusedAttnConfig. New fields may be appended at the end, and existing fields are never reordered, removed, or resized.

Values:

enumerator kNVTEFusedAttnConfigIsTraining#
enumerator kNVTEFusedAttnConfigDeterministic#
enumerator kNVTEFusedAttnConfigCudaGraph#
enumerator kNVTEFusedAttnConfigReturnMaxLogit#
enumerator kNVTEFusedAttnConfigAttnMaskType#
enumerator kNVTEFusedAttnConfigBiasType#
enumerator kNVTEFusedAttnConfigWindowSizeLeft#
enumerator kNVTEFusedAttnConfigWindowSizeRight#
enumerator kNVTEFusedAttnConfigBottomRightDiagonal#
enumerator kNVTEFusedAttnConfigSoftmaxType#
enumerator kNVTEFusedAttnConfigScalingMode#
enumerator kNVTEFusedAttnConfigDropout#
enumerator kNVTEFusedAttnConfigAttnScale#
enumerator kNVTEFusedAttnConfigQKVDtype#
enumerator kNVTEFusedAttnConfigODtype#
enumerator kNVTEFusedAttnConfigDODtype#
enumerator kNVTEFusedAttnConfigDQKVDtype#
enumerator kNVTEFusedAttnConfigQKVLayout#
enumerator kNVTEFusedAttnConfigOFormat#
enumerator kNVTEFusedAttnConfigDOFormat#
enumerator kNVTEFusedAttnConfigDQKVLayout#
enumerator kNVTEFusedAttnConfigQKVScaleInvFormat#
enumerator kNVTEFusedAttnConfigDOScaleInvFormat#
enumerator kNVTEFusedAttnConfigBatchSize#
enumerator kNVTEFusedAttnConfigNumAttnHeads#
enumerator kNVTEFusedAttnConfigNumGQAGroups#
enumerator kNVTEFusedAttnConfigHeadDimQK#
enumerator kNVTEFusedAttnConfigHeadDimV#
enumerator kNVTEFusedAttnConfigMaxSeqlenQ#
enumerator kNVTEFusedAttnConfigMaxSeqlenKV#
enumerator kNVTEFusedAttnConfigNumTokensQ#
enumerator kNVTEFusedAttnConfigNumTokensKV#
enumerator kNVTEFusedAttnConfigNumPagesK#
enumerator kNVTEFusedAttnConfigNumPagesV#
enumerator kNVTEFusedAttnConfigPageSizeK#
enumerator kNVTEFusedAttnConfigPageSizeV#
enumerator kNVTEFusedAttnConfigMaxPagesPerSeqK#
enumerator kNVTEFusedAttnConfigMaxPagesPerSeqV#
enumerator kNVTEFusedAttnConfigBiasBatchSize#
enumerator kNVTEFusedAttnConfigBiasNumHeads#
enumerator kNVTEFusedAttnConfigBiasSeqlenQ#
enumerator kNVTEFusedAttnConfigBiasSeqlenKV#
enumerator kNVTEFusedAttnConfigNumAttributes#
enum NVTEFusedAttnFwdParamsAttribute#

Attributes for NVTEFusedAttnFwdParams.

This enum is used to index the FusedAttnFwdParams struct. The order of its fields matches that of the declaration fields and the attr_sizes array of FusedAttnFwdParams. New fields may be appended at the end, and existing fields are never reordered, removed, or resized.

Values:

enumerator kNVTEFusedAttnFwdParamsQ#
enumerator kNVTEFusedAttnFwdParamsK#
enumerator kNVTEFusedAttnFwdParamsV#
enumerator kNVTEFusedAttnFwdParamsBias#
enumerator kNVTEFusedAttnFwdParamsSoftmaxOffset#
enumerator kNVTEFusedAttnFwdParamsS#
enumerator kNVTEFusedAttnFwdParamsO#
enumerator kNVTEFusedAttnFwdParamsAuxCtxTensors#
enumerator kNVTEFusedAttnFwdParamsCuSeqlensQ#
enumerator kNVTEFusedAttnFwdParamsCuSeqlensKV#
enumerator kNVTEFusedAttnFwdParamsCuSeqlensQPadded#
enumerator kNVTEFusedAttnFwdParamsCuSeqlensKVPadded#
enumerator kNVTEFusedAttnFwdParamsPageTableK#
enumerator kNVTEFusedAttnFwdParamsPageTableV#
enumerator kNVTEFusedAttnFwdParamsRngState#
enumerator kNVTEFusedAttnFwdParamsIsTraining#
enumerator kNVTEFusedAttnFwdParamsCudaGraph#
enumerator kNVTEFusedAttnFwdParamsReturnMaxLogit#
enumerator kNVTEFusedAttnFwdParamsAttnMaskType#
enumerator kNVTEFusedAttnFwdParamsBiasType#
enumerator kNVTEFusedAttnFwdParamsWindowSizeLeft#
enumerator kNVTEFusedAttnFwdParamsWindowSizeRight#
enumerator kNVTEFusedAttnFwdParamsBottomRightDiagonal#
enumerator kNVTEFusedAttnFwdParamsSoftmaxType#
enumerator kNVTEFusedAttnFwdParamsDropout#
enumerator kNVTEFusedAttnFwdParamsAttnScale#
enumerator kNVTEFusedAttnFwdParamsQKVLayout#
enumerator kNVTEFusedAttnFwdParamsOFormat#
enumerator kNVTEFusedAttnFwdParamsQKVScaleInvFormat#
enumerator kNVTEFusedAttnFwdParamsMaxSeqlenQ#
enumerator kNVTEFusedAttnFwdParamsMaxSeqlenKV#
enumerator kNVTEFusedAttnFwdParamsWorkspace#
enumerator kNVTEFusedAttnFwdParamsStream#
enumerator kNVTEFusedAttnFwdParamsNumAttributes#
enum NVTEFusedAttnBwdParamsAttribute#

Attributes for NVTEFusedAttnBwdParams.

This enum is used to index the FusedAttnBwdParams struct. The order of its fields matches that of the declaration fields and the attr_sizes array of FusedAttnBwdParams. New fields may be appended at the end, and existing fields are never reordered, removed, or resized.

Values:

enumerator kNVTEFusedAttnBwdParamsQ#
enumerator kNVTEFusedAttnBwdParamsK#
enumerator kNVTEFusedAttnBwdParamsV#
enumerator kNVTEFusedAttnBwdParamsO#
enumerator kNVTEFusedAttnBwdParamsDO#
enumerator kNVTEFusedAttnBwdParamsS#
enumerator kNVTEFusedAttnBwdParamsDP#
enumerator kNVTEFusedAttnBwdParamsAuxCtxTensors#
enumerator kNVTEFusedAttnBwdParamsDQ#
enumerator kNVTEFusedAttnBwdParamsDK#
enumerator kNVTEFusedAttnBwdParamsDV#
enumerator kNVTEFusedAttnBwdParamsDBias#
enumerator kNVTEFusedAttnBwdParamsDSoftmaxOffset#
enumerator kNVTEFusedAttnBwdParamsCuSeqlensQ#
enumerator kNVTEFusedAttnBwdParamsCuSeqlensKV#
enumerator kNVTEFusedAttnBwdParamsCuSeqlensQPadded#
enumerator kNVTEFusedAttnBwdParamsCuSeqlensKVPadded#
enumerator kNVTEFusedAttnBwdParamsDeterministic#
enumerator kNVTEFusedAttnBwdParamsCudaGraph#
enumerator kNVTEFusedAttnBwdParamsAttnMaskType#
enumerator kNVTEFusedAttnBwdParamsBiasType#
enumerator kNVTEFusedAttnBwdParamsWindowSizeLeft#
enumerator kNVTEFusedAttnBwdParamsWindowSizeRight#
enumerator kNVTEFusedAttnBwdParamsBottomRightDiagonal#
enumerator kNVTEFusedAttnBwdParamsSoftmaxType#
enumerator kNVTEFusedAttnBwdParamsDropout#
enumerator kNVTEFusedAttnBwdParamsAttnScale#
enumerator kNVTEFusedAttnBwdParamsQKVLayout#
enumerator kNVTEFusedAttnBwdParamsOFormat#
enumerator kNVTEFusedAttnBwdParamsDOFormat#
enumerator kNVTEFusedAttnBwdParamsDQKVLayout#
enumerator kNVTEFusedAttnBwdParamsQKVScaleInvFormat#
enumerator kNVTEFusedAttnBwdParamsDOScaleInvFormat#
enumerator kNVTEFusedAttnBwdParamsMaxSeqlenQ#
enumerator kNVTEFusedAttnBwdParamsMaxSeqlenKV#
enumerator kNVTEFusedAttnBwdParamsWorkspace#
enumerator kNVTEFusedAttnBwdParamsStream#
enumerator kNVTEFusedAttnBwdParamsNumAttributes#

Functions

NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(
NVTE_QKV_Layout qkv_layout,
)#

Get QKV layout group for a given QKV layout.

Parameters:

qkv_layout – [in] QKV layout, e.g. sbh3d.

Returns:

qkv layout group, e.g. h3d.

NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout)#

Get QKV format for a given QKV layout.

Parameters:

qkv_layout – [in] QKV layout, e.g. sbh3d.

Returns:

qkv format, e.g. sbhd.

NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout)#

Get Q format for a given QKV layout.

Parameters:

qkv_layout – [in] QKV layout, e.g. sbhd_bshd_bshd.

Returns:

q format, e.g. sbhd.

NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout)#

Get KV format for a given QKV layout.

Parameters:

qkv_layout – [in] QKV layout, e.g. sbhd_bshd_bshd.

Returns:

kv format, e.g. bshd.

NVTEFusedAttnConfig nvte_create_fused_attn_config(void)#

Create a fused-attention configuration.

void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config)#

Destroy a fused-attention configuration.

void nvte_get_fused_attn_config_attribute(
NVTEFusedAttnConfig config,
NVTEFusedAttnConfigAttribute attr,
void *buf,
size_t size_in_bytes,
size_t *size_written,
)#

Query an attribute in a fused-attention configuration.

void nvte_set_fused_attn_config_attribute(
NVTEFusedAttnConfig config,
NVTEFusedAttnConfigAttribute attr,
const void *buf,
size_t size_in_bytes,
)#

Set an attribute in a fused-attention configuration.

NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void)#

Create a fused-attention forward-parameter object.

void nvte_destroy_fused_attn_fwd_params(
NVTEFusedAttnFwdParams params,
)#

Destroy a fused-attention forward-parameter object.

void nvte_get_fused_attn_fwd_params_attribute(
NVTEFusedAttnFwdParams params,
NVTEFusedAttnFwdParamsAttribute attr,
void *buf,
size_t size_in_bytes,
size_t *size_written,
)#

Query an attribute in a fused-attention forward-parameter object.

void nvte_set_fused_attn_fwd_params_attribute(
NVTEFusedAttnFwdParams params,
NVTEFusedAttnFwdParamsAttribute attr,
const void *buf,
size_t size_in_bytes,
)#

Set an attribute in a fused-attention forward-parameter object.

NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void)#

Create a fused-attention backward-parameter object.

void nvte_destroy_fused_attn_bwd_params(
NVTEFusedAttnBwdParams params,
)#

Destroy a fused-attention backward-parameter object.

void nvte_get_fused_attn_bwd_params_attribute(
NVTEFusedAttnBwdParams params,
NVTEFusedAttnBwdParamsAttribute attr,
void *buf,
size_t size_in_bytes,
size_t *size_written,
)#

Query an attribute in a fused-attention backward-parameter object.

void nvte_set_fused_attn_bwd_params_attribute(
NVTEFusedAttnBwdParams params,
NVTEFusedAttnBwdParamsAttribute attr,
const void *buf,
size_t size_in_bytes,
)#

Set an attribute in a fused-attention backward-parameter object.

NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(
NVTEFusedAttnConfig cfg,
const char **message,
)#

Get fused-attention backend based on the user configuration.

Parameters:

  • cfg – [in] Fused-attention configuration created by nvte_create_fused_attn_config().

  • message – [out] If the configuration is supported, an empty string. If not supported, a diagnostic message explaining why the configuration is rejected. Pass NULL to skip the diagnostics. The library maintains a per-thread buffer, and callers do not need to allocate memory for the message. The buffer content stays valid until the next call of nvte_get_fused_attn_backend_v2.

Returns:

Fused-attention backend, NVTE_F16_arbitrary_seqlen or NVTE_FP8, if the configuration is supported; otherwise, NVTE_No_Backend.

NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
bool is_training,
NVTEDType q_dtype,
NVTEDType kv_dtype,
NVTE_QKV_Layout qkv_layout,
NVTE_Bias_Type bias_type,
NVTE_Mask_Type attn_mask_type,
NVTE_Softmax_Type softmax_type,
float dropout,
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,
int64_t window_size_left,
int64_t window_size_right,
bool return_max_logit,
bool cuda_graph,
bool deterministic,
)#

Get fused attention backend based on input parameters.

Deprecated:

This function has been deprecated in favor of nvte_get_fused_attn_backend_v2.

Note

nvte_get_fused_attn_backend has a narrower input signature compared to nvte_get_fused_attn_backend_v2. For the fields it cannot express, it fills them with the default values of nvte_get_fused_attn_backend_v2. This includes setting batch_size = 1, deriving output/gradient formats from qkv_layout, assuming a standard bias shape [b, h, sq, skv] for NVTE_POST_SCALE_BIAS, using delayed scaling for all FP8, and not supporting paged-KV attention checks. Users who need more precise control should use nvte_get_fused_attn_backend_v2 instead.

Parameters:

  • is_training – [in] Whether the model is in training mode.

  • q_dtype – [in] The data type of Tensor Q.

  • kv_dtype – [in] The data type of Tensors K, V.

  • qkv_layout – [in] The layout of Tensors Q, K, V.

  • bias_type – [in] The attention bias type.

  • attn_mask_type – [in] The attention mask type.

  • softmax_type – [in] The attention softmax type.

  • dropout – [in] The dropout probability.

  • num_attn_heads – [in] The number of heads in Q.

  • num_gqa_groups – [in] The number of heads in K, V.

  • max_seqlen_q – [in] The sequence length of Q.

  • max_seqlen_kv – [in] The sequence length of K, V.

  • head_dim_qk – [in] The head dimension of Q, K.

  • head_dim_v – [in] The head dimension of V.

  • window_size_left – [in] Sliding window size (the left half).

  • window_size_right – [in] Sliding window size (the right half).

  • return_max_logit – [in] Whether to produce Max along with Stats.

  • cuda_graph – [in] Whether cuda graph capture is enabled or not.

  • deterministic – [in] Whether determinism is required or not.

void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params)#

Compute dot product attention with Q, K, and V.

All inputs and outputs are carried by the opaque params handle. Create it with nvte_create_fused_attn_fwd_params(), populate it with nvte_set_fused_attn_fwd_params_attribute() or FusedAttnFwdParamsWrapper setters, and destroy it with nvte_destroy_fused_attn_fwd_params().

Parameters:

params – [inout] Fused-attention forward-parameter handle.

void nvte_fused_attn_fwd(
const NVTETensor Q,
const NVTETensor K,
const NVTETensor V,
const NVTETensor Bias,
const NVTETensor SoftmaxOffset,
NVTETensor S,
NVTETensor O,
NVTETensorPack *Aux_CTX_Tensors,
const NVTETensor cu_seqlens_q,
const NVTETensor cu_seqlens_kv,
const NVTETensor cu_seqlens_q_padded,
const NVTETensor cu_seqlens_kv_padded,
const NVTETensor page_table_k,
const NVTETensor page_table_v,
const NVTETensor rng_state,
size_t max_seqlen_q,
size_t max_seqlen_kv,
bool is_training,
bool return_max_logit,
bool cuda_graph,
float attn_scale,
float dropout,
NVTE_QKV_Layout qkv_layout,
NVTE_QKV_Format o_format,
NVTE_QKV_Format qkv_scale_inv_format,
NVTE_Bias_Type bias_type,
NVTE_Mask_Type attn_mask_type,
NVTE_Softmax_Type softmax_type,
int64_t window_size_left,
int64_t window_size_right,
bool bottom_right_diagonal,
NVTETensor workspace,
cudaStream_t stream,
)#

Compute dot product attention with separate Q, K and V.

Computes:

  • P = Q * Transpose(K) + Bias

  • S = ScaleMaskSoftmax(P)

  • D = Dropout(S)

  • O = D * Transpose(V)

Notes:

Tensors cu_seqlens_q_padded and cu_seqlens_kv_padded help identify the correct offsets of different sequences in tensors Q, K, V and O. When the QKV format (nvte_get_qkv_format(qkv_layout)) is bshd or sbhd, offset tensors are not used in the attention calculation and can be set to empty NVTETensors. When the QKV format is thd, these tensors should follow the following rules. When there is no padding between sequences, the offset tensors should be equal to cu_seqlens_q and cu_seqlens_kv respectively. When there is padding between sequences, users are responsible to adjust the offsets as needed. For example, a tensor of 4 sequences [a, PAD, b, b, c, PAD, PAD, d, d] should have cu_seqlens = [0, 1, 3, 4, 6] and cu_seqlens_padded= [0, 2, 4, 7, 9].

Deprecated:

This function has been deprecated in favor of nvte_fused_attn_fwd_v2.

Parameters:

  • Q – [in] The Q tensor.

  • K – [in] The K tensor.

  • V – [in] The V tensor.

  • Bias – [in] The Bias tensor.

  • SoftmaxOffset – [in] The SoftmaxOffset tensor.

  • S – [inout] The S tensor.

  • O – [out] The output O tensor.

  • Aux_CTX_Tensors – [out] Auxiliary output tensors when training, e.g. softmax stats, optional Max, rng_state.

  • cu_seqlens_q – [in] Cumulative sequence lengths for Q, [batch_size + 1].

  • cu_seqlens_kv – [in] Cumulative sequence lengths for K and V, [batch_size + 1].

  • cu_seqlens_q_padded – [in] Cumulative sequence offsets for Q, [batch_size + 1].

  • cu_seqlens_kv_padded – [in] Cumulative sequence offsets for KV, [batch_size + 1].

  • page_table_k – [in] Page table for K cache, [batch_size, max_pages_per_seq_k].

  • page_table_v – [in] Page table for V cache, [batch_size, max_pages_per_seq_v].

  • rng_state – [in] Seed and offset of CUDA random number generator.

  • max_seqlen_q – [in] Max sequence length used for computing for Q. it may be >= max(seqlen_q_i) for i=0,…batch_size-1.

  • max_seqlen_kv – [in] Max sequence length used for computing for K and V. it may be >= max(seqlen_kv_i) for i=0,…batch_size-1.

  • is_training – [in] Whether this is in training mode or inference.

  • return_max_logit – [in] Whether to produce Max along with Stats.

  • cuda_graph – [in] Whether cuda graph capture is enabled or not.

  • attn_scale – [in] Scaling factor for Q * K.T.

  • dropout – [in] Dropout probability.

  • qkv_layout – [in] QKV tensors’ layout.

  • o_format – [in] Output format.

  • qkv_scale_inv_format – [in] Format of scale-inverse tensors for QKV; if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout.

  • bias_type – [in] Bias type.

  • attn_mask_type – [in] Attention mask type.

  • softmax_type – [in] Attention softmax type.

  • window_size_left – [in] Sliding window size (the left half).

  • window_size_right – [in] Sliding window size (the right half).

  • bottom_right_diagonal – [in] Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix.

  • workspace – [in] Workspace tensor.

  • stream – [in] CUDA stream used for this operation.

void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params)#

Compute the backward of the dot product attention with Q, K and V.

All inputs and outputs are carried by the opaque params handle. Create it with nvte_create_fused_attn_bwd_params(), populate it with nvte_set_fused_attn_bwd_params_attribute() or FusedAttnBwdParamsWrapper setters, and destroy it with nvte_destroy_fused_attn_bwd_params().

Parameters:

params – [inout] Fused-attention backward-parameter handle.

void nvte_fused_attn_bwd(
const NVTETensor Q,
const NVTETensor K,
const NVTETensor V,
const NVTETensor O,
const NVTETensor dO,
const NVTETensor S,
NVTETensor dP,
const NVTETensorPack *Aux_CTX_Tensors,
NVTETensor dQ,
NVTETensor dK,
NVTETensor dV,
NVTETensor dBias,
NVTETensor dSoftmaxOffset,
const NVTETensor cu_seqlens_q,
const NVTETensor cu_seqlens_kv,
const NVTETensor cu_seqlens_q_padded,
const NVTETensor cu_seqlens_kv_padded,
size_t max_seqlen_q,
size_t max_seqlen_kv,
float attn_scale,
float dropout,
NVTE_QKV_Layout qkv_layout,
NVTE_QKV_Format o_format,
NVTE_QKV_Format do_format,
NVTE_QKV_Layout dqkv_layout,
NVTE_QKV_Format qkv_scale_inv_format,
NVTE_QKV_Format do_scale_inv_format,
NVTE_Bias_Type bias_type,
NVTE_Mask_Type attn_mask_type,
NVTE_Softmax_Type softmax_type,
int64_t window_size_left,
int64_t window_size_right,
bool bottom_right_diagonal,
bool deterministic,
bool cuda_graph,
NVTETensor workspace,
cudaStream_t stream,
)#

Compute the backward of the dot product attention with separate Q, K and V.

Notes:

Tensors cu_seqlens_q_padded and cu_seqlens_kv_padded help identify the correct offsets of different sequences in tensors Q, K, V and O. When the QKV format (nvte_get_qkv_format(qkv_layout)) is bshd or sbhd, offset tensors are not used in the attention calculation and can be set to empty NVTETensors. When the QKV format is thd, these tensors should follow the following rules. When there is no padding between sequences, the offset tensors should be equal to cu_seqlens_q and cu_seqlens_kv respectively. When there is padding between sequences, users are responsible to adjust the offsets as needed. For example, a tensor of 4 sequences [a, PAD, b, b, c, PAD, PAD, d, d] should have cu_seqlens = [0, 1, 3, 4, 6] and cu_seqlens_padded= [0, 2, 4, 7, 9].

Deprecated:

This function has been deprecated in favor of nvte_fused_attn_bwd_v2.

Parameters:

  • Q – [in] The Q tensor.

  • K – [in] The K tensor.

  • V – [in] The V tensor.

  • O – [in] The O tensor from forward.

  • dO – [in] The gradient of the O tensor.

  • S – [in] The S tensor.

  • dP – [inout] The gradient of the P tensor.

  • Aux_CTX_Tensors – [in] Auxiliary tensors from context when in training mode, e.g. softmax stats, optional Max, rng_state.

  • dQ – [out] The gradient of the Q tensor.

  • dK – [out] The gradient of the K tensor.

  • dV – [out] The gradient of the V tensor.

  • dBias – [out] The gradient of the Bias tensor.

  • dSoftmaxOffset – [out] The gradient of the SoftmaxOffset tensor.

  • cu_seqlens_q – [in] Cumulative sequence lengths for Q, [batch_size + 1].

  • cu_seqlens_kv – [in] Cumulative sequence lengths for K and V, [batch_size + 1].

  • cu_seqlens_q_padded – [in] Cumulative sequence offsets for Q, [batch_size + 1].

  • cu_seqlens_kv_padded – [in] Cumulative sequence offsets for KV, [batch_size + 1].

  • max_seqlen_q – [in] Max sequence length used for computing for Q. it may be >= max(seqlen_q_i) for i=0,…batch_size-1.

  • max_seqlen_kv – [in] Max sequence length used for computing for K and V. it may be >= max(seqlen_kv_i) for i=0,…batch_size-1.

  • attn_scale – [in] Scaling factor for Q * K.T.

  • dropout – [in] Dropout probability.

  • qkv_layout – [in] QKV tensors’ layout.

  • o_format – [in] Output format.

  • do_format – [in] Output gradient’s format.

  • dqkv_layout – [in] QKV gradient tensors’ layout.

  • qkv_scale_inv_format – [in] Format of scale-inverse tensors for QKV; if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout.

  • do_scale_inv_format – [in] Format of scale-inverse tensors for dO; if NVTE_QKV_Format_NOT_SET, inferred from the output layout.

  • bias_type – [in] Bias type.

  • attn_mask_type – [in] Attention mask type.

  • softmax_type – [in] Attention softmax type.

  • window_size_left – [in] Sliding window size (the left half).

  • window_size_right – [in] Sliding window size (the right half).

  • bottom_right_diagonal – [in] Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix.

  • deterministic – [in] Whether to execute with deterministic behaviours.

  • cuda_graph – [in] Whether cuda graph capture is enabled or not.

  • workspace – [in] Workspace tensor.

  • stream – [in] CUDA stream used for this operation.

void nvte_populate_rng_state_async(
NVTETensor rng_state_dst,
const NVTETensor seed,
size_t q_max_seqlen,
size_t kv_max_seqlen,
NVTE_Fused_Attn_Backend backend,
cudaStream_t stream,
)#

Update the RNG state with the seed and calculated offset.

Warning

This API is experimental and subject to change.

Parameters:

  • rng_state_dst – [in] RNG state to store seed and offset.

  • seed – [in] Seed for RNG state.

  • q_max_seqlen – [in] Max sequence length used for computing for Q. it may be >= max(seqlen_q_i) for i=0,…batch_size-1.

  • kv_max_seqlen – [in] Max sequence length used for computing for K and V. it may be >= max(seqlen_kv_i) for i=0,…batch_size-1.

  • backend – [in] Fused attention backend.

  • stream – [in] CUDA stream used for this operation.

uint32_t nvte_get_runtime_num_segments(
NVTETensor cu_seqlens,
NVTETensor workspace,
size_t len,
cudaStream_t stream,
)#

Get KV format for a given QKV layout.

Warning

This API is experimental and subject to change.

Parameters:

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • workspace – [in] Workspace tensor.

  • len – [in] batch_size x sequence_length.

  • stream – [in] CUDA stream used for this operation.

void nvte_extract_seed_and_offset(
int64_t *rng_state_ptr,
int captured,
int64_t *seed_ptr,
uint64_t seed_val,
int64_t *offset_ptr,
uint64_t offset_val,
uint32_t offset_intragraph,
cudaStream_t stream,
)#

Set the seed and offset for RNG state.

Warning

This API is experimental and subject to change.

Parameters:

  • rng_state_ptr – [out] A size 2 array storing the RNG’s seed and offset respectively.

  • captured – [in] Whether a CUDA graph is being captured.

  • seed_ptr – [in] Seed pointer.

  • seed_val – [in] Seed value.

  • offset_ptr – [in] Offset pointer.

  • offset_val – [in] Offset value.

  • offset_intragraph – [in] Intragraph offset in RNG states. For use with CUDA Graphs.

  • stream – [in] CUDA stream used for this operation.

void nvte_copy_to_kv_cache(
NVTETensor new_k,
NVTETensor new_v,
NVTETensor k_cache,
NVTETensor v_cache,
NVTETensor page_table,
NVTETensor cu_new_lens,
NVTETensor cu_cached_lens,
NVTE_QKV_Format qkv_format,
int b,
int max_ctx_len,
int max_seq_len,
int max_pages_per_seq,
int is_non_paged,
cudaStream_t stream,
)#

Copy keys and values into the KV cache.

Warning

This API is experimental and subject to change.

Parameters:

  • new_k – [in] Key tensor.

  • new_v – [in] Value tensor.

  • k_cache – [out] Key cache.

  • v_cache – [out] Value cache.

  • page_table – [in] Page table for K cache, [batch_size, max_pages_per_seq].

  • cu_new_lens – [in] Cumulative sequence lengths.

  • cu_cached_lens – [in] Cached cumulative sequence lengths.

  • qkv_format – [in] QKV format, e.g. sbhd.

  • b – [in] Batch size.

  • max_ctx_len – [in] Maximum context length.

  • max_seq_len – [in] Maximum sequence length.

  • max_pages_per_seq – [in] Maximum number of pages per sequence.

  • is_non_paged – [in] Whether the cache is paged or not.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_read_half_tensor(
const NVTETensor &tensor,
const NVTETensor &cu_seqlens,
NVTETensor half,
int half_idx,
cudaStream_t stream,
)#

Extract the first half (half_idx=0) or second half (half_idx=1) of a THD tensor.

Warning

This API is experimental and subject to change.

Parameters:

  • tensor – [in] Input tensor.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • half – [out] Output tensor.

  • half_idx – [in] Whether to read first or second half of input tensor.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_second_half_lse_correction(
NVTETensor lse,
const NVTETensor &lse_per_step,
const NVTETensor &cu_seqlens,
int lse_packed,
cudaStream_t stream,
)#

Correct the second half of the softmax LSE (LogSumExp) for context parallelism.

Warning

This API is experimental and subject to change.

Parameters:

  • lse – [out] Output tensor.

  • lse_per_step – [in] Input tensor.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • lse_packed – [in] Whether or not lse_per_step is packed.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_read_second_half_lse(
const NVTETensor &lse,
const NVTETensor &cu_seqlens,
NVTETensor half_lse,
int lse_packed,
int second_half_lse_seqlen,
cudaStream_t stream,
)#

Read the second half of the softmax LSE (LogSumExp) for context parallelism.

Warning

This API is experimental and subject to change.

Parameters:

  • lse – [in] Input tensor.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • half_lse – [out] Output tensor.

  • lse_packed – [in] Whether or the softmax LSE is in packed format.

  • second_half_lse_seqlen – [in] Sequence length.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_out_correction(
NVTETensor out,
const NVTETensor &out_per_step,
const NVTETensor &lse,
const NVTETensor &lse_per_step,
const NVTETensor &cu_seqlens,
int only_second_half,
int lse_packed,
cudaStream_t stream,
)#

Correct the THD format output of context parallelism in forward pass.

Warning

This API is experimental and subject to change.

Parameters:

  • out – [out] Output tensor.

  • out_per_step – [in] THD format output of context parallelism in forward pass.

  • lse – [in] Softmax LSE.

  • lse_per_step – [in] Softmax LSE per step.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • only_second_half – [in] Whether or not to correct only second half.

  • lse_packed – [in] Whether or the softmax LSE is in packed format.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_grad_correction(
NVTETensor grad,
const NVTETensor &grad_per_step,
const NVTETensor &cu_seqlens,
const char *first_half,
const char *second_half,
cudaStream_t stream,
)#

Update the two halves of each packed THD sequence during context-parallel backward.

first_half and second_half control how grad is updated from grad_per_step: “add” accumulates, “copy” replaces, “none” preserves, and “zero” clears. FP16, BF16, and FP32 gradients support (add, none), (none, add), (copy, none), (none, copy), (copy, zero), (zero, copy), (add, copy), and (copy, add). FP8 gradients are stored as raw encoded bytes, which this kernel cannot add numerically. They support only (copy, zero) and (zero, copy): copy preserves the FP8 values and zero clears the inactive sequence half.

Warning

This API is experimental and subject to change.

Parameters:

  • grad – [inout] Packed THD gradient to update.

  • grad_per_step – [in] Gradient from the current context-parallel step.

  • cu_seqlens – [in] Packed-sequence boundaries, [batch_size + 1].

  • first_half – [in] Operation for each sequence’s first half.

  • second_half – [in] Operation for each sequence’s second half.

  • stream – [in] CUDA stream used for this operation.

void nvte_cp_thd_get_partitioned_indices(
const NVTETensor &cu_seqlens,
NVTETensor output,
int total_tokens,
int world_size,
int rank,
cudaStream_t stream,
)#

Generate partitioned indices for inputs in THD format.

Warning

This API is experimental and subject to change.

Parameters:

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • output – [out] Output tensor.

  • total_tokens – [in] Total number of tokens.

  • world_size – [in] Total number of devices for context parallelism.

  • rank – [in] Device ID for current device.

  • stream – [in] CUDA stream used for this operation.

void nvte_thd_sequence_order_to_cp_rank_order(
const NVTETensor &inp,
const NVTETensor &cu_seqlens,
NVTETensor out,
int world_size,
int total_tokens,
cudaStream_t stream,
)#

Reorder THD tensor from sequence order to dual-chunk CP rank order.

Uses the padded THD sequence lengths to place each sequence’s two CP chunks in the order consumed by each CP rank.

Parameters:

  • inp – [in] Input THD tensor [total_tokens, …].

  • cu_seqlens – [in] Padded cumulative sequence lengths, [batch_size + 1], int32.

  • out – [out] Output tensor, same shape/dtype as inp.

  • world_size – [in] Context-parallel size.

  • total_tokens – [in] Total padded tokens (= inp.shape[0]).

  • stream – [in] CUDA stream used for this operation.

void nvte_thd_cp_rank_order_to_sequence_order(
const NVTETensor &inp,
const NVTETensor &cu_seqlens,
NVTETensor out,
int world_size,
int total_tokens,
cudaStream_t stream,
)#

Reorder THD tensor from dual-chunk CP rank order to sequence order.

Uses the padded THD sequence lengths to restore each sequence’s dual-chunk CP entries to sequence order.

Parameters:

  • inp – [in] Input THD tensor [total_tokens, …].

  • cu_seqlens – [in] Padded cumulative sequence lengths, [batch_size + 1], int32.

  • out – [out] Output tensor, same shape/dtype as inp.

  • world_size – [in] Context-parallel size.

  • total_tokens – [in] Total padded tokens (= inp.shape[0]).

  • stream – [in] CUDA stream used for this operation.

void nvte_thd_copy_valid_tokens_from_per_split_to_rank_local(
const NVTETensor &inp,
const NVTETensor &cu_seqlens_padded,
const NVTETensor &cu_seqlens,
NVTETensor out,
int total_tokens,
cudaStream_t stream,
)#

Copy valid token entries from a per-split THD tensor to a rank-local accumulator.

For each dual-chunk CP step/split, copies each sequence’s valid range at its padded THD token offsets and leaves padded entries untouched.

Parameters:

  • inp – [in] Per-split THD source tensor [total_tokens, …].

  • cu_seqlens_padded – [in] Padded cumulative sequence lengths, [batch_size + 1], int32.

  • cu_seqlens – [in] Valid cumulative sequence lengths, [batch_size + 1], int32.

  • out – [inout] Rank-local accumulator, same shape/dtype as inp.

  • total_tokens – [in] Total padded tokens (= inp.shape[0]).

  • stream – [in] CUDA stream used for this operation.

void nvte_convert_thd_to_bshd(
NVTETensor tensor,
NVTETensor cu_seqlens,
NVTETensor new_tensor,
int b,
int max_seq_len,
cudaStream_t stream,
)#

Convert tensor from THD to BSHD format.

Warning

This API is experimental and subject to change.

Parameters:

  • tensor – [in] Input tensor.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • new_tensor – [out] Output tensor.

  • b – [in] Batch size.

  • max_seq_len – [in] Maximum sequence length.

  • stream – [in] CUDA stream used for this operation.

void nvte_convert_bshd_to_thd(
NVTETensor tensor,
NVTETensor cu_seqlens,
NVTETensor new_tensor,
int t,
cudaStream_t stream,
)#

Convert tensor from BSHD to THD format.

Warning

This API is experimental and subject to change.

Parameters:

  • tensor – [in] Input tensor.

  • cu_seqlens – [in] Cumulative sequence lengths, [batch_size + 1].

  • new_tensor – [out] Output tensor.

  • t – [in] Packed sequence length.

  • stream – [in] CUDA stream used for this operation.

void nvte_prepare_flash_attn_fwd(
NVTETensor qkvi,
NVTETensor qkv,
cudaStream_t stream,
)#

Prepare QKV tensor for Flash Attention forward kernel.

Warning

This API is experimental and subject to change.

Parameters:

  • qkvi – [in] Input tensor.

  • qkv – [out] Output tensor.

  • stream – [in] CUDA stream used for this operation.

void nvte_prepare_flash_attn_bwd(
NVTETensor q,
NVTETensor k,
NVTETensor v,
NVTETensor qkv,
cudaStream_t stream,
)#

Prepare QKV tensor for Flash Attention backward kernel.

Warning

This API is experimental and subject to change.

Parameters:

  • q – [in] Input query tensor.

  • k – [in] Input key tensor.

  • v – [in] Input value tensor.

  • qkv – [out] Output tensor.

  • stream – [in] CUDA stream used for this operation.

void nvte_multi_tensor_transpose_to_bhsd(
NVTETensor *inputs,
NVTETensor *outputs,
size_t num_tensors,
NVTE_QKV_Format original_format,
cudaStream_t stream,
)#

Transpose multiple tensors from BSHD/SBHD to BHSD.

Each input tensor is 4D in BSHD or SBHD layout, and the corresponding output tensor is 4D in BHSD layout. Output tensors are pre-allocated and may have a larger last dimension.

Parameters:

  • inputs – [in] List of input tensors.

  • outputs – [inout] List of output tensors.

  • num_tensors – [in] Number of tensors in the list.

  • original_format – [in] Original QKV format (NVTE_BSHD or NVTE_SBHD).

  • stream – [in] CUDA stream.

void nvte_multi_tensor_pad_last_dim(
NVTETensor *inputs,
NVTETensor *outputs,
size_t num_tensors,
cudaStream_t stream,
)#

Pad the last dimension of multiple 2D tensors with zeros in one kernel launch.

Each tensor copies a row-major (rows, in_cols) input to a (rows, out_cols) output, zero-filling the region [in_cols, out_cols) in every row. Outputs must be pre-allocated with out_cols >= in_cols and matching dtype.

Parameters:

  • inputs – [in] List of input tensors.

  • outputs – [inout] List of output tensors.

  • num_tensors – [in] Number of tensors in the list.

  • stream – [in] CUDA stream.

class AttentionShape#
#include <fused_attn.h>

Parses a QKV tensor shape into canonical (b, h, s, d, t) dimensions and converts between QKV formats.

Public Functions

inline AttentionShape(NVTE_QKV_Format fmt, const size_t *shape)#
inline size_t b() const#
inline size_t h() const#
inline size_t s() const#
inline size_t d() const#
inline size_t t() const#
inline void to_format(
NVTE_QKV_Format dst_fmt,
size_t *dst_shape,
) const#

Private Members

size_t canonical_[5] = {}#

Private Static Functions

static inline std::pair<size_t, std::array<int, 4>> dim_order(
NVTE_QKV_Format fmt,
)#
class FusedAttnConfigWrapper#
#include <fused_attn.h>

C++ helper for constructing an NVTEFusedAttnConfig.

It owns an opaque NVTEFusedAttnConfig handle created by nvte_create_fused_attn_config(), and provides a convenient, chainable interface for setting every field in FusedAttnConfig.

Public Functions

inline FusedAttnConfigWrapper()#
FusedAttnConfigWrapper(const FusedAttnConfigWrapper&) = delete#
FusedAttnConfigWrapper &operator=(
const FusedAttnConfigWrapper&,
) = delete#
inline FusedAttnConfigWrapper(
FusedAttnConfigWrapper &&other,
) noexcept#
inline FusedAttnConfigWrapper &operator=(
FusedAttnConfigWrapper &&other,
) noexcept#
inline ~FusedAttnConfigWrapper()#
inline operator NVTEFusedAttnConfig() const noexcept#
inline NVTEFusedAttnConfig get() const noexcept#
inline FusedAttnConfigWrapper &set_is_training(bool val) noexcept#
inline FusedAttnConfigWrapper &set_deterministic(bool val) noexcept#
inline FusedAttnConfigWrapper &set_cuda_graph(bool val) noexcept#
inline FusedAttnConfigWrapper &set_return_max_logit(
bool val,
) noexcept#
inline FusedAttnConfigWrapper &set_attn_mask_type(
NVTE_Mask_Type val,
) noexcept#
inline FusedAttnConfigWrapper &set_bias_type(
NVTE_Bias_Type val,
) noexcept#
inline FusedAttnConfigWrapper &set_window_size_left(
int64_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_window_size_right(
int64_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_bottom_right_diagonal(
bool val,
) noexcept#
inline FusedAttnConfigWrapper &set_softmax_type(
NVTE_Softmax_Type val,
) noexcept#
inline FusedAttnConfigWrapper &set_scaling_mode(
NVTEScalingMode val,
) noexcept#
inline FusedAttnConfigWrapper &set_dropout(float val) noexcept#
inline FusedAttnConfigWrapper &set_attn_scale(float val) noexcept#
inline FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept#
inline FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept#
inline FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept#
inline FusedAttnConfigWrapper &set_dqkv_dtype(
NVTEDType val,
) noexcept#
inline FusedAttnConfigWrapper &set_qkv_layout(
NVTE_QKV_Layout val,
) noexcept#
inline FusedAttnConfigWrapper &set_o_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnConfigWrapper &set_do_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnConfigWrapper &set_dqkv_layout(
NVTE_QKV_Layout val,
) noexcept#
inline FusedAttnConfigWrapper &set_qkv_scale_inv_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnConfigWrapper &set_do_scale_inv_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnConfigWrapper &set_batch_size(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_num_attn_heads(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_num_gqa_groups(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_head_dim_v(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_max_seqlen_q(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_max_seqlen_kv(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_num_tokens_q(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_num_tokens_kv(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_num_pages_k(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_num_pages_v(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_page_size_k(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_page_size_v(size_t val) noexcept#
inline FusedAttnConfigWrapper &set_max_pages_per_seq_k(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_max_pages_per_seq_v(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_bias_batch_size(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_bias_num_heads(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_bias_seqlen_q(
size_t val,
) noexcept#
inline FusedAttnConfigWrapper &set_bias_seqlen_kv(
size_t val,
) noexcept#

Private Functions

template<typename T>
inline FusedAttnConfigWrapper &set_attr(
NVTEFusedAttnConfigAttribute attr,
T val,
) noexcept#

Private Members

NVTEFusedAttnConfig cfg_ = nullptr#
class FusedAttnFwdParamsWrapper#
#include <fused_attn.h>

C++ helper for constructing an NVTEFusedAttnFwdParams.

It owns an opaque NVTEFusedAttnFwdParams handle created by nvte_create_fused_attn_fwd_params(), and provides a convenient, chainable interface for setting every field in FusedAttnFwdParams.

Public Functions

inline FusedAttnFwdParamsWrapper()#
FusedAttnFwdParamsWrapper(const FusedAttnFwdParamsWrapper&) = delete#
FusedAttnFwdParamsWrapper &operator=(
const FusedAttnFwdParamsWrapper&,
) = delete#
inline FusedAttnFwdParamsWrapper(
FusedAttnFwdParamsWrapper &&other,
) noexcept#
inline FusedAttnFwdParamsWrapper &operator=(
FusedAttnFwdParamsWrapper &&other,
) noexcept#
inline ~FusedAttnFwdParamsWrapper()#
inline operator NVTEFusedAttnFwdParams() const noexcept#
inline NVTEFusedAttnFwdParams get() const noexcept#
inline FusedAttnFwdParamsWrapper &set_Q(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_K(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_V(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_SoftmaxOffset(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept#
inline FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(
NVTETensorPack *val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_cu_seqlens_q(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_page_table_k(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_page_table_v(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_rng_state(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept#
inline FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept#
inline FusedAttnFwdParamsWrapper &set_return_max_logit(
bool val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_attn_mask_type(
NVTE_Mask_Type val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_bias_type(
NVTE_Bias_Type val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_window_size_left(
int64_t val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_window_size_right(
int64_t val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(
bool val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_softmax_type(
NVTE_Softmax_Type val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept#
inline FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept#
inline FusedAttnFwdParamsWrapper &set_qkv_layout(
NVTE_QKV_Layout val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_o_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_max_seqlen_q(
size_t val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_max_seqlen_kv(
size_t val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_workspace(
NVTETensor val,
) noexcept#
inline FusedAttnFwdParamsWrapper &set_stream(
cudaStream_t val,
) noexcept#

Private Functions

template<typename T>
inline FusedAttnFwdParamsWrapper &set_attr(
NVTEFusedAttnFwdParamsAttribute attr,
T val,
) noexcept#

Private Members

NVTEFusedAttnFwdParams params_ = nullptr#
class FusedAttnBwdParamsWrapper#
#include <fused_attn.h>

C++ helper for constructing an NVTEFusedAttnBwdParams.

It owns an opaque NVTEFusedAttnBwdParams handle created by nvte_create_fused_attn_bwd_params(), and provides a convenient, chainable interface for setting every field in FusedAttnBwdParams.

Public Functions

inline FusedAttnBwdParamsWrapper()#
FusedAttnBwdParamsWrapper(const FusedAttnBwdParamsWrapper&) = delete#
FusedAttnBwdParamsWrapper &operator=(
const FusedAttnBwdParamsWrapper&,
) = delete#
inline FusedAttnBwdParamsWrapper(
FusedAttnBwdParamsWrapper &&other,
) noexcept#
inline FusedAttnBwdParamsWrapper &operator=(
FusedAttnBwdParamsWrapper &&other,
) noexcept#
inline ~FusedAttnBwdParamsWrapper()#
inline operator NVTEFusedAttnBwdParams() const noexcept#
inline NVTEFusedAttnBwdParams get() const noexcept#
inline FusedAttnBwdParamsWrapper &set_Q(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_K(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_V(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_O(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dO(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_S(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dP(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(
const NVTETensorPack *val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dK(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dV(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept#
inline FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_cu_seqlens_q(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_deterministic(
bool val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept#
inline FusedAttnBwdParamsWrapper &set_attn_mask_type(
NVTE_Mask_Type val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_bias_type(
NVTE_Bias_Type val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_window_size_left(
int64_t val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_window_size_right(
int64_t val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(
bool val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_softmax_type(
NVTE_Softmax_Type val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept#
inline FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept#
inline FusedAttnBwdParamsWrapper &set_qkv_layout(
NVTE_QKV_Layout val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_o_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_do_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_dqkv_layout(
NVTE_QKV_Layout val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_do_scale_inv_format(
NVTE_QKV_Format val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_max_seqlen_q(
size_t val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_max_seqlen_kv(
size_t val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_workspace(
NVTETensor val,
) noexcept#
inline FusedAttnBwdParamsWrapper &set_stream(
cudaStream_t val,
) noexcept#

Private Functions

template<typename T>
inline FusedAttnBwdParamsWrapper &set_attr(
NVTEFusedAttnBwdParamsAttribute attr,
T val,
) noexcept#

Private Members

NVTEFusedAttnBwdParams params_ = nullptr#
On this page