diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 2abd9824b6..a03f5ad9c2 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -83,6 +83,7 @@ def impl_test_self_attn( if not is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, QKVLayout.BS3HD, @@ -235,6 +236,7 @@ def test_cross_attn( if not is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, QKVLayout.BSHD_BS2HD, @@ -425,6 +427,7 @@ def impl_test_context_parallel_attn( def check_has_backend_for_mask(mask_type): return is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, qkv_layout, diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 0d1db0b9e1..1768e0227d 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -520,8 +520,9 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) - self.backend = FusedAttnHelper( + self.backend, message = FusedAttnHelper( self.is_training, + self.batch_size, self.dtype, self.dtype, self.qkv_layout, @@ -536,9 +537,10 @@ def _check_configs(self): self.head_dim_qk, self.head_dim_v, (-1, -1) if self.window_size is None else self.window_size, + self.attn_mask_type.is_bottom_right(), ).get_fused_attn_backend() if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - pytest.skip("Unsupported inputs combination or device compute capability.") + pytest.skip(message) if ( self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2dbf94fc20..a54ac1c962 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1789,12 +1789,23 @@ def test_dpa_fp8_extra_state(model, dtype): config = model_configs_fp8_extra_state[model] # Test backend availability is_training = True + fp8_recipe = recipe.DelayedScaling( + margin=0, + fp8_format=recipe.Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=True, + ) + fp8_meta = {} + fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, qkv_layout="sb3hd", is_training=is_training, deterministic=_deterministic, + fp8=True, + fp8_meta=fp8_meta, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported and not flash_attn_supported: @@ -2583,13 +2594,25 @@ def test_custom_mha_fp8_vs_f16(dtype, model): Both paths take F16 input and output. QKV layout is bs3hd""" config = model_configs_fp8[model] + os.environ["NVTE_UnfusedDPA_Emulate_FP8"] = "1" # Test backend availability is_training = True + fp8_meta = {} + fp8_recipe = recipe.DelayedScaling( + margin=0, + fp8_format=recipe.Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=True, + ) + fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, qkv_layout="bs3hd", + fp8=True, + fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, ) @@ -2667,6 +2690,7 @@ def _run_custom_mha_fp8(dtype, config, backend): fp8_format=recipe.Format.HYBRID, amax_history_len=1, amax_compute_algo="most_recent", + fp8_dpa=True, ) mha = Custom_MHA_FP8(config).to(dtype=dtype, device="cuda") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index f068f3581e..89d820ab86 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -311,6 +311,10 @@ def __init__( self.attn_type = "self" if (self.max_seqlen_q == self.max_seqlen_kv) else "cross" self.bias_shape = bias_shape self.window_size = check_set_window_size(self.attn_mask_type, window_size) + self.bottom_right_diagonal = self.attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + } self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type self.return_max_logit = return_max_logit @@ -389,6 +393,7 @@ def test(): head_dim_v=config.head_dim_v, attn_mask_type=config.attn_mask_type, window_size=config.window_size, + bottom_right_diagonal=config.bottom_right_diagonal, alibi_slopes_shape=alibi_slopes_shape, core_attention_bias_type=config.attn_bias_type, core_attention_bias_shape=core_attention_bias_shape, diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index be64fcb2be..af3fe25b61 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -181,6 +181,7 @@ list(APPEND transformer_engine_cpp_sources cudnn_utils.cpp transformer_engine.cpp fused_attn/fused_attn.cpp + fused_attn/config_and_params.cpp gemm/config.cpp normalization/common.cpp normalization/layernorm/ln_api.cpp diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp new file mode 100644 index 0000000000..c1e44e80af --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -0,0 +1,423 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "config_and_params.h" + +#include + +namespace { + +void bool_to_uint8(bool in, void *out) { + *reinterpret_cast(out) = static_cast(in); +} + +void uint8_to_bool(const void *in, bool &out) { + out = static_cast(*reinterpret_cast(in)); +} + +} // namespace + +namespace transformer_engine { + +namespace fused_attn { +// Forward declarations from fused_attn/utils.h. Declared here to avoid pulling the heavy +// cuDNN frontend header into this plain C++ translation unit. +size_t get_max_batch_size(size_t batch_size); +size_t get_max_tokens(size_t num_tokens); +} // namespace fused_attn + +void populate_fused_attn_config(FusedAttnConfig *cfg) { + NVTE_CHECK(cfg != nullptr, "FusedAttnConfig must not be NULL."); + + const int64_t b = static_cast(cfg->batch_size); + const int64_t h = static_cast(cfg->num_attn_heads); + const int64_t sq = static_cast(cfg->max_seqlen_q); + const int64_t skv = static_cast(cfg->max_seqlen_kv); + + const NVTE_QKV_Format q_format = nvte_get_q_format(cfg->qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(cfg->qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(cfg->qkv_layout); + const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + const bool has_bias = (cfg->bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + + const size_t num_tokens_q = + cfg->num_tokens_q != 0 ? cfg->num_tokens_q : static_cast(b * sq); + const size_t num_tokens_kv = + cfg->num_tokens_kv != 0 ? cfg->num_tokens_kv : static_cast(b * skv); + + // Bucket the THD (ragged) batch and token counts so the support probes and the runtime + // dispatch quantize into the same bucket, i.e. build and cache the same cuDNN graph. + const bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + const bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + cfg->bucketed_batch_size = + (is_ragged_q || is_ragged_kv) ? fused_attn::get_max_batch_size(cfg->batch_size) : 0; + cfg->bucketed_num_tokens_q = is_ragged_q ? fused_attn::get_max_tokens(num_tokens_q) : 0; + cfg->bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(num_tokens_kv) : 0; + + if (is_paged_kv) { + if (cfg->num_pages_k == 0) { + cfg->num_pages_k = static_cast(b); + } + if (cfg->num_pages_v == 0) { + cfg->num_pages_v = static_cast(b); + } + if (cfg->page_size_k == 0) { + cfg->page_size_k = static_cast(skv); + } + if (cfg->page_size_v == 0) { + cfg->page_size_v = static_cast(skv); + } + if (cfg->max_pages_per_seq_k == 0) { + cfg->max_pages_per_seq_k = 1; + } + if (cfg->max_pages_per_seq_v == 0) { + cfg->max_pages_per_seq_v = 1; + } + } + + if (has_bias) { + if (cfg->bias_batch_size == 0) { + cfg->bias_batch_size = static_cast(b); + } + if (cfg->bias_num_heads == 0) { + cfg->bias_num_heads = static_cast(h); + } + if (cfg->bias_seqlen_q == 0) { + cfg->bias_seqlen_q = static_cast(sq); + } + if (cfg->bias_seqlen_kv == 0) { + cfg->bias_seqlen_kv = static_cast(skv); + } + } +} + +} // namespace transformer_engine + +NVTEFusedAttnConfig nvte_create_fused_attn_config() { + return new transformer_engine::FusedAttnConfig( + transformer_engine::make_default_fused_attn_config()); +} + +void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config) { + delete transformer_engine::get_fused_attn_config_mutable(config); +} + +void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + + NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, + "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; + if (size_written != nullptr) { + *size_written = attr_size; + } + if (buf == nullptr) { + return; + } + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for fused attention config attribute (attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + + const auto &cfg = *get_fused_attn_config(config); + switch (attr) { + case kNVTEFusedAttnConfigIsTraining: + bool_to_uint8(cfg.is_training, buf); + break; + case kNVTEFusedAttnConfigDeterministic: + bool_to_uint8(cfg.deterministic, buf); + break; + case kNVTEFusedAttnConfigCudaGraph: + bool_to_uint8(cfg.cuda_graph, buf); + break; + case kNVTEFusedAttnConfigReturnMaxLogit: + bool_to_uint8(cfg.return_max_logit, buf); + break; + case kNVTEFusedAttnConfigQKVLayout: + std::memcpy(buf, &cfg.qkv_layout, attr_size); + break; + case kNVTEFusedAttnConfigOFormat: + std::memcpy(buf, &cfg.o_format, attr_size); + break; + case kNVTEFusedAttnConfigDOFormat: + std::memcpy(buf, &cfg.do_format, attr_size); + break; + case kNVTEFusedAttnConfigDQKVLayout: + std::memcpy(buf, &cfg.dqkv_layout, attr_size); + break; + case kNVTEFusedAttnConfigQKVScaleInvFormat: + std::memcpy(buf, &cfg.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnConfigDOScaleInvFormat: + std::memcpy(buf, &cfg.do_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnConfigBiasType: + std::memcpy(buf, &cfg.bias_type, attr_size); + break; + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(buf, &cfg.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnConfigSoftmaxType: + std::memcpy(buf, &cfg.softmax_type, attr_size); + break; + case kNVTEFusedAttnConfigScalingMode: + std::memcpy(buf, &cfg.scaling_mode, attr_size); + break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(buf, &cfg.attn_scale, attr_size); + break; + case kNVTEFusedAttnConfigDropout: + std::memcpy(buf, &cfg.dropout, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenQ: + std::memcpy(buf, &cfg.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenKV: + std::memcpy(buf, &cfg.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeLeft: + std::memcpy(buf, &cfg.window_size_left, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeRight: + std::memcpy(buf, &cfg.window_size_right, attr_size); + break; + case kNVTEFusedAttnConfigBottomRightDiagonal: + bool_to_uint8(cfg.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnConfigQKVDtype: + std::memcpy(buf, &cfg.qkv_dtype, attr_size); + break; + case kNVTEFusedAttnConfigODtype: + std::memcpy(buf, &cfg.o_dtype, attr_size); + break; + case kNVTEFusedAttnConfigDODtype: + std::memcpy(buf, &cfg.do_dtype, attr_size); + break; + case kNVTEFusedAttnConfigDQKVDtype: + std::memcpy(buf, &cfg.dqkv_dtype, attr_size); + break; + case kNVTEFusedAttnConfigBatchSize: + std::memcpy(buf, &cfg.batch_size, attr_size); + break; + case kNVTEFusedAttnConfigNumAttnHeads: + std::memcpy(buf, &cfg.num_attn_heads, attr_size); + break; + case kNVTEFusedAttnConfigNumGqaGroups: + std::memcpy(buf, &cfg.num_gqa_groups, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimQK: + std::memcpy(buf, &cfg.head_dim_qk, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimV: + std::memcpy(buf, &cfg.head_dim_v, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesK: + std::memcpy(buf, &cfg.num_pages_k, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesV: + std::memcpy(buf, &cfg.num_pages_v, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeK: + std::memcpy(buf, &cfg.page_size_k, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeV: + std::memcpy(buf, &cfg.page_size_v, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqK: + std::memcpy(buf, &cfg.max_pages_per_seq_k, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqV: + std::memcpy(buf, &cfg.max_pages_per_seq_v, attr_size); + break; + case kNVTEFusedAttnConfigBiasBatchSize: + std::memcpy(buf, &cfg.bias_batch_size, attr_size); + break; + case kNVTEFusedAttnConfigBiasNumHeads: + std::memcpy(buf, &cfg.bias_num_heads, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenQ: + std::memcpy(buf, &cfg.bias_seqlen_q, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenKV: + std::memcpy(buf, &cfg.bias_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensQ: + std::memcpy(buf, &cfg.num_tokens_q, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensKV: + std::memcpy(buf, &cfg.num_tokens_kv, attr_size); + break; + case kNVTEFusedAttnConfigBucketedBatchSize: + std::memcpy(buf, &cfg.bucketed_batch_size, attr_size); + break; + case kNVTEFusedAttnConfigBucketedNumTokensQ: + std::memcpy(buf, &cfg.bucketed_num_tokens_q, attr_size); + break; + case kNVTEFusedAttnConfigBucketedNumTokensKV: + std::memcpy(buf, &cfg.bucketed_num_tokens_kv, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + + NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, + "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for fused attention config attribute (attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + + auto &cfg = *get_fused_attn_config_mutable(config); + switch (attr) { + case kNVTEFusedAttnConfigIsTraining: + uint8_to_bool(buf, cfg.is_training); + break; + case kNVTEFusedAttnConfigDeterministic: + uint8_to_bool(buf, cfg.deterministic); + break; + case kNVTEFusedAttnConfigCudaGraph: + uint8_to_bool(buf, cfg.cuda_graph); + break; + case kNVTEFusedAttnConfigReturnMaxLogit: + uint8_to_bool(buf, cfg.return_max_logit); + break; + case kNVTEFusedAttnConfigQKVLayout: + std::memcpy(&cfg.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnConfigOFormat: + std::memcpy(&cfg.o_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDOFormat: + std::memcpy(&cfg.do_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDQKVLayout: + std::memcpy(&cfg.dqkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnConfigQKVScaleInvFormat: + std::memcpy(&cfg.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDOScaleInvFormat: + std::memcpy(&cfg.do_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasType: + std::memcpy(&cfg.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(&cfg.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigSoftmaxType: + std::memcpy(&cfg.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigScalingMode: + std::memcpy(&cfg.scaling_mode, buf, attr_size); + break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(&cfg.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnConfigDropout: + std::memcpy(&cfg.dropout, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenQ: + std::memcpy(&cfg.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenKV: + std::memcpy(&cfg.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeLeft: + std::memcpy(&cfg.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeRight: + std::memcpy(&cfg.window_size_right, buf, attr_size); + break; + case kNVTEFusedAttnConfigBottomRightDiagonal: + uint8_to_bool(buf, cfg.bottom_right_diagonal); + break; + case kNVTEFusedAttnConfigQKVDtype: + std::memcpy(&cfg.qkv_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigODtype: + std::memcpy(&cfg.o_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigDODtype: + std::memcpy(&cfg.do_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigDQKVDtype: + std::memcpy(&cfg.dqkv_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigBatchSize: + std::memcpy(&cfg.batch_size, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumAttnHeads: + std::memcpy(&cfg.num_attn_heads, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumGqaGroups: + std::memcpy(&cfg.num_gqa_groups, buf, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimQK: + std::memcpy(&cfg.head_dim_qk, buf, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimV: + std::memcpy(&cfg.head_dim_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesK: + std::memcpy(&cfg.num_pages_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesV: + std::memcpy(&cfg.num_pages_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeK: + std::memcpy(&cfg.page_size_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeV: + std::memcpy(&cfg.page_size_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqK: + std::memcpy(&cfg.max_pages_per_seq_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqV: + std::memcpy(&cfg.max_pages_per_seq_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasBatchSize: + std::memcpy(&cfg.bias_batch_size, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasNumHeads: + std::memcpy(&cfg.bias_num_heads, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenQ: + std::memcpy(&cfg.bias_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenKV: + std::memcpy(&cfg.bias_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensQ: + std::memcpy(&cfg.num_tokens_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensKV: + std::memcpy(&cfg.num_tokens_kv, buf, attr_size); + break; + case kNVTEFusedAttnConfigBucketedBatchSize: + std::memcpy(&cfg.bucketed_batch_size, buf, attr_size); + break; + case kNVTEFusedAttnConfigBucketedNumTokensQ: + std::memcpy(&cfg.bucketed_num_tokens_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigBucketedNumTokensKV: + std::memcpy(&cfg.bucketed_num_tokens_kv, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + } +} diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h new file mode 100644 index 0000000000..025d0166bf --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -0,0 +1,161 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file config_and_params.h + * \brief Internal backing objects for fused-attention config and parameter handles. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ + +#include "common/common.h" +#include "transformer_engine/fused_attn.h" + +#include + +namespace transformer_engine { + +struct FusedAttnConfig { + bool is_training = false; + bool deterministic = false; + bool cuda_graph = false; + bool return_max_logit = false; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + float attn_scale = 0.0f; + float dropout = 0.0f; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = false; + NVTEDType qkv_dtype = kNVTEFloat32; + NVTEDType o_dtype = kNVTEFloat32; + NVTEDType do_dtype = kNVTEFloat32; + NVTEDType dqkv_dtype = kNVTEFloat32; + size_t batch_size = 0; + size_t num_attn_heads = 0; + size_t num_gqa_groups = 0; + size_t head_dim_qk = 0; + size_t head_dim_v = 0; + size_t num_pages_k = 0; + size_t num_pages_v = 0; + size_t page_size_k = 0; + size_t page_size_v = 0; + size_t max_pages_per_seq_k = 0; + size_t max_pages_per_seq_v = 0; + size_t bias_batch_size = 0; + size_t bias_num_heads = 0; + size_t bias_seqlen_q = 0; + size_t bias_seqlen_kv = 0; + size_t num_tokens_q = 0; + size_t num_tokens_kv = 0; + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + + static constexpr size_t attr_sizes[] = { + sizeof(uint8_t), // is_training + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // do_format + sizeof(NVTE_QKV_Layout), // dqkv_layout + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_QKV_Format), // do_scale_inv_format + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(NVTEScalingMode), // scaling_mode + sizeof(float), // attn_scale + sizeof(float), // dropout + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTEDType), // qkv_dtype + sizeof(NVTEDType), // o_dtype + sizeof(NVTEDType), // do_dtype + sizeof(NVTEDType), // dqkv_dtype + sizeof(size_t), // batch_size + sizeof(size_t), // num_attn_heads + sizeof(size_t), // num_gqa_groups + sizeof(size_t), // head_dim_qk + sizeof(size_t), // head_dim_v + sizeof(size_t), // num_pages_k + sizeof(size_t), // num_pages_v + sizeof(size_t), // page_size_k + sizeof(size_t), // page_size_v + sizeof(size_t), // max_pages_per_seq_k + sizeof(size_t), // max_pages_per_seq_v + sizeof(size_t), // bias_batch_size + sizeof(size_t), // bias_num_heads + sizeof(size_t), // bias_seqlen_q + sizeof(size_t), // bias_seqlen_kv + sizeof(size_t), // num_tokens_q + sizeof(size_t), // num_tokens_kv + sizeof(size_t), // bucketed_batch_size + sizeof(size_t), // bucketed_num_tokens_q + sizeof(size_t), // bucketed_num_tokens_kv + }; + + bool operator<(const FusedAttnConfig &rhs) const { + return std::tie(is_training, deterministic, cuda_graph, return_max_logit, qkv_layout, o_format, + do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, + attn_mask_type, softmax_type, scaling_mode, attn_scale, dropout, max_seqlen_q, + max_seqlen_kv, window_size_left, window_size_right, bottom_right_diagonal, + qkv_dtype, o_dtype, do_dtype, dqkv_dtype, batch_size, num_attn_heads, + num_gqa_groups, head_dim_qk, head_dim_v, num_pages_k, num_pages_v, page_size_k, + page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_batch_size, + bias_num_heads, bias_seqlen_q, bias_seqlen_kv, num_tokens_q, num_tokens_kv, + bucketed_batch_size, bucketed_num_tokens_q, bucketed_num_tokens_kv) < + std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, + rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, + rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.bias_type, + rhs.attn_mask_type, rhs.softmax_type, rhs.scaling_mode, rhs.attn_scale, + rhs.dropout, rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.window_size_left, + rhs.window_size_right, rhs.bottom_right_diagonal, rhs.qkv_dtype, rhs.o_dtype, + rhs.do_dtype, rhs.dqkv_dtype, rhs.batch_size, rhs.num_attn_heads, + rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, rhs.num_pages_k, + rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, + rhs.max_pages_per_seq_v, rhs.bias_batch_size, rhs.bias_num_heads, + rhs.bias_seqlen_q, rhs.bias_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, + rhs.bucketed_batch_size, rhs.bucketed_num_tokens_q, rhs.bucketed_num_tokens_kv); + } +}; + +inline FusedAttnConfig make_default_fused_attn_config() { return FusedAttnConfig{}; } + +void populate_fused_attn_config(FusedAttnConfig *cfg); + +// Normalize cfg into the graph-cache key form used by cuDNN graph caching (ragged bucketing, +// bottom-right mask folding). Call after populate_fused_attn_config(). +FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg); + +inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { + NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); + return reinterpret_cast(config); +} + +inline FusedAttnConfig *get_fused_attn_config_mutable(NVTEFusedAttnConfig config) { + NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); + return reinterpret_cast(config); +} + +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..6e1bf518f6 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -10,6 +10,7 @@ #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" #include "../util/system.h" +#include "config_and_params.h" #include "fused_attn_f16_arbitrary_seqlen.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -225,308 +226,162 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { } } +namespace { + +// per-thread storage for the diagnostic string +// re-used (cleared + re-populated) on every call to nvte_get_fused_attn_backend_v2 on this thread +thread_local std::string fused_attn_backend_message_buffer; + +// Stash `reason` in the thread-local buffer and, if the caller asked for a diagnostic, +// publish a NUL-terminated pointer to it via `*message`. Safe to call with `message == nullptr`. +void set_message(const char **message, std::string reason) { + fused_attn_backend_message_buffer = std::move(reason); + if (message != nullptr) { + *message = fused_attn_backend_message_buffer.c_str(); + } +} + +} // namespace + // select a backend for fused attention -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) { +namespace { + +NVTE_Fused_Attn_Backend select_fused_attn_backend(const transformer_engine::FusedAttnConfig &cfg, + const char **message) { using namespace transformer_engine; - NVTE_Fused_Attn_Backend backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - auto cudnn_runtime_version = cudnnGetVersion(); + set_message(message, ""); + + cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(cfg.qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); + const auto cudnn_runtime_version = cudnnGetVersion(); - // For ragged offsets we only support 32-bit prior to cuDNN 9.5 - // Only used when THD format is requested. + // THD + 64-bit ragged offsets require cuDNN >= 9.5 const bool requires_64bit_ragged_offset = - (qkv_format == NVTE_THD && fused_attn::get_ragged_offset_dtype( - layout_group, num_attn_heads, num_gqa_groups, max_seqlen_q, - max_seqlen_kv, head_dim_qk, head_dim_v) == DType::kInt64); - const bool supported_ragged_offset_size = - (!requires_64bit_ragged_offset || cudnn_runtime_version >= 90500); - - if ((q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2) && - sm_arch_ >= 90 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - ( - // 9.2.1: {bshd, sbhd}, any seqlen, d=128, {no_mask, causal} - (cudnn_runtime_version >= 90201 && sm_arch_ < 100 && max_seqlen_q % 128 == 0 && - max_seqlen_kv % 128 == 0 && head_dim_qk == 128 && head_dim_v == 128 && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || - // 9.7: {bshd, sbhd}, any seqlen, d<=256 for sm90 and d<=128 for sm100, {padding, padding_causal} - (cudnn_runtime_version >= 90700 && - // TODO (cyang): add is_training to nvte_get_fused_attn_backend - // sm90: fwd d<=256, bwd d=128 only - // sm100: fwd d<=128, bwd d<=128 - ((sm_arch_ < 100 && (!is_training) && head_dim_qk <= 256 && head_dim_v <= 256) || - (sm_arch_ < 100 && is_training && head_dim_qk == 128 && head_dim_v == 128) || - (sm_arch_ >= 100 && head_dim_qk <= 128 && head_dim_v <= 128)) && - head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || - // 9.21: d_qk=192, d_v=128 - (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && - head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && - // pre-9.21: {bshd, sbhd}, {vanilla} - // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} - ((cudnn_runtime_version < 92100 && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || - (cudnn_runtime_version >= 92100 && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD))) && - !requires_64bit_ragged_offset && - // 9.10.0: known bugs with SDPA FP8 - (cudnn_runtime_version != 91000) && !return_max_logit) { - backend = NVTE_Fused_Attn_Backend::NVTE_FP8; - } else if ((q_dtype == NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16)) { - bool flag_arb = false; - if ( - // TODO(cyang): replace with cudnn-frontend check_support for cleaner logic and better error messaging - // architecture - ((cudnn_runtime_version < 8903 && (sm_arch_ == 80 || sm_arch_ == 90)) || - (cudnn_runtime_version >= 8903 && sm_arch_ >= 80 && sm_arch_ < 100) || - (cudnn_runtime_version >= 90700 && sm_arch_ >= 100)) && - // sequence length - ((cudnn_runtime_version < 90000 && max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0) || - (cudnn_runtime_version >= 90000)) && - // number of heads - ((cudnn_runtime_version < 8907 && num_attn_heads == num_gqa_groups) || - (cudnn_runtime_version >= 8907)) && - // head dimension - // multiples of 8 - (head_dim_qk % 8 == 0 && head_dim_v % 8 == 0 && - // <= 128 - ((head_dim_qk <= 128 && head_dim_v <= 128) || - // 9.1: <= 256 + Hopper + fprop - // 9.5: <= 256 + Hopper + bprop - (head_dim_qk <= 256 && head_dim_v <= 256 && - ((!is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90100) || - (is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90500))) || - // 9.9: any head_dim + Blackwell + fprop + non_paged + sq > 1 - (!is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 90900 && max_seqlen_q > 1 && - layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) || - // 9.10.2: any head_dim + any arch + fprop + paged - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq > 1 - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq = 1 + {no_mask, padding, BRCM, padding_BRCM} - (!is_training && cudnn_runtime_version >= 91002 && - (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD || max_seqlen_q > 1 || - (max_seqlen_q == 1 && attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) || - // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged - (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && - cudnn_runtime_version >= 91100)) && - // 9.11+ bug: 128 < d_qk <= 256, 128 < d_v <= 256 + Hopper + bprop + MLA - // Conditional to temporarily use blanket cudnn_runtime_version >= 9.11 until fixed - (!((cudnn_runtime_version >= 91100) && is_training && sm_arch_ == 90 && - head_dim_qk >= 128 && head_dim_v >= 128 && !(head_dim_qk == 192 && head_dim_v == 128) && - head_dim_qk != head_dim_v))) && - // bias type - ((cudnn_runtime_version < 8906 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS) || - (cudnn_runtime_version >= 8906 && - (bias_type == NVTE_Bias_Type::NVTE_NO_BIAS || - (bias_type == NVTE_Bias_Type::NVTE_ALIBI && - attn_mask_type != NVTE_Mask_Type::NVTE_NO_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - sm_arch_ >= 90) || - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS && sm_arch_ >= 90))) || - (cudnn_runtime_version >= 90000 && - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS && sm_arch_ >= 80))) && - // mask type - // pre-8.9.6: causal - ((cudnn_runtime_version < 8906 && attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - // 8.9.6: {bshd, sbhd} + {no_mask, causal, padding, padding_causal} - (cudnn_runtime_version >= 8906 && - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD) && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || - // 9.1: adds thd + {padding, padding_causal} - (cudnn_runtime_version >= 90100 && qkv_format == NVTE_QKV_Format::NVTE_THD && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || - // 9.3: adds {bshd, sbhd} + causal_bottom_right + self/cross-attn (sq <= skv) - (cudnn_runtime_version >= 90300 && - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD) && - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.5: adds {paged_kv_bshd, paged_kv_sbhd} + {padding, padding_causal, padding_causal_bottom_right} - (cudnn_runtime_version >= 90500 && - layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv)) && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.6: adds {bshd, sbhd, thd} + padding_causal_bottom_right + self/cross-attn (sq <= skv) - (cudnn_runtime_version >= 90600 && - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.7: removes s_q/s_kv % 64 = 0 for {causal_bottom_right, padding_causal_bottom_right} - // for any q_format/kv_format, and paged/non-paged - (cudnn_runtime_version >= 90700 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - ((attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - max_seqlen_q <= max_seqlen_kv)))) && - // bias + mask combination - (!(cudnn_runtime_version >= 8906 && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) && - bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS)) && - // qkv format - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD || - (qkv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90 && - ((cudnn_runtime_version >= 90100 && num_attn_heads == num_gqa_groups) || - cudnn_runtime_version >= 90600)) || - ((q_format == NVTE_QKV_Format::NVTE_SBHD || q_format == NVTE_QKV_Format::NVTE_BSHD || - q_format == NVTE_QKV_Format::NVTE_BHSD || - (q_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90) || - kv_format == NVTE_QKV_Format::NVTE_SBHD || kv_format == NVTE_QKV_Format::NVTE_BSHD || - kv_format == NVTE_QKV_Format::NVTE_BHSD || - (kv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90)) && - cudnn_runtime_version >= 90700)) && - // sliding window - // pre-9.2: full attn, causal - ((cudnn_runtime_version < 90200 && window_size_left == -1 && - (window_size_right == -1 || window_size_right == 0)) || - // 9.2: SWA (left, 0) + top-left diagonal + {bshd, sbhd} - (cudnn_runtime_version >= 90200 && - ((window_size_left == -1 && window_size_right == -1 && - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK) || - ((window_size_left == -1 || window_size_left >= 0) && window_size_right == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q == max_seqlen_kv)) && - max_seqlen_q <= max_seqlen_kv && dropout == 0.0 && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || - qkv_format == NVTE_QKV_Format::NVTE_SBHD)))) || - // 9.6: SWA (left, 0) + top-left/bottom-right diagonal + {bshd, sbhd, thd} - (cudnn_runtime_version >= 90600 && - ((window_size_left == -1 && (window_size_right == -1 || window_size_right == 0)) || - ((window_size_left >= 0 || window_size_left == -1) && - (window_size_right >= 0 || window_size_right == -1) && - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - // TODO(cyang): fix bug for BRCM + cross-attention on sm100 - (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && - cudnn_runtime_version <= 90700) || - cudnn_runtime_version > 90700)))) || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && - cudnn_runtime_version <= 90700) || - cudnn_runtime_version > 90700))))) && - max_seqlen_q <= max_seqlen_kv && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - dropout == 0.0)))) && - // check 64-bit ragged offset support - (supported_ragged_offset_size) && - // 9.10.0/9.10.1: known bugs with SDPA F16 - (cudnn_runtime_version != 91000) && (cudnn_runtime_version != 91001) && - // softmax type - // pre-9.13.1: vanilla - // 9.13.1+: vanilla, off-by-one, learnable - (cudnn_runtime_version >= 91301 || - (cudnn_runtime_version < 91301 && - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX)) && - // max_logit - // pre-9.21: no (the composite softmax node rejects the Stats + Max output combination) - // 9.21+: yes (Stats + Max via the unified softmax node) - (!return_max_logit || cudnn_runtime_version >= 92100) && - // determinism on Blackwell - // pre-9.18.1: fwd: deterministic; bwd: non-deterministic - // 9.18.1+: fwd: deterministic; bwd: non-deterministic/deterministic - (sm_arch_ < 100 || - (sm_arch_ >= 100 && (!is_training || - (is_training && !deterministic && - (dropout == 0.0 || bias_type == NVTE_Bias_Type::NVTE_NO_BIAS)) || - (is_training && deterministic && cudnn_runtime_version >= 91801 && - dropout == 0.0 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS))))) { - flag_arb = true; + (qkv_format == NVTE_THD && + fused_attn::get_ragged_offset_dtype(layout_group, cfg.num_attn_heads, cfg.num_gqa_groups, + cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.head_dim_qk, + cfg.head_dim_v) == DType::kInt64); + if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { + set_message(message, + "Configuration requires 64-bit ragged offsets, which require " + "cuDNN >= 9.5."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // THD requires padding-style mask + if (qkv_format == NVTE_QKV_Format::NVTE_THD && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + set_message(message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + const bool is_fp8 = (cfg.qkv_dtype == NVTEDType::kNVTEFloat8E4M3 || + cfg.qkv_dtype == NVTEDType::kNVTEFloat8E5M2); + const bool is_f16_or_bf16 = + (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); + + if (is_fp8) { + if (cfg.return_max_logit) { + set_message(message, "FP8 fused attention does not support return_max_logit=True."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (flag_arb) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + if (qkv_format != NVTE_QKV_Format::NVTE_BSHD && qkv_format != NVTE_QKV_Format::NVTE_SBHD && + qkv_format != NVTE_QKV_Format::NVTE_BHSD) { + set_message(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + "."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cudnn_runtime_version < 8900 && - backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP16/BF16 fused attention is supported by cuDNN 8.9.0+." - " Please upgrade your cuDNN version if possible." - << std::endl; + std::string fwd_reason = is_supported_fp8_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if ((cudnn_runtime_version == 91400) && (max_seqlen_kv > 1024) && (window_size_left != -1) && - (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK)) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of attention mask (non-causal) and " - "max_seqlen_kv (> 1024) does not support fused attention for cuDNN 9.14.0. " - " Please upgrade your cuDNN version if possible." - << std::endl; + if (cfg.is_training) { + std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); + if (!bwd_reason.empty()) { + set_message(message, std::move(bwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } } - if ((cudnn_runtime_version <= 91500) && is_training && + return NVTE_Fused_Attn_Backend::NVTE_FP8; + } + + if (is_f16_or_bf16) { + if (cudnn_runtime_version <= 91500 && cfg.is_training && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - (max_seqlen_kv % 128 != 0) && cuda_graph && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of attention mask (non-padding)," - " max_seqlen_kv (not divisible by 128), and qkv_format (BSHD/SBHD) for" - " backward fused attention with graph capture requires cuDNN 9.15.1+. " - "Please upgrade your cuDNN version if possible." - << std::endl; + (cfg.max_seqlen_kv % 128 != 0) && cfg.cuda_graph && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) { - if (cudnn_runtime_version < 91801) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " - "91801 is not supported. " - << " Please upgrade your cuDNN version if possible." << std::endl; - } else if (deterministic && is_training) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Deterministic fused attention on SM120 is not supported." - << std::endl; - } else { - // Known missing support for T3HD/TH3D layouts on SM120 - const bool is_t3hd_or_th3d = - (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D); - if (is_t3hd_or_th3d) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. " - << " Please consider using other THD layouts if possible." << std::endl; - } + std::string fwd_reason = is_supported_f16_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + if (cfg.is_training) { + std::string bwd_reason = is_supported_f16_bwd(cfg, handle); + if (!bwd_reason.empty()) { + set_message(message, std::move(bwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } - } else { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - return backend; + + set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; +} + +} // namespace + +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, + const char **message) { + using namespace transformer_engine; + return select_fused_attn_backend(*get_fused_attn_config(cfg), message); +} + +// Deprecated: thin wrapper preserving the historical narrow signature. New callers should +// construct an NVTEFusedAttnConfig and call nvte_get_fused_attn_backend_v2 directly to access +// the additional fields (attn_scale, format/layout fields, scaling_mode, paged-KV/bias shape, +// dO/dQKV dtypes, etc.) that this wrapper cannot express. +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) { + (void)is_training; + transformer_engine::FusedAttnConfig cfg = transformer_engine::make_default_fused_attn_config(); + cfg.qkv_layout = qkv_layout; + cfg.bias_type = bias_type; + cfg.attn_mask_type = attn_mask_type; + cfg.softmax_type = softmax_type; + cfg.attn_scale = 1.0f; // legacy default; matches the value pre-PR probes hardcoded + cfg.dropout = dropout; + cfg.max_seqlen_q = max_seqlen_q; + cfg.max_seqlen_kv = max_seqlen_kv; + cfg.window_size_left = window_size_left; + cfg.window_size_right = window_size_right; + cfg.cuda_graph = cuda_graph; + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + cfg.qkv_dtype = q_dtype; + cfg.o_dtype = q_dtype; // legacy: O dtype matches Q dtype + cfg.batch_size = 1; // legacy: pre-PR probes assumed batch=1 + cfg.num_attn_heads = num_attn_heads; + cfg.num_gqa_groups = num_gqa_groups; + cfg.head_dim_qk = head_dim_qk; + cfg.head_dim_v = head_dim_v; + cfg.is_training = false; // legacy wrapper cannot express dO/dQKV dtypes; skip bwd probe + cfg.return_max_logit = return_max_logit; + cfg.deterministic = deterministic; + return select_fused_attn_backend(cfg, /*message=*/nullptr); } // NVTE fused attention FWD with separate Q, K and V @@ -611,28 +466,72 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); const NVTEDType Q_type = static_cast(input_Q->data.dtype); const NVTEDType KV_type = static_cast(input_K->data.dtype); + NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); + const NVTEDType O_type = static_cast(output_O->data.dtype); + const NVTEScalingMode scaling_mode = input_Q->scaling_mode; + + size_t bias_b = 0, bias_h = 0, bias_sq = 0, bias_skv = 0; + if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && + input_Bias->data.dptr != nullptr && input_Bias->data.shape.size() >= 4) { + bias_b = input_Bias->data.shape[0]; + bias_h = input_Bias->data.shape[1]; + bias_sq = input_Bias->data.shape[2]; + bias_skv = input_Bias->data.shape[3]; + } - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, - return_max_logit, cuda_graph, false); + transformer_engine::FusedAttnConfig cfg = transformer_engine::make_default_fused_attn_config(); + cfg.is_training = false; // fwd-only probe; restored before dispatch + cfg.deterministic = false; + cfg.cuda_graph = cuda_graph; + cfg.return_max_logit = return_max_logit; + cfg.qkv_layout = qkv_layout; + cfg.o_format = o_format; + cfg.qkv_scale_inv_format = qkv_scale_inv_format; + cfg.bias_type = bias_type; + cfg.attn_mask_type = attn_mask_type; + cfg.softmax_type = softmax_type; + cfg.scaling_mode = scaling_mode; + cfg.attn_scale = attn_scale; + cfg.dropout = dropout; + cfg.max_seqlen_q = max_seqlen_q; + cfg.max_seqlen_kv = max_seqlen_kv; + cfg.window_size_left = window_size_left; + cfg.window_size_right = window_size_right; + cfg.bottom_right_diagonal = bottom_right_diagonal; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = O_type; + cfg.batch_size = b; + cfg.num_attn_heads = h_q; + cfg.num_gqa_groups = h_kv; + cfg.head_dim_qk = d_qk; + cfg.head_dim_v = d_v; + cfg.num_pages_k = static_cast(num_pages_k); + cfg.num_pages_v = static_cast(num_pages_v); + cfg.page_size_k = static_cast(page_size_k); + cfg.page_size_v = static_cast(page_size_v); + cfg.max_pages_per_seq_k = static_cast(max_pages_per_seq_k); + cfg.max_pages_per_seq_v = static_cast(max_pages_per_seq_v); + cfg.bias_batch_size = bias_b; + cfg.bias_num_heads = bias_h; + cfg.bias_seqlen_q = bias_sq; + cfg.bias_seqlen_kv = bias_skv; + cfg.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + NVTE_Fused_Attn_Backend fused_attention_backend = + select_fused_attn_backend(cfg, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - fused_attn_arbitrary_seqlen_fwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, num_pages_k, num_pages_v, - page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, is_training, - return_max_logit, attn_scale, dropout, qkv_layout, o_format, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, - input_V, input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, - input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); + cfg.is_training = is_training; + fused_attn_arbitrary_seqlen_fwd(cfg, input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, + output_O, Aux_CTX_Tensors, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, + input_cu_seqlens_kv_padded, input_page_table_k, + input_page_table_v, input_rng_state, wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training, - attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, - input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + cfg.is_training = is_training; + fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, + output_O, Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, + input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -692,11 +591,60 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); const NVTEDType Q_type = static_cast(input_Q->data.dtype); const NVTEDType KV_type = static_cast(input_K->data.dtype); + NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); + const NVTEDType O_type = static_cast(input_O->data.dtype); + const NVTEDType dO_type = static_cast(input_dO->data.dtype); + const NVTEDType dQKV_type = static_cast(output_dQ->data.dtype); + const NVTEScalingMode scaling_mode = input_Q->scaling_mode; + + size_t bias_b = 0, bias_h = 0, bias_sq = 0, bias_skv = 0; + if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && + output_dBias->data.shape.size() >= 4) { + bias_b = output_dBias->data.shape[0]; + bias_h = output_dBias->data.shape[1]; + bias_sq = output_dBias->data.shape[2]; + bias_skv = output_dBias->data.shape[3]; + } - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false, - cuda_graph, deterministic); + transformer_engine::FusedAttnConfig cfg = transformer_engine::make_default_fused_attn_config(); + cfg.is_training = true; + cfg.deterministic = deterministic; + cfg.cuda_graph = cuda_graph; + cfg.return_max_logit = false; + cfg.qkv_layout = qkv_layout; + cfg.o_format = o_format; + cfg.do_format = do_format; + cfg.dqkv_layout = dqkv_layout; + cfg.qkv_scale_inv_format = qkv_scale_inv_format; + cfg.do_scale_inv_format = do_scale_inv_format; + cfg.bias_type = bias_type; + cfg.attn_mask_type = attn_mask_type; + cfg.softmax_type = softmax_type; + cfg.scaling_mode = scaling_mode; + cfg.attn_scale = attn_scale; + cfg.dropout = dropout; + cfg.max_seqlen_q = max_seqlen_q; + cfg.max_seqlen_kv = max_seqlen_kv; + cfg.window_size_left = window_size_left; + cfg.window_size_right = window_size_right; + cfg.bottom_right_diagonal = bottom_right_diagonal; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = O_type; + cfg.do_dtype = dO_type; + cfg.dqkv_dtype = dQKV_type; + cfg.batch_size = b; + cfg.num_attn_heads = h_q; + cfg.num_gqa_groups = h_kv; + cfg.head_dim_qk = d_qk; + cfg.head_dim_v = d_v; + cfg.bias_batch_size = bias_b; + cfg.bias_num_heads = bias_h; + cfg.bias_seqlen_q = bias_sq; + cfg.bias_seqlen_kv = bias_skv; + cfg.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + NVTE_Fused_Attn_Backend fused_attention_backend = + select_fused_attn_backend(cfg, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; @@ -709,14 +657,12 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (softmax_type != NVTE_VANILLA_SOFTMAX) { input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_arbitrary_seqlen_bwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, input_Q, input_K, - input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQ, output_dK, - output_dV, output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, - handle); + fused_attn_arbitrary_seqlen_bwd(cfg, input_Q, input_K, input_V, input_O, input_dO, input_Bias, + input_SoftmaxOffset, output_S, output_dQ, output_dK, output_dV, + output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, + input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, + handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { size_t i = 0; const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); @@ -729,13 +675,9 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, - input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, - input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, - output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + fused_attn_fp8_bwd(cfg, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, + input_S, input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, + output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index 6df7ad35c8..11bccc09f4 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -47,22 +47,52 @@ namespace transformer_engine { namespace fused_attn { + void fused_attn_arbitrary_seqlen_fwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t num_pages_k, int64_t num_pages_v, - int64_t page_size_k, int64_t page_size_v, int64_t max_pages_per_seq_k, - int64_t max_pages_per_seq_v, int64_t bias_b, int64_t bias_h, int64_t bias_sq, int64_t bias_skv, - bool is_training, bool return_max_logit, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, 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, void *devPtrQ, void *devPtrK, - void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, - void *devPtrO, void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, + const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, + void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, void *devPtrO, + void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + + int64_t b = static_cast(cfg.batch_size); + int64_t h = static_cast(cfg.num_attn_heads); + int64_t hg = static_cast(cfg.num_gqa_groups); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.max_seqlen_kv); + int64_t d_qk = static_cast(cfg.head_dim_qk); + int64_t d_v = static_cast(cfg.head_dim_v); + int64_t bucketed_batch_size = static_cast(cfg.bucketed_batch_size); + int64_t bucketed_num_tokens_q = static_cast(cfg.bucketed_num_tokens_q); + int64_t bucketed_num_tokens_kv = static_cast(cfg.bucketed_num_tokens_kv); + int64_t num_pages_k = static_cast(cfg.num_pages_k); + int64_t num_pages_v = static_cast(cfg.num_pages_v); + int64_t page_size_k = static_cast(cfg.page_size_k); + int64_t page_size_v = static_cast(cfg.page_size_v); + int64_t max_pages_per_seq_k = static_cast(cfg.max_pages_per_seq_k); + int64_t max_pages_per_seq_v = static_cast(cfg.max_pages_per_seq_v); + int64_t bias_b = static_cast(cfg.bias_batch_size); + int64_t bias_h = static_cast(cfg.bias_num_heads); + int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const bool is_training = cfg.is_training; + const bool return_max_logit = cfg.return_max_logit; + const float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + bool bottom_right_diagonal = cfg.bottom_right_diagonal; + bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -104,56 +134,16 @@ void fused_attn_arbitrary_seqlen_fwd_impl( if (sm_arch_ != 120) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + b = bucketed_batch_size; + s_q = is_ragged_q ? bucketed_num_tokens_q : s_q; + s_kv = is_ragged_kv ? bucketed_num_tokens_kv : s_kv; } } const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; bool generate_stats = true; // Always return stats + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg); try { - FADescriptor_v1 descriptor{ - b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - num_pages_k, - num_pages_v, - page_size_k, - page_size_v, - max_pages_per_seq_k, - max_pages_per_seq_v, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - is_training, - dropout_probability, - qkv_layout, - o_format, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Layout_NOT_SET, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - true, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - return_max_logit, - }; - namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, @@ -178,11 +168,11 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; + using CacheType = std::map; static thread_local CacheType sdpa_f16_fprop_cache; // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType &cache, const FADescriptor_v1 &descriptor) -> graph_and_tensors { + auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); if (it != cache.end()) { @@ -432,7 +422,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto [mha_graph, Q, K, V, attn_scale, O, S1, S2, bias, softmax_offset, seq_q, seq_kv, page_table_k, page_table_v, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = get_graph(sdpa_f16_fprop_cache, descriptor); + dropout_seed, dropout_offset] = get_graph(sdpa_f16_fprop_cache, cache_cfg); // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. @@ -552,21 +542,46 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } void fused_attn_arbitrary_seqlen_bwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t bias_b, int64_t bias_h, - int64_t bias_sq, int64_t bias_skv, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_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, bool deterministic, void *devPtrQ, void *devPtrKTranspose, - void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, - void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, - void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, - void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, + void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, void *devPtrSoftmaxOffset, + void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, void *devPtrdBias, + void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, + void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, + void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, cudaStream_t stream, + cudnnHandle_t handle) { using namespace transformer_engine; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + + int64_t b = static_cast(cfg.batch_size); + int64_t h = static_cast(cfg.num_attn_heads); + int64_t hg = static_cast(cfg.num_gqa_groups); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.max_seqlen_kv); + int64_t d_qk = static_cast(cfg.head_dim_qk); + int64_t d_v = static_cast(cfg.head_dim_v); + int64_t bucketed_batch_size = static_cast(cfg.bucketed_batch_size); + int64_t bucketed_num_tokens_q = static_cast(cfg.bucketed_num_tokens_q); + int64_t bucketed_num_tokens_kv = static_cast(cfg.bucketed_num_tokens_kv); + int64_t bias_b = static_cast(cfg.bias_batch_size); + int64_t bias_h = static_cast(cfg.bias_num_heads); + int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_QKV_Format do_format = cfg.do_format; + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -606,57 +621,17 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (sm_arch_ != 120) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + b = bucketed_batch_size; + s_q = is_ragged_q ? bucketed_num_tokens_q : s_q; + s_kv = is_ragged_kv ? bucketed_num_tokens_kv : s_kv; } } // We choose between 32-bit and 64-bit offsets depending on need. // This allows us to support older cuDNN runtimes gracefully. const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg); try { - FADescriptor_v1 descriptor{ - b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - true, - dropout_probability, - qkv_layout, - o_format, - do_format, - dqkv_layout, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - deterministic, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - false, - }; - namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, @@ -684,11 +659,11 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; + using CacheType = std::map; static thread_local CacheType sdpa_f16_bprop_cache; // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType &cache, const FADescriptor_v1 &descriptor) -> graph_and_tensors { + auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); if (it != cache.end()) { @@ -946,7 +921,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto [mha_graph, q, k, v, o, dO, stats, attn_scale, dQ, dK, dV, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = get_graph(sdpa_f16_bprop_cache, descriptor); + dropout_seed, dropout_offset] = get_graph(sdpa_f16_bprop_cache, cache_cfg); // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. @@ -1072,24 +1047,29 @@ void fused_attn_arbitrary_seqlen_bwd_impl( using namespace transformer_engine::fused_attn; 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, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, - bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, 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) { + const FusedAttnConfig &cfg, 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) { using namespace transformer_engine; + const size_t batch = cfg.batch_size; + const size_t num_attn_heads = cfg.num_attn_heads; + const size_t num_gqa_groups = cfg.num_gqa_groups; + const size_t max_seqlen_q = cfg.max_seqlen_q; + const size_t max_seqlen_kv = cfg.max_seqlen_kv; + const size_t head_dim_qk = cfg.head_dim_qk; + const size_t head_dim_v = cfg.head_dim_v; + const size_t num_tokens_q = cfg.num_tokens_q; + const bool return_max_logit = cfg.return_max_logit; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const auto QKV_type = input_Q->data.dtype; NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1123,17 +1103,13 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrPageTableK = page_table_k ? page_table_k->data.dptr : nullptr; void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); + FusedAttnConfig graph_cfg = cfg; + populate_fused_attn_config(&graph_cfg); + if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { + graph_cfg.bias_batch_size = bias_b; + graph_cfg.bias_num_heads = bias_h; + graph_cfg.bias_seqlen_q = bias_sq; + graph_cfg.bias_seqlen_kv = bias_skv; } size_t i = 0; @@ -1210,14 +1186,9 @@ void fused_attn_arbitrary_seqlen_fwd( size_t workspace_size = 0; fused_attn_arbitrary_seqlen_fwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, num_pages_k, num_pages_v, page_size_k, - page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, bias_skv, - is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, o_format, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, - devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, - devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), + graph_cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, + devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, + devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { @@ -1236,21 +1207,18 @@ void fused_attn_arbitrary_seqlen_fwd( } void fused_attn_arbitrary_seqlen_bwd( - 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, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_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, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, + const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - const auto QKV_type = input_Q->data.dtype; + + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1271,19 +1239,13 @@ void fused_attn_arbitrary_seqlen_bwd( bias_skv = output_dBias->data.shape[3]; } - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); + FusedAttnConfig graph_cfg = cfg; + populate_fused_attn_config(&graph_cfg); + if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { + graph_cfg.bias_batch_size = bias_b; + graph_cfg.bias_num_heads = bias_h; + graph_cfg.bias_seqlen_q = bias_sq; + graph_cfg.bias_seqlen_kv = bias_skv; } void *devPtrdQ = output_dQ->data.dptr; @@ -1310,14 +1272,11 @@ void fused_attn_arbitrary_seqlen_bwd( size_t workspace_size = 0; fused_attn_arbitrary_seqlen_bwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, bias_b, bias_h, bias_sq, bias_skv, attn_scale, - p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, - devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, - devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, - devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); + graph_cfg, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, + devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, + devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, + devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, + &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1333,4 +1292,55 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_ERROR("Unexpected workspace_size."); } } + +std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + populate_fused_attn_config(&graph_cfg); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, + /*devPtrSoftmaxOffset=*/nullptr, /*devPtrS1=*/nullptr, /*devPtrS2=*/nullptr, + /*devPtrO=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, + /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, + /*devPtrPageTableK=*/nullptr, /*devPtrPageTableV=*/nullptr, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception &e) { + return e.what(); + } catch (...) { + return "is_supported_f16_fwd: unknown failure."; + } +} + +std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + populate_fused_attn_config(&graph_cfg); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrKTranspose=*/nullptr, + /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, + /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, + /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, /*devPtrdO=*/nullptr, + /*devPtrdBias=*/nullptr, /*devPtrdSoftmaxOffset=*/nullptr, + /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, + /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception &e) { + return e.what(); + } catch (...) { + return "is_supported_f16_bwd: unknown failure."; + } +} + } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h index 8f79b5bb4a..5065fbe93a 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h @@ -13,40 +13,40 @@ #include +#include + #include "common/common.h" +#include "config_and_params.h" #include "transformer_engine/fused_attn.h" namespace transformer_engine { 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, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, - bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, 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); + const FusedAttnConfig &cfg, 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( - 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, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_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, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, + const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); +// check if a given configuration is supported for F16/BF16 forward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message in the form of a string. +std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); + +// check if a given configuration is supported for F16/BF16 backward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message in the form of a string. +std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); + } // namespace transformer_engine #endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index eab1ae02e6..90d24d2b36 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -17,20 +17,41 @@ using namespace transformer_engine; // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, 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, void* devPtrQ, void* devPtrK, void* devPtrV, + const FusedAttnConfig &cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, - void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, + void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, - NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + void* workspace, size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); + + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + + int64_t b = static_cast(cfg.batch_size); + int64_t h = static_cast(cfg.num_attn_heads); + int64_t hg = static_cast(cfg.num_gqa_groups); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.max_seqlen_kv); + int64_t d_qk = static_cast(cfg.head_dim_qk); + int64_t d_v = static_cast(cfg.head_dim_v); + const bool is_training = cfg.is_training; + const float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const NVTEScalingMode scaling_mode = cfg.scaling_mode; + const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; + bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -60,46 +81,8 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg); try { - FADescriptor_v1 descriptor{b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - is_training, - dropout_probability, - qkv_layout, - o_format, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Layout_NOT_SET, - qkv_scale_inv_format, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - true, - qkv_tensor_type, - o_tensor_type, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - false}; - namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, @@ -124,11 +107,11 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; + using CacheType = std::map; static thread_local CacheType sdpa_fp8_fprop_cache; // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType& cache, const FADescriptor_v1& descriptor) -> graph_and_tensors { + auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); if (it != cache.end()) { @@ -375,7 +358,7 @@ void fused_attn_fp8_fwd_impl( auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, dropout_seed, - dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); + dropout_offset] = get_graph(sdpa_fp8_fprop_cache, cache_cfg); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -422,7 +405,7 @@ void fused_attn_fp8_fwd_impl( void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -447,27 +430,53 @@ void fused_attn_fp8_fwd_impl( // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_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, - bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, - void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, - void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, - void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, - void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, - void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, - void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, - void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, - cudnn_frontend::DataType_t do_tensor_type, cudnn_frontend::DataType_t dqkv_tensor_type, - NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { + const FusedAttnConfig &cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, + void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, + void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, + void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, + void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, + void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, + void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, + void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, + void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, + void* devPtrDropoutOffset, void* workspace, size_t* workspace_size, cudaStream_t stream, + cudnnHandle_t handle) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); + + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const cudnn_frontend::DataType_t do_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.do_dtype)); + const cudnn_frontend::DataType_t dqkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); + + int64_t b = static_cast(cfg.batch_size); + int64_t h = static_cast(cfg.num_attn_heads); + int64_t hg = static_cast(cfg.num_gqa_groups); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.max_seqlen_kv); + int64_t d_qk = static_cast(cfg.head_dim_qk); + int64_t d_v = static_cast(cfg.head_dim_v); + const float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_QKV_Format do_format = cfg.do_format; + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + const NVTEScalingMode scaling_mode = cfg.scaling_mode; + const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; + const NVTE_QKV_Format do_scale_inv_format = cfg.do_scale_inv_format; + bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -500,46 +509,8 @@ void fused_attn_fp8_bwd_impl( bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg); try { - FADescriptor_v1 descriptor{b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - true, - dropout_probability, - qkv_layout, - o_format, - do_format, - dqkv_layout, - qkv_scale_inv_format, - do_scale_inv_format, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - deterministic, - qkv_tensor_type, - o_tensor_type, - do_tensor_type, - dqkv_tensor_type, - false}; - namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, @@ -585,11 +556,11 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; + using CacheType = std::map; static thread_local CacheType sdpa_fp8_bprop_cache; // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType& cache, const FADescriptor_v1& descriptor) -> graph_and_tensors { + auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); if (it != cache.end()) { @@ -987,7 +958,7 @@ void fused_attn_fp8_bwd_impl( descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); + dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, cache_cfg); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1062,7 +1033,7 @@ void fused_attn_fp8_bwd_impl( void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1090,16 +1061,18 @@ void fused_attn_fp8_bwd_impl( // fused attention FWD FP8 with separate Q, K, V 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_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_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 mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + const FusedAttnConfig &cfg, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, const Tensor* input_SoftmaxOffset, 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) { using namespace transformer_engine; + + const size_t batch = cfg.batch_size; + const size_t num_attn_heads = cfg.num_attn_heads; + const size_t max_seqlen_q = cfg.max_seqlen_q; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; void *devPtrO = nullptr, *devPtrAmaxO = nullptr, *devPtrScaleO = nullptr; @@ -1166,22 +1139,16 @@ void fused_attn_fp8_fwd( void* devPtrDropoutOffset = reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - const DType QKV_type = input_Q->data.dtype; - const DType O_type = output_O->data.dtype; size_t workspace_size = 0; NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || (qkv_format == NVTE_QKV_Format::NVTE_BHSD)) { fused_attn::fused_attn_fp8_fwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, - devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, devPtrDescaleK, - devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, - devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, - get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, - qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); + cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, + devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, + devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, + workspace->data.dptr, &workspace_size, stream, handle); } else { NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, or BHSD.\n"); } @@ -1200,20 +1167,17 @@ void fused_attn_fp8_fwd( } // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( - 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, float attn_scale, float p_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 mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, - const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, - const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, const Tensor* output_dQ, - const Tensor* output_dK, const Tensor* output_dV, Tensor* output_dSoftmaxOffset, - const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, const Tensor* rng_state, - Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + const FusedAttnConfig &cfg, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + const Tensor* input_O, const Tensor* input_dO, const Tensor* input_dO_f16, const Tensor* input_M, + const Tensor* input_S, const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, + const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, + Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, + const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; void* devPtrV = input_V->data.dptr; @@ -1229,8 +1193,8 @@ void fused_attn_fp8_bwd( devPtrDescaleK_t = input_K->columnwise_scale_inv.dptr; } - void* devPtrO = input_O->data.dptr; const DType O_type = input_O->data.dtype; + void* devPtrO = input_O->data.dptr; void* devPtrDescaleO = nullptr; if (O_type == DType::kFloat8E4M3 || O_type == DType::kFloat8E5M2) { devPtrDescaleO = input_O->scale_inv.dptr; @@ -1286,28 +1250,19 @@ void fused_attn_fp8_bwd( void* devPtrDropoutOffset = reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - const DType QKV_type = input_Q->data.dtype; - const DType dO_type = input_dO->data.dtype; - const DType dQKV_type = output_dQ->data.dtype; size_t workspace_size = 0; NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || (dqkv_format == NVTE_QKV_Format::NVTE_BHSD)) { fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, - devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, + cfg, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, - devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, - devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, - devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, - devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, - devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), - get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), - input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, - &workspace_size, stream, handle); + devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, + devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, + devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, + devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, + devPtrDropoutOffset, workspace->data.dptr, &workspace_size, stream, handle); } else { NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, or BHSD.\n"); } @@ -1324,4 +1279,56 @@ void fused_attn_fp8_bwd( return; } } + +std::string is_supported_fp8_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_fwd_impl( + cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, + /*devPtrSoftmaxOffset=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, + /*devPtrDescaleQ=*/nullptr, /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, + /*devPtrDescaleS=*/nullptr, /*devPtrScaleS=*/nullptr, /*devPtrScaleO=*/nullptr, + /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, + /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, + /*devPtrDropoutOffset=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "is_supported_fp8_fwd: unknown failure."; + } +} + +std::string is_supported_fp8_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_bwd_impl( + cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, + /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, + /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, + /*devPtrdSoftmaxOffset=*/nullptr, /*devPtrDescaleQ=*/nullptr, + /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, /*devPtrDescaleO=*/nullptr, + /*devPtrDescaledO=*/nullptr, /*devPtrDescaleS=*/nullptr, /*devPtrDescaledP=*/nullptr, + /*devPtrScaleS=*/nullptr, /*devPtrScaledP=*/nullptr, /*devPtrScaledQ=*/nullptr, + /*devPtrScaledK=*/nullptr, /*devPtrScaledV=*/nullptr, /*devPtrAmaxdP=*/nullptr, + /*devPtrAmaxdQ=*/nullptr, /*devPtrAmaxdK=*/nullptr, /*devPtrAmaxdV=*/nullptr, + /*devPtrQ_t=*/nullptr, /*devPtrK_t=*/nullptr, /*devPtrdO_f16=*/nullptr, + /*devPtrdO_t=*/nullptr, /*devPtrDescaleQ_t=*/nullptr, /*devPtrDescaleK_t=*/nullptr, + /*devPtrDescaledO_t=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, + /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, + /*devPtrDropoutOffset=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "is_supported_fp8_bwd: unknown failure."; + } +} + } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index b9660128ca..1ede20c7f1 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,35 +8,36 @@ * \brief Functions for fused attention for FP8 */ +#include + +#include "config_and_params.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V 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_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_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 mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_SoftmaxOffset, 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); // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( - 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, float attn_scale, float p_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 mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_S, - const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, const Tensor *output_dQ, - const Tensor *output_dK, const Tensor *output_dV, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_O, const Tensor *input_dO, const Tensor *input_dO_f16, const Tensor *input_M, + const Tensor *input_S, const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, + const Tensor *output_dQ, const Tensor *output_dK, const Tensor *output_dV, + Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + +// check if a given configuration is supported for FP8 forward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message in the form of a string. +std::string is_supported_fp8_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); + +// check if a given configuration is supported for FP8 backward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message in the form of a string. +std::string is_supported_fp8_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 3e628b6581..c338f1a99d 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -9,6 +9,8 @@ #include "../common.h" #include "../cudnn_utils.h" +#include "../util/cuda_runtime.h" +#include "config_and_params.h" #include "transformer_engine/fused_attn.h" #include "utils.h" @@ -633,6 +635,44 @@ __global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, i } } // namespace fused_attn + +FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg) { + FusedAttnConfig cache_cfg = cfg; + + const int64_t s_q = static_cast(cache_cfg.max_seqlen_q); + const int64_t s_kv = static_cast(cache_cfg.max_seqlen_kv); + const bool is_padding = + (cache_cfg.attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || + (cache_cfg.attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (cache_cfg.attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + const bool is_bottom_right = + (cache_cfg.attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || + (cache_cfg.attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + if (is_bottom_right && s_q == s_kv && !is_padding) { + cache_cfg.bottom_right_diagonal = false; + } + + const NVTE_QKV_Format q_format = nvte_get_q_format(cache_cfg.qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(cache_cfg.qkv_layout); + const bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + const bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + const auto cudnn_runtime_version = cudnnGetVersion(); + const int device_id = cuda::current_device(); + const int sm_arch_ = cuda::sm_arch(device_id); + + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + cache_cfg.batch_size = cache_cfg.bucketed_batch_size; + if (is_ragged_q) { + cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q; + } + if (is_ragged_kv) { + cache_cfg.max_seqlen_kv = cache_cfg.bucketed_num_tokens_kv; + } + } + + return cache_cfg; +} + } // namespace transformer_engine void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 41656062a4..9bec83e157 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -273,66 +273,6 @@ struct FADescriptor { } }; -struct FADescriptor_v1 { - std::int64_t b; - std::int64_t h; - std::int64_t hg; - std::int64_t s_q; - std::int64_t s_kv; - std::int64_t d_qk; - std::int64_t d_v; - std::int64_t num_pages_k; - std::int64_t num_pages_v; - std::int64_t page_size_k; - std::int64_t page_size_v; - std::int64_t max_pages_per_seq_k; - std::int64_t max_pages_per_seq_v; - std::int64_t bias_b; - std::int64_t bias_h; - std::int64_t bias_sq; - std::int64_t bias_skv; - float attnScale; - bool isTraining; - float dropoutProbability; - 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 mask_type; - NVTE_Softmax_Type softmax_type; - std::int64_t window_size_left; - std::int64_t window_size_right; - bool bottom_right_diagonal; - bool deterministic; - cudnn_frontend::DataType_t qkv_tensor_type; - cudnn_frontend::DataType_t o_tensor_type; - cudnn_frontend::DataType_t do_tensor_type; - cudnn_frontend::DataType_t dqkv_tensor_type; - bool return_max_logit; - - bool operator<(const FADescriptor_v1 &rhs) const { - return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k, - page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, - bias_skv, attnScale, isTraining, dropoutProbability, qkv_layout, o_format, - do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, - deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, - dqkv_tensor_type, return_max_logit) < - std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k, - rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, - rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv, - rhs.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.qkv_layout, - rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, - rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, - rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic, - rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type, rhs.return_max_logit); - } -}; - __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 41e4b136bd..f26e03c5ad 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -11,6 +11,8 @@ #ifndef TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ #define TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ +#include + #include "stdint.h" #include "transformer_engine.h" @@ -194,7 +196,116 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); +/*! \brief Opaque fused-attention configuration handle. */ +typedef void *NVTEFusedAttnConfig; + +/*! \enum NVTEFusedAttnConfigAttribute + * \brief Attribute types for ``NVTEFusedAttnConfig``. + * + * New fields may only be appended at the end; existing fields are never + * reordered, removed, or resized. + */ +enum NVTEFusedAttnConfigAttribute { + kNVTEFusedAttnConfigIsTraining = 0, + kNVTEFusedAttnConfigDeterministic, + kNVTEFusedAttnConfigCudaGraph, + kNVTEFusedAttnConfigReturnMaxLogit, + kNVTEFusedAttnConfigQKVLayout, + kNVTEFusedAttnConfigOFormat, + kNVTEFusedAttnConfigDOFormat, + kNVTEFusedAttnConfigDQKVLayout, + kNVTEFusedAttnConfigQKVScaleInvFormat, + kNVTEFusedAttnConfigDOScaleInvFormat, + kNVTEFusedAttnConfigBiasType, + kNVTEFusedAttnConfigAttnMaskType, + kNVTEFusedAttnConfigSoftmaxType, + kNVTEFusedAttnConfigScalingMode, + kNVTEFusedAttnConfigAttnScale, + kNVTEFusedAttnConfigDropout, + kNVTEFusedAttnConfigMaxSeqlenQ, + kNVTEFusedAttnConfigMaxSeqlenKV, + kNVTEFusedAttnConfigWindowSizeLeft, + kNVTEFusedAttnConfigWindowSizeRight, + kNVTEFusedAttnConfigBottomRightDiagonal, + kNVTEFusedAttnConfigQKVDtype, + kNVTEFusedAttnConfigODtype, + kNVTEFusedAttnConfigDODtype, + kNVTEFusedAttnConfigDQKVDtype, + kNVTEFusedAttnConfigBatchSize, + kNVTEFusedAttnConfigNumAttnHeads, + kNVTEFusedAttnConfigNumGqaGroups, + kNVTEFusedAttnConfigHeadDimQK, + kNVTEFusedAttnConfigHeadDimV, + kNVTEFusedAttnConfigNumPagesK, + kNVTEFusedAttnConfigNumPagesV, + kNVTEFusedAttnConfigPageSizeK, + kNVTEFusedAttnConfigPageSizeV, + kNVTEFusedAttnConfigMaxPagesPerSeqK, + kNVTEFusedAttnConfigMaxPagesPerSeqV, + kNVTEFusedAttnConfigBiasBatchSize, + kNVTEFusedAttnConfigBiasNumHeads, + kNVTEFusedAttnConfigBiasSeqlenQ, + kNVTEFusedAttnConfigBiasSeqlenKV, + kNVTEFusedAttnConfigNumTokensQ, + kNVTEFusedAttnConfigNumTokensKV, + kNVTEFusedAttnConfigBucketedBatchSize, + kNVTEFusedAttnConfigBucketedNumTokensQ, + kNVTEFusedAttnConfigBucketedNumTokensKV, + kNVTEFusedAttnConfigNumAttributes +}; + +/*! \brief Create a default-initialized fused-attention configuration. + * + * Categorical fields (layouts, formats, masks, window sizes, scaling mode) are + * set to safe NOT_SET / no-op defaults. Numeric and tensor-derived fields, + * paged-KV shape, bias broadcast shape, and direction flags default to + * zero/false; callers must set the fields relevant to their query. + * + * \return A new configuration handle. Must be destroyed with + * ``nvte_destroy_fused_attn_config()``. + */ +NVTEFusedAttnConfig nvte_create_fused_attn_config(void); + +/*! \brief Destroy a fused-attention configuration handle. */ +void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config); + +/*! \brief Query an attribute in 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); + +/*! \brief Set 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); + +/*! \brief Get fused attention backend based on input parameters. + * + * This call exercises cudnn-frontend's support checks by building (and caching) + * the cuDNN execution graph for the supported configurations. The configuration + * parameters are a superset of those of ``nvte_fused_attn_fwd`` and + * ``nvte_fused_attn_bwd`` to maintain a consistent signature between graph + * building and runtime calls. + * + * \param[in] cfg Attention configuration created with + * ``nvte_create_fused_attn_config()`` (or the C++ + * ``FusedAttnConfigWrapper``). + * \param[out] message Empty on success, otherwise a diagnostic string describing + * why the configuration was rejected. The string pointer + * refers to a per-thread buffer owned by the library and + * remains valid only until the next call to + * ``nvte_get_fused_attn_backend_v2`` on the same thread; + * callers that need to retain the message across further + * calls must copy it. Pass NULL to skip diagnostics. + * + * \return Backend able to execute this configuration, or ``NVTE_No_Backend`` if none. + */ +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, + const char **message); + /*! \brief Get fused attention backend based on input parameters. + * + * \deprecated This function has been deprecated in favor of nvte_get_fused_attn_backend_v2. * * \param[in] is_training Whether the model is in training mode. * \param[in] q_dtype The data type of Tensor Q. @@ -708,6 +819,246 @@ class AttentionShape { size_t canonical_[5] = {}; }; +/*! \class FusedAttnConfigWrapper + * \brief C++ helper for constructing an ``NVTEFusedAttnConfig``. + * + * Owns an opaque ``NVTEFusedAttnConfig`` handle created via + * ``nvte_create_fused_attn_config()``. Provides typed, chainable setters for + * every field. + */ +class FusedAttnConfigWrapper { + public: + FusedAttnConfigWrapper() : cfg_{nvte_create_fused_attn_config()} {} + + FusedAttnConfigWrapper(const FusedAttnConfigWrapper &) = delete; + FusedAttnConfigWrapper &operator=(const FusedAttnConfigWrapper &) = delete; + + FusedAttnConfigWrapper(FusedAttnConfigWrapper &&other) noexcept : cfg_{other.cfg_} { + other.cfg_ = nullptr; + } + + FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { + if (this != &other) { + nvte_destroy_fused_attn_config(cfg_); + cfg_ = other.cfg_; + other.cfg_ = nullptr; + } + return *this; + } + + ~FusedAttnConfigWrapper() { + if (cfg_ != nullptr) { + nvte_destroy_fused_attn_config(cfg_); + } + } + + operator NVTEFusedAttnConfig() const noexcept { return cfg_; } + NVTEFusedAttnConfig get() const noexcept { return cfg_; } + + FusedAttnConfigWrapper &set_is_training(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigIsTraining, &u8_val, + sizeof(u8_val)); + return *this; + } + FusedAttnConfigWrapper &set_deterministic(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDeterministic, &u8_val, + sizeof(u8_val)); + return *this; + } + FusedAttnConfigWrapper &set_cuda_graph(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigCudaGraph, &u8_val, + sizeof(u8_val)); + return *this; + } + FusedAttnConfigWrapper &set_return_max_logit(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigReturnMaxLogit, &u8_val, + sizeof(u8_val)); + return *this; + } + FusedAttnConfigWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVLayout, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigOFormat, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDOFormat, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVLayout, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVScaleInvFormat, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDOScaleInvFormat, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasType, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigAttnMaskType, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigSoftmaxType, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_scaling_mode(NVTEScalingMode val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigScalingMode, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_attn_scale(float val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigAttnScale, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_dropout(float val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDropout, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_max_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxSeqlenQ, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_max_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxSeqlenKV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeLeft, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeRight, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bottom_right_diagonal(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBottomRightDiagonal, &u8_val, + sizeof(u8_val)); + return *this; + } + FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVDtype, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigODtype, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDODtype, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_dqkv_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVDtype, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_batch_size(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBatchSize, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_attn_heads(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumAttnHeads, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_gqa_groups(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumGqaGroups, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigHeadDimQK, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_head_dim_v(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigHeadDimV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_pages_k(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumPagesK, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_pages_v(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumPagesV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_page_size_k(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigPageSizeK, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_page_size_v(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigPageSizeV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_max_pages_per_seq_k(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxPagesPerSeqK, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_max_pages_per_seq_v(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxPagesPerSeqV, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bias_batch_size(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasBatchSize, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bias_num_heads(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasNumHeads, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bias_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasSeqlenQ, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bias_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasSeqlenKV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_tokens_q(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumTokensQ, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_num_tokens_kv(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumTokensKV, &val, sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bucketed_batch_size(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBucketedBatchSize, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bucketed_num_tokens_q(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBucketedNumTokensQ, &val, + sizeof(val)); + return *this; + } + FusedAttnConfigWrapper &set_bucketed_num_tokens_kv(size_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBucketedNumTokensKV, &val, + sizeof(val)); + return *this; + } + + private: + NVTEFusedAttnConfig cfg_ = nullptr; +}; + #endif // __cplusplus #endif diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index f7ffb5ad8d..540b38143d 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -91,6 +91,13 @@ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8) \ .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend); \ + pybind11::enum_(m, "NVTEScalingMode", pybind11::module_local()) \ + .value("NVTE_DELAYED_TENSOR_SCALING", NVTEScalingMode::NVTE_DELAYED_TENSOR_SCALING) \ + .value("NVTE_MXFP8_1D_SCALING", NVTEScalingMode::NVTE_MXFP8_1D_SCALING) \ + .value("NVTE_BLOCK_SCALING_1D", NVTEScalingMode::NVTE_BLOCK_SCALING_1D) \ + .value("NVTE_BLOCK_SCALING_2D", NVTEScalingMode::NVTE_BLOCK_SCALING_2D) \ + .value("NVTE_NVFP4_1D_SCALING", NVTEScalingMode::NVTE_NVFP4_1D_SCALING) \ + .value("NVTE_INVALID_SCALING", NVTEScalingMode::NVTE_INVALID_SCALING); \ pybind11::enum_( \ m, "Float8BlockScaleTensorFormat", pybind11::module_local()) \ .value("GEMM_READY", transformer_engine::Float8BlockScaleTensorFormat::GEMM_READY) \ diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index ecca4a3871..3cb3a1dd26 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -13,6 +13,7 @@ import jax.numpy as jnp from transformer_engine_jax import NVTE_Bias_Type +from transformer_engine_jax import NVTE_Fused_Attn_Backend from transformer_engine_jax import NVTE_Mask_Type from transformer_engine_jax import NVTE_QKV_Layout from transformer_engine_jax import NVTE_QKV_Format @@ -325,6 +326,7 @@ def canonicalize_attn_mask_type(attn_mask_type: str): def is_fused_attn_kernel_available( is_training, + batch_size, q_dtype, kv_dtype, qkv_layout, @@ -339,15 +341,26 @@ def is_fused_attn_kernel_available( head_dim_qk, head_dim_v, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, + return_reason: bool = False, ): """ - To check whether the fused attention kernel is supported + To check whether the fused attention kernel is supported. + + When ``return_reason`` is ``True``, returns ``(available, message)`` where ``message`` is + the diagnostic string for the reason why the fused attention kernel is not supported (empty on success). """ window_size_tuple = (-1, -1) if window_size is None else window_size def make_helper(attn_mask_type): + bottom_right = ( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ) return tex.FusedAttnHelper( is_training, + batch_size, q_dtype, kv_dtype, qkv_layout, @@ -362,9 +375,15 @@ def make_helper(attn_mask_type): head_dim_qk, head_dim_v, window_size_tuple, + bottom_right, ) - return make_helper(attn_mask_type).is_fused_attn_kernel_available() + helper = make_helper(attn_mask_type) + if return_reason: + backend, message = helper.get_fused_attn_backend() + available = backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend + return available, message + return helper.is_fused_attn_kernel_available() def _obtain_batch_and_max_seqlen(qkv, qkv_layout): diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 489bfde997..eaa9c8769a 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -16,7 +16,12 @@ from jax.experimental.custom_partitioning import SdyShardingRule import transformer_engine_jax -from transformer_engine_jax import NVTE_Fused_Attn_Backend +from transformer_engine_jax import ( + JAXX_Scaling_Mode, + NVTE_Fused_Attn_Backend, + NVTE_QKV_Format, + NVTE_QKV_Layout, +) from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -108,6 +113,7 @@ class FusedAttnHelper: """ is_training: bool + batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -122,21 +128,44 @@ class FusedAttnHelper: head_dim_qk: int head_dim_v: int window_size: Tuple[int, int] + bottom_right_diagonal: bool + attn_scale: float = 1.0 def is_fused_attn_kernel_available(self): - """Check if there is available fused attention kernel""" - return self.get_fused_attn_backend() != NVTE_Fused_Attn_Backend.NVTE_No_Backend + """Check if there is available fused attention kernel. + + Use ``get_fused_attn_backend()`` directly to also get the diagnostic message + explaining why a configuration was rejected. + """ + backend, _ = self.get_fused_attn_backend() + return backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend def get_fused_attn_backend(self): - """Get the fused attention kernel backend""" + """Get the fused attention kernel backend. + + Returns a ``(backend, message)`` tuple. ``message`` is empty on success, otherwise a + diagnostic string describing why the configuration was rejected when backend = NVTE_No_Backend. + """ + q_type = jax_dtype_to_te_dtype(self.q_dtype) return transformer_engine_jax.get_fused_attn_backend( self.is_training, - jax_dtype_to_te_dtype(self.q_dtype), + self.batch_size, + q_type, jax_dtype_to_te_dtype(self.kv_dtype), + q_type, + q_type, + q_type, + JAXX_Scaling_Mode.NO_SCALING, self.qkv_layout.value, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, self.attn_bias_type.value, self.attn_mask_type.value, self.softmax_type.value, + self.attn_scale, self.dropout_probability, self.q_num_heads, self.kv_num_heads, @@ -146,6 +175,7 @@ def get_fused_attn_backend(self): self.head_dim_v, self.window_size[0], self.window_size[1], + self.bottom_right_diagonal, not self.is_non_deterministic_allowed(), ) @@ -335,8 +365,10 @@ def abstract( out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) # backend determines the softmax buffer shape/dtype - backend = FusedAttnHelper( + input_batch = reduce(operator.mul, batch_shape) + backend, message = FusedAttnHelper( config.is_training, + input_batch, q_dtype, k_dtype, config.qkv_layout, @@ -351,6 +383,8 @@ def abstract( q_head_dim, v_head_dim, config.window_size, + config.bottom_right_diagonal, + attn_scale=float(config.scaling_factor), ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: @@ -369,7 +403,7 @@ def abstract( ) softmax_dtype = dtypes.canonicalize_dtype(jnp.float32) else: - raise ValueError(f"Unsupported {backend=}") + raise ValueError(f"Unsupported backend: {message}") softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index b9c7c849f2..5a6790793c 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include "common/common.h" @@ -151,12 +152,16 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, +std::tuple GetFusedAttnBackend( + bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + DType do_dtype, DType dqkv_dtype, JAXX_Scaling_Mode scaling_mode, 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 mask_type, NVTE_Softmax_Type softmax_type, - float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic); + float attn_scale, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 3fd6780d6d..c88f63a5e9 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -24,18 +24,62 @@ namespace transformer_engine { namespace jax { -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, +std::tuple GetFusedAttnBackend( + bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + DType do_dtype, DType dqkv_dtype, JAXX_Scaling_Mode scaling_mode, 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 mask_type, NVTE_Softmax_Type softmax_type, - float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic) { - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); - return backend; + float attn_scale, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic) { + if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { + o_format = nvte_get_q_format(qkv_layout); + } + if (do_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { + do_format = o_format; + } + if (dqkv_layout == NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET) { + dqkv_layout = qkv_layout; + } + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + + FusedAttnConfigWrapper cfg; + cfg.set_is_training(is_training) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .set_return_max_logit(false) + .set_qkv_layout(qkv_layout) + .set_o_format(o_format) + .set_do_format(do_format) + .set_dqkv_layout(dqkv_layout) + .set_qkv_scale_inv_format(qkv_scale_inv_format) + .set_do_scale_inv_format(do_scale_inv_format) + .set_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_scaling_mode(get_nvte_scaling_mode(scaling_mode)) + .set_attn_scale(attn_scale) + .set_dropout(dropout_probability) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_qkv_dtype(static_cast(q_dtype)) + .set_o_dtype(static_cast(o_dtype)) + .set_do_dtype(static_cast(do_dtype)) + .set_dqkv_dtype(static_cast(dqkv_dtype)) + .set_batch_size(batch_size) + .set_num_attn_heads(q_attn_heads) + .set_num_gqa_groups(kv_attn_heads) + .set_head_dim_qk(qk_head_dim) + .set_head_dim_v(v_head_dim); + + const char *message = nullptr; + auto backend = nvte_get_fused_attn_backend_v2(cfg, &message); + return {backend, message != nullptr ? std::string(message) : std::string()}; } /* @@ -274,11 +318,14 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(dtype), static_cast(dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + auto [backend, _fwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, JAXX_Scaling_Mode::NO_SCALING, + qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, scaling_factor, + dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, + v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, deterministic); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -550,11 +597,14 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(dtype), static_cast(dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + auto [backend, _bwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, JAXX_Scaling_Mode::NO_SCALING, + qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, scaling_factor, + dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, + v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, deterministic); PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, softmax_aux, rng_state, bias, softmax_offset); diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index db5468afe6..bfb7b2a826 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -200,12 +200,14 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVTE_BSHD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD) .value("NVTE_T3HD", NVTE_QKV_Layout::NVTE_T3HD) .value("NVTE_THD_T2HD", NVTE_QKV_Layout::NVTE_THD_T2HD) - .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD); + .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD) + .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); pybind11::enum_(m, "NVTE_QKV_Format", pybind11::module_local()) .value("NVTE_SBHD", NVTE_QKV_Format::NVTE_SBHD) .value("NVTE_BSHD", NVTE_QKV_Format::NVTE_BSHD) - .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD); + .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD) + .value("NVTE_QKV_Format_NOT_SET", NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET); pybind11::enum_(m, "NVTE_Softmax_Type", pybind11::module_local()) .value("NVTE_VANILLA_SOFTMAX", NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 76922d2b55..6b8107ee67 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -779,6 +779,8 @@ def __call__( enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) sequence_dim = 0 if self.transpose_batch_sequence else 1 + batch_dim = 1 - sequence_dim + batch_size = query.shape[batch_dim] seqlen_q = query.shape[sequence_dim] if qkv_layout == QKVLayout.BS3HD: seqlen_kv = seqlen_q @@ -795,10 +797,11 @@ def __call__( if not enable_fused_attn: raise ValueError("score_mod requires fused attention, but NVTE_FUSED_ATTN=0.") kernel_qkv_layout = qkv_layout.to_separate() if score_mod_requested else qkv_layout - has_fused_attn_kernel = is_fused_attn_kernel_available( + has_fused_attn_kernel, fused_attn_reject_reason = is_fused_attn_kernel_available( # This needs to be fixed: TE-Jax has historically correlated training mode # with deterministic mode. not deterministic, + batch_size, input_dtype, # self._assert_dtypes enforces Q, K, V, bias to have the same dtype, so # using input_dtype as kv dtype is sufficient. @@ -815,6 +818,7 @@ def __call__( head_dim_qk, head_dim_v, self.window_size, + return_reason=True, ) if score_mod_requested and not has_fused_attn_kernel: raise ValueError( @@ -824,13 +828,11 @@ def __call__( use_fused_attn = enable_fused_attn and has_fused_attn_kernel if enable_fused_attn and not has_fused_attn_kernel: + reason = fused_attn_reject_reason or "(no diagnostic message available)" warnings.warn( - "Fused attention is not enabled because there is no available kernel.\n" - "Fall back to the unfused attention.\n" - "Please try to update the cuDNN and TE to the latest version.\n" - f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n" - f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n" - f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n" + "Falling back to the unfused attention backend as fused attention does not" + f" support:\n{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\nReason" + f" for this rejection: {reason}\n" ) dropout_rng = None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index d3ee1a2e2c..f3078c1a9c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -447,6 +447,7 @@ def __init__( softmax_scale = 1.0 / math.sqrt( kv_channels if isinstance(kv_channels, int) else kv_channels[0] ) + self.softmax_scale = softmax_scale self.deterministic = ( not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) @@ -1653,6 +1654,7 @@ def forward( return_max_logit=self.return_max_logit, cuda_graph=is_graph_capturing(), num_splits=num_splits, + softmax_scale=self.softmax_scale, fp8_output=fp8_output, checkpoint_core_attention=checkpoint_core_attention, has_score_mod=score_mod is not None, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be94a6fa1..68bb0b199e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -23,6 +23,7 @@ import transformer_engine as te from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, + QKVFormat, AttnBiasType, AttnMaskType, SoftmaxType, @@ -259,6 +260,9 @@ class AttentionParams: Whether support for cuda graph capture is needed or not. num_splits : int, default = 1 The number of kernels to split attention to. + softmax_scale : float, default = 1.0 + Pre-softmax attention scale. Plumbed through to the cuDNN graph cache key so that the + backend probe builds the same execution graph the runtime call later reuses. fp8_output : bool, default = False Whether output is requested in FP8. checkpoint_core_attention : bool, default = False @@ -300,6 +304,7 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 + softmax_scale: float = 1.0 fp8_output: bool = False checkpoint_core_attention: bool = False has_score_mod: bool = False @@ -382,6 +387,7 @@ def get_attention_backend( return_max_logit = attention_params.return_max_logit cuda_graph = attention_params.cuda_graph num_splits = attention_params.num_splits + softmax_scale = attention_params.softmax_scale fp8_output = attention_params.fp8_output checkpoint_core_attention = attention_params.checkpoint_core_attention has_score_mod = attention_params.has_score_mod @@ -1360,17 +1366,56 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # on the C++ side, so pass it straight to the pybind function. q_type = TE_DType[qkv_dtype] kv_type = q_type + o_type = q_type + do_type = q_type + dqkv_type = q_type + scaling_mode = tex.NVTEScalingMode.NVTE_INVALID_SCALING + qkv_scale_inv_format = None + do_scale_inv_format = None if fp8 and fp8_meta["recipe"].fp8_dpa: - q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True) + recipe = fp8_meta["recipe"] + q_type = get_fp8_te_dtype(recipe, fprop_tensor=True) kv_type = q_type - fused_attention_backend = tex.get_fused_attn_backend( + cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" + if recipe.mxfp8(): + scaling_mode = tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + o_type = TE_DType[torch.bfloat16] + do_type = TE_DType[torch.bfloat16] + dqkv_type = TE_DType[torch.bfloat16] + qkv_scale_inv_format = "bhsd" + do_scale_inv_format = "bhsd" + elif recipe.float8_current_scaling() and cs_o_in_f16: + scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + o_type = TE_DType[torch.bfloat16] + do_type = TE_DType[torch.bfloat16] + dqkv_type = TE_DType[torch.bfloat16] + else: + scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + o_type = q_type + do_type = o_type + dqkv_type = q_type + o_format = q_format + do_format = o_format + dqkv_layout = qkv_layout + fused_attention_backend, reject_message = tex.get_fused_attn_backend( is_training, + batch_size, q_type, kv_type, + o_type, + do_type, + dqkv_type, + scaling_mode, QKVLayout[qkv_layout], + QKVFormat[o_format], + QKVFormat[do_format], + QKVLayout[dqkv_layout], + QKVFormat[qkv_scale_inv_format], + QKVFormat[do_scale_inv_format], AttnBiasType[fu_core_attention_bias_type], AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], + softmax_scale, attention_dropout, num_heads, num_gqa_groups, @@ -1380,12 +1425,16 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt head_dim_v, window_size[0], window_size[1], + bottom_right_diagonal, return_max_logit, cuda_graph, deterministic, ) if fused_attention_backend == FusedAttnBackend["No_Backend"]: - logger.debug("Disabling FusedAttention as no backend supports the provided input") + logger.debug( + "Disabling FusedAttention: %s", + reject_message, + ) use_fused_attention = False fused_attention_backend = None elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..982fbdb169 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -86,12 +86,18 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T * Attention **************************************************************************************************/ -NVTE_Fused_Attn_Backend get_fused_attn_backend( - bool is_training, const DType q_dtype, const DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float p_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); +// Returns (backend, reason). `reason` is empty on success, otherwise a diagnostic string +// describing why the configuration was rejected when backend = NVTE_No_Backend. +std::tuple get_fused_attn_backend( + bool is_training, size_t batch_size, const DType q_dtype, const DType kv_dtype, + const DType o_dtype, const DType do_dtype, const DType dqkv_dtype, NVTEScalingMode scaling_mode, + 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, float attn_scale, float p_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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic); std::vector fused_attn_fwd( size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index eb8813d4a0..706eb630e0 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -40,18 +40,53 @@ void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &s namespace transformer_engine::pytorch { // get the fused attention backend -NVTE_Fused_Attn_Backend get_fused_attn_backend( - bool is_training, const DType q_dtype, const DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float p_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) { - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, - bias_type, attn_mask_type, softmax_type, p_dropout, num_attn_heads, num_gqa_groups, - max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, - return_max_logit, cuda_graph, deterministic); - return fused_attention_backend; +std::tuple get_fused_attn_backend( + bool is_training, size_t batch_size, const DType q_dtype, const DType kv_dtype, + const DType o_dtype, const DType do_dtype, const DType dqkv_dtype, NVTEScalingMode scaling_mode, + 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, float attn_scale, float p_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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic) { + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + + FusedAttnConfigWrapper cfg; + cfg.set_is_training(is_training) + .set_deterministic(deterministic) + .set_cuda_graph(cuda_graph) + .set_return_max_logit(return_max_logit) + .set_qkv_layout(qkv_layout) + .set_o_format(o_format) + .set_do_format(do_format) + .set_dqkv_layout(dqkv_layout) + .set_qkv_scale_inv_format(qkv_scale_inv_format) + .set_do_scale_inv_format(do_scale_inv_format) + .set_bias_type(bias_type) + .set_attn_mask_type(attn_mask_type) + .set_softmax_type(softmax_type) + .set_scaling_mode(scaling_mode) + .set_attn_scale(attn_scale) + .set_dropout(p_dropout) + .set_max_seqlen_q(max_seqlen_q) + .set_max_seqlen_kv(max_seqlen_kv) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_qkv_dtype(static_cast(q_dtype)) + .set_o_dtype(static_cast(o_dtype)) + .set_do_dtype(static_cast(do_dtype)) + .set_dqkv_dtype(static_cast(dqkv_dtype)) + .set_batch_size(batch_size) + .set_num_attn_heads(num_attn_heads) + .set_num_gqa_groups(num_gqa_groups) + .set_head_dim_qk(head_dim_qk) + .set_head_dim_v(head_dim_v); + + const char *message = nullptr; + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2(cfg, &message); + return {fused_attention_backend, message != nullptr ? std::string(message) : std::string()}; } // helper function for S and dP quantizers